Simplifying Transformer Blocks 论文解读

前言¶

标准的Transformer Block并不简介，每个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¶

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

其中I代表的是一个Identity操作，A(X)表示原始注意力操作。这两个操作各自有一个可训练标量 $\alpha$ 和 $\beta$ ，初始化为 $\alpha=1$ , $\beta=0$ 。

这个设计的insight是每个token在训练前期更多的是关注自身相关性，类似的如Pre-LN操作，在Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks这项工作发现，Pre-LN相当于把 skip-branch 权重提高，降低residual-branch权重，以在较深的神经网络里仍然有良好的信号传播。

而The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit 该工作里提出了Shape Attention，也是收到信号传播理论的启发，将注意力公式更改为：

相比之下多了一个C矩阵，这是个常量矩阵（论文称其为centering matrix)，不参与训练。他的值被设置为当 $\frac{1}{\sqrt{d_k}}$ $X$ $W^Q$ $W^{K^T}$ $X^T$ querykey dot 为0时候，A(x)的值，那么我们回去看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]])

而新的可训练标量 $\gamma$ = $\beta$ ，以保证初始化时， $\beta A(x) - \gamma C = 0$

其中这些可训练标量如果改成headwise，即每个注意力头独立，则性能有部分提升。当然作者还是强调其中的一个重要的点是，显式的将MLP Block的系数降低：