Paper: https://arxiv.org/abs/2311.01906
GitHub: https://github.com/bobby-he/simplified_transformers
Abstract:
A simple design recipe for deep Transformers is to compose identical building blocks. But standard transformer blocks are far from simple, interweaving attention and MLP sub-blocks with skip connections & normalisation layers in precise arrangements. This complexity leads to brittle architectures, where seemingly minor changes can significantly reduce training speed, or render models untrainable.
In this work, we ask to what extent the standard transformer block can be simplified? Combining signal propagation theory and empirical observations, we motivate modifications that allow many block components to be removed with no loss of training speed, including skip connections, projection or value parameters, sequential sub-blocks and normalisation layers. In experiments on both autoregressive decoder-only and BERT encoder-only models, our simplified transformers emulate the per-update training speed and performance of standard transformers, while enjoying 15% faster training throughput, and using 15% fewer parameters.
I’m a bit confused by the motivation of removing the skip connections. I’ve read things from people in the mechanistic interpretability community refer to the “residual stream” as not just a training hack but as an important part of the reason why transformer models work well. That the fact that information can propagate through the layers past the nonlinearities is seen as a feature not a bug.