MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections
This work addresses efficiency and performance bottlenecks in Transformer architectures for language modeling, offering a method that yields significant gains with minimal parameter and computation overhead.
The paper tackles the limitations of residual connections in Transformers by proposing MUDD connections, which enhance cross-layer information flow with dynamic weights, resulting in models that achieve performance comparable to Transformers trained with 1.8X-2.4X more compute, such as MUDDPythia-2.8B matching Pythia-6.9B in pretraining perplexity and downstream tasks.
We propose MUltiway Dynamic Dense (MUDD) connections, a simple yet effective method to address the limitations of residual connections and enhance cross-layer information flow in Transformers. Unlike existing dense connection approaches with static and shared connection weights, MUDD generates connection weights dynamically depending on hidden states at each sequence position and for each decoupled input stream (the query, key, value or residual) of a Transformer block. MUDD connections can be seamlessly integrated into any Transformer architecture to create MUDDFormer. Extensive experiments show that MUDDFormer significantly outperforms Transformers across various model architectures and scales in language modeling, achieving the performance of Transformers trained with 1.8X-2.4X compute. Notably, MUDDPythia-2.8B matches Pythia-6.9B in pretraining ppl and downstream tasks and even rivals Pythia-12B in five-shot settings, while adding only 0.23% parameters and 0.4% computation. Code in JAX and PyTorch and pre-trained models are available at https://github.com/Caiyun-AI/MUDDFormer .