修正了一点typo和公式未正常显示的问题

Author: Jorwnpay

chapter2/README.md

@@ -103,7 +103,7 @@ if __name__ == "__main__":
 
 获得注意力分数后，我们还需要进行softmax归一化，softmax的公式如下。softmax归一化主要是为了保证每一个Query对所有Key的分数的总和为1.0，这样当我们将注意力分数作为加权系数乘Value求和后输出的向量整体分布稳定。如图1的例子，如果不做归一化，Query2女生与Key3小丽(女)和Key3小花(女)均匹配成功（即系数为1），那女生的平均身高就变成了166 x 1.0 + 170 x 1.0 = 336cm。如果说把匹配成功的分数改成0.5不就行了，这样的话Query1小明的身高不就是175 x 0.5 = 87.5cm了，所以还是需要归一化来决定什么时候系数时1.0，什么时候是0.5，这就是引入softmax非线性层的原因。
 
-$softmax(Z)_i = \frac{e^{z_i}}{\sum_{k=i}^n e^{z_k}}, Z=(z_0, z_1, ..., z_n)$
+ $ softmax(Z)_i = \frac{e^{z_i}}{\sum_{k=i}^n e^{z_k}}, Z=(z_0, z_1, ..., z_n) $ 
 
 而后归一化的注意力权重attn_weights就会乘以value得到注意力计算的输出。我们一般还会再添加一层线性层(Output Projection)，作为注意力模块的最终输出。线性层主要就是增加模型的表达能力并且存储知识，这就没什么可以解释的了。
 
@@ -119,7 +119,7 @@ $softmax(Z)_i = \frac{e^{z_i}}{\sum_{k=i}^n e^{z_k}}, Z=(z_0, z_1, ..., z_n)$
 
 而我图中多头注意力的呈现方式与大多数Transformer论文的配图略有不同，我认为这样更便于理解，而且数学上也完全是等价的。我们完全把Attention分成了独立的N份，每份的中间为度从hidden_dim降到了hidden_dim / N，这里的N就是头数(num_head)，这样保证了模型参数量不变的情况下，模型变成独立参数的N份，这样就实现Ensemble，Ensemble是神经网络中一种常见的提升模型能力的方式，或者说奇技淫巧(trick)，相当于三个臭皮匠胜过诸葛亮，三个独立的小模型的结果求平均会比一个大模型更好。因为这可以防止一个模型不小心陷入局部次优解。而在注意力中，多头还引入了一个新的意义，就是不同的头可以独立的关注到不同区域，也就是每个头softmax后的注意力图都不一样，增加了模型的学习和表达能力。
 
-注意：多头注意力的实现并不需要真的维护N个子网络，而是通过reshape将hidden_dim维度拆分成两个维度num_head和head_dim维度即可，而且最终先concate再过一个大的输出线性层和过N个小的输出线性层在求和在数学上也是等价的。具体实现代码如下：
+注意：多头注意力的实现并不需要真的维护N个子网络，而是通过reshape将hidden_dim维度拆分成两个维度num_head和head_dim维度即可，而且最终先concate再过一个大的输出线性层和过N个小的输出线性层再求和在数学上也是等价的。具体实现代码如下：
 
 ```
 class MultiHeadCausalSelfAttention(nn.Module):
@@ -223,7 +223,7 @@ if __name__ == "__main__":
 ```
 
 ## 三. KV Cache缓存机制
-假设我们有两个token在某一层的输入特征(hidden states): $[x_1, x_2]$, 经过三个线性层，我们得到他们分别的QKV: $[q_1, q_2], [k_1, k_2], [v_1, v_2]$, Attention后对应的输出为：$[y_1, y_2]$。在因果掩码(Causal Mask)下，Attention的计算公式应该如下，其中OutLinear只是对输出做映射的线性层：
+假设我们有两个token在某一层的输入特征(hidden states): $[x_1, x_2]$, 经过三个线性层，我们得到他们分别的QKV: $[q_1, q_2], [k_1, k_2], [v_1, v_2]$, Attention后对应的输出为： $[y_1, y_2]$ 。在因果掩码(Causal Mask)下，Attention的计算公式应该如下，其中OutLinear只是对输出做映射的线性层：
 
 $y_1 = OutLinear( softmax(q_1 @ [k_1]^T) @ [v_1] )$
 
@@ -237,11 +237,11 @@ $y_2 = OutLinear( softmax(q_2 @ [k_1,k_2]^T) @ [v_1,v_2] )$
 
 $y_3 = OutLinear( softmax(q_3 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$
 
-这时候我们会发现，当推理第三个token时，前两个token的计算完全没有变，也就是说$[y_1, y_2]$的结果和推理第二个token时一摸一样。也许有人要说大模型除了注意力还有前馈网络FFN呀，注意，FFN的计算中每个token都是独立的，不会相互影响，只要attention的输出结果一样，后续FFN的输出结果必然一样。这时候我们就会发现，一旦有了因果掩码，当推理第N个token的时候，再去计算前0到N-1个token就完全是浪费了，因为计算结果完全一样。我们再看如果推理第N个token的时候只计算$x_N$本身，他的$q_N,k_N,v_N$都可以直接通过$x_N$过线性层得到，其需要额外进入的就只有$[k_1,k_2,...,k_{N-1}],[v_1,v_2,...,v_{N-1}]$。而这些Key和Value完全没必要重走一边线性层的计算，因为推理第N-1个token的时候已经得到过一次了，所以当时只要存下来这是再读取拿来用就好了，这也就是KV Cache，用存储和加载的带宽增加换计算节省的一种优化方式。
+这时候我们会发现，当推理第三个token时，前两个token的计算完全没有变，也就是说 $[y_1, y_2]$ 的结果和推理第二个token时一摸一样。也许有人要说大模型除了注意力还有前馈网络FFN呀，注意，FFN的计算中每个token都是独立的，不会相互影响，只要attention的输出结果一样，后续FFN的输出结果必然一样。这时候我们就会发现，一旦有了因果掩码，当推理第N个token的时候，再去计算前0到N-1个token就完全是浪费了，因为计算结果完全一样。我们再看如果推理第N个token的时候只计算 $x_N$ 本身，他的 $q_N,k_N,v_N$ 都可以直接通过 $x_N$ 过线性层得到，其需要额外进入的就只有 $[k_1,k_2,...,k_{N-1}],[v_1,v_2,...,v_{N-1}]$ 。而这些Key和Value完全没必要重走一边线性层的计算，因为推理第N-1个token的时候已经得到过一次了，所以当时只要存下来这是再读取拿来用就好了，这也就是KV Cache，用存储和加载的带宽增加换计算节省的一种优化方式。
 
 ### 扩展知识1：为什么生成式语言模型一定要因果注意力（causal attention）？
 
-如果没有因果掩码，如下面公式所示，一旦推理到第3个token的时候，第1和第2个token可以看到后面的token，那么$[y_1,y_2,y_3]$三个值都会更新，这时候就需要重算所有的token了。以此类推，推理到第N个token的时候，前N-1个token还得全部输入重新计算attention。这就没法用KV Cache了，因为一旦attention的输出改变了，前N-1个token的FFN也需要重新计算。
+如果没有因果掩码，如下面公式所示，一旦推理到第3个token的时候，第1和第2个token可以看到后面的token，那么 $[y_1,y_2,y_3]$ 三个值都会更新，这时候就需要重算所有的token了。以此类推，推理到第N个token的时候，前N-1个token还得全部输入重新计算attention。这就没法用KV Cache了，因为一旦attention的输出改变了，前N-1个token的FFN也需要重新计算。
 
 $y_1 = OutLinear( softmax(q_1 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$
 
@@ -253,19 +253,19 @@ $y_3 = OutLinear( softmax(q_3 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$
 
 ### 扩展知识2：KV-Cache的优化能省下多少计算力？
 
-我们来估算下计算复杂度，我们假设特征的维度d是常数，仅考虑推理token数N的变化，那么注意力的计算复杂度其实就是来自于QKV计算中的两个矩阵乘，也就是$O(N^2)$，因为线性层的复杂度此时只有$O(N)$。而一旦引入了KV Cache，我们QKV的计算复杂度是要输入一个query，所以复杂度就缩减到了$O(N)$。一个KV-Cache就可以将计算复杂度从$O(N^2)$降低到$O(N)$还能保证数学等价性，简直血赚。
+我们来估算下计算复杂度，我们假设特征的维度d是常数，仅考虑推理token数N的变化，那么注意力的计算复杂度其实就是来自于QKV计算中的两个矩阵乘，也就是 $O(N^2)$ ，因为线性层的复杂度此时只有 $O(N)$ 。而一旦引入了KV Cache，我们QKV的计算复杂度是要输入一个query，所以复杂度就缩减到了 $O(N)$ 。一个KV-Cache就可以将计算复杂度从$O(N^2)$降低到$O(N)$还能保证数学等价性，简直血赚。
 
 ### 扩展知识3：如何简单有效地加速首轮(Prefilling)问答响应时间？
 
-首轮Prefilling回答速度之所以会慢，是因为首轮需要一次性把问题的prompt全部输入，同时推理所有token还是$O(N^2)$的复杂度，后续增量decoding每次只要推理一个token就快很多只要$O(N)$的复杂度。
+首轮Prefilling回答速度之所以会慢，是因为首轮需要一次性把问题的prompt全部输入，同时推理所有token还是 $O(N^2)$ 的复杂度，后续增量decoding每次只要推理一个token就快很多只要 $O(N)$ 的复杂度。
 
 优化首轮回答响应速度是工业界一个非常重要的研究课题，尤其是超长文本输入的时候，比如把整个几十万字的一本书全部作为输入来询问关于这本书的问题时。我这里提供一个简单的思路，也是我2023年末参加中国移动全球合作伙伴大会演示时想到的一个非常简单的技巧，就是把系统prompt提前计算好存成KV-Cache。因为大模型使用时除了问题本身，我们往往还会增加一个固定的系统prompt在问题前面，这部分的token是不会变的，所以我们完全可以提前离线计算好。这个小技巧让我们当时首轮平均响应速度从7秒降低到了3-4秒（具体提升比例受实际问题长短影响）。所以说KV-Cache真是个妙不可言的好东西。我2023年第一次理解了KV-Cache的原理时，深深的感受到了什么叫工程的美感，人类智慧的结晶。
 
 ### 扩展知识4：pytorch的动态长度推理怎么转换为需要静态张量形状的ONNX推理格式？
 
 我们2023年部署端侧大语言模型参展时另外遇到的一个问题就是动态转静态的问题。我们当时的安卓部署平台仅仅支持老版本的ONNX(一种工业界AI模型的存储格式，包含了模型结构与权重，可以直接拿去运行推理)，不支持动态轴，因为我们项目周期只有不足4个月，我们没有时间和人力去修改部署平台底层，因此我想到了一个取巧的办法，通过一定的冗余来将动态推理转换为静态。
 
-首先ONNX的静态推理意味着，模型的所有计算都必须有固定的张量形状(tensor shape)，可在大语言模型的推理中token数的N是变化的呀，因此我只能将token数的维度固定为2048，永远计算2048个token。好在有KV-Cache，计算复杂度只是线性提升，不是平方提升。然后实际推理第k个token的时候$(k < N)$，把k+1到N个token的注意力部分给mask掉，也就是赋予-inf的值，让其softmax后为0，防止其参与计算。
+首先ONNX的静态推理意味着，模型的所有计算都必须有固定的张量形状(tensor shape)，可在大语言模型的推理中token数的N是变化的呀，因此我只能将token数的维度固定为2048，永远计算2048个token。好在有KV-Cache，计算复杂度只是线性提升，不是平方提升。然后实际推理第k个token的时候 $(k < N)$ ，把k+1到N个token的注意力部分给mask掉，也就是赋予-inf的值，让其softmax后为0，防止其参与计算。
 
 这个部署策略我们后来也看到其他公司和其他人使用过，但我觉得这毕竟是无奈之举，会有计算浪费，还希望各个公司早日把推理库的底层改成支持动态轴(有一些张量维度的长度可以变)，2025年了，别再拖着了。