技术摘要：用于深度仅解码器 Transformer 的多分辨率残差路由

问题陈述

标准的仅解码器 Transformer 依赖于 PreNorm 残差连接，其中子层更新通过固定的单位权重进行聚合。虽然这种结构能够实现稳定的训练，但它施加了一种僵化的累积规则，随着网络深度的增加，可能会稀释单个层的贡献。最近的进展，如 Attention Residuals 和 Block Attention Residuals (Block AttnRes)，通过使用基于内容的、学习到的对先前输出的路由，取代了固定累积，解决了这一问题。Block AttnRes 特别通过将多个层压缩为块级残差摘要（$C_b$）来提高效率。

然而，作者指出 Block AttnRes 存在一个关键局限性：单个块摘要仅代表了块内残差轨迹的低频（DC）分量。通过仅存储更新的总和 ($\sum u_{b,i}$)，该机制丢弃了宝贵的方向结构，例如：

1. 子层之间的不平衡： 该块是由 Attention 更新主导还是由 MLP 更新主导。
2. 时间动态： 块内的早期更新是否指向与后期更新不同的方向。

论文假设，随着残差轨迹长度的增长，这种丢失的方向信息对于有效的深度路由变得愈发至关重要。

方法论：WAV v1

作者提出了 WAV v1 (Waveform Attention v1)，它是 Block AttnRes 的一种轻量级、多分辨率扩展。WAV v1 不仅保留了原始的 Attention 和 MLP 模块，还通过两个零和方向细节基底（zero-sum directional detail bases）增强了每个块的表示，而不仅仅是用总和来表示一个块。

1. 多分辨率表示

对于包含 $m$ 个子层更新 $\{u_{b,i}\}_{i=1}^m$ 的块 $b$，WAV v1 计算：

• 状态基底 ($C_b$)： 标准的块总和，代表低频状态。
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• 相位基底 ($D^{phase}_b$)： 一个对比 Attention 与 MLP 更新的细节源。它为 Attention 更新分配 $+1$ 的符号，为 MLP 更新分配 $-1$ 的符号。
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{其中 } a_i \in \{+1, -1\}$$
• 分裂基底 ($D^{split}_b$)： 一个对比块的前半部分与后半部分的细节源。它为前半部分分配 $+1$，为后半部分分配 $-1$。
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{其中 } r_i \in \{+1, -1\}$$

这些基底在产生更新时进行在线累积，产生的计算开销微乎其微。

2. 路由机制

深度 Softmax 混合器接收一个扩展后的源池。对于 Block AttnRes，该池为 $\{e, C_0, \dots, C_{b-1}, PC\}$。对于 WAV v1，它扩展到包含每个已完成块及其当前部分块的细节基底：
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. 稳定性机制

为了防止训练初期出现不稳定，WAV v1 采用了两种保守策略：

• 负向细节偏差： 细节源被初始化为一个强烈的负偏差 ($\beta_D = -2.0$)，确保模型在初始化时表现得像标准的 Block AttnRes，并且仅在有益时才利用细节。
• 脱离式 RMS 匹配： 细节张量相对于其对应的块摘要 ($C$) 进行缩放，使用脱离（detached）的 RMS 比率，并限制最大缩放因子 $\rho$。这防止了细节仅仅因为较大的原始量级而占据主导地位。
• 保守读取： 在默认情况下，最终的预测混合器仅读取嵌入（embedding）和 $C$ 源，而不读取细节源，从而进一步确保稳定性。

核心贡献

1. 公式化： 论文将块残差路由框架化为一个多分辨率表示问题，区分了低频状态信息 ($C_b$) 和块内轨迹信息（方向细节）。
2. 算法： 引入了 WAV v1，这是一种掉入式（drop-in）扩展，它通过增加相位和分裂细节基底来增强块摘要，而无需修改核心的 Attention 或 MLP 函数。
3. 实证证据： 初步的缩放实验表明，多分辨率路由的效用是依赖于深度的，仅在更深的架构中才会显著受益。

实验结果

作者在字符级数据集 TinyStories 和 Text8 上评估了 12 层、24 层和 48 层的 GPT 仅解码器模型。

• 深度依赖趋势：
  • 12 层： WAV v1 的表现逊于 Block AttnRes，表明在浅层网络中方向细节是冗余或带有噪声的。
  • 24 层： WAV v1 具有竞争力，在 TinyStories 上显示出轻微改进，并在 Text8 上表现相当。
  • 48 层： WAV v1 在两个数据集上均实现了最佳验证损失，显著优于 Block AttnRes、ReZero 和 LayerScale。
• 定量收益 (48 层)：
  • TinyStories： 验证损失从 0.4960 (Block AttnRes) 降低到 0.4738 (WAV v1)，困惑度 (PPL) 降低了 2.19%。
  • Text8： 验证损失从 0.9363 降低到 0.9305，PPL 降低了 0.57%。
• 训练动态： 在 48 层时，WAV v1 在训练早期就与 Block AttnRes 分离，并保持持续优势，表明这种改进不是后期阶段的人为产物，而是优化轨迹的基础性转变。

意义与主张

论文声称，深度残差路由不仅受益于选择哪个块进行读取，还受益于路由每个块的内部方向结构。结果表明，随着残差轨迹变长，更新轨迹的“形状”（例如，Attention 主导地位或早期与晚期动态）成为了深度混合器的一个有效信号。

作者将 WAV v1 定位为现有方法的补充方法：

• 不同于 ReZero 或 LayerScale（它们调节更新的幅度），WAV v1 改变了路由器可用的表示。
• 不同于 Delta Attention Residuals（它们对单个增量进行路由，成本很高），WAV v1 提供了一种块级、多分辨率的近似，在效率与结构丰富性之间取得了平衡。

作者总结道，未来对残差路由方法的评估应考虑深度，因为多分辨率信息的益处可能与残差轨迹的长度和结构有关。他们承认目前的局限性，包括缺乏不同随机种子下的标准差、仅限于小型字符级模型以及使用人工设计的细节基底，并建议未来的工作应探索可学习的基底和更大规模的 Token 级模型。