Technical Report of Attention Residuals

作者/机构: Kimi Team https://github.com/MoonshotAI/Attention-Residuals

A1 主要贡献

本文旨在解决现代大型语言模型(LLM)中普遍存在的残差连接问题。标准的残差连接与 PreNorm 结合使用时,会以固定的单位权重累积所有层的输出,导致隐藏状态随网络深度不受控制地增长,并逐渐稀释每一层的贡献。为解决此问题,本文提出了注意力残差(Attention Residuals, AttnRes)

核心思想:AttnRes 提议用对先前所有层输出的 softmax 注意力机制来取代固定的累加操作。这使得每一层能够通过学习到的、依赖于输入的权重,选择性地聚合更早期的表征。AttnRes 的核心机制是为每一层学习一个伪查询向量(pseudo-query),通过这个向量与先前所有层的输出(作为键和值)计算注意力权重,从而实现对信息的动态、内容感知的深度检索。

可扩展性设计:由于在大型模型训练中,对所有先前层输出进行注意力计算会带来巨大的内存和通信开销,本文进一步提出了块注意力残差(Block AttnRes)。该方法将网络层划分为多个块(block),在块内部使用标准的残差连接进行信息累积,形成块级表征,然后在块之间应用注意力机制。这种设计将内存和通信开销从 $O(Ld)$ 降低到 $O(Nd)$(L为层数,N为块数,d为隐藏维度),同时保留了 Full AttnRes 的大部分性能优势。

系统优化:为了使 Block AttnRes 在大规模训练和推理中实用高效,本文开发了相应的系统优化。
1. 训练优化:在流水线并行训练中,引入了基于缓存的跨阶段通信机制,避免了在不同虚拟阶段间重复传输历史块表征,显著降低了通信开销。
2. 推理优化:采用了一种两阶段计算策略,利用在线 softmax【索引31,Online normalizer calculation for softmax,2018,http://cs.PF】技术,将跨块注意力的计算成本分摊,使得推理延迟开销极小。

主要贡献总结
- 提出 Attention Residuals:用学习到的深度 softmax 注意力取代固定的残差累积,并提出了其可扩展变体 Block AttnRes,将内存和通信开销从 $O(Ld)$ 降至 $O(Nd)$。通过统一的结构化矩阵分析,证明了标准残差及其变体对应于深度方向的线性注意力,而 AttnRes 实现了深度方向的 softmax 注意力。
- 大规模基础设施优化:开发了跨阶段缓存和两阶段计算策略等系统优化,使 Block AttnRes 成为标准残差连接的实用替代方案,训练开销极小,推理延迟开销低于2%。
- 全面的评估与分析:通过扩展定律实验、组件消融实验以及在一个1.4T tokens上预训练的48B参数模型上的下游基准测试,验证了 AttnRes 的有效性。分析表明,AttnRes 缓解了 PreNorm 的稀释问题,使隐藏状态幅度有界,并实现了更均匀的梯度分布。


图 1: Attention Residuals 概览。(a) 标准残差:具有统一加性累积的标准残差连接。(b) Full AttnRes:每一层通过学习到的注意力权重选择性地聚合所有先前层的输出。(c) Block AttnRes:层被分组为块,将内存从 $O(Ld)$ 减少到 $O(Nd)$。

A3 背景知识/关键Observation/设计原则

2 动机

符号表示。考虑一批输入序列,形状为 $B \times T \times d$,其中 $B$ 是批量大小,$T$ 是序列长度,$d$ 是隐藏维度。为清晰起见,我们为单个 token 编写公式:$h_l \in R^d$ 表示进入第 $l$ 层的隐藏状态,其中 $l \in \{1, . . . , L\}$ 是层索引,$L$ 是总层数。词元嵌入是 $h_1$。函数 $f_l$ 表示第 $l$ 层应用的变换。在 Transformer 模型中,我们将每个自注意力或 MLP 视为一个独立的层。

2.1 通过残差训练深度网络

残差学习。残差学习【索引12,Deep Residual Learning for Image Recognition,2015,http://cs.CV】被证明是训练深度网络的一项关键技术,因为它允许梯度绕过变换。具体来说,每一层按如下方式更新隐藏状态:

$$ \boldsymbol{h}_l = \boldsymbol{h}_{l-1} + f_{l-1}(\boldsymbol{h}_{l-1}) $$

展开这个递归式,第 $l$ 层的隐藏状态是嵌入和所有先前层输出的总和:$h_l = h_1 + \sum_{i=1}^{l-1} f_i(h_i)$。残差连接背后的关键洞见是恒等映射:每一层都保留了一条直接路径,让信息和梯度都能无变化地流动。在反向传播过程中,关于中间隐藏状态的梯度是:

$$\frac{\partial \mathcal{L}}{\partial \boldsymbol{h}_{l}}=\frac{\partial \mathcal{L}}{\partial \boldsymbol{h}_{L}} \cdot \prod_{j=l}^{L-1}\left(\mathbf{I}+\frac{\partial f_{j}}{\partial \boldsymbol{h}_{j}}\right)$$

展开这个乘积会得到单位矩阵 $I$ 以及包含层雅可比矩阵 $\partial f_j / \partial h_j$ 的高阶项。单位项始终被保留,为梯度从损失函数到任何层提供了一条直接路径,无论深度如何。

泛化残差。虽然有效,但残差更新中的固定单位系数统一对待每一层的贡献,没有提供一种机制来根据深度调整混合方式。Highway networks【索引45,Highway Networks,2015,cs.LG】通过引入学习到的逐元素门控来放宽这一限制:

$$\boldsymbol{h}_l = (1 - \boldsymbol{g}_l) \odot \boldsymbol{h}_{l-1} + \boldsymbol{g}_l \odot f_{l-1}(\boldsymbol{h}_{l-1})$$

其中 $g_l \in [0, 1]^d$ 在变换和恒等路径之间进行插值。更一般地,两者都是加权递归 $h_l = \alpha_l \cdot h_{l-1} + \beta_l \cdot f_{l-1}(h_{l-1})$ 的实例,残差连接设置 $\alpha_l=\beta_l=1$,而 Highway 网络设置 $\alpha_l=1-g_l, \beta_l=g_l$。

局限性。无论是固定的还是门控的,这两种方法都有一个根本性的约束:每一层只能访问其直接输入 $h_{l-1}$,这是一个单一的压缩状态,它混合了所有早期层的输出,而不是各个独立的输出。这带来了几个局限性:(1) 无选择性访问:不同类型的层(例如,注意力层与MLP层)接收相同的聚合状态,尽管它们可能从不同的加权方式中受益;(2) 不可逆的损失:通过聚合丢失的信息无法在更深层被选择性地恢复;(3) 输出增长:后层学习越来越大的输出来影响累积的残差,这可能会破坏训练的稳定性。这些局限性促使我们寻找一种机制,让每一层能够从所有先前层中选择性地聚合信息。

A2 方法细节

3 注意力残差:时间和深度的统一视角

时间与深度的对偶性。以上讨论的局限性让人联想到序列建模中的类似瓶颈,这表明我们应为深度维度寻求类似的解决方案。与 RNN 在时间维度上一样,残差连接将所有先前信息压缩到深度维度上的单一状态 $h_l$ 中。对于序列建模,Transformer 通过用注意力机制取代循环【索引3,Neural Machine Translation by Jointly Learning to Align and Translate,2016,http://cs.CL;索引52,Attention is All you Need,2017,Advances in NeurIPS】,改进了 RNN,允许每个位置以数据依赖的权重选择性地访问所有先前位置。我们为深度维度提出同样的方法:

$$\boldsymbol{h}_l = \alpha_{0\to l} \cdot \boldsymbol{h}_1 + \sum_{i=1}^{l-1} \alpha_{i\to l} \cdot f_i(\boldsymbol{h}_i)$$

其中 $\alpha_{i \to l}$ 是特定于层的注意力权重,满足 $\sum_{i=0}^{l-1} \alpha_{i \to l} = 1$。与序列长度(可以达到数百万个 token)不同,网络深度通常是适度的($L < 1000$),这使得在深度上进行 $O(L^2)$ 的注意力计算在计算上是可行的。我们将这种方法称为注意力残差(Attention Residuals),简称 AttnRes。

3.1 完全注意力残差

注意力权重计算。注意力权重可以写为 $\alpha_{i \to l} = \phi(q_l, k_i)$,其中 $\phi : R^d \times R^d \to R_{\ge 0}$ 是一个核函数,而 $q_l$ 和 $k_i$ 是查询和键向量【索引23,Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention,2020,Proceedings of ICML;索引70,Understanding Transformer from the Perspective of Associative Memory,2025,cs.LG】。不同的 $\phi$ 选择可以恢复不同的残差变体(§6.2);我们采用带有归一化的 $\phi(q, k) = \exp(q^\top \text{RMSNorm}(k))$【索引66,Root mean square layer normalization,2019,Advances in NeurIPS】,从而在深度上产生 softmax 注意力:

$$\alpha_{i \to l}=\frac{\phi\left(\boldsymbol{q}_{l}, \boldsymbol{k}_{i}\right)}{\sum_{j=0}^{l-1} \phi\left(\boldsymbol{q}_{l}, \boldsymbol{k}_{j}\right)}$$

对于每一层 $l$,我们定义:

$$\begin{aligned} \boldsymbol{q}_l = \boldsymbol{w}_l, \quad \boldsymbol{k}_i = \boldsymbol{v}_i = \begin{cases} \boldsymbol{h}_1 & i = 0 \\ f_i(\boldsymbol{h}_i) & 1 \le i \le l - 1 \end{cases} \end{aligned}$$

其中查询 $q_l = w_l$ 是一个特定于层的可学习向量,属于 $R^d$。$\phi$ 内部的 RMSNorm 防止具有大范数输出的层主导注意力权重。

Full AttnRes 的输入。第 $l$ 层的输入随后是:

$$ \boldsymbol{h}_{l}=\sum_{i=0}^{l-1} \alpha_{i \rightarrow l} \cdot \boldsymbol{v}_{i} $$

我们称这种形式为完全注意力残差(Full AttnRes)。对于每个 token,Full AttnRes 需要 $O(L^2d)$ 的算术运算和 $O(Ld)$ 的内存来存储层输出。由于深度远小于序列长度,算术成本是适度的。

开销。在普通训练中,$O(Ld)$ 的内存与为反向传播保留的激活值完全重叠,因此 Full AttnRes 不会引入额外的内存开销。然而,在大规模训练中,激活重计算和流水线并行被广泛采用:那些本可以被释放和重新计算的层输出现在必须为所有后续层保留,并且在流水线并行下,每个输出还必须跨越阶段边界传输。此时,内存和通信开销都会增长到 $O(Ld)$。

分块优化。Full AttnRes 中一个刻意的设计选择是,伪查询 $w_l$ 是一个与层的前向计算解耦的可学习参数。这种独立性意味着任何一组层的注意力权重都可以并行计算,而无需等待它们的顺序输出,特别是允许将 $L$ 层分组为 $N$ 个块,每块 $S$ 层,并在每个块内批量处理注意力计算,将每层的内存 I/O 从 $O(Ld)$ 减少到 $O((S+N)d)$(详细的两阶段策略我们推迟到§4)。然而,在当前的分布式训练方案中,主要成本不是本地内存带宽,而是流水线并行下的跨阶段通信:每个层输出仍然必须在阶段之间传输,而这种 $O(Ld)$ 的通信开销无法通过本地批处理来缓解。这促使我们引入了下面介绍的 Block AttnRes 变体,它将跨阶段表示的数量从 $L$ 减少到 $N$。我们预计,未来的互连技术改进将使完整的 $O(Ld)$ 通信变得可行,从而完全实现 Full AttnRes 的潜力。

3.2 块注意力残差

核心思想。我们提出了块注意力残差(Block Attention Residuals),它将 $L$ 层划分为 $N$ 个块:在每个块内,层输出通过求和被简化为单一表示;在块之间,我们仅对 $N$ 个块级表示和词元嵌入应用完全注意力。这使得内存和通信开销都从 $O(Ld)$ 减少到 $O(Nd)$。

块内累积。具体来说,我们将 $L$ 层划分为 $N$ 个大小为 $S = L/N$ 的块,假设 $L$ 可以被 $N$ 整除;否则,最后一个块包含剩余的 $L \pmod N$ 层。令 $B_n$ 表示块 $n$($n = 1, . . . , N$)中的层索引集合。为了形成一个块,我们对该块中所有层的输出求和:

$$\boldsymbol{b}_n = \sum_{j \in \mathcal{B}_n} f_j(\boldsymbol{h}_j)$$

我们进一步将 $b_n^i$ 表示为 $B_n$ 中前 $i$ 层的部分和,因此 $b_n = b_n^S$。当 $L$ 不能被 $N$ 整除时,最终的部分和被视为最后一个块的表示。与 Full AttnRes 一样,$\phi$ 内部的 RMSNorm 防止了完整块和部分和之间的幅度差异对注意力权重产生偏见。


图 2: Block Attention Residuals 的 PyTorch 风格伪代码。block_attn_res 使用一个学习到的伪查询 wl 对块表示计算 softmax 注意力;forward 是一个单层的前向传播过程,它维护 partial_block($b_n^i$,块内残差)和 blocks([$b_0, . . . , b_{n-1}$],块间历史)。

块间注意力。在 Full AttnRes 中,第 $l$ 层的输入是通过对直到 $f_{l-1}(h_{l-1})$ 的所有输出进行注意力计算得到的。块级变体用块表示替换了这些单独的输出,定义 $b_0 = h_1$,以便词元嵌入始终作为源包含在内。对于块 $n$ 中的第 $i$ 层,值矩阵为:

$$\begin{aligned} \mathbf{V}= \begin{cases}{\left[\boldsymbol{b}_{0}, \boldsymbol{b}_{1}, \ldots, \boldsymbol{b}_{n-1}\right]^{\top}} & \text { if } i=1 \text { (first layer of block } n \text { ) } \\ {\left[\boldsymbol{b}_{0}, \boldsymbol{b}_{1}, \ldots, \boldsymbol{b}_{n-1}, \boldsymbol{b}_{n}^{i-1}\right]^{\top}} & \text { if } i \geq 2 \text { (subsequent layers) }\end{cases} \end{aligned}$$

键和注意力权重遵循公式3和公式2。网络第一层的输入是词元嵌入,即 $b_0 = h_1$。在每个块中,第一层接收先前的块表示和词元嵌入,而后续层还会额外关注部分和 $b_n^{i-1}$。最终的输出层聚合所有 $N$ 个块表示。图2提供了 Block AttnRes 的 PyTorch 风格伪代码。

效率。由于现在每层关注 $N$ 个块表示而不是 $L$ 个单独的输出,内存从 $O(L)$ 减少到 $O(N)$,计算从 $O(L^2)$ 减少到 $O(N^2)$。块数 $N$ 在两个极端之间插值:$N = L$ 恢复为 Full AttnRes,而 $N = 1$ 简化为标准残差连接,其中嵌入作为 $b_0$ 被隔离。根据经验,我们发现 $N \approx 8$ 可以在不同模型规模上恢复大部分的收益,每个 token 只需要存储八个隐藏状态(见§5)。除了内存和计算,块结构也有利于推理延迟:块边界定义了§3中描述的块级优化的分发粒度,固定的块数 $N$ 限制了 KV 缓存的大小。并行的块间结果通过在线 softmax【索引31,Online normalizer calculation for softmax,2018,http://cs.PF】与顺序的块内部分和合并,保持了精确的等价性(§4)。

4 基础设施设计

与标准残差连接相比,Block AttnRes 带来了额外的系统挑战。对于大规模模型训练,块表示必须在流水线阶段之间传播,这在朴素实现中会导致大量的通信。在推理期间,对累积块表示的重复访问会增加延迟,而长上下文预填充则放大了缓存块表示的内存成本。我们通过训练中的跨阶段缓存,以及推理中的两阶段计算策略和内存高效的预填充方案来解决这些挑战。


图 3: 基于缓存的流水线通信示例,有4个物理 rank 和每个 rank 2个虚拟阶段,其中带阴影的框表示 AttnRes 块的末端。数字表示微批次索引。每个 rank 缓存先前接收到的块;阶段转换仅传输增量块(+[b1, b2]),而不是完整的历史记录。

4.1 训练

小规模训练。对于小规模训练,AttnRes 增加的计算开销微乎其微,并且没有额外的内存使用,因为无论如何激活值都需要为反向传播而保存。在大型分布式训练中,流水线并行是 AttnRes 面临的主要基础设施挑战。Full AttnRes 需要所有 $L$ 个层输出在各阶段间传输;Block AttnRes 将其减少到 $N$ 个块表示,而下面的优化进一步最小化了剩余的开销。

流水线通信。对于标准残差连接,流水线并行【索引18,GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism,2019,Advances in NeurIPS】在相邻阶段之间传输一个固定大小的隐藏状态,与流水线深度无关。Block AttnRes 在每个阶段都需要所有累积的块表示来进行块间注意力计算,朴素地在每次转换时传输完整的历史记录会产生冗余通信。考虑一个交错式流水线调度【索引33,Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM,2021,http://cs.CL】,有 $P$ 个物理阶段,每个物理阶段有 $V$ 个虚拟阶段。为简单起见,假设每个物理阶段平均为每个 token 产生 $N_p$ 个维度为 $d$ 的块表示。总共有 $C = P \cdot V$ 个块(每个物理阶段在每个虚拟阶段中),第 $j$ 个块累积了 $jN_p$ 个块。朴素地在每次转换时传输所有累积的块会产生每个 token 的通信成本:

$$\text{Comm}_{\text{na\"ive}} = \sum_{j=1}^{C-1} j N_p \cdot d = \frac{C(C-1)}{2} N_p d.$$

跨阶段缓存。由于每个物理阶段连续处理多个虚拟阶段,我们可以通过在本地缓存块来消除这种冗余:在早期虚拟阶段接收到的块保留在本地内存中,无需重新传输。第一个虚拟阶段($v=1$)没有缓存,正常累积;对于 $v \ge 2$,每次转换仅传送自接收方在上一虚拟阶段的相应块以来累积的约 $P N_p$ 个增量块。总通信量减少为:

$$\text{Comm}_{\text{cached}} = \underbrace{\frac{P(P-1)}{2} N_p d}_{\text{first virtual stage}} + \underbrace{(V-1) P^2 N_p d}_{\text{subsequent virtual stages}} .$$

缓存将每次转换的峰值成本从 $O(C)$ 降低到 $O(P)$,这是一个 $V$ 倍的改进,使得在稳态的 1F1B 调度中可以与计算完全重叠。反向传播过程也受益于同样的方案。图3用 $P=4$ 和 $V=2$ 的例子说明了这种优化:对于第二个虚拟阶段,缓存消除了6次冗余的块传输。

算法 1: 块n的两阶段计算

内存开销。通过跨阶段缓存,每个块在所有 $V$ 个虚拟阶段中只存储一次,这相对于标准的每层激活缓存来说是微不足道的。关键是,每层的激活占用空间与标准架构完全相同,因为激活检查点消除了所有块间注意力的中间产物,并且检查点的输入 $p_l$ 与其替换的隐藏状态 $h_l$ 的内存大小相匹配。就挂钟时间而言,当不启用流水线并行时,Block AttnRes 增加的训练开销可以忽略不计;在流水线并行下,测得的端到端开销小于4%。

4.2 推理

两阶段计算策略。下面描述的两阶段计算策略适用于 Full 和 Block AttnRes:在任何一种情况下,层都被分组为大小为 $S$ 的块,其中阶段1批量处理块间查询,阶段2处理顺序的块内回溯。对于 Full AttnRes,这将每层的 I/O 从 $O(Ld)$ 减少到 $O((S+N)d)$(详细推导见附录B);Block AttnRes 进一步将存储的表示从 $L$ 个减少到 $N$ 个,因为每个块被压缩成一个向量。接下来,我们专注于 Block AttnRes,并详细介绍两阶段计算策略以及用于长上下文输入的序列分片预填充方案。该策略利用了伪查询向量与前向计算解耦的特性(§3),可以将一个块内的所有 $S = L/N$ 个查询批处理成单个矩阵乘法,从而将内存访问从 $S$ 次读取摊销为1次。

算法1实例化了一个利用此特性的两阶段计算策略:
- 阶段1:通过对缓存的块表示进行一次批处理查询,同时计算所有 $S$ 层的块间注意力,返回输出和 softmax 统计信息(最大值和对数-和-指数)。这摊销了内存访问成本,将每个块的读取次数从 $S$ 次减少到仅一次。
- 阶段2:使用不断演变的部分和,为每层顺序计算块内注意力,然后通过在线 softmax【索引31,Online normalizer calculation for softmax,2018,http://cs.PF】与阶段1的输出合并。由于在线 softmax 的合并是逐元素的,此阶段自然可以与周围的操作进行核函数融合,进一步减少I/O开销。

效率分析。通过两阶段设计,阶段2的I/O占用与标准残差连接相似,而主要的额外成本来自阶段1的块间注意力。由于这些块间读取通过批处理被分摊到块中的所有层,每层的总内存访问成本仅为 $(\frac{N}{S} + 3)d$ 次读取和 $2d$ 次写入(表1)。这远低于先前残差泛化方法(如 (m)HC)在典型设置下的残差流 I/O。在实践中,阶段1还可以与块中第一层的计算部分重叠,进一步减少其挂钟时间影响。因此,在典型的推理工作负载上,端到端的推理延迟开销小于2%。

表 1: 每种方案下残差机制导致的每 token 每层内存访问成本。层函数 $f_l$ 的内部 I/O 不包括在内。对于 AttnRes,Full 和 Block 变体都使用附录B中描述的两阶段推理调度;摊销成本是在一个块内的 $N$ 层上平均的。典型值:$L=128, N=8, S=L/N=16, m=4$。

内存高效的预填充。在预填充期间存储块表示需要 $N \cdot T \cdot d$ 个元素,对于一个128K-token序列和8个块,会产生15 GB的内存。我们通过在 $P$ 个张量并行设备上沿序列维度对这些表示进行分片来缓解这个问题,允许阶段1在本地序列分片上独立执行。阶段2的在线-softmax合并随后集成到标准的TP all-reduce通信路径中:输出被reduce-scatter,本地合并,然后通过all-gather重构,自然地允许与RMSNorm等操作进行核函数融合。这将每个设备的内存占用减少到 $N \cdot (T /P ) \cdot d$——将128K上下文的例子从15 GB降低到每个设备约1.9 GB。结合分块预填充(例如,16K块大小),开销进一步减少到每个设备低于0.3 GB。

A4 实验环境

架构细节:我们的架构与 Kimi Linear [【索引69,Kimi Linear: An Expressive, Efficient Attention Architecture,2025,http://cs.CL】] 相同,这是一个遵循 Moonlight [【索引28,Muon is Scalable for LLM Training,2025,cs.LG】] / DeepSeek-V3 [【索引9,DeepSeek-V3 Technical Report,2025,http://cs.CL】] 设计的专家混合(MoE)Transformer。该架构以3:1的比例交错使用 Kimi Delta Attention (KDA) 和 Multi-Head Latent Attention (MLA) 层,每层之后跟着一个 MoE 前馈层。唯一的修改是在残差连接中加入了 AttnRes;所有其他组件(模型深度、隐藏维度、专家路由和 MLP 结构)保持不变。AttnRes 每层只引入一个 RMSNorm 和一个伪查询向量 $w_l \in R^d$,占总参数量的比例可以忽略不计。关键的是,所有伪查询向量必须初始化为零。这确保了初始的注意力权重 $\alpha_{i \to l}$ 在源层之间是均匀的,从而在训练开始时将 AttnRes 简化为等权重平均,防止了训练波动,这一点我们通过经验得到了验证。

A4 实验结果

5.1 扩展定律

我们对五种模型尺寸(表2)进行了扫描,并为每种尺寸训练了三个变体:一个 PreNorm 基线模型、一个 Full AttnRes 模型和一个带有约8个块的 Block AttnRes 模型。它们使用8192个 token 的上下文窗口和一个余弦学习率调度进行训练。在每个扩展定律尺寸组内,所有变体共享相同的超参数,这些超参数是在基线模型下选择的,以确保公平比较;这种设置有意偏向基线,因此使得比较结果更为保守。遵循标准实践,我们拟合了形式为 $L = A \times C^{-\alpha}$ 的幂律曲线【索引22,Scaling Laws for Neural Language Models,2020,cs.LG;索引15,Training Compute-Optimal Large Language Models,2022,http://cs.CL】,其中 $L$ 是验证损失,$C$ 是以 PFLOP/s-天为单位的计算量。

扩展行为:图4展示了拟合的扩展曲线。基线模型遵循 $L = 1.891 \times C^{-0.057}$,而 Block AttnRes 拟合为 $L = 1.870 \times C^{-0.058}$,Full AttnRes 拟合为 $L = 1.865 \times C^{-0.057}$。所有三个变体都表现出相似的斜率,但 AttnRes 在整个计算范围内始终实现了更低的损失。根据拟合曲线,在5.6 PFLOP/s-天的计算量下,Block AttnRes 的损失达到1.692,而基线为1.714,相当于1.25倍的计算优势。Full AttnRes 和 Block AttnRes 之间的差距随着规模的扩大而缩小,在最大尺寸时仅为0.001。我们还在表2中列出了 mHC(-lite)【索引64,mHC-lite: You Don’t Need 20 Sinkhorn-Knopp Iterations,2026,cs.LG】以供参考。Full AttnRes 的性能优于 mHC,而 Block AttnRes 在每层内存I/O更低的情况下与之持平:AttnRes 为5.5d,而 mHC(m=4个流)为34d(表1)。

表 2: Baseline vs Block AttnRes (N = 8) vs Full AttnRes vs mHC(-lite) [64]:模型配置、超参数和验证损失。

† 表示我们的 MoE 模型中激活的参数数量,不包括嵌入。‡ 所有模型均使用8192的上下文长度进行训练。⋆ Lb = L/2 表示 Transformer 块的数量。


图 4: 注意力残差的扩展定律曲线。Full 和 Block AttnRes 在所有尺度上都持续优于基线。Block AttnRes 紧随 Full AttnRes,在最大尺度上恢复了大部分增益。

5.2 主要结果

训练方案:我们研究的最大模型基于完整的 Kimi Linear 48B 配置:27个 Transformer 块(54层),每个块有256个专家中的8个被路由加上1个共享专家,总计48B参数,激活参数为3B。该模型应用 Block AttnRes,每块6层,产生9个块加上词元嵌入,总共10个深度源。我们遵循与 Kimi Linear 1.4T-token 运行相同的训练方案【索引69,Kimi Linear: An Expressive, Efficient Attention Architecture,2025,http://cs.CL】:所有模型都使用4096-token的上下文窗口、Muon 优化器【索引28,Muon is Scalable for LLM Training,2025,cs.LG】和 WSD(Warmup–Stable–Decay)学习率调度【索引16,MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies,2024,http://cs.CL】进行预训练,全局批量大小为8M tokens。最终模型的训练分为两个阶段:(i) 在1T tokens上进行WSD预训练,(ii) 在约400B高质量 tokens上进行中期训练,遵循 Moonlight【索引28】的退火方案。中期训练后,我们继续使用逐步增长的32K tokens序列长度进行训练。

训练动态:我们在图5中比较了最终的 Baseline 和 Block AttnRes 模型在1T tokens上的训练动态。
- 验证损失:AttnRes 在整个训练过程中始终获得较低的验证损失,差距在衰减阶段扩大,并导致最终损失显著降低。
- 输出幅度:Baseline 遭受 PreNorm 稀释问题【索引60,On Layer Normalization in the Transformer Architecture,2020,cs.LG;索引27,SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm,2026,cs.LG】:随着隐藏状态幅度随深度单调增长,更深的层被迫从固定尺度的归一化输入中学习越来越大的输出以保持影响力。Block AttnRes 将这种增长限制在每个块内,因为块边界处的选择性聚合重置了累积,产生了一个有界的周期性模式。
- 梯度幅度:由于所有残差权重固定为1,Baseline 没有提供调节深度梯度流的方法,导致最早的层梯度过大。Block AttnRes 中可学习的 softmax 权重(图8)在源之间引入了对概率质量的竞争,从而导致梯度分布更加均匀。


图 5: 基线和 Block AttnRes 的训练动态。(a) 训练期间的验证损失。(b) 训练结束时每个 transformer 块的输出幅度。(c) 每个 transformer 块的梯度幅度。

下游性能:我们评估了两个模型在三个领域的表现(表3):
- 语言理解与推理:MMLU, MMLU-Pro Hard, GPQA-Diamond, BBH, ARC-Challenge, HellaSwag, and TriviaQA。
- 推理(代码与数学):GSM8K, MGSM, Math, CMath, HumanEval, and MBPP。
- 中文语言理解:CMMLU and C-Eval。

如表3所示,Block AttnRes 在所有基准测试中均持平或优于基线。在多步推理任务如 GPQA-Diamond (+7.5) 和 Minerva Math (+3.6),以及代码生成任务如 HumanEval (+3.1) 上的提升尤为显著,而知识导向的基准如 MMLU (+1.1) 和 TriviaQA (+1.9) 也显示出稳固的增益。这一模式与以下假设一致:改进的深度信息流有利于组合性任务,其中后层可以有选择地检索和构建早期表示。

表 3: AttnRes 与基线的性能比较,两者均经过相同的预训练方案。每行最佳结果加粗显示。

5.3 消融研究

我们对表2中的16头模型进行了消融研究,以验证 AttnRes 的关键设计选择(表4)。所有模型共享相同的超参数和计算预算。

与先前方法的比较:我们将 AttnRes 与 PreNorm 基线(损失1.766)和两种代表性的残差连接泛化方法进行比较。DenseFormer【索引36,DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging,2024,http://cs.CL】允许每层访问所有先前的输出,但使用固定的、与输入无关的标量系数进行组合;它相对于基线没有显示出增益(1.767),凸显了输入依赖加权的重要性。mHC【索引59,mHC: Manifold-Constrained Hyper-Connections,2026,http://cs.CL】通过 $m$ 个并行流和学习的混合矩阵引入了输入依赖性,将损失提高到1.747。AttnRes 通过 softmax 注意力实现了更明确的内容依赖选择:Full AttnRes 达到1.737,Block AttnRes 达到1.746,仅用每层一个查询向量就优于这两种方法。

跨层访问:我们比较了三种粒度的跨层访问。Full AttnRes 直接源于时间-深度对偶性(§3),对所有先前层应用注意力,并取得了最低的损失(1.737)。一种简单的降低其内存成本的方法是滑动窗口聚合(SWA),它只保留最近的 $W=8$ 个层输出加上词元嵌入;它比基线有所改进(1.764),但远不及 Full 和 Block AttnRes,这表明选择性地访问远距离层比关注许多邻近层更重要。Block AttnRes 提供了更好的权衡:当块大小 $S=4$ 时,它达到了1.746的损失,同时保持了每层的内存开销恒定。图6扫描了 $S$ 从 $S=1$(即 Full AttnRes)到越来越粗的 groupings 的整个范围。损失随着 $S$ 的增长而平稳下降,$S=2, 4, 8$ 都接近1.746,而更大的块($S=16, 32$)则趋向于基线。

表 4: AttnRes 关键组件的消融实验(16层模型)。


图 6: 块大小对验证损失的影响(16层模型)。

组件设计
- 输入依赖的查询:一个自然的扩展是通过从当前隐藏状态投影来使查询依赖于输入。这进一步将损失降低到1.731,但每层引入了一个 $d \times d$ 的投影,并在解码期间需要顺序内存访问,因此我们默认使用学习到的查询。
- 与输入无关的混合:我们移除了查询和键,并用可学习的、与输入无关的标量来加权先前层,这损害了性能(1.749 vs. 1.737)。
- softmax vs. sigmoid:用 sigmoid 替换 softmax 会降低性能(1.741)。我们将其归因于 softmax 的竞争性归一化,它强制在源之间进行更清晰的选择。
- 多头注意力:我们在 Block AttnRes 上测试了每头(H=16)的深度聚合,允许不同的通道组关注不同的源层。这损害了性能(1.752 vs. 1.746),表明最优的深度混合在通道间是基本一致的:当一个层的输出相关时,它是作为一个整体相关的。
- 键上的 RMSNorm:移除 RMSNorm 会降低 Full AttnRes(1.743)和 Block AttnRes(1.750)的性能。对于 Full AttnRes,它防止了具有天然较大输出的单个层主导 softmax。这对于 Block AttnRes 变得更加关键,因为块级表示在更多层上累积,可能会产生大的幅度差异;RMSNorm 防止这些差异偏向注意力权重。

5.4 分析

5.4.1 最佳架构

架构扫描实验。为了理解 AttnRes 如何重塑最佳架构的缩放规律,我们在固定的计算和参数预算下进行了一项受控的能力重新分配研究。我们的核心问题是 AttnRes 是否改变了首选的深度-宽度-注意力权衡,特别是,鉴于其在深度维度上的潜在优势,它是否比传统的 Transformer 设计启发式更偏爱更深的模型。为了隔离与深度直接相关的结构性因素,我们根据内部经验观察固定了每个专家的 MLP 扩展比($d_{ff}/d_{model} \approx 0.45$)。我们进一步固定了总训练计算量(FLOPs $\approx 6.5 \times 10^{19}$)和活动参数($\approx 2.3 \times 10^8$),确保任何性能变化纯粹来自架构重新分配,而非整体容量差异。在此约束预算下,我们枚举了在 $d_{model}/L_b \in \{15, 30, 45, 60, 75\}$ 和 $H/L_b \in \{0.3, 0.4, 0.5, 0.6, 0.7\}$ 的 $5 \times 5$ 网格上的25种配置,其中 $L_b = L/2$ 是 Transformer 块的数量,$H$ 是注意力头的数量。结果如图7所示。

结果分析。两个热图都呈现出共同的模式:损失随着 $d_{model}/L_b$ 的增长和 $H/L_b$ 的缩小而降低,两种方法都在 $H/L_b \approx 0.3$ 时达到最优。尽管有这种共同趋势,AttnRes 在所有25种配置中都比基线取得了更低的损失,差距在0.019-0.063之间。最明显的区别在于最优值的位置:基线在 $d_{model}/L_b \approx 60$ 时达到最低损失(1.847),而 AttnRes 将其移至 $d_{model}/L_b \approx 45$(1.802)。在固定的参数预算下,较低的 $d_{model}/L_b$ 对应于更深、更窄的网络,这表明 AttnRes 能更有效地利用额外的深度。我们注意到,这种对深度的偏好并不直接转化为部署建议,因为更深的模型通常由于其顺序计算而产生更高的推理延迟【索引39,Efficiently Scaling Transformer Inference,2022,cs.LG】。相反,这次扫描作为一个诊断工具,揭示了 AttnRes 在何处受益最大,并且这种深度偏好可以在架构选择时与推理成本一起考虑。


图 7: 固定计算量下的架构扫描(≈ 6.5 × 10^19 FLOPs,≈ 2.3 × 10^8 激活参数)。每个单元格报告了 (dmodel/Lb, H/Lb) 配置的验证损失,其中 Lb = L/2 是 Transformer 块的数量;星号标记了最优值。

5.4.2 分析学习到的 AttnRes 模式

可视化学习到的权重。我们在图8中可视化了16头模型(来自表2)中学习到的权重 $\alpha_{i \to l}$,包括 full 和 block (N=8) AttnRes 两种情况。每个热图显示了第 $l$ 个注意力或 MLP 层(行)如何在先前的源(列)上分配其注意力,其中 pre-attention 和 pre-MLP 层分开显示。我们强调三个关键观察:

  • 保留的局部性。每一层都最强烈地关注其直接前驱,但选择性的非对角线集中区域也出现了(例如,第4层关注早期源,第15-16层在块设置下回溯),这表明除了标准残差路径外,还学习到了跳跃连接。
  • 层专业化。嵌入 $h_1$ 在整个过程中都保持着不可忽略的权重,尤其是在 pre-attention 层中。pre-MLP 的输入对近期表示显示出更强的对角线依赖,而 pre-attention 的输入则保持更宽的感受野,这与注意力在层间路由信息而 MLP 在局部操作的观点一致。
  • Block AttnRes 保留结构。对角线主导、嵌入持久性和层专业化都从 full 变体转移到了 block 变体,这表明块级压缩在保留基本信息通路的同时,起到了隐式正则化的作用。


图 8: 16头模型中 full(上)和 block(下)Attention Residuals 的深度注意力权重分布,按 token 平均。该模型有16个注意力和16个 MLP 层。每行显示第 l 个注意力(左)或 MLP(右)层如何在先前的源上分配权重。对角线主导表明局部性仍然是主要的信息通路,而对源 0(嵌入)的持续权重和偶尔的非对角线集中揭示了学习到的跳跃连接。块注意力(N = 8)以更清晰、更果断的权重分布恢复了基本结构。

A7 补充细节

6 讨论

表 5: 残差更新机制的比较。权重:混合系数是架构固定的、学习后静态的,还是输入依赖的(动态的)。来源:第 l 层可以访问哪些早期表示。为清晰起见,大多数公式省略了归一化。

1 ConvPool:池化操作后跟卷积(通道投影)。2 $\phi(q, k) = \exp(q^\top \text{RMSNorm}(k))$; $k_i = v_i$; $v_0 = h_1$, $v_{i\ge1} = f_i(h_i)$。softmax 在所有源上联合归一化。3 与 Full 相同的 $\phi$ 和归一化;$v_i = b_i$, $v_n^j = b_n^j$。

6.1 序列-深度对偶性

类比RNNs。残差连接通过一个固定的递归 $h_l = h_{l-1} + f_{l-1}(h_{l-1})$ 在深度上传播信息,这与 RNNs 在时间上传播信息的方式非常相似。Test-Time Training (TTT)【索引46,Learning to (Learn at Test Time): RNNs with Expressive Hidden States,2024,ArXiv】正式化了这一类比的序列侧(可参考 Fast Weight Programmers【索引43,Learning to control fast-weight memories: An alternative to dynamic recurrent networks,1992,Neural Computation;索引32,Metalearned Neural Memory,2019,ArXiv】),将每个递归步骤视为在自监督损失上的梯度下降:

$$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta \nabla \ell(\mathbf{W}_{t-1}; \boldsymbol{x}_t),$$

其中一个慢网络参数化了 $\ell$,状态 $W$ 每个 token 更新一次。当 $f$ 是线性时,这简化为普通的线性注意力 $S_t = S_{t-1} + k_t v_t^\top$。标准残差在深度维度上表现出相同的加性形式,其中 $h_l$ 作为状态,每个层 $f_l$ 充当一次“梯度步骤”。

扩展对偶性。正如【索引4】所指出的,这种对偶性可以扩展到更丰富的变体(表5)。序列侧的数据依赖门控【索引47,Retentive Network: A Successor to Transformer for Large Language Models,2023,http://cs.CL;索引63,Gated Linear Attention Transformers with Hardware-Efficient Training,2024,Proceedings of ICML】对应于深度侧的 Highway networks【索引45,Highway Networks,2015,cs.LG】;delta 规则【索引42,Linear Transformers Are Secretly Fast Weight Programmers,2021,Proceedings of ICML;索引62,Gated Delta Networks: Improving Mamba2 with Delta Rule,2025,Proceedings of ICLR;索引69,Kimi Linear: An Expressive, Efficient Attention Architecture,2025,http://cs.CL】对应于 DDL【索引67,Deep Delta Learning,2026,cs.LG】;MRLA【索引10,Cross-Layer Retrospective Retrieving via Layer Attention,2023,http://cs.CV】则镜像了 GLA 的【索引63】门控线性注意力。这些方法都改进了递归更新,但仍停留在递归范式内。AttnRes 更进一步,用直接的跨层注意力取代了深度递归,就像 Transformers 用自注意力取代了时间递归一样。由于当前架构的层数仍在 softmax 注意力的实际可行范围内,我们采用普通的深度注意力。结合更具表现力但内存效率更高的(例如,线性复杂度)替代方案是未来工作的自然方向。

6.2 残差连接作为结构化矩阵

深度混合矩阵。上述讨论的残差变体都可以被视为对先前层输出的加权聚合。我们用一个深度混合矩阵 $M \in R^{L \times L}$ 来形式化这一点,其中 $M_{i \to l}$ 是第 $l$ 层分配给第 $i$ 层输出的权重。这些变体的不同之处在于这些权重是如何产生的(固定的、学习的或输入依赖的),以及 $M$ 是被约束为低秩还是允许是稠密的。$M$ 的半可分秩(semiseparable rank)【索引8,Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality,2024,CoRR】为比较它们提供了一个统一的视角。具体来说,第 $l$ 层的输入是 $h_l = \sum_{i=0}^{l-1} M_{i \to l} v_i$,其中 $v_0 = h_1$(嵌入),$v_i = f_i(h_i)$ for $i \ge 1$。图9可视化了代表性方法的 $M$;我们下面推导每一种。

(m)HC

$$\begin{aligned} \left[\begin{array}{cccc} 1 & & & \\ \gamma_{1 \rightarrow 2}^{\times} & g_{2} & & \\ \gamma_{1 \rightarrow 3}^{\times} & g_{2} \gamma_{2 \rightarrow 3}^{\times} & g_{3} & \\ \gamma_{1 \rightarrow 4}^{\times} & g_{2} \gamma_{2 \rightarrow 4}^{\times} & g_{3} \gamma_{3 \rightarrow 4}^{\times} & g_{4} \end{array}\right]\left[\begin{array}{cccc} \boldsymbol{\beta}_{0}^{\top} \boldsymbol{\alpha}_{1} & & & \\ \boldsymbol{\beta}_{0}^{\top} \mathbf{A}_{1 \rightarrow 2}^{\times} \boldsymbol{\alpha}_{2} & \boldsymbol{\beta}_{1}^{\top} \boldsymbol{\alpha}_{2} & & \\ \boldsymbol{\beta}_{0}^{\top} \mathbf{A}_{1 \rightarrow 3}^{\times} \boldsymbol{\alpha}_{3} & \boldsymbol{\beta}_{1}^{\top} \mathbf{A}_{2 \rightarrow 3}^{\times} \boldsymbol{\alpha}_{3} & \boldsymbol{\beta}_{2}^{\top} \boldsymbol{\alpha}_{3} & \\ \boldsymbol{\beta}_{0}^{\top} \mathbf{A}_{1 \rightarrow 4}^{\times} \boldsymbol{\alpha}_{4} & \boldsymbol{\beta}_{1}^{\top} \mathbf{A}_{2 \rightarrow 4}^{\times} \boldsymbol{\alpha}_{4} & \boldsymbol{\beta}_{2}^{\top} \mathbf{A}_{3 \rightarrow 4}^{\times} \boldsymbol{\alpha}_{4} & \boldsymbol{\beta}_{3}^{\top} \boldsymbol{\alpha}_{4} \end{array}\right] \end{aligned}$$


Highway

Block AttnRes

$$\begin{aligned} \begin{bmatrix} \phi\left(\boldsymbol{w}_{1}, \boldsymbol{k}_{0}\right) & & & \\ \phi\left(\boldsymbol{w}_{2}, \boldsymbol{k}_{0}\right) & \phi\left(\boldsymbol{w}_{2}, \boldsymbol{k}_{1}\right) & & \\ \phi\left(\boldsymbol{w}_{3}, \boldsymbol{k}_{0}\right) & \phi\left(\boldsymbol{w}_{3}, \boldsymbol{k}_{1}\right) & \phi\left(\boldsymbol{w}_{3}, \boldsymbol{k}_{2}\right) & \\ \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{0}\right) & \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{1}\right) & \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{2}\right) & \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{3}\right) \end{bmatrix} \begin{bmatrix} \phi\left(\boldsymbol{w}_{1}, \boldsymbol{k}_{0}\right) & & \\ \phi\left(\boldsymbol{w}_{2}, \boldsymbol{k}_{0}\right) & \phi\left(\boldsymbol{w}_{2}, \boldsymbol{k}_{1}\right) & \\ \phi\left(\boldsymbol{w}_{3}, \boldsymbol{k}_{0}\right) & \multicolumn{2}{c}{\phi\left(\boldsymbol{w}_{3}, \boldsymbol{k}_{1}+\boldsymbol{k}_{2}\right)} & \\ \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{0}\right) & \multicolumn{2}{c}{\phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{1}+\boldsymbol{k}_{2}\right)} & \phi\left(\boldsymbol{w}_{4}, \boldsymbol{k}_{3}\right) \end{bmatrix} \end{aligned}$$


图 9: 四种残差变体的深度混合矩阵 M (L=4; Block AttnRes 使用块大小 S=2)。为清晰起见,Highway 显示为标量门。AttnRes 面板显示未归一化的 ϕ 分数;背景色将共享相同源(Full AttnRes)或相同源块(Block AttnRes)的条目分组。

  • 标准残差【索引12,Deep Residual Learning for Image Recognition,2015,http://cs.CV】, $h_l = h_{l-1} + f_{l-1}(h_{l-1})$。展开得到 $h_l = \sum_{i=0}^{l-1} v_i$,所以对于所有 $i < l$,$M_{i \to l} = 1$,$M$ 是一个全为1的下三角矩阵:
$$\begin{aligned} \begin{bmatrix} \boldsymbol{h}_1 \\ \boldsymbol{h}_2 \\ \vdots \\ \boldsymbol{h}_L \end{bmatrix} = \begin{bmatrix} 1 & & & \\ 1 & 1 & & \\ \vdots & \vdots & \ddots & \\ 1 & 1 & \cdots & 1 \end{bmatrix} \begin{bmatrix} \boldsymbol{v}_0 \\ \boldsymbol{v}_1 \\ \vdots \\ \boldsymbol{v}_{L-1} \end{bmatrix} \end{aligned}$$
  • Highway【索引45,Highway Networks,2015,cs.LG】, $h_l = (1-g_l) h_{l-1} + g_l f_{l-1}(h_{l-1})$(为清晰起见,这里用标量门书写;逐元素的扩展是直接的)。定义进位乘积 $\gamma^\times_{i \to l} := \prod_{j=i+1}^l (1 - g_j)$,权重对于嵌入是 $M_{0 \to l} = \gamma^\times_{1 \to l}$,对于 $i \ge 1$ 是 $M_{i \to l} = g_{i+1} \gamma^\times_{i+1 \to l}$。由于累积乘积通过标量门分解,$M$ 是1-半可分的【索引8,Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality,2024,CoRR】,与标准残差具有相同的秩,但具有输入依赖的权重。权重通过构造总和为1,使得 Highway 成为一种无 softmax 的深度方向的 stick-breaking 注意力【索引49,Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study,2025,Proceedings of ICLR】。

  • (m)HC【索引72,Hyper-Connections,2025,cs.LG;索引59,mHC: Manifold-Constrained Hyper-Connections,2026,http://cs.CL】维护 $m$ 个并行流 $H_l \in R^{d \times m}$,通过以下方式更新:

$$\mathbf{H}_{l}=\mathbf{H}_{l-1} \mathbf{A}_{l}+f_{l-1}\left(\mathbf{H}_{l-1} \boldsymbol{\alpha}_{l-1}\right) \boldsymbol{\beta}_{l-1}^{\top}$$

其中 $A_l \in R^{m \times m}$ 是一个学习的转移矩阵,$\alpha_{l-1} \in R^m$ 将流混合成 $f_{l-1}$ 的单个输入,$\beta_{l-1} \in R^m$ 将输出分布回流中。展开递归得到有效权重:

$$\mathbf{M}_{i \rightarrow l}=\beta_{i}^{\top} \mathbf{A}_{i+1 \rightarrow l}^{\times} \boldsymbol{\alpha}_{l},$$

其中 $A^\times_{i \to j} := \prod_{k=i+1}^j A_k$。$m \times m$ 的转移使得 $M$ 是 $m$-半可分的【索引8,Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality,2024,CoRR】。mHC【索引59,mHC: Manifold-Constrained Hyper-Connections,2026,http://cs.CL;索引64,mHC-lite: You Don’t Need 20 Sinkhorn-Knopp Iterations,2026,cs.LG】进一步约束每个 $A_l$ 为双随机矩阵,以稳定跨深度的累积乘积。

  • Full AttnRes 通过 $\phi(w_l, k_i) = \exp(w_l^\top \text{RMSNorm}(k_i))$ 和归一化计算 $M_{i \to l} = \alpha_{i \to l}$,其中 $k_i = v_i$ 是输入依赖的层输出,产生一个稠密的、秩为 $L$ 的 $M$。

  • Block AttnRes 将层划分为 $N$ 个块 $B_1, . . . , B_N$。对于一个已完成的早期块 $B_n$ 中的源 $i$,所有源共享块级键/值 $b_n$,因此对于每个 $i \in B_n$,$M_{i \to l} = \alpha_{n \to l}$。在当前块内,每层额外关注演变中的部分和 $b_n^{i-1}$,为每个块内位置引入一个额外的不同源。因此,$M$ 的有效秩介于 $N$ 和 $N + S$ 之间(其中 $S$ 是块大小),在标准残差($N=1$)和 Full AttnRes($N=L$)之间插值。

实用性。结构化矩阵的视角有两个目的。首先,它能够提供仅从递归形式中无法显现的分析性见解。例如,AttnRes 的输入依赖的 $M$ 揭示了深度方向的注意力汇点(attention sinks)(§5.4.2),即某些层无论输入如何都持续吸引高权重,这反映了序列方向注意力中的相同现象【索引57,Efficient streaming language models with attention sinks,2023,arXiv】。其次,它通过揭示核函数 $\phi$ 的哪些属性是重要的来为新设计提供信息。例如,当 $\phi$ 分解为 $\phi(q, k) = \varphi(q)^\top\varphi(k)$(对于某个特征映射 $\varphi$【索引23,Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention,2020,Proceedings of ICML】)时,深度方向的注意力会塌缩成一个递归——这正是前面提到的 MRLA-GLA 和 DDL-DeltaNet 对应关系背后的结构。

先前残差作为深度线性注意力。结构化矩阵的视角进一步通过展示现有残差变体实际上是深度轴上的线性注意力实例,与序列-深度对偶性相关联。例如,展开的 (m)HC 权重 $M_{i \to l} = \beta_i^\top A^\times_{i+1 \to l} \alpha_l$(公式10)可以自然地解释为注意力,其中 $\alpha_l$ 扮演第 $l$ 层发出的查询角色,$\beta_i$ 作为总结第 $i$ 层贡献的键,而累积转移 $A^\times_{i+1 \to l}$ 则充当控制中间层查询-键交互的深度相对位置算子【索引69,Kimi Linear: An Expressive, Efficient Attention Architecture,2025,http://cs.CL】。值得注意的是,$m$ 个并行流对应于深度轴上的状态扩展【索引40,HGRN2: Gated Linear RNNs with State Expansion,2024,http://cs.CL;索引29,Residual Matrix Transformers: Scaling the Size of the Residual Stream,2025,cs.LG】,将递归状态从 $d$ 扩展到 $d \times m$,从而增加了 $M$ 的半可分秩。【索引58,Your DeepSeek mHC Might Not Need the “m”,2026,Zhihu blog post】表明,用单位矩阵替换 $A^\times_{i+1 \to l}$ 仍能获得有竞争力的性能,突显了状态扩展的作用。通过这个视角,像 (m)HC 这样的方法充当了带有矩阵值状态的深度线性注意力,而 AttnRes 则充当了深度 softmax 注意力。

7 相关工作

归一化、缩放与深度稳定性。标准残差更新 $h_{l+1} = h_l + f_l(h_l)$【索引12,Deep Residual Learning for Image Recognition,2015,http://cs.CV】在归一化位置和梯度传播之间存在根本性的矛盾。PostNorm【索引52,Attention is All you Need,2017,Advances in NeurIPS】能保持有界的幅度但会扭曲梯度,因为残差路径上的重复归一化会在深度上累积导致梯度消失【索引60,On Layer Normalization in the Transformer Architecture,2020,cs.LG】。PreNorm【索引34,Transformers without Tears: Improving the Normalization of SelfAttention,2019,Proceedings of IWSLT;索引60】恢复了一条干净的恒等路径,但引入了无界的幅度增长:由于 $\|h_l\|$ 以 $O(L)$ 的速度增长,每层的相对贡献会缩小,迫使更深的层产生越来越大的输出,从而限制了有效深度【索引27,SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm,2026,cs.LG】。后续工作通过缩放残差路径【索引54,DeepNet: Scaling Transformers to 1,000 Layers,2022,http://cs.CL】、混合归一化【索引73,HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization,2025,http://cs.CL】、放大的跳跃连接【索引4,Post-LayerNorm Is Back: Stable, ExpressivE, and Deep,2026,cs.LG】或学习的逐元素门控【索引45,Highway Networks,2015,cs.LG】来调和这两个要求(见表5)。AttnRes 通过用对单个早期层输出的选择性聚合取代加性递归来回避这种矛盾,既避免了 PreNorm 的累积幅度增长,也避免了 PostNorm 的重复尺度收缩。

多状态循环。以上所有单状态方法都只将第 $l$ 层以 $h_{l-1}$ 为条件,从中无法选择性地检索单个早期层的贡献。有几种方法通过将递归扩展到多个并行流来解决这个问题:Hyper-Connections【索引72,Hyper-Connections,2025,cs.LG】及其稳定变体 mHC【索引59,mHC: Manifold-Constrained Hyper-Connections,2026,http://cs.CL】维护 $m$ 个带有学习混合矩阵的流;DDL【索引67,Deep Delta Learning,2026,cs.LG】维护一个通过 delta 规则的擦除-写入机制更新的矩阵状态;SiameseNorm【索引27,SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm,2026,cs.LG】维护两个参数共享的流——一个 PreNorm 和一个 PostNorm——以保留恒等梯度和有界表示。虽然这些方法缓解了信息压缩问题,但它们仍然以直接前驱的状态为条件;AttnRes 是正交的,它提供了对单个早期层输出的选择性访问,同时与任何归一化或门控方案兼容。我们在§6.2中讨论了与 Hyper-Connections 的形式化联系。

跨层连接性。另一条研究路线通过让每一层直接访问单个早期层的输出来绕过单状态瓶颈。最简单的方法是使用静态权重:DenseNet【索引17,Densely Connected Convolutional Networks,2018,http://cs.CV】连接所有先前的特征图;ELMo【索引38,Deep Contextualized Word Representations,2018,Proceedings of NAACL】用学习的标量权重计算层表示的 softmax 加权和;DenseFormer【索引36,DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging,2024,http://cs.CL】和 ANCRe【索引68,ANCRe: Adaptive Neural Connection Reassignment for Efficient Depth Scaling,2026,cs.LG】分配训练后固定的每层学习标量系数。对于输入依赖的聚合,MUDDFormer【索引56,MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections,2025,Proceedings of ICML】通过一个小型 MLP 在四个解耦的流上生成位置依赖的权重;MRLA【索引10,Cross-Layer Retrospective Retrieving via Layer Attention,2023,http://cs.CV】对所有先前层应用逐元素的 sigmoid 门控,尽管其可分离的查询-键乘积更接近于线性注意力而非基于 softmax 的检索。其他方法则用更有针对性的设计来换取完整的跨层访问:Value Residual Learning【索引71,Value Residual Learning,2025,Proceedings of ACL】只访问单个早期层;LAuReL【索引30,LAuReL: Learned Augmented Residual Layer,2025,cs.LG】用对前 $k$ 个激活的低秩投影来增强残差;Dreamer【索引24,Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves,2026,http://cs.AI】将序列注意力与深度注意力和稀疏专家结合起来。AttnRes 通过每层单个 $d$ 维伪查询,将 softmax 归一化、输入依赖的权重与对所有先前层的选择性访问结合起来,并引入了将成本从 $O(L^2)$ 降低到 $O(LN)$ 的块结构。基于缓存的流水线通信和两阶段计算策略(§4)使得 Block AttnRes 在大规模应用中具有可忽略的开销,变得切实可行。

A5 结论

受序列与深度之间对偶性的启发,我们引入了 AttnRes,它用学习到的、依赖于输入的深度注意力取代了固定的、统一的残差累积。我们通过消融研究和扩展定律实验验证了该方法,表明其增益在不同规模下持续存在。由于 Full AttnRes 必须在每一层访问所有先前的层输出,跨层聚合的内存占用以 $O(Ld)$ 的速度增长,这对于当前硬件上的大型模型来说是不可行的。因此,我们引入了 Block AttnRes,它将层划分为 $N$ 个块,并对块级表示进行注意力计算。经验表明,使用大约8个块可以恢复 Full AttnRes 的大部分增益,而随着未来硬件限制的放宽,采用更细粒度的分块仍然是一个有前途的方向。结合跨阶段缓存和两阶段计算策略,Block AttnRes 在大规模应用中是实用的,仅产生边际的训练开销和最小的推理开销。

A6 附录

B Full Attention Residuals 的优化推理 I/O

朴素实现的局限性。Full AttnRes 的一个朴素实现会在每一层扫描所有之前的层输出,因此内存流量会随着深度线性扩展。然而,正如§4.2所述,伪查询 $w_l$ 是一个与输入和隐藏状态都无关的可学习参数。因此,我们可以通过一个两阶段的调度方案,跨层批量处理块间的访问,从而将总 I/O 远低于朴素实现的上限。

优化方案设置。请注意,下面引入的块划分纯粹是一种推理调度手段。与 Block AttnRes 不同,它不改变模型架构,也不用块摘要替换每层的源;它只是为了使摊销论证具体化。假设模型有 $L$ 层和隐藏维度 $d$,划分为 $N$ 个大小为 $S = L/N$ 的连续块。推理一次处理一个块:阶段1联合计算块中所有 $S$ 层的块间注意力,针对所有之前的块;阶段2顺序处理块内的依赖关系。

阶段1:批处理块间注意力。考虑块 $n$ 及其 $S$ 层。查询 $\{w_l\}_{l \in B_n}$ 在执行开始前都是已知的,因此之前的 $(n-1)S$ 个键值对只需从 HBM 中读取一次,并在所有 $S$ 个查询中重用。因此,块 $n$ 的读取成本是

$$\text{Read}_{\text{inter}}^{(n)} = 2(n-1)Sd,$$

其中因子2解释了键和值。对所有 $N$ 个块求和并使用 $SN = L$:

$$\text{Read}_{\text{inter}} = \sum_{n=1}^{N} 2(n-1)Sd = 2Sd \cdot \frac{N(N-1)}{2} = dL(N-1).$$

总共写入成本为:

$$Write_{inter} = Ld$$

阶段2:顺序块内注意力。阶段1覆盖了当前块之前的所有源。然而,在块内,每一层都依赖于它之前的层,所以这些必须按顺序处理。第 $t$ 层($1 \le t \le S$)读取 $t-1$ 个块内键值对,成本为 $2(t-1)d$。对一个块求和:

$$\text{Read}_{\text{intra}}^{(n)} = \sum_{t=1}^{S} 2(t-1)d = S(S-1)d.$$

每层总摊销 I/O。对所有 $N$ 个块的两个阶段求和:

$$\text{Read}_{\text{total}} = dL(N - 1) + N \cdot S(S - 1)d, \qquad \text{Write}_{\text{total}} = 2Ld.$$

除以 $L$ 并使用 $SN=L$:

$$\text{Read per layer} = (N - 1)d + (S - 1)d = (S + N - 2)d, \quad \text{Write per layer} = 2d,$$

总的I/O为:

$$\text{Total I/O per layer} = (S + N) \, d.$$

因此,批量处理块间读取将每层的 I/O 从 $O(L)$ 降低到 $O(S+N)$。该调度遵循与 Block AttnRes 相同的两阶段划分:块间注意力占了大部分流量,而顺序计算则局限在每个块内部。