大语言模型的核心算法——简要解析
Transformer架构的数学本质与演进
自注意力机制的核心原理
Transformer架构的灵魂在于自注意力机制,它允许模型在处理序列中的每个元素时,动态地关注序列中的所有其他位置。从数学角度看,自注意力的计算过程可以表达为:
$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
当输入序列$X \in \mathbb{R}^{n \times d_{model}}$进入自注意力层时,首先通过三个不同的线性变换生成查询(Query)、键(Key)和值(Value)矩阵:
$$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$$
其中$W^Q, W^K, W^V \in \mathbb{R}^{d_{model} \times d_k}$是可学习的参数矩阵。查询矩阵$Q$代表当前位置想要获取的信息,键矩阵$K$代表每个位置能够提供的信息特征,而值矩阵$V$则是实际要传递的信息内容。通过计算$Q$和$K$的点积,模型得到了一个注意力分数矩阵,表示每个位置对其他所有位置的关注程度。
这里的缩放因子$\sqrt{d_k}$起着关键作用。随着维度$d_k$的增加,点积的方差会线性增长,可能导致softmax函数进入饱和区域,梯度变得极小。具体来说,假设$q$和$k$的分量是独立同分布的随机变量,均值为0,方差为1,则它们点积的方差为$d_k$。通过除以$\sqrt{d_k}$进行缩放,确保了注意力分数保持在合理的范围内,维持了训练的稳定性。
多头注意力(Multi-Head Attention)进一步扩展了这一机制。通过并行运行多个注意力头,每个头关注不同的表示子空间,模型能够同时捕获多种类型的依赖关系:
$$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,...,\text{head}_h)W^O$$
其中每个注意力头计算为:
$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$
这里$Wi^Q \in \mathbb{R}^{d{model} \times d_k}$,$Wi^K \in \mathbb{R}^{d{model} \times d_k}$,$Wi^V \in \mathbb{R}^{d{model} \times d_v}$,而$W^O \in \mathbb{R}^{hdv \times d{model}}$将所有头的输出整合。这种设计让模型能够从多个角度理解输入序列的结构和语义。
位置编码的革命性演进
原始Transformer使用正弦位置编码来注入序列的顺序信息:
$$PE_{(pos,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
$$PE_{(pos,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
这种编码方式虽然简单有效,但在处理超长序列时存在局限性。
2024年,旋转位置嵌入(RoPE)已经成为事实标准,被LLaMA、GPT-NeoX等主流模型广泛采用。RoPE的核心创新在于通过旋转矩阵编码绝对位置信息,同时自然地融入相对位置依赖。对于位置$m$的token embedding $\mathbf{x}_m$,RoPE应用旋转变换:
$$f(\mathbf{x}_m, m) = \mathbf{R}_m \mathbf{x}_m$$
其中旋转矩阵$\mathbf{R}_m$定义为:
$$\mathbf{R}_m = \begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \cdots \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \cdots \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \cdots \\ 0 & 0 & \sin m\theta_1 & \cos m\theta_1 & \cdots \\ \vdots & \vdots & \vdots & \vdots & \ddots \end{pmatrix}$$
频率参数$\theta_i = 10000^{-2i/d}$。RoPE的优雅之处在于,当计算两个位置$m$和$n$的注意力分数时:
$$\langle f(\mathbf{q}, m), f(\mathbf{k}, n) \rangle = \mathbf{q}^T \mathbf{R}_m^T \mathbf{R}_n \mathbf{k} = \mathbf{q}^T \mathbf{R}_{n-m} \mathbf{k}$$
点积结果自然包含了相对位置信息$(n-m)$。这使得模型能够处理训练时未见过的序列长度,实现了真正的长度外推能力。
相比之下,ALiBi(Attention with Linear Biases)采用了更加直接的方法,通过在注意力分数中添加线性偏置来编码位置信息:
$$\text{Attention}_{ALiBi}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} - m \cdot |i-j|\right)V$$
其中$m$是头特定的斜率参数,$|i-j|$是位置差的绝对值。
Flash Attention与硬件协同优化
2024年7月发布的Flash Attention v3代表了算法与硬件协同设计的巅峰。标准注意力计算的时间复杂度为$O(N^2d)$,空间复杂度为$O(N^2)$,其中$N$是序列长度,$d$是特征维度。Flash Attention通过分块计算和内存优化,将空间复杂度降低到$O(N)$。
其核心思想是将注意力矩阵分块计算,避免存储完整的$N \times N$注意力矩阵。对于输入序列,Flash Attention将其分为大小为$B_r \times B_c$的块,并使用以下更新规则:
$$m_{ij}^{new} = \max(m_{ij}^{old}, \tilde{m}_{ij})$$
$$l_{ij}^{new} = e^{m_{ij}^{old} - m_{ij}^{new}} l_{ij}^{old} + e^{\tilde{m}_{ij} - m_{ij}^{new}} \tilde{l}_{ij}$$
$$O_{ij}^{new} = \frac{l_{ij}^{old}}{l_{ij}^{new}} e^{m_{ij}^{old} - m_{ij}^{new}} O_{ij}^{old} + \frac{\tilde{l}_{ij}}{l_{ij}^{new}} e^{\tilde{m}_{ij} - m_{ij}^{new}} \tilde{O}_{ij}$$
其中$m$是最大值,$l$是指数和,$O$是输出。这种增量更新方式允许逐块处理,显著减少了内存访问。
在FP16精度下,Flash Attention v3达到了740 TFLOPs/s的性能,相当于H100理论峰值的75%利用率。它还支持FP8低精度计算,在保持数值稳定性的同时达到了1.2 PFLOPs/s的性能。
分组查询注意力(GQA)作为另一项重要创新,通过将查询头分组共享键值对,在效率与质量之间找到了完美平衡。假设有$H$个查询头和$G$个键值组($G < H$),则第$i$个查询头使用第$\lfloor i \cdot G/H \rfloor$组的键值对:
$$\text{GQA}_i = \text{Attention}(Q_i, K_{\lfloor i \cdot G/H \rfloor}, V_{\lfloor i \cdot G/H \rfloor})$$
以Llama 2 70B为例,使用8个组可以将KV缓存减少87.5%,而模型性能几乎不受影响。
主流模型的架构创新与技术突破
GPT-4的稀疏专家混合架构
根据已知信息,GPT-4采用了Mixture of Experts(MoE)架构,总参数约1.8万亿,包含16个专家网络,每个约1110亿参数。关键的是,每次前向传播仅激活2个专家,实际使用约2200亿参数。
MoE架构的核心在于路由机制。对于输入token $x$,路由网络计算每个专家的选择概率:
$$G(x) = \text{TopK}(\text{softmax}(W_g \cdot x + \epsilon), k)$$
其中$W_g$是路由权重,$\epsilon$是噪声项用于探索,$\text{TopK}$选择概率最高的$k$个专家。最终输出为:
$$y = \sum_{i=1}^k G_i(x) \cdot E_i(x)$$
这里$E_i(x)$是第$i$个专家的输出,$G_i(x)$是对应的门控权重。
负载均衡是MoE训练的关键挑战。DeepSeek-V3通过auxiliary-loss-free策略优化了这一过程,定义负载均衡损失为:
$$\mathcal{L}_{balance} = \alpha \cdot \sum_{i=1}^N \left( f_i - \frac{1}{N} \right)^2$$
其中$f_i$是第$i$个专家的平均激活频率,$N$是专家总数,$\alpha$是平衡系数。
Claude的Constitutional AI训练范式
Anthropic的Constitutional AI代表了一种全新的模型对齐方法。不同于传统的RLHF需要大量人工标注,Constitutional AI通过两个阶段实现自监督学习。
在第一阶段,模型生成响应后,基于预定义的宪法原则进行自我批评和修订。这个过程可以形式化为条件生成:$$p\!\left(y_{\text{revised}} \mid x,\, y_{\text{initial}},\, \mathcal{C}\right) = \prod_{t=1}^{T} p\!\left(y_t \mid y_{
其中 $\mathcal{C}$ 表示宪法原则集合,$y{\text{initial}}$ 是初始响应,$y{\text{revised}}$ 是修订后的响应。
在第二阶段,使用AI反馈(RLAIF)训练偏好模型。偏好模型的目标是学习条件分布:
$$P(y_1 \succ y_2 | x, \mathcal{C}) = \sigma(r_\theta(x, y_1, \mathcal{C}) - r_\theta(x, y_2, \mathcal{C}))$$
其中$r_\theta$是参数化的奖励函数,$\sigma$是sigmoid函数,$y_1 \succ y_2$表示$y_1$优于$y_2$。
LLaMA 3.1的密集架构优化
Meta的LLaMA 3.1 405B虽然采用密集架构而非MoE,但通过一系列技术创新实现了卓越性能。模型使用RMSNorm替代LayerNorm:
$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma$$
其中:
$$\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2}$$
这种简化不仅提高了训练稳定性,还减少了约7%的计算开销。
SwiGLU激活函数是另一项关键创新:
$$\text{SwiGLU}(x) = \text{Swish}(xW_1) \otimes (xW_2)$$
其中Swish函数定义为:
$$\text{Swish}(x) = x \cdot \sigma(\beta x) = \frac{x}{1 + e^{-\beta x}}$$
当$\beta = 1$时,SwiGLU提供了比ReLU更平滑的梯度,其导数为:
$$\frac{d\text{Swish}}{dx} = \sigma(\beta x) + \beta x \cdot \sigma(\beta x)(1 - \sigma(\beta x))$$
这种平滑性有助于模型学习更复杂的表示,特别是在深层网络中。
训练技术的范式转变
从RLHF到DPO的演进
强化学习从人类反馈(RLHF)曾是大模型对齐的黄金标准。RLHF的目标是最大化期望奖励,同时限制与参考策略的偏离:
$$\max_{\pi_\theta} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_\theta(y|x)} [R_\phi(x,y)] - \beta \mathbb{D}_{KL}[\pi_\theta(y|x) || \pi_{ref}(y|x)]$$
通过拉格朗日乘数法,可以得到最优策略的闭式解:
$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{ref}(y|x) \exp\left(\frac{R(x,y)}{\beta}\right)$$
其中$Z(x)$是配分函数:
$$Z(x) = \sum_y \pi_{ref}(y|x) \exp\left(\frac{R(x,y)}{\beta}\right)$$
直接偏好优化(DPO)的关键洞察是,可以从最优策略的形式反推出奖励函数:
$$R(x,y) = \beta \log \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} + \beta \log Z(x)$$
基于Bradley-Terry模型,人类偏好的概率可以表示为:
$$P(y_w \succ y_l | x) = \sigma(R(x, y_w) - R(x, y_l))$$
代入奖励函数的表达式并化简(注意$Z(x)$项会抵消),得到DPO的损失函数:
$$\mathcal{L}_{DPO}(\pi_\theta; \pi_{ref}) = -\mathbb{E}_{(x,y_w,y_l) \sim \mathcal{D}} \left[ \log \sigma \left( \beta \log \frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)} \right) \right]$$
这种方法消除了奖励模型训练阶段,直接优化策略参数,计算效率提升11%,内存使用减少11%。
长上下文处理的技术突破
Ring Attention通过分布式计算突破了序列长度限制。假设有$P$个设备,序列长度为$N$,每个设备处理$N/P$长度的子序列。注意力计算分解为:
$$\text{Attention}_{global} = \bigcup_{p=1}^P \text{LocalAttention}_p + \text{CrossAttention}_{p \rightarrow (p+1) \mod P}$$
在环形拓扑中,设备$p$在第$t$轮接收来自设备$(p-t) \mod P$的键值对,计算交叉注意力:
$$A_{p,t} = \text{softmax}\left(\frac{Q_p K_{(p-t) \mod P}^T}{\sqrt{d_k}}\right) V_{(p-t) \mod P}$$
最终通过$P$轮通信完成全局注意力计算,通信复杂度为$O(N)$而非$O(N^2)$。
StreamingLLM基于"注意力汇聚"现象,保留初始$k_{sink}$个tokens作为锚点,配合大小为$w$的滑动窗口:
$$\text{KVCache} = \{\text{KV}_{1:k_{sink}}\} \cup \{\text{KV}_{t-w+1:t}\}$$
注意力计算时,初始tokens获得的注意力权重满足:
$$\sum_{i=1}^{k_{sink}} \alpha_i \approx 0.2 \sim 0.3$$
即使这些tokens在语义上并不重要。这种设计使得模型能够稳定处理400万tokens的序列,速度比基线方法快22.2倍。
附录:核心算法
A. 自注意力机制的梯度推导与优化
考虑自注意力的前向传播:
$$A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right), \quad Y = AV$$
我们需要推导损失函数$\mathcal{L}$关于$Q$、$K$、$V$的梯度。首先定义中间变量:
$$S = \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{n \times n}$$
对于softmax函数,其Jacobian矩阵为:
$$\frac{\partial A_{ij}}{\partial S_{ik}} = A_{ij}(\delta_{jk} - A_{ik})$$
其中$\delta_{jk}$是Kronecker delta函数。
通过链式法则,损失关于$S$的梯度为:
$$\frac{\partial \mathcal{L}}{\partial S} = \frac{\partial \mathcal{L}}{\partial Y} \cdot \frac{\partial Y}{\partial A} \cdot \frac{\partial A}{\partial S}$$
展开计算:
$$\frac{\partial \mathcal{L}}{\partial S_{ij}} = \sum_{k} \frac{\partial \mathcal{L}}{\partial Y_{ik}} V_{jk} A_{ij} - \sum_{k,l} \frac{\partial \mathcal{L}}{\partial Y_{ik}} V_{lk} A_{il} A_{ij}$$
简化为矩阵形式:
$$\frac{\partial \mathcal{L}}{\partial S} = A \odot \left(\frac{\partial \mathcal{L}}{\partial Y} V^T - \text{diag}\left(A \cdot \left(\frac{\partial \mathcal{L}}{\partial Y} V^T\right)\right)\right)$$
进而得到关于$Q$和$K$的梯度:
$$\frac{\partial \mathcal{L}}{\partial Q} = \frac{1}{\sqrt{d_k}} \frac{\partial \mathcal{L}}{\partial S} K$$
$$\frac{\partial \mathcal{L}}{\partial K} = \frac{1}{\sqrt{d_k}} \frac{\partial \mathcal{L}}{\partial S}^T Q$$
值矩阵的梯度更直接:
$$\frac{\partial \mathcal{L}}{\partial V} = A^T \frac{\partial \mathcal{L}}{\partial Y}$$
B. 旋转位置编码(RoPE)的数学性质
定理1:RoPE保持内积的相对位置依赖性。
证明:对于两个位置$m$和$n$的向量$\mathbf{q}$和$\mathbf{k}$,应用RoPE后的内积为:
$$\langle f_q(\mathbf{q}, m), f_k(\mathbf{k}, n) \rangle = \mathbf{q}^T \mathbf{R}_\Theta^T(m) \mathbf{R}_\Theta(n) \mathbf{k}$$
由于旋转矩阵的性质$\mathbf{R}\Theta^T(m) = \mathbf{R}\Theta(-m)$,我们有:
$$\mathbf{R}_\Theta^T(m) \mathbf{R}_\Theta(n) = \mathbf{R}_\Theta(n-m)$$
将向量$\mathbf{q}$和$\mathbf{k}$分解为2维子空间:
$$\mathbf{q} = [\mathbf{q}_0, \mathbf{q}_1, ..., \mathbf{q}_{d/2-1}]^T$$
对于每个2维子空间$i$,内积计算为:
$$\langle f_q(\mathbf{q}_i, m), f_k(\mathbf{k}_i, n) \rangle = \mathbf{q}_i^T \begin{pmatrix} \cos((n-m)\theta_i) & -\sin((n-m)\theta_i) \\ \sin((n-m)\theta_i) & \cos((n-m)\theta_i) \end{pmatrix} \mathbf{k}_i$$
展开得:
$$= q_{i,0}k_{i,0}\cos((n-m)\theta_i) + q_{i,0}k_{i,1}\sin((n-m)\theta_i) - q_{i,1}k_{i,0}\sin((n-m)\theta_i) + q_{i,1}k_{i,1}\cos((n-m)\theta_i)$$
$$= (q_{i,0}k_{i,0} + q_{i,1}k_{i,1})\cos((n-m)\theta_i) + (q_{i,0}k_{i,1} - q_{i,1}k_{i,0})\sin((n-m)\theta_i)$$
这表明内积仅依赖于相对位置$(n-m)$,证毕。
C. DPO损失函数的变分推导
从RLHF的目标函数出发:
$$J(\pi_\theta) = \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_\theta} [R(x,y)] - \beta \mathbb{D}_{KL}[\pi_\theta || \pi_{ref}]$$
KL散度可以展开为:
$$\mathbb{D}_{KL}[\pi_\theta || \pi_{ref}] = \mathbb{E}_{y \sim \pi_\theta} \left[\log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)}\right]$$
将目标函数重写为:
$$J(\pi_\theta) = \mathbb{E}_{y \sim \pi_\theta} \left[R(x,y) - \beta \log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)}\right]$$
对$\pi_\theta$求变分导数并令其为零:
$$\frac{\delta J}{\delta \pi_\theta(y|x)} = R(x,y) - \beta \log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)} - \beta - \lambda(x) = 0$$
其中$\lambda(x)$是保证$\sumy \pi\theta(y|x) = 1$的拉格朗日乘数。
解得最优策略:
$$\pi^*(y|x) = \pi_{ref}(y|x) \exp\left(\frac{R(x,y) - \lambda(x)}{\beta}\right)$$
归一化条件给出:
$$\sum_y \pi^*(y|x) = 1 \Rightarrow e^{\lambda(x)/\beta} = \sum_y \pi_{ref}(y|x) \exp\left(\frac{R(x,y)}{\beta}\right) = Z(x)$$
因此:
$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{ref}(y|x) \exp\left(\frac{R(x,y)}{\beta}\right)$$
反解奖励函数:
$$R(x,y) = \beta \log \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} + \beta \log Z(x)$$
对于偏好对$(y_w, y_l)$,Bradley-Terry模型给出:
$$P(y_w \succ y_l) = \frac{\exp(R(x,y_w))}{\exp(R(x,y_w)) + \exp(R(x,y_l))} = \sigma(R(x,y_w) - R(x,y_l))$$
代入奖励函数表达式($Z(x)$项抵消):
$$P(y_w \succ y_l) = \sigma\left(\beta \log \frac{\pi^*(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi^*(y_l|x)}{\pi_{ref}(y_l|x)}\right)$$
最大似然估计给出DPO损失函数:
$$\mathcal{L}_{DPO} = -\mathbb{E} \left[\log \sigma\left(\beta \log \frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)}\right)\right]$$
D. Mixture of Experts的负载均衡优化
考虑MoE层的前向传播,输入$x \in \mathbb{R}^d$,有$N$个专家${Ei}{i=1}^N$,门控网络$G$。
标准MoE的输出为:
$$y = \sum_{i=1}^N G_i(x) E_i(x)$$
其中门控权重通过softmax计算:
$$G(x) = \text{softmax}(W_g x + b_g)$$
为了实现稀疏激活,使用Top-K操作:
$$\tilde{G}(x) = \text{KeepTopK}(G(x), k)$$
负载均衡的目标是使每个专家的期望负载相等。定义专家$i$的负载为:
$$L_i = \mathbb{E}_{x \sim \mathcal{D}} [\mathbb{1}[i \in \text{TopK}(G(x), k)]]$$
理想情况下,$L_i = k/N$ 对所有$i$成立。
引入辅助损失促进均衡:
$$\mathcal{L}_{aux} = \alpha \cdot N \cdot \sum_{i=1}^N f_i \cdot P_i$$
其中$f_i$是批次中路由到专家$i$的token比例:
$$f_i = \frac{1}{|B|} \sum_{x \in B} \mathbb{1}[i \in \text{TopK}(G(x), k)]$$
$P_i$是专家$i$的平均路由概率:
$$P_i = \frac{1}{|B|} \sum_{x \in B} G_i(x)$$
这个损失函数鼓励$f_i$和$P_i$都接近$1/N$,从而实现负载均衡。
E. Flash Attention的数值稳定性分析
Flash Attention使用在线softmax算法避免数值溢出。对于输入序列分块$Q_i, K_j, V_j$,计算过程如下:
定义局部最大值和指数和:
$$m_{ij} = \max(Q_i K_j^T / \sqrt{d_k})$$
$$\ell_{ij} = \sum \exp(Q_i K_j^T / \sqrt{d_k} - m_{ij})$$
增量更新规则保证数值稳定性:
$$m_i^{new} = \max(m_i^{old}, m_{ij})$$
$$\ell_i^{new} = e^{m_i^{old} - m_i^{new}} \ell_i^{old} + e^{m_{ij} - m_i^{new}} \ell_{ij}$$
$$O_i^{new} = \frac{e^{m_i^{old} - m_i^{new}} \ell_i^{old}}{\ell_i^{new}} O_i^{old} + \frac{e^{m_{ij} - m_i^{new}} \ell_{ij}}{\ell_i^{new}} \tilde{O}_{ij}$$
定理2:Flash Attention的输出与标准注意力在数值上等价。
证明:考虑完整的注意力计算:
$$A = \text{softmax}(S) = \frac{\exp(S)}{\sum_j \exp(S_j)}$$
使用log-sum-exp技巧:
$$\log \sum_j \exp(S_j) = m + \log \sum_j \exp(S_j - m)$$
其中$m = \max_j S_j$。Flash Attention通过分块计算维护:
- 全局最大值:$m{global} = \max{i,j} m_{ij}$
- 全局指数和:$\ell{global} = \sum{i,j} \exp(m{ij} - m{global}) \ell_{ij}$
- 加权输出:$O{global} = \frac{1}{\ell{global}} \sum{i,j} \exp(m{ij} - m{global}) \ell{ij} O_{ij}$
通过归纳法可证明这与标准计算等价,且避免了指数溢出。
F. 梯度累积与混合精度训练的理论分析
在大模型训练中,由于内存限制,常使用梯度累积技术。设批量大小为$B$,累积步数为$K$,则有效批量大小为$B_{eff} = K \cdot B$。
标准SGD更新:
$$\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t; \mathcal{B}_{eff})$$
梯度累积更新:
$$g_k = \frac{1}{K} \nabla_\theta \mathcal{L}(\theta_t; \mathcal{B}_k), \quad k = 1, ..., K$$
$$\theta_{t+1} = \theta_t - \eta \sum_{k=1}^K g_k$$
定理3:在凸优化条件下,梯度累积与标准SGD收敛速度相同。
证明:假设损失函数$\mathcal{L}$是$L$-光滑的,即:
$$||\nabla \mathcal{L}(\theta_1) - \nabla \mathcal{L}(\theta_2)|| \leq L ||\theta_1 - \theta_2||$$
对于梯度累积,期望梯度为:
$$\mathbb{E}[\sum_{k=1}^K g_k] = \mathbb{E}[\nabla_\theta \mathcal{L}(\theta_t; \mathcal{B}_{eff})]$$
方差分析:
$$\text{Var}[\sum_{k=1}^K g_k] = \frac{1}{K} \text{Var}[\nabla_\theta \mathcal{L}(\theta_t; \mathcal{B})]$$
这与使用大批量$B_{eff}$的方差相同,因此收敛速度一致。
混合精度训练使用FP16计算和FP32主权重。损失缩放防止梯度下溢:
$$\mathcal{L}_{scaled} = s \cdot \mathcal{L}$$
$$g_{FP16} = \nabla_\theta \mathcal{L}_{scaled} = s \cdot \nabla_\theta \mathcal{L}$$
更新规则:
$$\theta_{FP32} = \theta_{FP32} - \eta \cdot \frac{g_{FP16}}{s}$$
动态损失缩放通过监控梯度溢出自适应调整$s$:
$$s_{t+1} = \begin{cases} s_t \cdot \rho_{up} & \text{if no overflow for } N \text{ steps} \\ s_t / \rho_{down} & \text{if overflow detected} \end{cases}$$
典型参数:$\rho{up} = 2$,$\rho{down} = 2$,$N = 2000$。
G. 量化技术的信息论分析
权重量化可以从率失真理论角度分析。设原始权重$W \in \mathbb{R}^{m \times n}$,量化后权重$\hat{W}$,量化函数$Q$。
均方量化误差:
$$D = \mathbb{E}[||W - \hat{W}||_F^2] = \sum_{i,j} (W_{ij} - \hat{W}_{ij})^2$$
对于$k$-bit均匀量化:
$$\hat{W}_{ij} = s \cdot \text{round}\left(\frac{W_{ij}}{s}\right)$$
其中量化步长:
$$s = \frac{\max(W) - \min(W)}{2^k - 1}$$
量化信噪比(SQNR):
$$\text{SQNR} = 10 \log_{10} \left(\frac{\mathbb{E}[W^2]}{\mathbb{E}[(W - \hat{W})^2]}\right) \approx 6.02k + 4.77 - 20\log_{10}\left(\frac{\sigma_W}{\mu_W}\right)$$
AWQ(Activation-aware Weight Quantization)通过重要性加权优化量化:
$$\min_{\hat{W}} \sum_{x \in \mathcal{D}} ||f(x; W) - f(x; \hat{W})||^2$$
定义权重重要性:
$$I_{ij} = \mathbb{E}_{x \sim \mathcal{D}} \left[\left|\frac{\partial f(x; W)}{\partial W_{ij}}\right|\right]$$
AWQ保持top-$p\%$重要权重为全精度:
$$\hat{W}_{ij} = \begin{cases} W_{ij} & \text{if } I_{ij} > \tau_p \\ Q_k(W_{ij}) & \text{otherwise} \end{cases}$$
实验表明,保持0.1%的关键权重可以将量化损失减少75%。
H. 长序列注意力的计算复杂度优化
标准自注意力的计算和内存复杂度分析:
- 计算复杂度:$O(n^2 \cdot d)$
- 内存复杂度:$O(n^2 + n \cdot d)$
其中$n$是序列长度,$d$是隐藏维度。
线性注意力近似:
通过核函数近似,可以将复杂度降至线性:
$$\text{Attention}(Q,K,V) \approx \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)\sum_j \phi(K_j)}$$
其中$\phi: \mathbb{R}^d \rightarrow \mathbb{R}^r$是特征映射。
使用随机特征:
$$\phi(x) = \frac{1}{\sqrt{r}} [\cos(w_1^T x), \sin(w_1^T x), ..., \cos(w_r^T x), \sin(w_r^T x)]$$
复杂度降至$O(n \cdot r \cdot d)$,其中$r \ll n$。
稀疏注意力模式:
定义注意力掩码$M \in {0,1}^{n \times n}$,稀疏度$\rho = ||M||_0 / n^2$。
局部窗口注意力:
$$M_{ij} = \mathbb{1}[|i - j| \leq w]$$
复杂度:$O(n \cdot w \cdot d)$
跨步注意力:
$$M_{ij} = \mathbb{1}[i \mod s = 0 \text{ or } j \mod s = 0]$$
复杂度:$O(n^2 / s \cdot d)$
组合模式通过并集实现:
$$M = M_{local} \cup M_{strided} \cup M_{random}$$
BigBird证明了这种组合可以保持完整注意力的表达能力。