Transformer design improves performance with fine-grained layer mixing
M3Former: Fine-Grained Cross-Layer Mixing for Transformers
Machine Learning
Summary
Transformers are popular models that process information in steps called layers, connecting each new step to all previous ones through a simple addition. The authors found that this simple connection limits how well different parts of the model can mix and use information from earlier steps. They created a new method called M3Former that lets each part mix information from all earlier layers in a more flexible and detailed way. This small change improves model performance and training stability without much extra cost.
What this means in practice
- •For deep learning engineers: Build transformer models that achieve better accuracy and training stability with only a slight increase in computation and parameters.
- •For machine learning infrastructure teams: Optimize large-scale transformer training pipelines by using fine-grained layer mixing to improve model efficiency and stability for complex models.
Authors
Jiasong Han, Yunkai Chai, Tong Zhu, Hengyu Zhang, Wei Fang, Jintao Wang, Guoqi Li, Yu Cheng
Abstract
Residual connections give each Transformer block access to all earlier layers only through one fixed sum, broadcast to all of the block's inputs. This broadcast bottleneck imposes three constraints: the query, key, value and residual inputs of a block, which we call its roles, read the same vector; all hidden dimensions within a role share one set of weights over earlier layers; and every earlier layer has the same fixed weight for every token. We propose the Multi-Role Multi-Group Mixing Residual (M3), which removes all three: every role, and every group of hidden dimensions within a role, reads its own token-dependent mixture of all earlier layers. M3 starts as an exact residual connection and adds as little as 1.4% parameters and 1.8% FLOPs. From 100M to 1.4B parameters, M3Former outperforms the Transformer and coarse-grained cross-layer connections at every size, matching a Transformer 1.19-1.34x the size of the one it modifies, trained with 1.4-1.8x the compute. Relaxing only some of the constraints recovers only part of this gain. M3Former also improves sparse mixture-of-experts models of up to 4.6B parameters and trains more stably in deeper networks. Code: https://github.com/Kingchet10/m3former