摘要： 残差连接配合 PreNorm 已成为现代大语言模型的标准范式，但它们以固定的单位权重累加所有层的输出。这种均匀的聚合方式导致隐藏状态的幅值随深度不受控地增长，逐步稀释每一层的贡献。我们提出 Attention Residuals (AttnRes) ，用对前序层输出的 softmax attention 替代这种固定累加，使每一层能够以学习得到的、依赖输入的权重选择性地聚合先前的表示。为了解决在大规模模型训练中对所有前序层输出做 attention 带来的显存和通信开销，我们引入 Block AttnRes ，将各层划分为若干块，仅在块级表示之间做 attention，在保留 Full AttnRes 大部分增益的同时降低了显存占用。结合基于缓存的流水线通信策略和两阶段计算方案，Block AttnRes 成为标准残差连接的即插即用替代方案，仅带来极小的额外开销。Scaling law 实验证实这一改进在不同模型规模下保持一致，消融实验验证了内容依赖的深度选择机制的有效性。我们进一步将 AttnRes 集成到 Kimi Linear 架构（总参数量 48B / 激活参数量 3B）中，在 1.4T token 上进行预训练。AttnRes 缓解了 PreNorm 稀释问题，使输出幅值和梯度分布在各深度上更加均匀，并在所有评估任务上提升了下游性能。

1. Intro

残差连接通过恒等映射让梯度绕过变换层，实现了深度网络的稳定训练。然而，如果我们观察残差连接的公式：

如果将上面这个公式展开，可以看到每一层接收到的其实是前序层的所有输出的均匀相加求和，没有任何机制来选择性地增强或者抑制个别层的贡献。

实践中，PreNorm基本成为主流范式，但是上面这种均匀加和导致隐藏状态的幅值随着深度不断增长，逐渐稀释了每一层的相对贡献。早期层的信息逐渐被淹没。

本文观察到，深度方向的累加与RNN中的递推之间存在形式化的对偶关系，于是提出Attention Residuals，将固定的累加$h_l={\sum }_i{v}_i$替换为$h_l={\sum }_i{a}_{i\to l}\cdot v_i$，其中$a_{i\to l}$是每层学到的一个pseudo-query伪查询计算得到的额softmax attention权重。这一轻量级机制仅用每层一个 d 维向量即可实现选择性的、内容感知的跨深度检索。

事实上，标准残差连接及先前基于递推的变体都可以被证明执行的是深度方向的linear attention ；AttnRes 将它们推广为深度方向的 softmax attention ，从而在深度维度上完成了与序列维度上同样具有变革性的线性到 softmax 的转变。

理想情况下，AttnRes的训练几乎不会引入任何额外开销，因为它所需要的每层的输出在做bwd的时候本来就会保留。然而，在现在的large scale训练下，activation的重计算和pipeline parallel被广泛使用，这些内容实际上没有被保存，Ful AttnRes的引入在这种情况下引入了巨大的额外占用。

为了解决如上情况，引入Block AttnRes来保持效率，将各个层划分为N个块，每个块内的层输出通过标准的Residual Addition合并为单一表示，从而降低显存和通信的开销。

Scaling law 实验证实 AttnRes 在各计算预算下始终优于基线，其中 Block AttnRes 的 loss 相当于使用 1.25× 更多计算量训练的基线。我们进一步将 AttnRes 集成到 Kimi Linear 架构（总参数量 48B / 激活参数量 3B）中，在 1.4T token 上进行预训练。对训练动态的分析表明，AttnRes 缓解了 PreNorm 稀释问题，输出幅值在各深度上保持有界，梯度范数在各层间分布更加均匀。在下游基准测试中，我们的最终模型在所有评估任务上均优于基线。

本文的Contribution如下：

• Attention Residuals。 我们提出 AttnRes，用学习到的深度方向 softmax attention 替代固定的残差累加，及其可扩展变体 Block AttnRes，将显存和通信从$O(Ld)$降低到$O(Nd)$。通过统一的结构化矩阵分析，我们证明标准残差及先前基于递推的变体对应于深度方向的linear attention，而 AttnRes 执行的是深度方向的 softmax attention。
• 面向规模化的基础设施。 我们开发了使 Block AttnRes 在大规模训练中实用且高效的系统优化，包括在流水线并行下消除冗余传输的跨阶段缓存，以及通过 online softmax 摊销跨块 attention 的两阶段推理策略。由此带来的训练开销极小，在典型推理工作负载上延迟开销不到 2%。
• 全面的评估与分析。 我们通过 scaling law 实验、组件消融和在 1.4T token 上预训练的 48B 参数模型的下游基准测试验证了 AttnRes，展示了相对于标准残差连接的一致改进。训练动态分析进一步表明，AttnRes 缓解了 PreNorm 稀释，使隐藏状态幅值保持有界，梯度分布在各深度上更加均匀。

2. Motivation

2.1. Training Deep Networks via Residuals

Residual Learning##

考虑残差连接：

展开可以得到：

残差连接背后的关键洞察是恒等映射 ：每一层都保留了一条直接路径，使得信息和梯度都能不受干扰地流动。

考虑梯度：

展开该乘积得到$I$加上涉及层 Jacobian 矩阵 $\partial f_j/\partial h_j$ 的高阶项。恒等项始终保留，提供了从损失到任何层的直接梯度路径，不受深度影响。

Generalizing Residuals

Highway Network引入可学习的Elementwise gating：

Limitations

然而，上述的方法中都存在一个根本性的局限性：每一层只能访问其直接输入$h_{l-1}$，而它是一个将所有早期层输出混合在一起的单一压缩状态，这导致：

1. 访问变得没有选择性，不同类型的层（attention vs. MLP）接收到相同的剧和信息，即使它们可能受益于不同的加权方式；
2. 不可逆的信息损失；
3. 输出的数值大小不断增大，较深的层学习到用越来越大的输出来在residual addition中获得影响力。

这些局限性促使我们设计一种机制，让每一层能够选择性地从所有前序层聚合信息。

3. Attention Residuals: A Unified View of Time and Depth

上面提到的所有Limitation让人回想起在序列建模/时序预测领域曾经遇到的问题。

The Duality of Time and Depth

正如RNN在时间维度上所作的，Residual将历史信息压缩为单一状态$h_l$，而Transformer能够以Data Dependent的方式访问所有的prefix。对于这次的问题，可以提出相同的方法论：

虽然计算的难度变成了$O(L^2)$，但网络的深度一般不会像sequence问题那样成千上万，往往只有有限的层。这种方法称为AttnRes。

3.1. Ful Attention Residuals

attention weight可以写作${\alpha }_{i\to l}=\varphi \left({\mathbf{q}}_l,{\mathbf{k}}_i\right),\varphi :{\mathbb{R}}^d\times {\mathbb{R}}^d\to {\mathbb{R}}_{\ge 0}$，不同的$\varphi$可以得到不同的残差变体，此处采用$\varphi (\mathbf{q},\mathbf{k})=\exp \left({\mathbf{q}}^{\top }\text{ RMSNorm }(\mathbf{k})\right)$，得到深度方向上的softmax attention：

对每层，定义：

其中$q_l$是一个每层的可学习向量，第$l$层的输入：

我们将此形式称为 full attention residuals 。对每个 token，Full AttnRes 需要 $O(L^{2}d)$ 的计算量和 $O(Ld)$的显存来存储层输出。由于深度远小于序列长度，计算开销是适度的。

Overhead

$O(Ld)$ 的显存与反向传播中已经需要保留的激活完全重叠，因此 Full AttnRes 在常规训练中不引入额外的显存开销。然而在大规模训练中，激活重计算和流水线并行被广泛采用：原本可以释放并重新计算的层输出现在必须为所有后续层保留，在流水线并行下还需在阶段边界之间传输。此时显存和通信开销均增长为 $O(Ld)$。

Blockwise optimization

Full AttnRes 中一个刻意的设计选择是伪查询 $w_l$ 是与层的前向计算解耦的可学习参数。这种独立性意味着，任意一组层的 attention 权重可以并行计算，无需等待它们的顺序输出；特别地，可以将 $L$ 层分组为 $N$ 个各含 $S$ 层的块，并在每个块内批量化 attention 计算，将每层的显存 I/O 从 $O(Ld)$ 降低到 $O((S+N)d)$。然而在当前的分布式训练体制下，主要成本不是本地显存带宽而是流水线并行下的跨阶段通信：每个层输出仍然必须在阶段间传输，这个 $O(Ld)$ 的通信开销无法通过本地批量化来缓解。这促使了下面介绍的 Block AttnRes 变体，它将跨阶段表示的数量从 $L$ 减少到 $N$。我们预期未来互连技术的改进将使完整的 $O(Ld)$ 通信变得可行，充分释放 Full AttnRes 的潜力。

3.2. Block Attention Residuals

我们将 $L$ 层划分为 $N$ 个各含 $S=L/N$ 层的块，假设 $L$ 能被 $N$ 整除；否则，最后一个块包含剩余的 $L\mathrm{mod}N$ 层。令 $B_n$ 表示第 $n$ 个块 $(n=1,...,N)$ 中的层索引集合。为形成一个块，我们将其所有层输出求和：

Inter-Block Attention

在 Full AttnRes 中，第 l 层的输入通过对直到 $f_{l-1}(h_{l-1})$ 的所有输出做 attention 来计算。块级变体用块表示替代这些独立输出，定义 $b_0=h_1$以确保 token embedding 始终作为源包含在内。对于第 $n$ 个块中的第 $i$ 层，值矩阵为：

网络最前面一层的输入就是 token embedding，即 $b_0=h_1$。在每个块中，第一层接收前序块表示和 token embedding，后续各层还额外 attend 到部分和 $b_n^{i-1}$。最终的输出层聚合所有$N$ 个块表示。

4. Infrastructure Design

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

4.1. Training

对于小规模训练，AttnRes 仅增加微小的计算开销且不增加额外显存使用，因为无论如何都需要为反向传播保存激活。在大规模分布式训练下，流水线并行构成了 AttnRes 的主要基础设施挑战。

Pipeline communication

在标准残差连接下，流水线并行在相邻阶段之间传输固定大小的隐藏状态，与流水线深度无关。Block AttnRes 需要在每个阶段保留所有累积的块表示以进行跨块 attention，朴素地在每次转换时传输完整历史会导致冗余通信。

考虑具有 $P$ 个物理阶段和每物理阶段 $V$ 个虚拟阶段的交错流水线调度。为简单起见，假设每个物理阶段平均产生 $N_p$ 个维度为 $d$ 的块表示。总共 $C=PV$ 个分块（每个物理阶段在每个虚拟阶段中各一个），第 $j$ 个分块累积 $j{N}_p$ 个块。朴素地在每次转换时传输所有累积块带来的每 token 通信成本为：

Cross-stage caching

由于每个物理阶段依次处理多个虚拟阶段，我们可以通过本地缓存来消除这种冗余：在较早虚拟阶段期间接收的块保留在本地显存中，无需重新传输。第一个虚拟阶段 $(v=1)$没有缓存，按正常方式累积；对于 $v\ge 2$，每次转换仅传送自接收方在前一虚拟阶段中对应分块以来累积的约 $P{N}_p$ 个增量块。总通信量减少为：

4.2. Inference

下面描述的两阶段计算策略同时适用于 Full 和 Block AttnRes：在两种情况下，层都被分组为大小为 S 的块，Phase 1 批量化跨块查询，Phase 2 处理顺序的块内回溯。对于 Full AttnRes，这将每层 I/O 从 $O(Ld)$ 降低到 $O((S+N)d)$；Block AttnRes 进一步将存储的表示从 L 减少到 N，因为每个块被压缩为单一向量。

Tow-phase computation strategy

Block AttnRes 的逐层 attention 计算类似于自回归解码，其中块表示充当跨层复用的共享 KV cache。朴素实现在每一层计算 attention 残差，每次都需要遍历所有前序块，导致 $O(L\cdot N)$ 次显存访问。由于伪查询向量与前向计算解耦，一个块内的所有 $S=L/N$ 个查询可以批量化为单次矩阵乘法，将显存访问从 $S$ 次读取摊销为 1 次。

• Phase 1 通过对缓存的块表示进行单次批量查询，同时为所有 $S$ 层计算跨块 attention，返回输出和 softmax 统计量（最大值和 log-sum-exp）。这摊销了显存访问成本，将读取次数从每块$S$ 次减少到仅 1 次。
• Phase 2 使用不断演进的部分和为每一层顺序计算块内 attention，然后通过 online softmax 与 Phase 1 的输出合并。由于 online-softmax 合并是逐元素的，这一阶段天然支持与周围操作的 kernel 融合，进一步减少 I/O 开销。

通过两阶段设计，Phase 2 保持了与标准残差连接相似的 I/O 占用，而主要的额外成本来自 Phase 1 的跨块 attention。由于这些跨块读取通过批量化被摊销到块内的所有层，每层的总显存访问成本仅为 $(N/S+3)d$ 次读取和 $2d$ 次写入。这远低于先前残差推广方案如 (m)HC 在典型设置下的残差流 I/O。在实践中，Phase 1 还可以与块中第一层的计算部分重叠，进一步降低其挂钟时间影响。因此，在典型推理工作负载上，端到端推理延迟开销不到 2%。

Memory-efficient prefilling

预填充期间存储块表示需要 $N\cdot T\cdot d$ 个元素，对于 128K token 序列和 8 个块，这将产生 15 GB 的显存占用。我们通过沿序列维度在 $P$ 个张量并行设备上分片这些表示来缓解此问题，使 Phase 1 能够在本地序列分片上独立执行。Phase 2 的 online-softmax 合并随后集成到标准的 TP all-reduce 通信路径中：输出先 reduce-scatter，本地合并，再通过 all-gather 重建，天然支持与 RMSNorm 等操作的 kernel 融合。这将每设备显存占用降低到 $N\cdot (T/P)\cdot d$。

以 128K 上下文为例，从 15 GB 降至约每设备 1.9 GB。结合分块预填充（如 16K 块大小），开销进一步降至每设备不到 0.3 GB。

5. Experiments

我们的架构与 Kimi Linear 完全一致，这是一个遵循 Moonlight / DeepSeek-V3 设计的MoE Transformer，以 3:1 的比例交错使用 Kimi Delta Attention 和 Multi-Head Latent Attention 层，每层后跟一个 MoE 前馈层。唯一的修改是在残差连接中添加了 AttnRes；所有其他组件（模型深度、隐藏维度、专家路由和 MLP 结构）保持不变。AttnRes 每层仅引入一个 RMSNorm 和一个伪查询向量 $w_l\in R^d$，占总参数量的可忽略部分。至关重要的是，所有伪查询向量必须初始化为零。这确保初始 attention 权重 ${\alpha }_{i\to l}$ 在各源层上是均匀的，从而在训练开始时将 AttnRes 退化为等权平均，并防止训练波动，我们已通过实验验证了这一点。

5.1. Scaling Laws

扫描五个模型大小，每个大小训练三个变体：PreNorm 基线、Full AttnRes 和使用约 8 个块的 Block AttnRes。它们使用 8192 token 的上下文窗口和余弦学习率调度进行训练。在每个 scaling law 大小组内，所有变体共享在基线下选择的相同超参数，以确保公平比较；这种设置有意偏向基线，因此使比较更加保守。遵循标准做法，我们拟合形如 $\mathcal{L}=A\times C^{-\alpha }$的幂律曲线，其中 $\mathcal{L}$ 是验证 loss，$C$ 是以 PFLOP/s-days 衡量的计算量。

三个变体展现出相似的斜率，但 AttnRes 在整个计算范围内始终取得更低的 loss。根据拟合曲线，在 5.6 PFLOP/s-days 时，Block AttnRes 达到 1.692 而 Baseline 为 1.714，等效于 1.25× 的计算优势。Full 和 Block AttnRes 之间的差距随规模缩小，在最大规模时仅为 0.001。我们还在表 2 中列出了 mHC(-lite) 作为参考。Full AttnRes 优于 mHC，而 Block AttnRes 在更低的每层显存 I/O 下与之持平：5.5d vs. mHC 的 34d。

5.2. Main Results

Training Recipes

我们研究的最大模型基于完整的 Kimi Linear 48B 配置：27 个 Transformer block（54 层），256 个路由专家中选择 8 个加 1 个共享专家，产生总共 48B 和 3B 激活参数。该模型应用 Block AttnRes，每块 6 层，产生 9 个块加上 token embedding，总共 10 个深度方向的源。

我们遵循与 Kimi Linear 1.4T token 运行相同的数据和训练方案：所有模型使用 4096 token 的上下文窗口、Muon 优化器和 WSD（Warmup–Stable–Decay）学习率调度进行预训练，全局批大小为 8M token。最终模型的训练分两个阶段进行：

1. 在 1T token 上的 WSD 预训练阶段，之后是
2. 在约 400B 高质量 token 上的中期训练 (mid-training) 阶段，遵循 Moonlight的退火方案。

中期训练之后，我们继续以逐步增加的序列长度训练至 32K token。由于我们的架构使用混合 KDA/MLA attention，其中 MLA 不使用位置编码 (NoPE)，上下文扩展无需 YaRN或 attention 温度重缩放等修改。

Training dynamics

• 验证 loss： AttnRes 在整个训练过程中始终取得更低的验证 loss，差距在衰减阶段扩大，最终 loss 显著更低。
• 输出幅值： Baseline 受到 PreNorm 稀释问题的困扰：由于隐藏状态幅值随深度单调增长，更深的层被迫从固定尺度的归一化输入中学习越来越大的输出以保持影响力。Block AttnRes 将这种增长限制在每个块内，因为块边界处的选择性聚合重置了累加，产生有界的周期性模式。
• 梯度幅值： 由于所有残差权重固定为 1，Baseline 没有调节跨深度梯度流的手段，导致最早层的梯度不成比例地大。Block AttnRes 中可学习的 softmax 权重在各源之间引入了概率质量的竞争，产生了明显更均匀的梯度分布。

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

5.3. Ablation Study

Comparison with prior methods

DenseFormer 赋予每层对所有前序输出的访问权，但使用固定的、不依赖输入的标量系数来组合它们；它相比基线没有增益 (1.767)，突显了依赖输入的加权方式的重要性。

mHC 通过 m 个并行流和学习到的混合矩阵引入了输入依赖性，改进到 1.747。

AttnRes 通过 softmax attention 的显式内容依赖选择进一步推进：Full AttnRes 达到 1.737，Block AttnRes 达到 1.746，两者仅用每层一个查询向量即优于上述方法。

Cross-layer access

我们比较了三种粒度的跨层访问。Full AttnRes 直接源于时间-深度对偶，对所有前序层做 attention，取得最低 loss (1.737)。一种减少其显存开销的简单方式是滑动窗口聚合 (SWA)，仅保留最近 W=8 个层输出加上 token embedding；它优于基线 (1.764) 但远不及 Full 和 Block AttnRes，表明选择性地访问远距离层比关注许多邻近层更为重要。

Block AttnRes 提供了更好的权衡：块大小 S=4 时达到 1.746，同时保持每层恒定的显存开销。图 6 将 S 从 S=1（即 Full AttnRes）扫描到越来越粗的分组。Loss 随 S 增大而优雅退化，S=2, 4, 8 均接近 1.746，而更大的块 (S=16, 32) 则趋向基线。在实践中，出于基础设施效率考虑，我们将块数固定为约 8 (§4)。随着未来硬件缓解显存容量限制，采用更细粒度的块大小或 Full AttnRes 是进一步提升性能的自然路径。

Component design

我们进一步消融了 attention 机制的各个组件：

• 依赖输入的查询。 一种自然的扩展是通过从当前隐藏状态投影来使查询依赖输入。这进一步将 loss 降低到 1.731，但每层引入了一个 $d\times d$ 的投影，并且在解码期间需要顺序显存访问，因此我们默认使用学习到的查询。
• 不依赖输入的混合。 我们移除了查询和键，替换为可学习的、不依赖输入的标量来加权前序层，这损害了性能 (1.749 vs. 1.737)。
• softmax vs. sigmoid。 用 sigmoid 替换 softmax 降低了性能 (1.741)。我们将此归因于 softmax 的竞争性归一化，它迫使在各源之间做出更锐利的选择。
• 多头 attention。 我们在 Block AttnRes 上测试了逐头的深度聚合 (H=16)，允许不同的通道组 attend 到不同的源层。这损害了性能 (1.752 vs. 1.746)，表明最优的深度方向混合在各通道上大体一致：当一层的输出是相关的，它作为整体都是相关的。
• 键上的 RMSNorm。 移除 RMSNorm 降低了 Full AttnRes (1.743) 和 Block AttnRes (1.750) 的性能。对于 Full AttnRes，它防止自然输出幅值较大的个别层主导 softmax。对于 Block AttnRes 这一点更为关键，因为块级表示在更多层上累积，可能产生较大的幅值差异；RMSNorm 防止这些差异偏置 attention 权重。

5.4. Analysis

5.4.1. Optimal Architecture

为了理解 AttnRes 如何重塑最优架构 scaling，我们在固定计算和参数预算下进行了受控的容量重分配研究。

我们的核心问题是 AttnRes 是否改变了首选的深度-宽度-attention 权衡，特别是鉴于其在深度维度上的潜在优势，它是否相比传统 Transformer 设计启发式更偏好更深的模型。

为了隔离与深度直接耦合的结构因素，我们基于内部经验观察固定了每专家 MLP 扩展比 ($d_{ff}/d_{model}\approx 0.45$)。我们进一步固定了总训练计算量 (FLOPs ≈ 6.5 × 10¹⁹) 和激活参数量 (≈ 2.3 × 10⁸)，确保任何性能变化纯粹来自架构重分配而非总体容量差异。

在此受限预算下，我们在 $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×5 网格上枚举了 25 种配置，其中 $L_b=L/2$ 是 Transformer block 数，H 是 attention head 数。

两个热力图展现了共同的模式：loss 随 $d_{model}/L_b$ 增大和 $H/L_b$ 减小而降低，两种方法的最优点均在 $H/L_b\approx 0.3$ 处。

尽管有这种共同趋势，AttnRes 在全部 25 种配置中均取得了低于基线的 loss，差异为 0.019–0.063。

最明显的差异在于最优点的位置：基线在 $d_{model}/L_b\approx 60$ 处取得最低 loss (1.847)，而 AttnRes 将其移至 d_model/L_b ≈ 45 (1.802)。

在固定参数预算下，较低的 $d_{model}/L_b$ 对应更深更窄的网络，表明 AttnRes 能更有效地利用额外的深度。我们注意到，这种对深度的偏好不能直接转化为部署建议，因为更深的模型由于其顺序计算通常会产生更高的推理延迟。相反，这一扫描作为诊断工具揭示了 AttnRes 最受益之处，这种深度偏好可以在架构选择中与推理成本一并考虑。

5.4.2. Analyzing Learned AttnRes Patterns

我们强调三个关键观察：

• 保留的局部性。 每一层最强烈地 attend 到其直接前驱，但选择性的非对角集中也会出现（如第 4 层 attend 到早期源，第 15-16 层在 block 设置下回溯到远处），表明在标准残差路径之外学习到了 skip connection。
• 层间特化。 embedding ** h**₁ 在整个过程中保持非平凡的权重，特别是在 pre-attention 层中。Pre-MLP 输入显示出对近期表示更尖锐的对角依赖，而 pre-attention 输入维持更广的感受野，这与 attention 在层间路由信息而 MLP 在局部操作的特性一致。
• Block AttnRes 保留了结构。 对角主导性、embedding 持久性和层间特化都从 full 变体转移到了 block 变体，表明块级压缩起到了隐式正则化的作用，同时保留了基本的信息通路。

6. Disscussion

6.1. Sequence-Depth Duality

残差连接通过固定递推 $h_l=h_{l-1}+f_{l-1}(h_{l-1})$ 在深度方向传播信息，这与 RNN 在时间方向传播信息的方式非常相似。Test-Time Training (TTT) 从序列侧形式化了这一类比（参见 Fast Weight Programmers），将每个递推步骤视为对自监督损失的一次梯度下降：

其中，一个慢速网络对 $\ell$ 进行参数化，状态 $\mathbf{W}$ 在每个 token 上更新一次。当 $f$ 为线性函数时，该过程退化为标准的线性 attention：

标准残差在深度方向呈现出同样的加性形式：$h_l$作为状态，每层 $f_l$ 相当于一次“梯度步”。

这种对偶性还能扩展到更丰富的变体。序列侧的数据依赖门控对应于深度侧的 Highway 网络；delta 规则对应于 DDL；MRLA 则镜像了 GLA 的门控线性 attention。这些方法都仍在递推范式内对更新规则进行细化。AttnRes 更进一步，用显式的跨层 attention 替代深度方向的递推，正如 Transformer 用自注意力替代时间方向的递推。由于当前架构的层数仍处于 softmax attention 的可行范围内，我们采用常规的深度方向 attention。引入更具表达力、但更显存高效（例如线性复杂度）的替代方案，是未来工作的自然方向。

6.2. Residual Connections as Structured Matrices

上面讨论的残差变体都可以被视为对前序层输出的加权聚合。我们用一个深度混合矩阵 M ∈ ${\mathbb{R}}^{L\times L}$ 来形式化这一点，其中 $M_{i\to l}$ 是第 $l$ 层赋予第 $i$ 层输出的权重。

各变体的区别在于这些权重如何产生（固定的、学习到的或依赖输入的）以及 M 是否被约束为低秩或允许为稠密的。** M** 的半可分 (semiseparable) 秩提供了统一的比较视角。

具体而言，第 $l$ 层的输入为 $h_l={\sum }_{i=0}^{l-1}{M}_{i\to l}{v}_i$，其中 $v_0=h_1$（embedding），对 $i\ge 1$ 有 $v_i=f_i(h_i)$。

我们逐一推导如下。

• 标准残差
  • $h_l=h_{l-1}+f_{l-1}(h_{l-1})$
  • 展开得 $h_l={\sum }_{i=0}^{l-1}{v}_i$
  • 因此对所有 $i<l$ 有 $M_{itol}=1$，M 是全 1 的下三角矩阵。
• Highway
  • $h_l=(1-g_l),h_{l-1}+g_l,f_{l-1}(h_{l-1})$（此处为简明起见使用标量门控；逐元素扩展是直接的）
  • 定义传递乘积 ${\gamma }_{i\to l}^{\times }:={\prod }_{j=i+1}^l(1-g_j)$
  • 权重为 $M_{0\to l}={\gamma }^{\times }1\to l$（对 embedding）和 M{i\to l}=g_{i+1},\gamma^{times}_{i+1\to l}$$（对 i\ge 1$）
  • 由于累积乘积通过标量门控因式分解，M 是 1-半可分的 [8]，与标准残差相同的秩但具有依赖输入的权重。
  • 权重的和自动为 1，使 Highway 成为无 softmax 的深度方向 stick-breaking attention [49] 的实例。
• (m)HC
  • 维护 $m$ 个并行流 $H_l\in {\mathbb{R}}^{d\times m}$，通过以下方式更新：$H_l=H_{l-1}{A}_l+f_{l-1}(H_{l-1}{\alpha }_{l-1}){\beta }_{l-1}^{\top }$
  • 其中 $A_l\in {\mathbb{R}}^{m\times m}$ 是学习到的转移矩阵，${\alpha }_{l-1}\in {\mathbb{R}}^m$ 将流混合为 $f_{l-1}$ 的单一输入，${\beta }_{l-1}\in {\mathbb{R}}^m$ 将输出分配回各流。
  • 展开递推得到有效权重：
    • 其中 $A_{i\to j}^{\times }:={\prod }_{k=i+1}^j{A}_k$
    • $m\times m$ 的转移使 M 成为 m-半可分的
    • mHC [59, 64] 进一步约束每个 $A_l$ 为双随机矩阵，稳定跨深度的累积乘积。$M_{i\to l}={\beta }_i^{\top }{A}_{i+1\to l}^{\times }{\alpha }_l$
• Full AttnRes
  • 通过 $\varphi (w_l,k_i)=\exp (w_l^{\top }\mathrm{RMSNorm}(k_i))$ 并归一化计算 $M_{i\to l}={\alpha }_{i\to l}$
  • 其中 $k_i=v_i$ 是依赖输入的层输出
  • 产生稠密的、秩为 $L$ 的 M 。
• Block AttnRes
  • 将层划分为 N 个块 B₁, …, B_N。
  • 对于已完成的早期块 $B_n$ 中的源 $i$，它们共享块级键/值 $b_n$，因此对 $B_n$ 中的每个 $i$ 有 $Mi\to l={\alpha }_{n\to l}$。
  • 在当前块内，每一层还额外 attend 到不断演进的部分和 $b_n^{i-1}$，为每个块内位置引入一个额外的独立源。
  • M 的有效秩因此介于 $N$ 和 $N+S$ 之间（其中 $S$ 是块大小），在标准残差 (N=1) 和 Full AttnRes (N=L) 之间插值。

Prior Residuals as Depth-Wise Linear Attention

结构化矩阵的视角进一步通过表明现有残差变体实际上是深度轴上的线性 attention 实例来与序列-深度对偶相联系。例如，展开的 (m)HC 权重 $M_{i\to l}=\beta i\top A{i+1\to l}^{\times }{\alpha }_l$ 允许自然的 attention 解释：α l 扮演第 l 层发出的查询角色，β i 作为概括第 i 层贡献的键，累积转移 $A_{i+1\to l}^{\times }$ 作为控制跨中间层的查询-键交互的深度相对位置算子 。

值得注意的是，m 个并行流对应于深度轴上的状态扩展 ，将递推状态从 d 扩展到 $d\times m$，从而增加 M 的半可分秩。用恒等矩阵替换 $A_{i+1\to l}^{\times }$ 仍能产生有竞争力的性能，突出了状态扩展的作用。通过这一视角，(m)HC 等方法充当具有矩阵值状态的深度方向线性 attention，而 AttnRes 充当深度方向的 softmax attention。

7. Related Work

Normalization, Scaling, and Depth Stablility

标准残差更新 $h_{l+1}=h_l+f_l(h_l)$ 呈现了归一化位置与梯度传播之间的根本张力。PostNorm 维持了有界的幅值但扭曲了梯度，因为对残差路径的反复归一化在深度上复合为梯度消失。PreNorm 恢复了干净的恒等路径，但引入了无界的幅值增长：由于 $\Vert h_l\Vert$ 以 $O(L)$ 增长，每层的相对贡献缩小，迫使更深的层产生越来越大的输出并限制了有效深度。后续工作通过缩放残差路径、混合归一化、放大 skip connection或学习到的逐元素门控来调和这两个目标。

AttnRes 通过用对个别早期层输出的选择性聚合替代加性递推来回避这一张力，既避免了 PreNorm 的累积幅值增长，也避免了 PostNorm 的反复尺度收缩。

Multi-State Recurrence

上述所有单状态方法都仅以 $h_{l-1}$ 为条件，无法从中选择性地检索个别早期层的贡献。若干方法通过将递推扩展到多个并行流来解决这一问题：Hyper-Connections 及其稳定变体 mHC 维护 m 个流和学习到的混合矩阵；DDL 维护通过 delta 规则的擦写 (erase-and-write) 机制更新的矩阵状态；SiameseNorm 维护两个参数共享的流——一个 PreNorm 和一个 PostNorm——以保留恒等梯度和有界表示。虽然这些方法减轻了信息压缩，但它们仍以直接前驱的状态为条件；AttnRes 是正交的，提供对个别早期层输出的选择性访问，同时与任何归一化或门控方案兼容。

Cross-Layer Connectivity

另一条研究路线通过赋予每层对个别早期层输出的直接访问来绕过单状态瓶颈。最简单的方法使用静态权重：DenseNet 拼接所有前序特征图；ELMo 用学习到的标量权重计算层表示的 softmax 加权和；DenseFormer 和 ANCRe 分配训练后固定的学习到的逐层标量系数。对于依赖输入的聚合，MUDDFormer 过小型 MLP 在四个解耦流上生成位置依赖的权重；MRLA 在所有前序层上应用逐元素 sigmoid 门控，尽管其可分离的查询-键乘积更接近线性 attention 而非基于 softmax 的检索。其他方法以更有针对性的设计换取了完整的跨层访问：Value Residual Learning 仅访问单一较早层；LAuReL 用前 k 个激活的低秩投影增强残差；Dreamer 将序列 attention 与深度 attention 和稀疏专家结合。

AttnRes 通过每层一个 d 维伪查询将 softmax 归一化的、依赖输入的权重与对所有前序层的选择性访问相结合，并引入块结构将成本从 $O(L^2)$ 降至 $O(LN)$。

Conclusion

受序列与深度之间对偶性的启发，我们引入了 AttnRes，用学习到的、依赖输入的深度方向 attention 替代固定的、均匀的残差累加。我们通过消融研究和 scaling law 实验验证了该方法，表明其增益在各规模下持续存在。由于 Full AttnRes 必须在每一层访问所有前序层的输出，跨层聚合的显存占用增长为 $O(Ld)$，在当前硬件上对大规模模型来说过于昂贵。因此我们引入 Block AttnRes，将层划分为 $N$ 个块并在块级表示上做 attention。实验表明，使用约 8 个块即可恢复 Full AttnRes 的大部分增益，而更细粒度的分块仍是未来硬件约束放松后的一个有前景的方向。配合跨阶段缓存和两阶段计算策略，Block AttnRes 在大规模下是实用的，仅带来极小的训练开销和极低的推理开销。

相当于添加了一个纵向的SoftmaxAttention，几个Infra优化恰好就是Flash Attention的优化，很有趣！听说一作是高三的学生。