Transformer Block精简方式：理论与实践结合的全面深入探索_AI热点日报

结合信号传播理论，研究人员移除了TransformerBlock中的跳跃连接、输出投影、值投影、层归一化及串行结构，采用并行简化设计。在解码器和编码器模型上，该方案减少15%可训练参数并提升15%训练速度。

在深度学习领域，Transformer 已成为最基础的神经网络架构之一。然而，标准的 Transformer Block 结构并不简洁——它包含了 Attention、MLP、Skip Connection、Normalization 等多个子模块。一些看似微小的改动，轻则导致模型训练速度下降，重则使模型完全无法收敛。

本教程将带您深入探索 Transformer Block 的极致精简方案。研究者结合信号传播理论与实验观察，在不牺牲训练速度的前提下，成功移除了 Skip Connection、Out Project、Value Project、Normalization 以及 串行组织 Block 的形式。在 Decoder-only 和 Encoder-only 两类模型上，该方案减少了 15% 的可训练参数，并提升了 15% 的训练速度。

官方仓库：bobby-he/simplified_transformers
论文：Simplifying Transformer Blocks。

一些标记注解

每个 Transformer Block 由以下子模块组成，每个子模块都配备了一个标量系数（后续会用到）。

Removing Skip Connection

Value-SkipInit 方法

作者此前的工作 Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation 删除了残差连接，提出了 Value-SkipInit 操作，将自注意力相关操作修改为：

其中 I 代表 Identity 操作，A(X) 表示原始注意力操作。这两个操作各自配备了一个可训练标量 α 和 β，初始化为 0,0。

设计的 核心思想 是：每个 token 在训练初期更多地关注自身相关性。类似的操作如 Pre-LN（Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks）发现，Pre-LN 相当于提高了 skip-branch 权重，降低了 residual-branch 权重，从而在较深的神经网络中依然保持良好的信号传播。

Shape Attention 与 Centering Matrix

另一项工作 The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit 提出了 Shape Attention，同样受信号传播理论启发，将注意力公式修改为：

相比原始公式，这里多了一个 C 矩阵（称为 Centering Matrix），这是一个常量矩阵，不参与训练。它的取值被设定为当 query-key dot 为 0 时 A(x) 的值，即仅剩 mask 值时的结果。代码实现如下：

# Centered attention, from https://arxiv.org/abs/2306.17759
uniform_causal_attn_mat = torch.ones(
    (max_positions, max_positions), dtype=torch.float32
) / torch.arange(1, max_positions + 1).view(-1, 1)
self.register_buffer(
    "uniform_causal_attn_mat",
    torch.tril(
        uniform_causal_attn_mat,
    ).view(1, 1, max_positions, max_positions),
    persistent=False,
)

对于 CausalLM 来说，MASK 是一个下三角矩阵，形状为 (S, S)。第 i 行只有前 i 个位置有值（未被掩盖），经过 softmax 后，1.0 概率被均匀分配到这些有值的位置。这就是代码中执行 ones / arange 操作的原因。示例代码如下：

import torch
max_positions = 32
mask = torch.tril(torch.ones(max_positions, max_positions)) + torch.triu(torch.ones(max_positions, max_positions), 1) * -65536
print(torch.softmax(mask, -1))
# 输出：
tensor([[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.3333, 0.3333, 0.3333,  ..., 0.0000, 0.0000, 0.0000],
        ...,
        [0.0333, 0.0333, 0.0333,  ..., 0.0333, 0.0000, 0.0000],
        [0.0323, 0.0323, 0.0323,  ..., 0.0323, 0.0323, 0.0000],
        [0.0312, 0.0312, 0.0312,  ..., 0.0312, 0.0312, 0.0312]])

可训练标量与 MLP 系数调整

新的可训练标量 α = 1.0，以保证初始化时 Attention 先以 Identity 方式工作。如果将这些可训练标量改为 headwise（每个注意力头独立），则性能会有所提升。作者特别强调：需要显式地将 MLP Block 的系数调低：

论文中针对 18 层 Transformer，将 MLP 系数设置为 0.1。