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

2025-02-13

AI Summary

To address the limited cross-layer information flow caused by static residual connections in Transformers, this paper proposes the Multi-path Unified Dynamic Dense (MUDD) connection mechanism. MUDD dynamically generates sparse, position-aware connection weights—separately for Q/K/V projections and residual streams—enabling fine-grained, low-overhead positional information reuse and departing from conventional shared-weight paradigms. The method employs a lightweight dynamic weight module compatible with both JAX and PyTorch, introducing only a 0.23% parameter overhead and enabling plug-and-play integration into arbitrary Transformer architectures. Experiments demonstrate that MUDD-enhanced Pythia-2.8B matches the pretraining perplexity and downstream task performance of Pythia-6.9B, achieves five-shot learning accuracy comparable to Pythia-12B, and delivers 1.8–2.4× computational efficiency gains.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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 .

Problem

• Enhance cross-layer information flow
• Dynamic connection weights generation
• Improve Transformer model efficiency

Innovation

• Dynamic connection weights
• Enhanced cross-layer flow
• Seamless Transformer integration