探索Transformer Block精简的方式

2023-12-06

前言,标准的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精简的方式 (https://ic.work/) AI 人工智能第1张

每个transformer block如上述公式组成，每个子模块都配备了一个系数，这个后续会使用到,Removing Skip Connection,作者先前的一项工作Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation 删除了残差连接，提出的操作Value-SkipInit，将自注意力相关操作修改为>
探索Transformer Block精简的方式 (https://ic.work/) AI 人工智能第3张
其中I代表的是一个Identity操作，A表示原始注意力操作。这两个操作各自有一个可训练标量和，初始化为 , 。
这个设计的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，也是收到信号传播理论的启发，将注意力公式更改为>
探索Transformer Block精简的方式 (https://ic.work/) AI 人工智能第4张
相比之下多了一个C矩阵，这是个常量矩阵已经通过前面的一系列修改实现了，因此可以直接删除Norm层,
当然还是得看实验效果，回到这张图，可以看到移除了Norm对收敛还是有一定影响的。作者猜测在信号传播理论范围之外，Norm层能加速训练收敛，如Scaling Vision Transformers to 22 Billion Parameters,
引入了更多LayerNorm层，将ViT缩放至22B参数量上,因此作者还是主张保留PreLN结构>

最后实验,作者也补充了一些训练速度benchmark，模型准确率，以及收敛趋势的实验>

总结,作者对Transformer Block移除了各种参数，减少了15%参数量，提高了15%的训练速度，各个环节都有做充分的实验，但一些经验性得到的结论也并没有直接回答一些问题。
实验规模并不大，而标准的TransformerBlock还是在各个Scale里得到广泛验证的，期待有人进一步试验,你说的对，但我还是套LLAMA结构