Mixture-of-Depths Attention

Lianghui Zhu; Yuxin Fang; Bencheng Liao; Shijie Wang; Tianheng Cheng; Zilong Huang; Chen Chen; Lai Wei

cs.CLMar 16, 2026

Mixture-of-Depths Attention

Authors: Lianghui Zhu, Yuxin Fang, Bencheng Liao, Shijie Wang, Tianheng Cheng, Zilong Huang, Chen Chen, Lai Wei, +5 more

Organizations: School of EIC, Huazhong University of Science & Technology · ByteDance Seed

Abstract

Scaling depth is a key driver for large language models (LLMs). Yet, as LLMs become deeper, they often suffer from signal degradation: informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers. We introduce mixture-of-depths attention (MoDA), a mechanism that allows each attention head to attend to sequence KV pairs at the current layer and depth KV pairs from preceding layers. We further describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2's efficiency at a sequence length of 64K. Experiments on 1.5B-parameter models demonstrate that MoDA consistently outperforms strong baselines. Notably, it improves average perplexity by 0.2 across 10 validation benchmarks and increases average performance by 2.11% on 10 downstream tasks, with a negligible 3.7% FLOPs computational overhead. We also find that combining MoDA with post-norm yields better performance than using it with pre-norm. These results suggest that MoDA is a promising primitive for depth scaling. Code is released at https://github.com/hustvl/MoDA .

Figures & tables

Figure 1: We propose mixture-of-depths attention (MoDA) to address the modern LLM’s information dilution problem in a dynamic and hardware-efficient way. Compared with vanilla causal sequence attention, MoDA additionally allows query to attend to depth memories, i.e., depth KV pairs $\{K_{i},V_{i}\}_{i=0}^{l-1}$ at the same query position from preceding layers.

Figure 2: Comparing MoDA and strong open-sourced baseline, i.e., OLMo2 [ 27 ] , with validation loss and downstream performance under the 1.5B-parameter setting. Models using MoDA achieve lower C4 [ 30 ] validation loss and better downstream performance, i.e., HellaSwag [ 48 ] , WinoGrande [ 32 ] , and ARC-Challenge [ 10 ] , than OLMo2.

Figure 3: Conceptual comparison of mechanisms that utilize the depth stream. (a) Depth Residual [ 16 ] is the standard residual connection along depth: it reads the current representation and writes back by addition. (b) Depth Dense [ 20 , 28 ] reads a set of historical representations and linearly projects them back to width $D$ ; it writes back by concatenation along depth, preserving all intermediate states. (c) We introduce Depth Attention as an intermediate formulation, which uses attention to read historical depth KV pairs in a data-dependent way. It writes back by concatenating the current layer’s keys and values along depth. (d) We propose the upgraded version of Depth Attention, i.e., Mixture-of-Depths Attention (MoDA), which combines depth attention with standard sequence attention. It writes both the current layer’s output and its KV pairs to depth streams for subsequent layers.

Methods	Depth Dense	Depth Attention	Mixture-of-Depths Attention
Is data-dependent?	✘	✔	✔
Is unified softmax?	✘	✘	✔
Parameters	$\frac{1}{2}G^{2}L^{2}H_{k}^{2}d^{2}+\frac{1}{2}G^{2}LH_{k}^{2}d^{2}$ $O(L^{2}D^{2})$	$G^{2}LH_{k}^{2}d^{2}+2GLH_{k}^{2}d^{2}$ $O(LD^{2})$	$2GLH_{k}^{2}d^{2}$ $O(L\frac{D^{2}}{G})$
Decoding Cache	$GLH_{k}d$ $O(LD)$	$2LH_{k}d$ $O(L\frac{D}{G})$	$2LH_{k}d$ $O(L\frac{D}{G})$
Prefilling Cache	$GTLH_{k}d$ $O(TLD)$	$2TLH_{k}d$ $O(TL\frac{D}{G})$	$2TLH_{k}d$ $O(TL\frac{D}{G})$
Decoding FLOPs	$G^{2}LH_{k}^{2}d^{2}+G^{2}L^{2}H_{k}^{2}d^{2}$ $O(L^{2}D^{2})$	$2GL^{2}H_{k}d+2GLH_{k}d$ $O(L^{2}D)$	$2GL^{2}H_{k}d+2GLH_{k}d$ $O(L^{2}D)$

Table 1: Asymptotic complexity of depth-stream mechanisms. Here,

T

is the sequence length,

D

is the model width,

G

is the group size of Group Query Attention (GQA) [ 2 ] ,

H_{k}

is the number of key heads (equal to value heads

H_{v}

H_{q}

is the number of query heads (=

GH_{k}

d

is the head dimension, and

L

is the number of layers. We report dominant terms and omit constant factors.

Figure 4: Hardware view of MoDA depth-cache access. Left: flash-compatible hardware-efficient MoDA keeps a depth KV cache of length $T\times L$ for each sequence, so each query potentially scans a long concatenated depth KV. Right: chunk-aware MoDA groups queries by chunk size $C$ and reorganizes depth KV by chunk, reducing the effective depth span from $T\times L$ to $(C\times L)/G$ per chunk, where $G$ is the GQA group number. This layout improves depth KV calculation efficiency and reduces memory access overhead.

No.	$T$	$G$	$H_{q}$	$H_{k}$	$L$	FA2-triton (ms)	MoDA-triton (ms)	Depth Utilization ( $\eta_{\text{depth}}$ )	Extra Time Percentage
Scaling Sequence Length $T$ ( $B{=}1$ , $d{=}64$ , $C{=}64$ )
(1)	4096	8	64	8	64	7.970	10.750	12.50%	25.86%
(2)	8192	8	64	8	64	28.700	35.427	12.50%	18.99%
(3)	16384	8	64	8	64	116.700	127.661	12.50%	8.59%
(4)	32768	8	64	8	64	459.854	480.914	12.50%	4.38%
(5)	65536	8	64	8	64	1831.668	1883.026	12.50%	2.73%

Table 2: Efficiency comparison of hardware-efficient MoDA and FlashAttention-2 Triton kernels under “forward&backward” setting. We report runtime (ms), depth utilization (

\eta_{\text{depth}}

), and relative extra time across three scaling settings. Here,

B

denotes batch size,

d

denotes head dimension, and

C

denotes chunk size. We launch all experiments on A100 GPU with bfloat16 data type.

Model	Layer	Mixture-of-Depths Attention (MoDA)				Params (M)	FLOPs (T)	Train PPL	C4 Val PPL	Downstream Average
Model		Sequence KV	Depth KV	Extra FFN KV Proj.	Extra Attn KV Proj.
Baseline Models
(1) OLMo2		36	✔								669.0	8.01	14.49	18.59	56.93
(2) OLMo2	38	✔				700.5	8.41	14.27	18.31	57.11
Our Models
(3) Ours	36	✔	✔			669.0	8.02	14.08	18.48	58.10

Table 3: Performance of different mixture-of-depths attention (MoDA) variants on the training set, C4 validation set, and downstream benchmarks. We train the 700M models on 400B tokens. For MoDA settings: ‘Sequence KV’ means the each token only attends to the sequence keys/values, can be regarded as the vanilla attention mechanism. ‘Depth KV’ means the each token attends to its depth keys/values. ‘Extra FFN KV Proj.’ means further project the FFN’s input

X

to the depth keys/values, which are then used in subsequent attention operations. ‘Extra Attn KV Proj.’ means set individual depth key/value projections rather than reuse the original key/value projections of sequence attention. The width

D

, GQA group size

G

, sequence length

T

are set to 1024, 2, and 4096, respectively. We further report the parameters and FLOPs of the models.

Model	PIQA	Hella- Swag	Wino- Grade	OpenBook- QA	BoolQA	SciQ	ARC-E	ARC-C	COPA	MMLU	Average
700M Models
(1) OLMo2	73.72	58.77	55.33	35.60	56.24	89.50	66.84	33.44	77.00	24.69	57.11
(2) Ours	73.39	59.19	60.22	37.20	59.33	89.60	67.37	34.78	82.00	25.61	58.87
1.5B Models
(3) OLMo2	76.55	65.86	63.22	38.80	63.61	90.60	72.98	42.47	81.00	27.73	62.28
(4) Ours	76.82	66.24	65.59	41.60	67.34	92.10	72.81	46.82	85.00	29.59	64.39

Table 4: Performance of the proposed MoDA models with varying model sizes on the downstream benchmarks. We train the 700M and 1.5B models on the 400B tokens of OLMo2 dataset. The width

D

, GQA group size

G

, sequence length

T

are set to 1024, 2, and 4096, respectively. We mark the best performance with the bold font.

Model	C4	ICE	m2d2-s2orc	Pile	Wiki-text	Books	CC	peS2o	Reddit	Stack	Average
700M Models
(1) OLMo2	18.32	17.43	24.37	9.53	12.26	16.78	20.53	9.17	23.84	3.93	15.61
(2) Ours	18.29	17.24	23.64	9.48	12.06	16.58	20.52	9.14	23.75	3.90	15.46
1.5B Models
(3) OLMo2	16.16	15.37	21.10	8.45	10.41	14.19	18.13	8.19	21.21	3.57	13.67
(4) Ours	15.97	15.08	20.92	8.33	10.16	13.95	17.88	8.09	20.85	3.52	13.47

Table 5: Per-domain validation perplexity of the proposed MoDA models with varying model sizes. We train the 700M and 1.5B models on the 400B tokens of OLMo2 dataset. The width

D

, GQA group size

G

, sequence length

T

are set to 1024, 2, and 4096, respectively. Lower perplexity indicates better performance and is marked with the bold font.

Model	Layer	Norm	Mixture-of-Depths Attention (MoDA)			Params (M)	FLOPs (G)	FineWeb-Edu Val Loss
Model	Layer	Norm	Sequence KV	Depth KV	Extra FFN KV Proj.	Params (M)	FLOPs (G)	FineWeb-Edu Val Loss
Experiments with Deeper Models (48 Layers)
(1) OLMo2	48	prenorm	✔			123.38	136.61	3.3800
(2) OLMo2	48	postnorm	✔			123.38	136.61	3.4062
(3) Ours	48	prenorm	✔	✔		123.38	137.89	3.3759
(4) Ours	48	postnorm	✔	✔		123.38	137.89	3.3653

Table 6: Layer-number analysis of MoDA under deeper (48-layer) and shallower (24-layer) model settings. We compare vanilla attention (OLMo2) and MoDA variants with different MoDA choices, under both pre-norm and post-norm configurations. Models are trained with the same data recipe, and we report parameter count, FLOPs, and FineWeb-Edu validation loss. Across both depth regimes, introducing Depth KV consistently improves validation loss, and adding Extra FFN KV Projection yields further gains at moderate compute overhead.

Figure 5: Mixture-of-Depths Attention (MoDA) heatmaps with the combined-softmax formulation. Columns correspond to uniformly sampled layers $\{0,11,23,35\}$ , and rows correspond to randomly selected heads in each layer. The first column shows attention over sequence KV only, while the other columns show concatenated Sequence KV | Depth KV ; the red dashed line marks the boundary between the two KV blocks. Across layers and heads, substantial attention mass is consistently assigned to the depth-KV block, indicating that MoDA effectively leverages depth information in addition to standard sequence attention.

No.	Naive Torch	Flash-Compatible	Chunk-Aware	Group-Aware	Time (ms)
$B{=}1,\,T{=}1024,\,G{=}8,\,H_{q}{=}64,\,H_{k}{=}8,\,d{=}64,\,L{=}64,\,C{=}64$
(1)	✔				2128.900
(2)		✔			13.102
(3)		✔	✔		6.286
(4)		✔	✔	✔	1.460

Table 7: Ablation of the proposed kernel implementation strategies under a fixed configuration. Each row incrementally enables one optimization component: (1) naive PyTorch baseline, (2) Flash-compatible depth-KV layout, (3) Flash-compatible & chunk-aware depth-KV layout, and (4) Flash-compatible & chunk-aware & group-aware indexing. We report end-to-end “forward&backward” runtime in milliseconds (ms), where lower is better and the best performance is marked with the bold font. Experiments are run on a single A100 GPU with bfloat16 under fixed setting

B{=}1,\,T{=}1024,\,G{=}8,\,H_{q}{=}64,\,H_{k}{=}8,\,d{=}64,\,L{=}64,\,C{=}64

Explore similar work

Jun 3, 2026

Depth-Attention: Cross-Layer Value Mixing for Language Models

Self-attention selects information freely across the sequence, but across depth, Transformers merely add each layer's output to the residual stream, so later layers cannot selectively reuse earlier-layer representations. Recent cross-layer methods improve this flow but operate on hidden states outside attention, adding state beyond the key-value cache at inference--a cost that becomes increasingly salient as modern LLMs compress the cache with grouped-query and multi-head latent attention. We introduce Depth-Attention, which performs this selection inside the attention module itself: before a layer attends over the sequence, its query attends over the keys of earlier layers at the same token position and mixes their values into the value that self-attention then reads. Because Depth-Attention reuses the standard attention queries, keys, and value-cache slots, storing depth-mixed values in place of the original values, it adds no parameters and introduces no persistent inference state beyond the standard key-value cache--the same cache size as a vanilla decoder and less than hidden-state-based cross-layer methods. On Qwen3-style decoders at 1.5B and 3B parameters, Depth-Attention attains the lowest perplexity and the highest average downstream accuracy, improving over the vanilla Transformer by up to 2.3 accuracy points and surpassing strong cross-layer baselines in perplexity and average accuracy, while adding under 0.01% extra arithmetic FLOPs and no additional persistent inference state. The gains hold from 360M to 3B parameters and extend to looped Transformers.

Boyi Zeng, Yiqin Hao, Zitong Wang +7

LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2

Efficient Transformer Inference Decoder-Only Language Models Transformer Attention

Ask PDF

Oct 13, 2026

Depth Memory

Residual connections propagate information through a network by additively accumulating layer outputs. Recently proposed alternatives, such as Manifold-Constrained Hyper-Connections (mHC) and Attention Residuals (AttnRes), instead use more expressive depth aggregation mechanisms that improve language modeling performance at scale. In this work, we consider the mixing of representations over depth as a problem of storing and retrieving information across layers. Formalizing this perspective, we introduce Depth Memory (DM), a unified framework that describes the standard pre-layer-norm (PreLN) transformer, mHC, and AttnRes in terms of streams that carry representations through the network. One such stream, for example, could be an explicit memory stream that is queried to enrich the input to every layer, and whose output is subsequently written back into the memory for use by later layers. For AttnRes, this memory stream is a growing cache of explicit layer representations. For mHC, it is an expanded residual stream. Building on this view, we propose Linear Depth Memory (LDM), which formulates information retrieval across depth as a regression problem solved using recursive least squares. The resulting architecture propagates three streams across depth: the standard residual stream, a linear associative memory stream, and an inverse covariance stream. For high-throughput training and inference, we propose a hardware-aware implementation of LDM that reduces peak activation memory and memory traffic. LDM consistently achieves the lowest perplexities among PreLN, mHC and AttnRes baselines at the 400M and 1.5B parameter scales, alongside competitive zero-shot accuracy.

David González-Martínez, Philipp Nazari, T. Konstantin Rusch +1

ELLIS Institute Tübingen · Max Planck Institute for Intelligent Systems, Tübingen · Tübingen AI Center +2

Memory-Augmented Language Models Associative Memory Attention Residual Connections

Ask PDF

Oct 2, 2025

Accelerating Attention with Basis Decomposition

Attention is a core operation in large language models (LLMs). We present BD Attention (BDA), a lossless algorithmic reformulation of attention. BDA is enabled by a simple matrix identity from Basis Decomposition (BD), which restructures multi-head projections into a compact form while preserving exact outputs. Unlike I/O-aware system optimizations such as FlashAttention, BDA provides a mathematically guaranteed acceleration that is architecture-agnostic. On DeepSeek-V2-Lite (16B, FP16), BDA requires only 4s of offline preparation with no retraining required and, on modern GPUs, achieves 34% faster key/value projections and 25% smaller weights, while increasing perplexity (PPL) by just 0.02% (FP16) or 0.0004% (FP32), a negligible effect on model performance. These results position BDA as a theoretically exact method for lossless attention acceleration that is complementary to existing engineering-level optimizations. Our code is available at https://github.com/abcbdf/basis-decomposition-official.

Jialin Zhao

Department of Computer Science Tsinghua University

Self-Attention Transformer Attention LLM Inference Acceleration Efficient Attention Attention Mechanisms Multi-Head Attention

Ask PDF