MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections

TL;DR

MUDDFormer introduces Multiway Dynamic Dense connections to relax residual bottlenecks in Transformers by generating per-position, per-stream dense weights for cross-layer aggregation. By decoupling the Transformer inputs into four streams (Q, K, V, R) and applying dynamic, stream-specific dense connections, the approach significantly enhances information flow across layers and scales efficiently across language and vision tasks. Empirical results show substantial gains over strong baselines, including matching much larger models compute-wise on pretraining perplexity and downstream tasks, with only a small parameter and compute overhead. The method also offers stability via normalization variants and demonstrates favorable efficiency trade-offs, making it a practical enhancement for future foundation models.

Abstract

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 .

Figures (13)

• Figure 1: Downstream average accuracy of Pythia and MUDDPythia with different sizes.
• Figure 2: Architecture of Multiway Dynamic Dense Connections.
• Figure 3: Scaling curves of MUDDFormer and baseline models.
• Figure 4: Depth scaling of MUDDFormer and Transformer++.
• Figure 5: MUDDFormer with dense vs MoE models.