MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections
Authors: Da Xiao
Qingye Meng
Shengping Li
Xingyuan Yuan
Venue: ICML 2025 poster (2025)
Keywords: model architecture
Transformer
residual connections
dense skip connections
LLMs
pre-training
Primary area: deep_learning->foundation_models
TL;DR: We propose MUltiway Dynamic Dense (MUDD) connections to significantly improve Transformer by enhancing cross-layer information flow. MUDDFormer matches performance of Transformer trained with ~1.8x-2.4x compute in language modeling and scales well.

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 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.

FoDDFormer:
Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections
Da Xiao
Qingye Meng?
Shengping Li Xingyuan Yuan
Abstract
58
Downstream Avg. Acc (%)
4.2x
57
We propose MUltwayy 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) a Transformer block. MUDD connections can be seamlessly integrated into any Transformer ar- chitecture create MUDDFormer. Extensive experiments show that MUDDFormer signifi- cantly outperforms Transformers across various model architectures and scales in language mod- eling, achieving the performance of Transform- ers trained with -1.4x compute. Notably, MUDDPythia-2.55 matches Pythia-6.9B in pre- training ppl and downstream tasks and even rivals Pythia- 1ythia-12BBPo in five-shot settings, while adding only 0.23% parameters and 0.4% computation. Code in JAX and PyTorch and pre-trained mod- els are available at https:// /ttl::////ttttttt wiggghht..coo com/ Caiyun- Beiyyn-AA/MMUDA AI/MUDDE ormer.
2. 2x4x
55
53
Pythia-O-shat
5-shot 5-shot
51
MUDDPychia-O-shot
50
1.4B+300B
2.88-300B
6.98
128*300B
Compute
Figure I. Downstream average accuracy of Pythia and MUD- DPythia with different sizes.
common by prevailing Transformer-based LLMs.
On one hand, although theoretical (Merrill et al., 2022) and experimental (Tay et al., 2021b) work have suggested that adding layers increases the expressive capacity and gener- alization performance of Transformers, it is observed that increasing depth beyond a certain point yields diminishing returns (Petty al., 2023). The common practice of using Pre-Norm stabilize training leads to the issue of represen- tation collapse (Liu et al., 2020b), where hidden features in deeper layers become highly similar, and for popular fam- ilies of open-weight LLMs, a large fraction of the layers can be removed with minimal degradation of performance (Gromov al., 2024).
1. Introduction
On the other hand, mechanistic interpretability studies reveal that Transformers do in-context learning tasks by compos- ing model components (attention heads and MLPs) across different layers to form circuits, (Elhage et al., 2020; Wang et al., 2023; Merullo et al., 2024; Ni et al., 2025), where lay- ers communicate with each other by writing to and reading from different subspaces of residual stream. The resid- ual stream as the shared communication channel may be overloaded and become the bottleneck for very deep models, hindering the formation of sophisticated circuits spanning distant layers necessary for complex tasks.
Residual connections (He et al., 2016), which help mitigate the vanish gradient problem, have become indispensable to training deep learning architectures from CNNs to Trans- formers, the latter becoming the de facto backbone for foun- dation models. Though being very simple and effective, residual connections still have limitations be solved, espe- cially with deep Transformers with dozens of layers made
Beijing University of Posts and Telecommunications Bei- jing. China *ColorfuiClouds Technology Co., Ltd., Beijing, China. Correspondence to: Da Xiao <xiaoda99 @ @bupt.edu.cn..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Dense connections (Huang et al., 2017) was proposed as a promising solution to the above issues, by allowing subse- quent layers to directly access outputs of preceding layers

---

Multiway Dynamic Dense Connections
(Figure 2 (a)). It has shown effectiveness for CNNs with DenseNet (Huang al., 2017), for encoder-decoder Trans- formers with Deep Transformers (Wang et al. 2019) and for decoder-only Transformers with DenseFormer (Pagliardini et al., 2024). However, these dense connection approaches use static (either fixed (Huang et al., 2017) or learnable (Wang et al., 2019; Pagliardini et al., 2024)) dense connec- tion weights that are shared across sequence positions and different input streams of a Transformer block. As will be shown, this rigidity severely limits their expressive capacity in Transformers.
connected ferd-forward network (FFN), both wrapped with Pre-LayerNorm (LN) residual connections:
XA = MIAI(L B(X)= = FFN(LN(XA)) + X
+
(2)
In this architecture, the output X; R.TxD (D is model dim) of layer is used as input layer + 1. With dense connections, the input to layer + 1 is aggregation X; E RT*x of outputs of all i + preceding layers, from the embedding till layer i: X == {Xo, X.} (Figure 2 (a)). In this way, the cross-layer communication bandwidth is significantly increased compared to residual connections. To obtain Transformer with dense connections, we just add Depth-wise Aggregate (DA) module after each layer to provide input for the next layer (ef. Eq. 1):
In this work, we propose MUltiway Dynamic Dense (MUDD) connections, a simple yet effective approach to address the shortcomings of residual connections. 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 (the query, key, value or residual) of a Transformer block. These weights are used by depth-wise wise aggregate modules to combine outputs from all preceding layers, creating multiple input streams for the current layer. MUDD connections can be seen as depth-wise multi-head attention (Vaswani al., 2017) and the cross- layer communication bandwidth is expanded far beyond the restriction of the residual stream.
X X, = EmbeddinngX)
X; B, X, (X.,,, ie [1, (3)
TenneeTransfomeet = X
In the following subsections, we progressively derive Trans- former with MUDD Connections (MUDDFormer), focusing on the computation inside the DA module after layer
MUDD connections can be seamlessly integrated into any Transformer architecture to create MUDDFormer models. We conduct extensive experiments focusing on language model pretraining to evaluate MUDDFormer's effective- ness, efficiency and scalability. MUDDFormer significantly outperforms Transformer across various model architectures and scales (from 405M model on 7B tokens to 2.8B models on 300B tokens), achieving performance of Transformers trained with ~1.8 1..xx compute. Notably, MUDDPythia- pretraining perplexity and
2.1. Static Dense Connections
In the simplest case, the DA module aggregates previous lay- ers' outputs by taking a weighted sum of them (Figure 2 (b). also equivalent to DenseFormer (Pagliardini al., 2024)):
i+1 (++1)xTX
X, = DA Mtatiic
stati (X,;; = -wsum 0:
X1
(4)
1 TxD
2.8B matches Pythia-6.9B downstream tasks and even rivals Pythia- 12B in five-shot in-context learning settings (Figure 1), while adding only 0.23% parameters and 0.4% computation. We also evaluate MUDD connections on vision Transformers and analyze the trained models to elucidate why MUDD connections work.
1=0
where wsum(-) is the weighted sum function taking the sequences of weights and values as inputs. Scalar ag is the jth value of the dense connection weight vector as E Ri+1 which trainable parameters, i.e., 0f {az}-
2. Method
2.2. Dynamic Dense Connections
We begin with a standard Transformer decoder with L layers on input sequence X {2o, .... æT): ):
In Transformer-like sequence models, each layer processes information from multiple sequence positions and may ben- efit from differentiated and input-dependent dense connec- tivity for each position. Dynamic dense connections expand the connection weight for X, from static scalar a,j to vector Aij E RRT, allowing X; to contribute differentially to each position t E [1, 71 of X, based on the hidden state X,[t] [t] E R at that position. The i + 1 weight vectors stack into a matrix A, E (+++1) which is generated dynami- cally by function A4(-) depending on X; (Figure 2 (c), cf.
Xo- Embedding(X)
X; (,( [1
(1)
Transformer X
where B.(-) is the ith Transformer block! composed of multi-head attention (MHA) module followed by a fully
In this paper we use "layer" and "block" interchangeably.
2

---

Multiway Dynamic Dense Connections
Blockg
Block3
Block's
LN & FFN
X3
Block
Blocks
X2
X2
X
X2
R22
$20
DA
DA
X
X2
X
Block2
Blockg
Block2
Block's
X1
LN
X
X
X1
X1
Block
Block
Block;
Block,
Block'1
MHA
%
Xo
%
X
T [seq, longth)
X
X
X
X
(a) Dense
(b) Static dense connection
(c) Dynamic dense
(d) Multiway dynamic dense connection
(e) Transformer block w/ single/muitiple inputs
connectivity
connection
Figure 2. Architecture of Multiway Dynamic Dense Connections.
Eq. (4)):
streams for query, key, value and residual, respectively (Eq. (7), Figure 2 (e) bottom), and then instantiate four DA mod- ules, each specializing in one stream's 's dense connectivity (Eq. (8), Figure (d))
X; DA,
dynamic
(X,;;
Tx(++1)
TXD (++1)XTXD
wsum Ag = -A, );
X2
(5)
X = MHA(LL(AX )) B'(XO,x* .VV (x = =FF((LN(X)) XX
x (7)
Tx1
TXD
Σ
Aij o Xj
(with broadcasting)
X X -x² = =x = = -mmbeeding (X)
where A. is the jth column of A; (with a slight abuse of notation). We instantiate A; RDD - Ré+1 with MLP parameterized by W1 and W2 which computes connection weights position-wise:
X; =B((X 1- x xx,,,,, x
(8)
X 2 X
DA i [1,
MUDD
A:(X,) GELU((RMNNmm(X))1 + (6)
By making the dense connections multiway, the cross-layer communication bandwidth is further increased significantly. Multiway dynamic dense connections can be seen as depth- wise multi(4)-head attention. This vertical cross-layer at- tention can be composed with the horizontal cross-token attention Transformer to form pathways adaptively, en- hancing information flow across the whole model when performing in-context learning tasks.* At this point, finally obtain MUDDFormer by integrating static, dynamic and multiway dense connections. Complete pseudo-code for MUDDFormer is given in Appendix B.
We apply RMSNorm to X, before MLP to stabilize training. We also add static weight vector a; acting as learnable prior for dense connectivity. The trainable parameters are 0r = (Wi E (+++1),, Wa E IR (+++1)x(++1), a, E 1R++)..
Overall, dynamic dense connections can be viewed as form of depth-wise single-headed self-attention (Vaswani et al., 2017) with query X; and keys X.,. and the attention weights are computed from the query side?.
2.3. Multiway Dynamic Dense Connections
In a Transformer block, single input is reused simulta- neously as the query, key, value and residual of the MHA module (Figure (e) top). These input streams play diver- gent roles and we hypothesize that they benefit from dif- ferentiated dense connectivity.
2.4. Parameter Re-allocation
Due to the dense connections, MUDDFormer's upper layers have the opportunity to process more information than lower layers and thus may need more parameters. We re-allocate the parameters of standard Transformers to make the size
enable this, we first turn normal Transformer block (XX) into a multi-input one B' B((xQ, x RR) XR) by decoupling its input into four FFor further explanation, see Appendix
"This is a logical view for clarity. In practice, these DSs can he combined for efficiency. See pseudocode in Appendix B. *For an illustrative example of this, see Figure
3

---

Multiway Dynamic Dense Connections
of FFN sub-layers grows with depth. Specifically, let Dj be the hidden dim of the original FFN, we compute (i) the hidden dim of FFN layer for MUDDFormer using linear interpolation:
3. Experiments
Implementation Details We implement MUDDFormer model and training JAX. We initialize the MUDD connec- tion weight generating parameters W1 and W2 with N(0, D) and 0 respectively, and initialize the static weight vector a; with 1 at dzi and 0 elsewhere. This reduces MUDDFormer to Transformer at the beginning of training, which is found to be critical for good performance. If PrePostDANorm is used, we initialize the scale parameters Pre-DA and Post- DA RMSNorms with and le-3, respectively, and initialize a; because X; is added as the residual after DA modules (Eq. 10). For the other parameters outside DA modules, use Xavier normal initializer.
D1 0.5(L - 1.5(i Dj
(9)
--1
i.e. the FFN hidden dim D grows linearly from y y to y. The total number of parameters remains unchanged.
2.5. Optional Normalization
To stabilize training models with large depth/width ratios, we propose variant of MUDDFormer by applying RM- SNorm before and after DA module, and adding residual connection to DA module after the post-RMMNorm:
Organization Our evaluation focuses on language model- ing with decoder-only Transformer architecture, analyzing both pretraining loss scaling laws (Section 3.1 and Sec- tion 3.2) and downstream task performance (Section 3.3) with large scale training on the Pile dataset (Gao al., 2020). Section 3.4 elucidates why MUDDFormer works through analyzing trained models, followed by efficiency analysis (Section 3.5) and ablations (Section 3.6). Extended vision experiments are provided Appendix F.
X (Norm(Xoo,,oo ..... Norm((X)) (PreDANor))
(10)
X; = (X,,);
(PostDANorm)
It is similar to the hybrid-norm strategy used by recent models such as Gemma 2 (Team et al., 2024b) and Grok-1 (xai org, 2024), though we apply it to DA modules instead of MHA/MLP modules. We use this PrePostDANorm variant to train the DeepNarrow models in scaling law experiments in Section 3.1 and the MUDDViT model in Appendix F,
3.1. Scaling Laws
Settings Table 2 (top half) specifies the model sizes and hyperparameters for scaling experiments, which are mostly taken from GPT-3 specifications (Brown et al., 2020). We untie input and output embedding matrices. We train with context length 2048 and set the number of training tokens to roughly match Chinchilla scaling laws (Hoffmann et al., 2022). The other hyperparameters are in Appendix D.I. In another set of scaling experiments we trade off some width for depth to train DeepNarrow models (Table 2 (bottom half)) to see if MUDD connections offer any benefits for this style of scaling.
2.6. Complexity Analysis
Table shows the ratios of extra parameters and computa- tion introduced by MUDD connections with both analytical results and typical concrete values. The derivations are in Appendix C. The ratio of extra parameters, i.e. parameter count of W1 and W2 of DA modules divided by that of the whole model, proportional to the rectified depth/width ratio 7 = L+3 The ratio of extra computation, i.c. FLOPs of generating MUDD connection weights and cross-layer aggregation divided by FLOPs of the whole forward pass, besides proportional to 1). decreases with p = T- Both ratios are negligible for commonly used settings.
Table 2, Model sizes and hyperparameters for scaling experiments.
learning rate
batch size (in tokens)
params
Dlayers
dmodel
Thends
tokens
Table Ratios extra parameters and computation introduced by MUDD connections: (last row) analytical results and (upper rows) concrete values for typical model architectures and hyperparame- ters. =. number of layers, T = sequence length.
405M
24
1024
16
3e-4
0.5M
7B
834M
24
1536
24
2.5e-4
0.5M
15B
1.4B
24
2048
32
2c-4
0.5M
26B
Scaling by depth
797M
34
1280
20
2.5e-4
0.5M
15B
Model Size RAparams
RArLors
L
=
T
1.34B
1536
24
2e-4
0.5M
26B
D
1.4B
0.22%
0.38%
24 2048
4096
0.0132
2
1.34B
0.49%
0.8%
42 1536
4096
0.0293
2.67
Baselines We compare MUDD with two recently proposed approaches to enhancing residual connections in Trans- formers: Fonserrrmer (Pagliardini et al., 2024) (same as the static dense connections described in Section 2.2) and Hyper-Connections (Zhu et al., 2025). We also compare to
2,8B
0.23%
0.4%
2560
4096
0.0137
1.6
6.9B
0.14%
0.26% 1/
32 4096
4096
0.0085
1
Formula
6
3 + p/4

---

Multiway Dynamic Dense Connections
2:30
Orig. Transformer
Transforme++
2.5
Transformer++
(aur) MUDDFormer
DenseFormer
2.25
ALaane
Hyper-Connections (our) DDFormer
2.4
(our) MUDDFormer
Loss
2.20
9
2.3
2.15
Alwuno -0.0641
2.2-
Ml..
2.10
=.042
2.1-
Alaupo
405M
1.88-M45MM
0.062
# Paramss
405M778
834M1158 Compute
1.48*268
Figure 5. MUDDFormer with dense vs MoE models.
Figure Scaling curves MUDDFormer and baseline models,
major architectural improvements (RoPE, SwiGLU MLP, etc.) over original Transformer since its invention.
2.6
Orig. Transformer Transformer++
Figure 4 demonstrates MUDDFormer's enhanced depth uti- lization: while Transformer+4 shows diminishing returns beyond 24 layers (almost coincident scaling curves). MUD- DFormer DeepNarrow maintains gains up 42 layers. This validates that MUDD connections alleviate depth-induced bottlenecks by enhancing cross-layer information flow.
2.5
MUDDFormer
deep Transfomer+++
2.4
deep MUDDFormer
2.3
2.2
2.08x
2.1-
3.2. MoE Models
405M~78
634/797M*158
1.4/1.38+268
Settings Transformer with Sparse Mixture-of-Experts (MoE) and MUDDFormer both architectures with dy- namic weights. MoE uses these weights to select experts within layer while MUDDFormer uses the weights to ag- gregate outputs across layers. They are complementary approaches and can be combined. To empirically compare MoE and MUDDFormer, we train a Transforemr++ MoE model 1.8B- BB-A405M with the same activated parameters as the 405M dense model, using the same training settings in the scaling law experiments. For MoE specific settings, largely follow OLMoE (Muennighoff et al., 2025), using dropless token-choice routing to choose 2 experts out of 16 for each token with expert hidden dim 1408. The FLOPs of the MoE model is nearly identical to that of the dense model. We apply MUDD connections to both the dense and MoE models. All the models are trained on 26B tokens.
Compute
Figure 4. Depth scaling of MUDDFormer and Transformer++.
Transformer with dynamic dense connections (DDFormer) as described in Section 2.2. All these approaches are ap- plied to an improved and now widely adopted Transformer architecture (Touvron al., 2023) with rotary positional encoding (RoPH) (Su al., 2024), SwiGLU MLP (Shazeer, 2020), etc. (often called Transformer+++) We also include the plot for the original Transformer architecture used by GPT-3 as a comparison. The details for these baselines are in Appendix D.2.
Results Figure 3 plots Pile validation loss scaling curves of the models. While DenseFormer and Cypne--Conesttions show clear advantage over Transformer++, DDFormer out- performs them by adding dynamicity to dense connections. MUDDFormer further improves upon DDFormer by making the dense connections multiway, significantly outperform- ing all baselines on models ranging from 405M to 1.4B. It can be estimated that MUDDFormer-S00 matches the loss of Transformer++ trained with 1.89 compute.
Results As shown in Figure 5, compared to the same Transformer++44053 500 loss reduction of T A--05 MoE, which is 4.4x larger. MoE works by decoupling param- eters and computation and relies on significantly expand- ing parameters for good performance, in contrast MUD- DFormer's efficient utilization of both parameters and com- putation. More notably, MUDD demonstrates more benefits for the MoE model (Aloss -0.0641) I) than for the dense model (Aloss -0.0596), a result of particular practical sig- nificance given the increasing adoption MoE architec- tures in frontier LLMs (Team et al., 2024a; Liu et al., 2024b; Team, 2025; Yang et al., 2025).
MUDDFommer405M achieves
Figure 3 also shows that architectural improvement, MUDDFormer's gain over Transformer++ (A LMUDD) re- mains stable while scaling, exceeding Transformer+** 's own gain over Transformer (AL44) beyond 834M parameters. This shows the favorable scalability of MUDDFormer, par- ticularly considering that Transformert+ has incorporated
5

---

Multiway Dynamic Dense Connections
Table 3. Zero-shot and five-shot downstream evaluations results.
Pile ppl.
FLAN ppll
LAM BADA
Wino Grande
ARC
ARC -C
Logi QA
Hella Swag
RACE -M
RACE
Avg
Model
PIQA
SciQ
BoolQ
-E
acct /Lacc
0-shot
Pythia-
7.29
9.30
61.6
71.0
57.2
60.5
26.1
86.6
21.4
63.3
40.5
37.3
33,9
50.8
MUDDPythia-IL4h
6.92
8.54
63.9
71.8
57.4
61.6
26.2
87.2
23.0
62.0
42.6
38.7
34.7
51.7/40.9
Pythia-2.8B
6.63
8.16
64.7
73.9
59.4
64.4
29.5
88.2
21.2
64.5
45.4
38.1
34.9
53.1
MUDDPythia-2.88
6.29
7.50
68.5
74.6
61.4
66.5
31.9
90.4
21.5
68.1
46.8
39.0
36.7
55.0/+1.9
Pythia-6.9B
6.29
7.85
67.3
75.2
60.9
67.3
31.3
89.7
25.3
63.7
48.0
40.6
37.0
55.1
Pythia-12B
6.01
7.26
70.5
76.0
63.9
70.2
31.8
90.2
22.4
67.4
50.3
40.6
38.3
56.5
MUDDFM-2.8B
6.01
7.08
70.7
75.7
63.4
70.4
34.2
91.8
24.0
67.4
49.5
40.6
38.1
56.9
5-shot
Pythia-
54.5
71.0
57.5
63.1
28.9
92.2
22.9
63.0
40.5
35.4
34.6
51.2
MUDDPythia-LL4B
58.2
73.0
59.0
64.1
28.2
94.0
23.8
61.5
42.6
37.9
35.2
52.5/+1.3
Pythia-2.8B
60.5
73.6
60.6
67.3
32.3
94.3
21.7
65.6
45.1
38.4
35.6
54.1
MUDDPythia-2.88
63.6
75.5
63.6
70.3
34.0
95.5
28.1
67.5
47.1
44.5
37.3
57.0/+2.9
Pythia-6.9B
63.8
75.5
63.7
70.2
35.6
95.1
27.0
65.7
48.1
39.0
36.5
56.4
Pythia-12B
67.3
76.0
64.2
71.0
36.5
95.3
21.8
68.0
50.3
40.1
38.8
57.2
MUDDFM-2.8B
65.6
76.4
66.8
73.0
39.2
95.6
25.2
70.9
49.8
41.4
38.0
58.4
3.3. Large Scaling Training and Downstream Evaluations
LM evaluation harness (Gao et al., 2023).
Results As shown in Table and Figure 1, besides lower Pile validation ppl. MUDDPythia also significantly outper- forms Pythia at 1.4B and 2.8B scales on downstream task accuracies. Notably, MUDDPythia-2.88 matches Pythia- 6.9B (2.46x compute) on both pretraining ppl and down- stream evaluation. Augmented with better Transformer++ architecture and training recipe, MUDDFormer-2.. even outperforms Pythia-12h,
Settings We compare MUDDFormer with the open source Pythia model suit (Biderman et al., 2023) large scale training on 300B tokens of Pile. Specifically, we train two models, MUDDPythia-1445 and MUDDPythia-2.8B. and compare them with Pythia models ranging from 1.4B to 12B. For fair comparison and clear quantification of the gain brought by MUDD, except adding MUDD connections as described in Section 2, MUDDPythia uses exactly the same architecture choices (e.g. parallel attention and MLP, rotary embedding with 1/4 head dim) and training hyper- parameters (e.g. optimizer settings, learning rate schedule, batch size, context length, initialization methods) as Pythia (refer Biderman et al. (2023) Appendix E for details). To evaluate if MUDD connections also work well with more advanced Transformer++ architecture and training recipe at this large scale, we also train MUDDFormer-2.0 based on Transformer++ instead of Pythia architecture with a larger learning rate of 3.2c-4 cosine decayed to 3.2e-6. Except these two changes, the other architectural and training hy- perparameters are kept the same as MUDDPythia-2.899
Table also reports target span perplexities randomly sampled subset of the FLAN Collection dataset (Longpre et al., 2023), which features data of instructing following, chain-of-though in-context few-shot learning, etc. The advantage of MUDDPythia on FLAN is even larger than on Pile with MUDDPythia-2.85 significantly outperforming Pythia-6.99 6.9B, showing that MUDD connections have more advantage in improving these valued emergent abilities of LLMs (Wei et al., 2022). The enhanced in-context learning ability is also manifested by the larger Aaccuracies of 5- shot results compared to 0-shot (e.g. 2.9% vs. 1.9%), where MUDDPythia-22. 2.88 is on par with Pythia- (4.29x com- pute). The multiway dense connections applied to the query, key, and value streams of MHA modules are likely to en- hance the functionality of attention heads, which are crucial for in-context learning (also see analysis in Section 3.4).
Evaluation Datasets Besides the datasets used by Pythia for downstream evaluation (LAMBADA (Paperno et al., 2016), PIQA (Bisk al., 2020), WinoGrande (Sakaguchi al., 2021), ARC (Clark et al., 2018), SciQ (Welbl al., 2017), LogiQA (Liu al., 2020a)), we also include BoolQ (Clark et al., 2019) and HellaSwag (Zellers et al., 2019) for commonsense reasoning, RACE (Lai et al., 2017) for reading comprehension, all of which are widely used bench- marks. We evaluate zero-shot and five-shot results using
Finally, the gains of MUDD connections with the 2.8B model are larger than those with the 1.4B model in both zero-shot and five-shot evaluations, further demonstrating the scalability of MUDD connections.
6

---

Multiway Dynamic Dense Connections
3.4. Analyzing and Understanding
We analyze and compare Pythia-2.8B and MUDDPy 2.8B trained in Section 3.3 to elucidate why MUDD connec- tions work. The analysis is done 1024 randomly sampled sequences of length 2048 from Pile validation set.
Representation Collapse Figure quantifies representa- tion collapse through cosine similarity between the inputs of adjacent layers. While Pythia exhibits progressive col- lapse with >0.97 similarity in later layers, MDDDPythia maintains more distinct input representations, particularly in the value stream. Thanks to input stream decoupling and stream-specing aggregation, DA modules can freely aggregate distinet value input streams for MHA modules of each layer at each sequence position. These values will then be moved to other positions by MHA at this layer without polluting the residual stream of the current position. The rel- ative importance of dense connections for the value stream is also evidenced by ablation studies (Section 3.6). We com- pare illustrative V-composition circuits' in Transformer and MUDDFormer in Figure 7 to highlight the benefit of MUDD connections and input stream decoupling, which results in more direct and cleaner information pathways. MUDD has similar effect on e/K-composition circuits.
:
Transformer, stream medisted by and coupled with residual stream polluted other flow-in
In MUDDFomeer, VO stream straiggt-through and decoupled from residual stream
Figure 7. Illustrative V-composition circuits in Transformer vs. MUDDFormer. Colored circles are MHA's inputs of the query (yellow), key (red), value (green) and residual (black) streams and oulput (blue). LayerNorms and MLPs are omitttd..
deeper layers. 6 This vitalization of attention heads stems from the multiway dense connections on the Q/K/V streams of MHA modules, ultimately improving in-context learning.
1.0
Pythia MUDDPythia
Head activation ratio
0.8
0.4
1.00
0.2
Cosine similarity
0.75
15
20
30
0.50
Layer index
0.25
Figure 8. Attention head activation ratio by layers.
0.00
Pythia
MUDDPythia Q
3.5. Training and Inference Efficiency
0.25
MUDDPythia-K
MUDDPythia-V
Besides theoretical complexity analysis in Section 2.6, we assess training and inference efficiency of MUDDFormer compared with Transformer in real-world settings.?
-0.50
MUDDPythia-R
0
10
15
20
25
30
Layer index
Figure 6. Cosine similarity between the inputs of the current layer and the preceding layer.
Settings Though we use Pythia's architecture in large scale training, the evaluation in this section is done on untrained models with Transformer++ (Llama) architecture, which of more practical relevance due to better performance. We measure on three model sizes: 1.3B, 2.8B and 6.9B, which are "Llama-ed" version of Pythia 1.4B, 2.8B and 6.9B re- spectively. We train Google Cloud TPU v5p-128 pods with context length of 2048, batch size of 2M tokens and measure training throughput. We do inference on NVIDIA A 100 80G GPU with prompt length of 4096, batch size of I and measure the speed to generate 128 tokens. We repeat the measurements 3 times and take the average. We implement training and inference in pure JAX and PyTorch respectively
Attention Head Activation Transformer models often ex- hibit null attention (Vig& & Belinkov, 2019) where attention heads focus on the initial tokens or some special tokens as default "attention sinks" (Xiao et al., 2024b) when relevant tokens are found. We define head be active position if its maximum attention weight does not fall on the first two positions of the sequence or on tokens "<bos>", "" and "\n"". We compute the activation ratio of attention heads for each layer by averaging over all heads of that layer and all sequence positions and plot the results in Figure 8. In Pythia, most heads remain inactive beyond the first few lay- ers, limiting their contribution. MUDDPythia demonstrates ~2.4 x higher activation ratio across layers, particularly in
VVisualization of sampled attention patterns for heads Pythia and MUDDPythia are shown Appendix G.
FFor additional results of analytical and measured memory us- age and training wall-clock time of MUDDFormer Transformer, see Appendix E.
"common and important in many tasks, e.g. Wang et al. (2023); Ni et al. (2025). For introduction, refer Elhage et al. (2020).

---

Multiway Dynamic Dense Connections
Table 4. Training throughput and inference speed comparison be-
Transformert +
tween Transfomer++ and MUDDFormer.
11.6
MUDDFormer-SW6 MUDDFormer-kxp
Model Size
Training (K tokens/s)
Inference (tokens/s)
Perplexity
11.4
TFM++
MUDDFM
TFM++
MUDDFM
11.2
1.3B
1147
1030
x89.8%
325
286 188.1%
2.8B
684
575
x84.0%
181
163
x90.0%
6.9B
332
318
x95.6%
95.5
89.7
x94.0%
11.0
15WW
* *1x
WSW8
fxx2
*
2x1
10.8
1x1
1xx1
without writing any Pallas/Triton-biee custom kernels. We
0.90
0.95
1.00
0.90
0.95
1.00
Rel. Training Speed
Rel. Inference Speed
use forch.complll (PyTorch 2.5.1) 1) to accelerate both Trans- former++ and MUDDFormer.
Figure 9, PPL vs. relative training and inference speed of MUD- DFomer variants.
Results As shown in Table 4, the training and inference overheads, while larger than the theoretical estimates in Ta- ble I and not negligible, are entirely acceptable considering the significant performance gain. The overheads primarily stem from the series of small operations and additional I/O introduced by DA modules. We believe that kernel fusion techniques offer potential for further acceleration and leave it for future work.
Variants with Sparse Connectivity We design MUD- DFormer variants by approximating its dense connections with two sparse connectivity patterns: 1) dilation and peri- odicity (MUDDFormer-kky) also used in (Pagliardini et al., 2024)): each DA module aggregates the outputs of every k-th block, and the DA modules are inserted after every p blocks. 2) sliding window (MUDDForme:=SWw)) each DA module accesses to the outputs from only previous blocks plus the embeddings. Figure 9 shows the Pile validation perplexities and relative training and inference speed (com- pared to Transformer++)) of these MUDDFormer variants. We measure perplexities with 405M models to align with ab- lation results in Table 5, while training and inference speeds are measured using 1.3B untrained models as in Table 4, the size of which is of more practical value.
3.6. Ablations and Variants
We conduct ablation studies with the 405M Trans- forner+t/MDDDFommer models in scaling law experiments for language modeling in Section 3.1.
Ablation settings We do two groups of experiments and report the perplexity results in Table 5. In the first group, we progressively add the four components, i.e. static (Sec- tion 2.1), dynamic (Section 2.2), multiway (Section 2.3) dense connections and parameter re-allocation (2.4)) to Transformer++ to finally obtain MUDDFormer to compare the contribution of each component. In the second group, we focus on the multiway aspect and study the effect of dense connections for the four decoupled inputs streams by replacing each of them with normal residual connection respectively.
While fully dense (1x1) connectivity achieves the best performance, these sparse variants provide a spectrum of performance-efficicn trade-offs. For example, switch- ing from MUDDFormer-Ixx 1 to MUDDFormer-22x in- creases relative training/infferenee speed from 89.8%/88.11 1%% to 97. 97.8%//93.49 4% with only 0.18 increase in ppl. In compar- ison, MUDDFormeer-SW is inferior with a lower speed ppl ratio, highlighting the indispensability of long-range (>8) cross-layer interactions.
Results All three ingredients of dense connections, i.e. static, dynamic and multiway, make contributions. While parameter re-allocation is effective on MUDDFormer, it deteriorates Transformer+++. Removing dense connections for each of the four streams hurts performance and the value stream benefits most from dense connections.
4. Related Work
Enhancing Residual Connections Despite the pervasive use of residual connections (He et al., 2016) in modern deep architectures, various approaches have been proposed to address their issues such as representational collapse and diminishing return for deeper models by strengthening cross-layer communication. Huang et al. (2017) introduced densely connected convolutional networks (DenseNet) for image classification. Inspired by DenseNet, Pagliardini et al. (2024) proposed DenseFormer for Decoder-only Trans- formers, which uses Depth Weighted Averaging modules to aggregate outputs from all preceding layers with static, learnable weights Similarly, Wang et al. (2019) proposed Dynamic Linear Combination of Layers (DLCL) in deep
Table 5. Ablations of MUDDFormer's components.
Config
ppl
Config
ppl
Transformer++
11.68
MUDDFormer
10.77
Static Dense
11.44
dense
10.89
+Dynamic Dense
11.09
dense
10.90
+Multiway Static Dense
11.27
V dense
11.05
+ Mulltiay Dynamic Dense
10.83
R dense
11.14
Mul. Dyn. Dense Re-alloc
10.77
Re-alloc
11.93
8

---

Multiway Dynamic Dense Connections
encoder-decoder Transformers for neural machine transla- tion, which also used static and learnable dense connection weights. MUDD connections also have some linkage with HighwayyNetworks (HN) (Srivastava et al., 2015). Both take inspiration from sequence model architectures and apply depthwise (HN from LSTM, MUDD from attention). HN also the first to propose the critical concept of input de- pendent gating when mixing outputs between layers. Most recently, Zhu al. (2025) proposed Hlype--Connectinn (HC), an alternative to residual connections that uses both static and dynamic weights to adjust inter-layer dependen- cies. Other research has explored different forms of cross- layer attention (EINokrashy et al., 2022; Fang et al., 2023; Wang et al., 2024) which retrieve or update representations across different layers in a more flexible manner.
attention. The computation overhead, both theoretical and practical, is much lower.
Intra-Layer Architectural Innovations Many other stud- ies attempt to enhance the performance or efficiency of foun- dational sequence models with individual layers, including attention mechanisms (Liu et al., 2024a; Xiao et al., 2024a; Ye al., 2024; Leviathan et al., 2024), sub-quadratic lin- ear attention or RNN/SSM architectures (Gu & Dao, 2024; Dao & Gu, 2024; Peng et al., 2024; Yang al., 2024) and sparse Mixture-of-Experts (MoE) (Fedus et al., 2022; Dai et al., 2024). By contrast, MUDD connections focus on cross-layer communication, making it orthogonal and com- plementary to these approaches. We leave the exploration of combining MUDD connections with these within-layer optimizations for future work.
MUDDFormer is closely related to DenseFormer and HC but differs in critical ways. First, unlike DenseFormer, our MUDD connections dynamically compute per-posiie posstisiionn weights conditioned on the hidden states. Second, although HC uses a combination of static and dynamic weights to ex- pand the hidden states, it does not employ explicit all-to-all cross-layer dense connectivity. Moreover, none of existing approaches consider decoupling the four input streams of Transformer block by multiway design, which is shown to bring significant performance gain in MUDDFormer.
5. Conclusion
We introduced Multiway Dynamic Dense connections to address the limitations of residual connections and enhance cross-layer information flow in Transformers. Experimental results showed that MUDDForner effective, efficient and scalable. It significantly outperforms Transformer baselines in language and vision tasks and improves emergent abilities such as in in-context learning with minimal overhead. MUDD connections have the potential to become an indispensable component of next-generation foundation models.
Mechanistic Interpretability Research in this field em- ploys various attribution methods (Conmy et al., 2023; Hanna et al., 2024) to uncover the circuits within Trans- formers that underlie specific capabilities (Elhage et al., 2020; Wang et al., 2024; Ni al., 2025). These studies reveal the critical role of cross-layer interactions between attention heads and MLPs in enabling complex reasoning a key insight motivating MUDD connections' design, which explicitly facilitates such interactions.
Acknowledgements
We are grateful to Google Cloud for providing the compute for model training, and to Shun Wang and Tingting Zhang for their technical support and help in troubleshooting TPU resource allocation and training.
Cross-Layer KV Cache Optimization Brandon et al. (2024) proposes Cross-Layer Attention (CLA) to reduce Transformer KV cache size by sharing keys and values be- tween adjacent layers, trading expressiveness for efficiency. Our MUDD connections enable cross-layer information flow between KV caches via dense connections on key and value streams. This enhances KV cache expressiveness and utility, improving in-context learning as evidenced by ex- periments. OmniNet (Tay et al., 2021a) achieves a fully global KV Cache receptive field by allowing each token to attend to all tokens in all layers. The authors reported that the additional omni (i.e. all-to-all) attention in OmniNet is computationally expensive and proposed to mitigate it by efficient attention variants. In contrast, MUDD connections much more efficient because: 1) Only depth-wise aggre- gation (DA) introduced, since the other omni attention paths can be formed by composition of DA and within-layer MHA; 2) DA is implemented as lightweight query-wise
Impact Statement
Our work introduces Multiway Dynamic Dense (MUDD) connections, lightweight architectural innovation that fun- damentally enhances Transformer-basee foundation mod- els. MUDD connections achieve up to 2.4x training effi- ciency gains while improving in-context learning capabili- ties, a critical requirement for real-world LLM applications. This breakthrough directly addresses the escalating com- putational and environmental costs of training ever-larger models, offering a sustainable pathway to pigh-performance AI without exponential parameter growth.
The open-source release of MUDDFormer architectures and pre-trained models will democratize access to state- of-the-art model efficiency techniques, particularly bene- fiting resource-constraine researchers and organizations. Enhanced attention mechanisms from MUDD's multiway design could advance mechanistic interpretability research

---

Multiway Dynamic Dense Connections
by producing more structured attention patterns and circuits.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, In Proceedings of the Forty-First Interna- tional Conference on Machine Learning (ICML), 2024.
While our method itself's architecture-focused, acknowl- edge that more capable language models could potentially be misused for harmful content generation. However, MUD- DFormer's efficiency gains may paradoxically mitigate this risk by reducing the energy barrier to developing safer, smaller models that match larger counterparts' performance. We commit implementing rigorous model release proto- cols aligned with responsible AI practices.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T.,, Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 10XXi:22100111900 2020.
References
Elhage, N., Neel, N., Olsson, C., et al. A math- ematical framework for transformer circuits. Transformer Circuits Thread, 2020. attps//transfoomer-ciiico 2021/ (021/framework/indee html.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' D'Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. Pythia: A suite for analyz- ing large language models across training and scaling. In International Conference on Machine Learning (ICML), pp. 2397-2430. PMLR, 2023.
URL
EINokrashy, M., AlKhamissi, B., and Diab, M. Depth-wise attention (dwatt): A layer fusion method for data-efficient classification. arXiv preprint arXiv:2209. 22162220..1168,, 2022.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Fang, Y., Cai, Y., Chen. J., Zhao, J., Tian, G., and Li. G. Cross-layer retrospective retrieving via layer attention. arXiv preprint arXiv:2302.088:
2023.
Fedus, W., Zoph, B., and Shazeer, N. Switch transform- ers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-522 5270, 2022.
Brandon, W., Mishra, M., Nrusimha, A., Panda, R., and Kelly, J. R. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint 2rXiv:2405.12981, 2024.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0002 2020.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac h, A., Li,
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.1000 10044 2019.
H., McDonell, K., Muennighoff, N., Ociepa, C., Phang,
J.. Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L.., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A framework for few-shot language model evaluation, 12 2023, URL https: /ttp///eoooooo///o///// org/ records/ 10256836.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv:1803 05557 2018.
Gromov, A., Tirumala, K., Shapourian, H., Glorioso, P., and Roberts, D. A. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403. 1788,, 2024.
Conmy, A., Mavor-Parker, A., Lynch, A., Heimersheim, S., and Garriga-Alonso A. Towards automated circuit discovery for mechanistic interpretability. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 16318-16352, 2023.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. In Proceedings of the First Conference on Language Modeling, 2024.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, Zeng, W., Yu, X., Wu, Y., et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts lan- guage models. arXiv preprint arXiv:2401.06066, 2024.
Hanna, M., Pezzelle, S., and Belinkov, Y. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms. In Proceedings of Conference on Language Modeling (COLM), 2024.
10

---

Multiway Dynamic Dense Connections
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition (CVPR), Pp. 770-778, 2016.
Merrill, W. Sabharwal, A., and Smith, N. A. Saturated transformers constant-depth threshold circuits. Trans- actions of the Association for Computational Linguistics, 2022.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., al. An empirical analysis of compute-optimat large language model training. Ad- vances in Neural Information Processing Systems, 35: 30016-30030, 2022.
Merullo, J., Eickhoff, C., and Pavlick, Talking heads: Un- derstanding inter-layer communication in transformer lan- guage models. arXiv preprint arXiv:2406.0m 09519, 2024.
Muennighoff, N., Soldaini, L.., Groeneveld, D., Lo, K., Mor- rison, J., Min, S., Shi, W., Walsh, P., Tafjord, O., Lambert, N., al. Olmoe: Open mixture-of-experts language models. 2025.
Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR), pp. 4700-4708, 2017.
Ni, R., Xiao, D., Meng, Q., Li, Li, X., Zheng, S., and Liang, H. Benchmarking and understanding compositional re- lational reasoning of Ilms. In Proceedings of the Thirty- Ninth AAAI Conference Artificial Intelligence (AAAI), 2025.
Korthikanti, V. A., Casper, J., Lym, S., McAfee, L.., Ander- sch, M., Shoeybi, M., and Catanzaro, B. Reducing acti- vation recomputation in large transformer models. Pro- ceedings of Machine Learning and Systems, 5:341-353, 2023.
Pagliardini, M., Mohtashami, A., Fleuret, F., and Jaggi, M. Denseformer: Enhancing information flow in transform- ers via depth weighted averaging. In Proceedings of the Thirly-Eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024.
Lai, G., Xie, Q.. Liu, H., Yang. Y., and Hovy, E. Race: Large-scale reading comprehension dataset from exami- nations. arXiv preprint arXiv:100.00668, 2017.
Paperno, D.,, Kruszewski, G., Lazaridou, A., Pham, Q. N.,
Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint 2016.
Leviathan, Y., Kalman, M., and Matias, Y. Selec- tive attention improves transformer. arXiv preprint arXiv:2410.02703, 2024.
Liu, A., Feng. B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo, D., et al. Deepseck-v2: A strong, economical, and efficient mixture-of-expers language model. arXiv preprint arXiv:2405.0444 2024a.
Peng, B., Goldstein, D., Anthony, Q., Albalak, A., AI- caide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., et al. Eagle and finch: Rwkv with matrix- valued states and dynamic recurrence. arXiv preprint arXiv:2404.0589) 058892 2024.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C.. Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek- v3 technical report. arXiv preprint 17X1v22412.19437, 2024b.
Petty, J., van Steenkiste, S., Dasgupta, I., Sha, F., Garrette, D., and Linzen, T. The impact of depth and width trans- former language model generalization. arXiv preprint arXiv:2310. 19956, 2023.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Liu, J., Cui, L., Liu, H., Huang, D., Wang. Y., and Zhang, Y. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08124 2020a.
Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99 106. 2021.
Liu, I.., Liu, x., Gao, J., Chen, W., and Han, J. Understand- ing the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), 2020b.
Shazeer, N. Glu variants improve transformer. arXiv
preprint arXiv:2002.05202. 2020.
Srivastava, R. K., Greff, K., and Schmidhuber, J. Training
very deep networks. Advances in neural information processing systems, 28, 2015.
Longpre, S., Hou, I.., Vu, T., Webson, A., Chung, H. W., Tay, Y., Zhou, D., Le. Q. V., Zoph, B., Wei, J., et al. The flan collection: Designing data and methods for effec- tive instruction tuning. arXiv preprint arXiv:2301.13688, 2023.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.

---

Multiway Dynamic Dense Connections
Tay, Y., Dehghani, M., Aribandi, V., Gupta, J., Pham, P. M., Qin, Z.,, Bahri, D., Juan, D.-C., and Metzler, D. Omninet: Omnidirectional representations from transformers. In International Conference on Machine Learning (ICML), pp. 10193-10202. PMLR, 2021a.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206. 07682, 2022.
Welbl, J., Liu, N. F., and Gardner, M. Crowdsoure- ing multiple choice science questions. arXiv preprint arXiv: 2017.
Tay, Y., Dehghani, M., Rao, J., Fedus, W., Abnar, S., Chung, H. W. Narang, S., Yogatama, D., Vaswani, A., and Metzler, D. Scale efficiently: Insights from pre- training and fine-tuning transformers. arXiv preprint arXiv:2109. 10686, 2021b.
xai org. Grok-1. 2024. URL htttps://gtthub...... xai- orooorrgggool grok-1. 1.
Team, G., Georgiev, P.. Lei, V. I., Burnell, R., Bai, L., Gulati, A., Tanzer, G., Vincent, D., Pan, Z., Wang, S., et al. Gemini 1.5: Unlocking multimodal understand- ing across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024a.
Xiao, D., Meng, Q., Li, S., and Yuan, X. Improving trans- formers with dynamically composable multi-heat atten- tion. 2024a.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations (ICLR), 2024b.
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, I.,, Mesnard, T,, Shahri- ari, B., Ramé, A., et al. Gemma Improving open language models at practical size. arXiv preprint arXiv:2408.00116 2024b.
Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Gao, C., Huang, C., Lv, C.,, et al. Qwen3 technical report. arXiv preprint 2Xiv:25505003399 2025.
Team, L. The Ilaaa 4 herd: The beginning of new of natively multimodal ai innova- tion. 2025. URL https: http:/:::/aa..aaa...oemmmm ctn.t/og/c///o///////o/o/ 1lama 4 www.mmmmmmmaaammmmmmmmmmmmmmmmma -nnteeelggeelleeee////e
Yang. S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Par- allelizing linear transformers with the delta rule over sequence length. In Proceedings of the Thirty-Eight Annual Conference on Neural Information Processing Systems, 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv: 2023.
Ye, T., Dong, I.., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. 2024.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: 1905.0780000 2019.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Zhu, D., Huang, H., Huang, Z., Zeng. Y., Mao, Y., Wu, B., Min, Q., and Zhou, X. Hyper-connections. In Pro- ceedings of the Thirteenth International Conference Learning Representations (ICLR), 2025.
Vig, J. and Belinkov, Y. Analyzing the structure of atten- tion in transformer language model. arXiv preprint 120.::9906.04284,, 2019.
Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the wild: a circuit for in- direct object identification in gpt-2 small. In Proceedings of the Eleventh International Conference on Learning Representations (JCLR), 2023.
Wang, K., Xia, X., Liu, J., Yi, Z., and He, T. Strengthen- ing layer interaction via dynamic layer attention. arXiv preprint ar arxiv:2406. 13392, 2024.
Wang, Q. Li, B., Xiao, T., Zhu, J., Li, C., Wong, D. F. and Chao, L. S. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACI.), 2019.
12

---

Multiway Dynamic Dense Connections
A. Dynamic Dense Connections as Depth-wise Self-Attention
Given a sequence & RTxD, the output of dot-product self-attention at position is:
softmax(1; (x.., J.;WWV
(1I)
where E,W W @f(+l)xdd (d is head dim) are + imput-dependent keys. Combining Eq. (5) and Eq. (6) in Section 2, the output of dynamic DA at layer i for position t is (operating depthwise across layers):
(GEL,U( a.()..((
(12)
Comparing Eq. (11) and Eq. (12), W1 plays the role of query projection (W²) and W E R(i+1)xd inner dim of MLP, also the head dim) + 1 keys as parameters independent of input. So dynamic DA can be seen as lightweight self-attention except:
the
Keys are independent of input;
A learnable positional bias a; is used;
Softmax is removed. Instead, GELU activation is applied to query (more like linear attention);
WV transformation is used.
While theoretically these simplifications may impact the representation capacity, empirically found that adding more sophisticated ingredients in DA (e.g. input dependent keys, softmax) does not bring improvement and slow down training.
B. PyTorch Style Pseudo-code for MUDDFormer
aatch_sizer reg_Iene D model_din dim
L layer_index; nam_ways 4; diddidee__di C+ (1+1)
def generate_dwls (x, thdd_theta): # x: BXTxD
wl, w2, a mudd theta # wl: DxK, w2: Xx (Cx (1+1)), at (x+1)
dw GRTU (RMSNorm (x) wl) 0 w2 t
dw rearrange (dw, (C CBTE', C=4) dw
def DA (Xs, mudd_theta) # Xs: List (1+1)x(BxTx0] 11
dw generate_dw [Xs[-1] mudd_theta)
12
xs = 11
13
for c, way in enumerate (C'0', 'K', WV, 'R'I):
x sum [dw[c, :, :, iti+1] * Xsli] BTI, BTD-BBDD for j range (1en [X
append return xs
18
19
def muddformer (x, model):
20
x model. enreedlembddingg (x)
Xs [x]
xq, xk, xv, blocks:
x, x, x,
for black in model attn (xq) (N() LN (xv)) +
attn_out bltn
block. iin (IN (ittantn + attn_out
36
append (x)
17
xq, xk, xv, xr DA (Xs, block. muddk_
return xr
13

---

Multiway Dynamic Dense Connections
C. Details of Complexity Analysis
Compared to Transformer++, extra compute and parameters MUDDFormer are introduced by DA modules, increasing from bottom layer to top layer. due to varied hidden dim K; of DA layer i. To estimate the overhead, we calculate Raporams and RAFLOPs. the ratios of extra parameters and computation by analyzing an average middle layer in MUDDFormer. For this layer, the average hidden dimension of DA is K 4(I + 1) 4(-+1 1) 2(L 3). In typical Transformer architecture we assume D) > K and define p negligible in terms of parameters and computation.
L+3 T 1 D
We omit RMSNorm because
Ratio of extra parameters
In Transformer++ of I. layers and D model dimension, the number of parameters is 12L (4LD²² for Attention and SLD for FFN). In MUDDFormer, the layer adds DK; parameters in Wi and Kद in W2 of DA, so the ratio of extra parameters as follows.
DK;
K?
(13)
R Aparams
12LD2
Approximating K2 with K and assuming D> K:
DK + KK²
Rapparama
12D2
DK
(assume D >>>>> K and ignore K²)
12D
K
12D
(14)
+(L+3)
(recall K 2(L + 3))
12D
L+3
6D
L 3
(denote 17
Thus, the ratio of extra parameters scales linearly with the rectified depth/width ratio 1) L+3
Ratio of extra FLOPs.
In a Transformer++ model, the pretraining FLOPs 2LDT(12D + T) for sequence of length T. At layer i of MUDDFormer, the extra FLOPs includes DTDK,+27T K? for generating dynamic dense weights Ai, (Eq, (6)) and 2T, for depthwise aggregate (Eq. (5)).
generate dense weight A.g
Depthwise Aguregate Code
Code Line
EL. 2TDK, + 2TK? + 2TDK:
(15)
RAFLOPs
2LDT(12D +T)
Similarly, we consider an average layer of MUDDFormer for simplification, then the extra FLOPS becomes 4TDK + +2TK"
14

---

Multiway Dynamic Dense Connections
4TDK + 2TK
RAFLOPs
2DT(12D +T)
2DK K
(divide 2l')
D(12D + +T)
2DK
(assume 2D > K and ignore K
D(12D +T)
(16)
4D(L + 3)
(recall K = +
12D +DT
(L+ +3/D
(divide 4D²2)
+ T//4D
T L +
3 + p/4 (denote
D D
1]
Therefore, the ratio is approximately 34p/1' decreasing with the rectified depth/width ratio.
D. Hyperparameters and Baselines for Scaling Law Experiments
D.1. Hyperparameters
We use the AdamW optimizer with 31 0.9, A2 = 0.95, gradient clip value 1.0, weight decay of 0.1, 1% learning rate warmup steps followed by cosine decay to 10% its maximal value, and no dropout. These hyperparameters are mostly taken from the GPT-3 paper (Brown et al., 2020) and are also used by the baseline models listed below.
D.2. Baseline Models
Transformer: The standard Transformer based on GPT-3.
Transformer++: An improved Transformer architecture adopted by Llama (Touvron et al., 2023) etc. with rotary positional encoding (RoPE) (Su et al., 2024), SwiGLU MLP (Shazeer, 2020), RMSNorm instead LayerNorm and no linear bias.
DenseFormer:
The DenseFormer model from Pagliardini al. (2024) without dilation, which has the best performance according the paper. We implemented the model in JAX based on the PyTorch code released by the authors
The dynamic hyper-connections with expansion rate n 4 (DHCx4) from Zhu al. (2025), which achieves superior results language model pre- pre-training and is the recommended configuration in the paper. We implemented the model in JAX based on the PyTorch Implementation given in Appendix J of the paper.
Typee--onnecttion:
DDFormer: Transformer with dynamic dense connections but without multiway splitting as described in Section 2.2.
E. Memory Usage and Wall-Clock Time for Main Experiments
Theoretical analysis of memory usage In theory, peak activation memory usage for training transformer in float1 6 with L layers, N heads, sequence length T, batch size B and model dim D occurs at the beginning of backpropagation and is composed of two parts:
1. hidden states for L layers: 2LBTD (gradient checkpointing)
2. activation memory for a layer: BTD(34 + GNT/D) (outlined in (Korthikanti al., 2023), recomputation of layer when backpropagate
15

---

Multiway Dynamic Dense Connections
We compare theoretical activation memory usages in Table 6. DenseFormer adds 2LBTD store gradients for each layer's hidden state when backpropagating DA after layer L. Based this, MUDDFormer adds another 6BTD for recomputation of layer L's multi-way hidden states Q.K,V when backpropagate it (they are not stored for all layers but are recomputed (L+3)/
during backpropagation). The extra memory ratio for MUDD is (L+ 3)/(Z (L + 17 + ANT/D), and typical values are less
than 30%. During inference, the activation memory usage dominated by the KV cache, which not impacted by the extra memory brought by MUDD.
Table 6. Comparison of theoretical activation memory usages.
model
activation memory
memory ratio
TFM+
2LBTD + BTD 6NT/D)
DenseFormer
2LBTD + BTD(31 6NT/DD + 2LBTD
17 3NT/D) 2LBTD BTD 6NT/D) 6BTD 2LBTD (L 177 3NT/D)
MUDDEM
17+
Measured memory usage and wall-clock time Table 7, we report actual memory usage (measured using jax. profile) and wall-clock time for both the main (Figure 3, Table 3) and efficiency (Table 4) experiments. Besides activation memory, the actual memory usage also includes model parameters, gradients and optimizer states, and affected by JAX compiler optimizations. For all model architectures and sizes, the relative training speed is 80%-90%, which could be further improved by custom kernel implementation. The extra memory ratio of MUDD is 20%-30% comparable to that of HyperConnections (see Table in their paper). As noted in our paper, the results for efficiency experiments (row 3-5) are of more practical relevance because they represent more commonly used architecture (Transformer++)) and model sizes.
Table 7. Comparison of measured activation memory usages and wall-clock time.
model
model size
wall-cloch time (hour)
rel. speed
mem (GB)
mem ratio
v5p pod size
tokens
batch size
TFM++/MUDDNN
405M
5.7/7
81%
86/111
29%
16
7B
0.5M
834M
20/25.1
79%
225/273
21%
16
15B
0.5M
1.4B
29.5/32.5
91% %
301/386
28%
32
26B
0.5M
2.8B
122/145
84%
1352/1648
22%
128
300B
2M
6.9B
251/262
96%
1887/2222
17%
128
300B
2M
Pythia/ MUDDPythia
1.4B
163/183
1296/1655
28%
300B
2M
2.8B
124/154
81%
1318/1655
25%
128
300B
2M
F. Image Classification with ViT
Besides decoder-only transformer for language modeling, we apply MUDD connections to Vision Transformer (ViT, an encoder-only Transformer) (Dosovitskiy et al., 2020) for image classification on the ImageNet- Ik dataset (ILSVRC-2012) Implementation and experimental settings (e.g. the use of LandAAugmenionnMiixUp fixed 2D sincos position embedding and global average pooling) are based the Big Vision codebase". We use AdamW with 31 0.9, 82 0.999, gradient clip value of 1.0, weight decay of 0.3, learning rate of 0.003 with 10000 warmup steps followed by cosine decay to O, batch size of 4096, RandAugment of leve 10, Mixup of 0.2 and dropout rate of 0.1. We use VIT-S/16 the baseline model and equip it with MUDD connections to obtain MUDDViT-S/IE We also compare with a 1.72x larger model ViT-M/16 (Table 8). We report validation loss and top-1 accuracy results on 90 and 300 epochs in Table 9. As can be seen, the gain from MUDD connections decreases a bit during the training progress, probably because many epochs of repeated passes over the same dataset diminish the additional expressive capacity brought by MUDD connections. Despite this, MUDDVTT-//6 still outperforms ViT-S/16 by 2% on epoch 300, also surpassing ViT-M/16.
Table VIT Model architectures ImageNet- Ik classification.
Model
naayers
dmodel
dpip
Theaad
params
(MUDDVTT-S/L ViT-M/16
12
384
1536
6
22M
12
512
2048
8
39M
hhans///ghhhhccmmmogggerrrrreer vision
16

---

Multiway Dynamic Dense Connections
classffcci Table ViT for ImageNet- 1k classification results.
Model
val. loss acc@e90
acc@=300
Rel. size
VIT-S/16
0.993
53,4
76.0
MUDDViT-S/16
0.871
56.9
78.1
1.007
ViT-M/16 0.890
55.2
77.9
1.72
G. Visualization
Head activation from attention patterns In Section 3.4, we show that the ratio head activations in MUDDPythia-2.88 larger than that in Pythia-2.88. Here we draw the actual attention patterns on randomly sampled sequence of length 32 from Pile validation set for the 32 heads in layer 25 of these two models Figure 10 and 11. It clear that attentions in Pythia mainly concentrate the sink token (inactive) while attentions in MUDDPythia disperse various tokens (active).
Cross-layer dynamic weights To better understand MUDDPythia, we visualize dynamic dense connection weights. Due to the high variance of the norm of hidden states X,, we scale those weights by the average norm of each layer, rectifying the importance of weights. The rectified mean and standard deviation of dynamic connection weights in MUDDPythia-2.85 are shown in Figure 12 and 13, respectively. It evident that the patterns of the four streams (query, key, value, residual) differ from each other, validating the necessity of separating them from the standard residual stream. is noteworthy that in the value-stream connections, most layers have salient and more dynamic weight on output of the first layer, thus forming long-range channel transport bottom information for attention heads in upper layers.
Head
Head
Head 2
Head
Hoadd
Hoad
Head 6
Head
0.8
0.8
Head
Head
Head
Head
Head
Hocdd
Head
Head 24
Head
Head 26
Head 27
Head
Head 29
Head 30
Head
0.2
Figure 10. Attention patterns for the 32 heads in the 25th layer of Pythia-2.8D
17

---

Multiway Dynamic Dense Connections
1.0
Head
Meadd
Head 10
Head 11
He.ad
Head 13
Head 14
Head 15
Head 16
Head
16
Head 19
Head
Head 21
Head 22
Head
00
Head
Head
Head
Hoad
Head
Figure 11. Attention patterns for the 32 heads in the 25th layer of MUDDPychia-2.8B
DA index
DA
V
10.
inden
Index
20
20
25
25
30
30
Figure 12. Mean of dynamic dense connections of MUDDPythia-22
18

---

Multiway Dynamic Dense Connections
200
DR indeon
index
10
10
10
10
DA index
DA index
35
50
Figure Standard deviation dynamic dense connections MUDDPythia-2.88
19


Reviews:
[{"review_id": "QP7zLkmUAg", "reviewer": "Reviewer_9FeX", "summary": "The paper proposes Multiway Dynamic Dense (MUDD) connections to enhance cross-layer information flow in Transformers by augmenting residual connections with dynamic, multi-headed dense aggregations. Key innovations include decoupling the query, key, value, and residual streams of each Transformer block, enabling dynamic weighting of connections based on position-dependent hidden states. This improves expressiveness over static/dense-only approaches (e.g., DenseFormer) and boosts scalability. MUDDFormer matches performance of larger models (e.g., Pythia-6.9B) with 2.4× less compute, achieving state-of-the-art results in language modeling while adding minimal overhead (~0.2% params, ~0.4% FLOPs). Ablations and analyses (e.g., attention activation, representation collapse) validate the method’s effectiveness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks!\n\n>Limited discussion on other dynamic weight methods (e.g., Llama 3’s routing) or global memory approaches (OmniNet).\n\nAlthough we don't know what you mean by \"Llama 3's routing\", we do a comparison with another representative dynamic weight method MoE, and a global memory approach OmniNet (which in fact has already been discussed briefly in Related Work). For detail please see reply to 5mRp.\n\n>Limited theoretical understanding of dynamic connections’ role.\n\nIn order to get a better understanding of dynamic connections’ role, we extend the discussions in the paper on connections between MUDD and self-attention.\n\n**dynamic dense connections as single-head self-attention**\n\nGiven a sequence $x\\in R^{T\\times D}$, the output of dot-product self-attention at position i is:\n\n$softmax(x_iW^Q(x_{:i}W^K)^T)x_{:i}W^V$ (a)\n\nwhere $x_{:i}W^K\\in R^{(i+1)\\times d}$ (d is head dim) are i+1 input-dependent keys.\n\nCombining (5) and (6) in the paper, the output of dynamic DA at layer i for position t is (operating depthwise across layers):\n\n$(GELU(X_i[t]W_1)W_2+a_i)X_{:i}[t]$ (b)\n\nComparing (a) and (b), $W_1$ plays the role of query projection ($W^Q$) and $W_2^T\\in R^{(i+1)\\times d}$ (d=i+1 is the inner dim of MLP and is also the head dim) are i+1 keys as parameters independent of input. So dynamic DA can be seen as lightweight self-attention except: \n- keys are independent of input;\n- a learnable positional bias $a_i$ is used;\n- softmax is removed. Instead, GELU activation is applied to query (more like linear attention);\n- $W^V$ transformation is not used.\n\nWhile theoretically these simplifications may impact the representation capacity, we empirically found that adding more sophisticated ingredients in DA (e.g. input dependent keys, softmax) do not bring improvement and slow down training.\n\n**multiway dynamic dense connections as multi-head self-attention**\n\nFor aggregating each of the four input streams (Q, K, V, R) there is a head. In our implementation (see Pseudo-code in Appendix A), the four heads share the query projection $W^Q$ while each head has its own set of keys $W_2^T$. The head dim is $d$ expanded four times.\n\nWhile there is no more theoretical analysis, there is an intuitive understanding of dynamic connections’ role (just as pointed out in the original Transformer paper (Vaswani et, al., 2017) in \"Why Self-Attention\"): as the role the horizontal attention over sequence in MHA plays, the depth-wise attention over layers in MUDD creates a *direct* path between any pair of layers, besides *mediated* through the shared residual stream which may get crowded when the model gets deeper.\n\n>Overheads could be prohibitive for edge/hardware-constrained scenarios.\n\nConsidering its significant improvement over Transformer as shown in extensive experiments, one can in principle use a MUDDFormer model to replace a twice bigger Transformer model and still get similar performance, so the resource requirement is actually *lower*, making MUDDFormer suitable for edge/hardware-constrained scenarios. Nevertheless, we admit that deploying LLMs on edge devices is beyond our expertise and the overhead on them needs further verification. We'd like to know how you make this speculation. We could discuss further if you have specific concerns.\n\n>Consider releasing pretrained models for reproducibility.\n\nSure. We'll release training and inference code in both JAX and PyTorch and the pretrained models once the paper gets accepted.\n\n>How does MUDD handle long contexts (e.g., 8k tokens)? Does it exacerbate memory issues?\n\nFor training a model with L layers, N heads, seq length T, batch size B and model dim D, the extra activation memory ratio for MUDDFormer over Transformer is (L+3)/(L+17+3NT/D) (please see reply to RYaY for derivation), which decreases as context length T increases. We measure and compare the actual memory usage for training MUDDFormer-6.9B and Transformer++-6.9B on 8K tokens of contexts, using the same batch size of 2M tokens as the efficiency experiment in Table 4. The extra memory usage ratio of MUDDFormer is ~10%, which would not become the bottleneck or cause memory issues under normal circumstances. For inference, the activation mem usage is dominated by the KV cache which grows with context length, so the relative impact of MUDD connections on memory consumption diminishes as the context gets longer.", "claims_and_evidence": "The claims include:\n\nMUDD improves Transformer scalability and efficiency by addressing residual stream bottlenecks.\nSupported by experiments showing better perplexity curves vs. baselines (e.g., 834M MUDD matches 1.8× compute-to-loss of Transformer++).\nMultiway dense connections enhance in-context learning (ICL).\nEvidenced by 5-shot performance rivaling 12B models and analysis of attention head activation.\nMinimal overhead (0.2% params, 0.4% FLOPs).\nVerified in Section 2.6 and Table 1, though real-world overheads (Table 4) are slightly higher.\nPotential gaps:\n\nClaims about \"circuit composability\" rely on indirect analyses (e.g., head activation) rather than explicit mechanistic studies.\nEfficiency claims could vary with hardware/implementation (real-world overheads are ~10-15% for training).", "methods_and_evaluation_criteria": "The dense connection mechanism (dynamic, multiway) is well-motivated for addressing depth-related issues. Evaluation combines scaling laws, downstream tasks, and structural analyses. Benchmarks (Pile, FLAN) and baselines (Pythia, DenseFormer) are appropriate. However, results could benefit from more diverse tasks (e.g., structured vs. semi-structured data).", "theoretical_claims": "No strict Theoretical Claims. This paper relies on empirical validation.", "experimental_designs_or_analyses": "Scaling laws span 405M to 1.4B models, with fair comparisons to Transformer++/Pythia.\n\nAblations isolate dynamic/multiway contributions.\n\nDownstream evaluations cover both generative and ICL tasks, though tests on vision models are limited.", "relation_to_broader_scientific_literature": "Builds on DenseFormer/Hyper-Connections but innovates via dynamic weights and multiway streams.", "essential_references_not_discussed": "Limited discussion on other dynamic weight methods (e.g., Llama 3’s routing) or global memory approaches (OmniNet).", "other_strengths_and_weaknesses": "Strengths:\n\n1. Simple yet effective method with minimal overhead.\n2. Strong empirical results across scales.\n3. Insightful analyses on representation collapse and head activation.\n4. Extend experiments on imagenet. \n\nWeaknesses:\n\n1. Limited theoretical understanding of dynamic connections’ role.\n2. Overheads could be prohibitive for edge/hardware-constrained scenarios.", "other_comments_or_suggestions": "1. Consider releasing pretrained models for reproducibility.", "questions_for_authors": "1. How does MUDD handle long contexts (e.g., 8k tokens)? Does it exacerbate memory issues?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8UnzZAn47c", "reviewer": "Reviewer_RYaY", "summary": "The paper proposes a new method for aggregating information from previous layers that improves information flow in transformers. Specifically, the authors combine two methods for the aggregation: gated, token-specific weighting and separate streams for inputs to the attention (queries, keys, values, and residuals). The authors demonstrate improvements on the transformer models with fixed FLOPs per training, when compared to models not incorporating the MUDD framework.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks! (mem = memory, backprop = backpropagating)\n\n>mem usage and wall-clock time for main experiments\n\nPeak activation mem usage for training a transformer in float16 with L layers, N heads, seq length T, batch size B and model dim D occurs at the beginning of backprop and is composed of two parts:\n- hidden states for L layers: 2LBTD  (grad checkpointing)\n- activation mem for a layer: BTD(34+6NT/D)  (outlined in [1], recomputation of layer L when backprop it)\n\nComparison of activation mem usages:\n|model|activation mem|mem $\\Delta$ ratio|\n|-|-|-|\n|TFM++|2LBTD+BTD(34+6NT/D)|1|\n|DenseFormer|2LBTD+BTD(34+6NT/D)+**2LBTD**|L/(L+17+3NT/D)|\n|MUDDFM|2LBTD+BTD(34+6NT/D)+**6BTD**+**2LBTD**|(L+3)/(L+17+3NT/D)|\n\nDenseFormer adds 2LBTD to store gradients for each layer's hidden state when backprop DA after layer L. Based on this, MUDDFormer adds another 6BTD for recomputation of layer L's multi-way hidden states Q,K,V when backprop it (they are not stored for all layers but are recomputed during backprop). The extra memory ratio for MUDD is (L+3)/(L+17+3NT/D) and typical values are less than 30%. During inference, the activation mem usage is dominated by the KV cache, which is not impacted by the extra mem brought by MUDD.\n\nWe report actual mem usage (measured using jax.profile) and wall-clock time for both the main (Fig 2, Tab 3) and efficiency (Tab 4) experiments below. Beside activation mem, the actual mem usage also includes model params, grads and optimizer states, and is affected by JAX compiler optimizations.\n|model|model size|wall-clock time (hour)| rel. speed |mem (GB)|mem $\\Delta$ ratio|v5p pod size|tokens|batch size|\n|-|-|-|-|-|-|-|-|-|\n|TFM++ / MUDDFM|405M|5.7/7|81%|86/111|29%|16|7B|.5M|\n||834M|20/25.1|79%|225/273|21%|16|15B|.5M|\n||1.4B|29.5/32.5|91%|301/386|28%|32|26B|.5M|\n||2.8B|122/145|84%|1352/1648|22%|128|300B|2M|\n||6.9B|251/262|96%|1887/2222|17%|128|300B|2M|\n|Pythia / MUDDPythia|1.4B|163/183|89%|1296/1655|28%|64|300B|2M|\n||2.8B|124/154|81%|1318/1655|25%|128|300B|2M|\n\nFor all model architectures and sizes, the relative training speed is ~80%-90%, which could be further improved by custom kernel implementation. The extra memory ratio of MUDD is ~20%-30%, comparable to that of HyperConnections (Tab 9 in their paper). As noted in our paper, the results for efficiency experiments (row 3-5) are of more practical relevance because they represent more commonly used architecture (Transformer++) and model sizes.\n\n[1] Korthikanti et al. (2022). Reducing activation recomputation in large transformer models\n\n> the analysis (eq. 11 and 12) could be clarified further with a more detailed derivation.\n\nDue to space limit, we'll refine the derivation and add explanation in the revised version. \n\n>citation to HighwayNetworks\n\nAs the precursor to residual connections, HighwayNetworks (HN) indeed has some fundamental linkage with MUDD connections, e.g. both take inspiration from sequence model architectures and apply depthwise (HN from LSTM, MUDD from attention). HN is also the first to propose the critical concept of *input dependent* gating when mixing outputs between layers. We'll cite HN in the revised version and discuss the linkage with MUDD in more detail.\n\n>The discussion of how Pre-Norm stabilizes training by preventing representation collapse (Liu et al., 2020b)\n\nThere may be a misunderstanding. What Liu et al. propose is that pre-ln *causes*, rather than prevents, representation collapse (RC) (see the paragraph before Section 4.2 in their paper). Some other papers (e.g. Hyper-Connections, ResiDual) also cite Liu et al. when discussing RC. Since RC in deep pre-ln Transformers is one of the issues tackled by MUDD, we cite this paper too.\n\n>dynamic dense connections as depth-wise self-attention needs further justification.\n\nGiven a sequence $x\\in R^{T\\times D}$, the output of dot-product self-attention at position i is:\n\n$softmax(x_iW^Q(x_{:i}W^K)^T)x_{:i}W^V$ (a)\n\nwhere $x_{:i}W^K\\in R^{(i+1)\\times d}$ (d is head dim) are i+1 input-dependent keys.\n\nCombining (5) and (6) in the paper, the output of dynamic DA at layer i for position t is (operating depthwise across layers):\n\n$(GELU(X_i[t]W_1)W_2+a_i)X_{:i}[t]$ (b)\n\nComparing (a) and (b), $W_1$ plays the role of query projection ($W^Q$) and $W_2^T\\in R^{(i+1)\\times d}$ (d=i+1 is the inner dim of MLP, also the head dim) are i+1 keys as parameters independent of input. So dynamic DA can be seen as lightweight self-attention except: \n- keys are independent of input;\n- a learnable positional bias $a_i$ is used;\n- softmax is removed. Instead, GELU activation is applied to query (more like linear attention);\n- $W^V$ transformation is not used.\n\nWhile theoretically these simplifications may impact the representation capacity, we empirically found that adding more sophisticated ingredients in DA (e.g. input dependent keys, softmax) does not bring improvement and slow down training.\n\n>grammatical errors and wording issues\n\nThank you! We'll fix them in the revised version.", "claims_and_evidence": "The authors claim that adding multiway dynamic dense (MUDD) connections improves the performance of the model. They demonstrate this on the example of the transformer architecture by integrating MUDD and showing that it can achieve the performance of models trained with twice the theoretical compute.\n\nWhile the evidence in terms of theoretical compute is very strong, the paper would benefit from reporting additional empirical measurements of performance from the main training, such as memory usage and wall-clock time for the main experiments. Aggregating multiple layers might lead to increased memory consumption, so it is important to consider this while designing a scalable architecture. Moreover, there are discrepancies between theoretical FLOP computations and practical wall-clock time. While the paper reports tokens/s in Table 4, this metric reflects latency for different models than those compared in the main claim with scaling curves. The paper would be improved if the timing was shown for the main experiment or if both were provided.", "methods_and_evaluation_criteria": "The selected benchmarking datasets and evaluation criteria are sensible but should be extended to include memory and wall-clock time measurements for the main experiments.", "theoretical_claims": "The only theoretical claim is the complexity analysis presented in Appendix B. However, the transitions in the analysis (equations 11 and 12) are extremely informal and could be clarified further with a more detailed derivation.", "experimental_designs_or_analyses": "The experimental comparison in the fixed FLOP scenario, particularly in relation to Pythia, makes sense. There are additional analyses that further give intuition about the underlying mechanisms in MUDD.", "relation_to_broader_scientific_literature": "The paper effectively positions itself as a continuation of research on DenseNet and DenseFormer by introducing dynamic connections to improve selectiveness in aggregation. It covers relevant recent literature, and it may be beneficial to also cite HighwayNetworks [1], which introduced a gated version of the residual connection earlier. This citation could further enrich the paper’s motivation by linking it to foundational ideas that combine aspects of DenseNet with the original HighwayNetworks concepts.\n\n[1] Srivastava, Rupesh Kumar, Klaus Greff, and Jürgen Schmidhuber. \"Highway networks.\" arXiv preprint arXiv:1505.00387 (2015).", "essential_references_not_discussed": "I did not find any essential references not discussed in the paper.", "other_strengths_and_weaknesses": "The paper would benefit from additional clarity and motivation in some sections, and figures could be described more thoroughly. Some claims might be refined or supported with further justification. For instance:\n\n- l14-l17: The discussion of how Pre-Norm stabilizes training by preventing representation collapse (Liu et al., 2020b) - the citation does not support this claim and the claim itself is not related to the paper.\n- l105-108: presenting dynamic dense connections as depth-wise single-headed self-attention needs further justification.\n\nAdditionally, the paper is full of grammatical errors and wording issues:\n\n- l078: \"benefit from differentiated\" -> \"benefit from different.\", also 'different' as a word here is redundant because of what follows directly after.\n- l088: \"differentially\" -> \"differently.\"\n- l138: \"grows\" -> \"grow.\", l640: \"calculates\" -> \"calculate\" (this error appears in multiple locations).\n- l643: \"negible\" \netc.", "other_comments_or_suggestions": "-", "questions_for_authors": "In summary, while the paper contains several minor issues that, if addressed, could further improve its clarity and impact, the main result is both strong and potentially impactful, particularly from a scalability perspective. I am willing to reconsider my evaluation once these issues have been addressed.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U13SZBqtxC", "reviewer": "Reviewer_5mRp", "summary": "The paper proposes Multiway Dynamic Dense (MUDD) connections to improve cross-layer information flow in Transformers. By dynamically computing per-position weights for query, key, value, and residual streams, MUDD mitigates representation collapse and enhances depth efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks!\n\n> it lacks comparisons to alternative cross-layer architectures like OmniNet (Tay et al., 2021a) or MoE-based models (Fedus et al., 2022).\n\n**vs. OmniNet**\n\nWe've already compared MUDD connections with OmniNet briefly in Related Work. While both aim at promoting information flow across the whole model, they have major differences which we elaborate below.\n\nThe additional Attend operation in OmniNet for omni (i.e. all-to-all) attention is computationally expensive (overhead generally larger than the whole original forward pass because of much longer seq length). The authors propose to mitigate it by efficient attention variants. MUDD connections is much more efficient by: 1) Only depth-wise aggregation (DA) is introduced, because the other omni attn paths can be formed by composition of DA and within-layer MHA; 2) DA is implemented as lightweight query-wise attention (pls see reply to RYaY). The computation overhead (both theoretical and practical) is much lower.\n\nThere is no official open source implementation of OmniNet and we do not implement and evaluate it due to time constraints. Instead here we do a rough comparison based on the C4 language modeling results in Table 1 of their paper, because it is most similar to our main experiments in terms of task, data and training compute. They compare models with roughly the same number of parameters, as also done in our experiments.\nMUDDPythia-1.4B/2.8B's relative reduction on ppl (5.1%/5.1%, see Table 2) is larger than OmniNet_T (4.2%) with the most expensive Transformer self-attention, which further degrades to OmniNet_P's 2.6% and OmniNet_B's 0.4% when more practical efficient attention is used to mitigate the prohibitive computation cost of OmniNet_T.\n\nThe distinctive advantage of OmniNet over MUDD is allowing early layers to attend to late layers (i.e. back attention). This recurrent processing can increase the effective depth. How to enable back attention while maintaining efficiency is a very valuable research direction which we'd like to explore in the future.\n\n**vs. MoE**\n\nMoE and MUDDFormer are both architectures with dynamic weights. However, MoE uses these weights to select experts *within* a layer while MUDDFormer uses the weights to aggregate outputs *across* layers. They are complementary approaches and can be combined. To empirically compare MoE and MUDDFormer, we train two MoE models MoE-1.8B-A405M and MoE-4B-A834M with the same activated parameters as 405M and 834M dense models, using the same training settings in the scaling law experiments. For MoE specific settings, we largely follow OLMoE, choose 2 experts out of 16, use token choice policy with expert capacity 1.5 and load balance loss coefficient 0.01. Although the FLOPs of MoE is larger than that of MUDDFormer because of expert capacity greater than 1, and the number of total parameters is 4.4x - 4.8x larger, MUDDFormer is slightly better than MoE in perplexity (table below), showing its efficient utilization of *both* parameter and computation. In contrast, MoE works by *decoupling* parameter and computation and relies on significantly expanding parameters for good performance.\n    \n|total params|activated params|TFM++|TFM++ MoE|MUDDFormer\n|-|-|-|-|-|\n|405M / 1.8B / 405M |405M|11.69 |10.83|**10.77** \n|834M / 4B / 834M |834M|9.47 |8.88|**8.85** \n\n>Inference slowdown (~10–15%) is measured but not optimized.\n\nAs discussed in the paper, inference slowdown primarily stems from the series of small operations. Currently we implement inference in pure PyTorch and optimize speed mainly relying on torch.compile. We believe it can be further optimized by writing custom kernels which is left for future work. We'd like to note that even without further optimization, the 10-15% slowdown is entirely worth paying in practice considering the 1.8x-2.4x gain in computation and parameter utilization.\n\n> Limited comparison to alternative architectures beyond DenseFormer and Hyper-Connections.\n\nAs most recent cross-layer architectures, the authors of DenseFormer and Hyper-Connections (HC) already compare them with several related approaches and show their advantage over these approaches, e.g. DenseFormer > Depthwise Attention (ElNokrashy et al., 2022), HC > ResiDual (Xie et al., 2023), Altup (Baykal et al., 2024). Therefore, we do not compare with these cross-layer approaches anymore by showing MUDD connections > HC > Denseformer.\n\n> No real-world deployment discussion, which would strengthen practical relevance.\n\nAs a general component that significantly enhance Transformer architecture, MUDD connections have broad real-world applications, especially for compute or memory resource-constrained settings, e.g. running LLMs on mobile, considering MUDDFormer's much improved parameter and computation utilization over Transformer. For example, one can use a ~2B MUDDFormer model to replace a twice bigger Transformer model and still get similar performance, so that the resource requirement can be lowered,", "claims_and_evidence": "Most claims are well-supported, but efficiency trade-offs and broader comparisons need discussion.", "methods_and_evaluation_criteria": "Methods and evaluation are well-aligned with the problem, but efficiency trade-offs and broader comparisons could strengthen the study.  \n\nExperiments confirm efficiency gains, depth utilization, and representation diversity.  However,  Inference slowdown (~10–15%) is measured (Table 4) but not addressed.  Focuses on DenseFormer and Hyper-Connections, lacking comparisons to OmniNet or MoE models.", "theoretical_claims": "The paper does not present formal theoretical claims or proofs. Its contributions focus on empirical improvements through MUDD connections rather than theoretical guarantees. The complexity analysis (Section 2.6) is straightforward and aligns with standard FLOP and parameter counting methods. No issues were found.", "experimental_designs_or_analyses": "The experimental design is rigorous and supports the claims, though broader architectural comparisons and efficiency optimizations would strengthen the analysis.", "relation_to_broader_scientific_literature": "The paper builds on residual connections (He et al., 2016) and DenseFormer (Pagliardini et al., 2024), addressing Transformer depth inefficiencies (Merrill et al., 2022) and representation collapse (Liu et al., 2020b) with dynamic multiway dense connections, though it lacks comparisons to alternative cross-layer architectures like OmniNet (Tay et al., 2021a) or MoE-based models (Fedus et al., 2022).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:  \n- The paper presents a novel combination of dense connections and dynamic weight generation, improving Transformer efficiency with minimal overhead.  \n- Strong empirical validation with extensive scaling laws, ablation studies, and interpretability analysis.  \n- Clear and well-organized writing with useful visualizations and pseudo-code for reproducibility.  \n\nWeaknesses:  \n- Inference slowdown (~10–15%) is measured but not optimized.  \n- Limited comparison to alternative architectures beyond DenseFormer and Hyper-Connections.  \n- No real-world deployment discussion, which would strengthen practical relevance.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_flag": "Flag this paper for an ethics review."}}]

Author rebuttal:

Meta-review: I recommend accepting this paper. The proposed MUDD connections enhance Transformers by dynamically weighting cross-layer connections for each input stream. All three reviewers gave weak accept ratings. The method shows good results, matching Transformers trained with 1.8x-2.4x compute while adding only very few parameters. Despite some efficiency tradeoffs, the significant performance improvements and strong scaling properties across model sizes make this a valuable contribution with practical applications, particularly for resource-constrained settings.
Decision: Accept (poster)