--- author: - Lianghui Zhu - Yuxin Fang - Bencheng Liao - Shijie Wang - | \ Tianheng Cheng - Zilong Huang - Chen Chen - Lai Wei - Yutao Zeng - Ya Wang - Yi Lin - Yu Li - Xinggang Wang bibliography: - main.bib title: 'Mixture-of-Depths Attention' --- ```{=latex} \def\ie{\textit{i.e.}} ``` ```{=latex} \def\eg{\textit{e.g.}} ``` ```{=latex} \def\etc{\textit{etc}} ``` ```{=latex} \def\etal{{\em et al.~}} ``` ```{=latex} \def\vs{{\em vs.}} ``` ```{=latex} \newcommand{\myPara}[1]{\noindent\textbf{#1.}} ``` ```{=latex} \newcommand{\corr}{\faEnvelope} ``` ```{=latex} \newcommand{\name}{MoDA} ``` ```{=latex} \newcommand{\TODO}[1]{\textcolor{red}{[TODO: #1]}} ``` ```{=latex} \maketitle ``` ![ We propose mixture-of-depths attention (`\name{}`{=latex}) to address the modern LLM's information dilution problem in a dynamic and hardware-efficient way. Compared with vanilla causal sequence attention, `\name{}`{=latex} 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. ](figures/dsa_pipeline_v6.2.png){#fig:dsa_pipeline width="1.\\textwidth"} ```{=latex} \newcommand{\customfootnote}[2]{% \begingroup \renewcommand\thefootnote{#1}% \footnotetext{#2}% \addtocounter{footnote}{-1}% \endgroup } ``` ```{=latex} \customfootnote{*}{This work was done when Lianghui Zhu and Bencheng Liao was interning at ByteDance Seed.} ``` ```{=latex} \begin{tcolorbox}[boxsep=0pt, left=10pt, right=10pt, top=5pt, bottom=5pt]\setlength{\parindent}{0cm} \setlength{\parskip}{0cm} \abstractlist\par { \setlength{\parskip}{0cm} \ifdefempty{\checkdatalist}{}{\checkdatalist\par} } \end{tcolorbox} ``` ```{=latex} \tcbset{reset} ``` ```{=latex} \FloatBarrier ``` Introduction ============ Recent progress in large language models (LLMs) [@team2023gemini; @achiam2023gpt4; @guo2025deepseek; @liu2024deepseek] has been driven by scaling along four major dimensions: context length [@child2019sparse; @yuan2025nsa; @dai2019transformerxl], training data [@achiam2023gpt4; @team2023gemini], model width [@touvron2023llama; @bai2023qwen], and model depth [@wang2024deepnet; @chen2026post]. Although these dimensions remain effective, incremental gains are becoming increasingly costly, motivating interest in complementary architectural scaling strategies. In current LLM practice, scaling is often realized more through data, context, and especially width, whose optimization behavior and system efficiency are generally easier to realize at scale. Depth, by contrast, remains comparatively under-exploited despite its strong representational appeal. In principle, deeper stacks can support richer hierarchical computation. Yet modern Transformers often fail to convert additional layers into proportional benefits due to the optimization problem [@he2016resnet] and information dilution [@huang2017densenet; @pagliardini2024denseformer]. The resulting question is central to the architecture design: how can a model scale depth while maintaining optimization stability and preventing information dilution? ![ Comparing `\name{}`{=latex} and strong open-sourced baseline, i.e., OLMo2 [@olmo2024olmo2], with validation loss and downstream performance under the 1.5B-parameter setting. Models using `\name{}`{=latex} achieve lower C4 [@raffel2020c4] validation loss and better downstream performance, i.e., HellaSwag [@zellers2019hellaswag], WinoGrande [@sakaguchi2021winogrande], and ARC-Challenge [@clark2018arc], than OLMo2. ](figures/overall_results_v1p2.png){#fig:overall_results width="100%"} The standard residual pathway (ResNet-style) improves optimization stability in deep networks [@he2016resnet], but it still compresses depth history into a single hidden-state trajectory, leaving information dilution largely unresolved. Many methods [@zhu2024hc; @xie2025mhc; @li2025vwn] have been tried to address this problem by upgrading the residual connection. Dense cross-layer connections (DenseNet-style) preserve richer layer-wise history and thus mitigate information dilution [@huang2017densenet; @pagliardini2024denseformer; @chen2017dpn], but their parameter growth is substantial at LLM scale, which has limited their adoption as a mainstream architecture. The success of attention [@vaswani2017attention] in sequence modeling suggests a broader principle: data-dependent dynamic mixing can preserve and retrieve historical information more effectively than fixed-pattern aggregation. This motivates extending the same principle from sequence modeling to depth modeling, i.e., enabling each layer to adaptively read useful states from earlier layers. Adaptive cross-layer retrieval is therefore promising, yet practical designs still require a better balance among expressivity, efficiency, and hardware friendliness. In this work, we introduce *mixture-of-depths attention* (`\name{}`{=latex}), a unified attention mechanism in which each head jointly attends to sequence KV of the current layer and depth KV from all preceding layers. Methodologically, we analyze Transformer stacking through a \`\`read, operate, write" lens, comparing depth residual, depth dense, and depth attention in a common design space. `\name{}`{=latex} occupies an efficient point that preserves data-dependent depth retrieval without dense cross-layer overhead. To make `\name{}`{=latex} practical at scale, we develop a hardware-aware implementation [@dao2022flashattention; @dao2023flashattentionv2; @yang2024fla] that fuses sequence and depth attention in one forward pass with shared online-softmax states. Besides, the proposed chunk-aware depth-KV layout and group-aware indexing significantly improve memory access efficiency. This fused kernel reaches 97.3% of FlashAttention-2 efficiency at 64K sequence length, showing that depth-aware aggregating can be integrated without sacrificing modern GPU efficiency. We validate `\name{}`{=latex} on decoder-only language models trained with the 400B-token OLMo2 recipe [@olmo2024olmo2] at 700M and 1.5B scales. In our main 1.5B setting, `\name{}`{=latex} improves average perplexity by 0.2 across 10 validation benchmarks and increases average downstream performance by 2.11% on 10 tasks. We also find that combining `\name{}`{=latex} with post-norm yields better performance than using it with pre-norm. Additional analyzes, i.e., model-size scaling, attention visualization, and layer-number studies, show robust gains and reduced attention-sink [@xiao2023attnsink] behavior via better probability allocation to informative sequence and depth KV. The contributions of this paper are summarized as: - We propose `\name{}`{=latex}, a unified attention formulation for dynamic mixtures of sequence and depth, which improves the aggregation of depth-wise information and addresses the information dilution problem of modern LLMs in a data-dependent way. - We present a hardware-efficient fused algorithm that makes `\name{}`{=latex} practical for long-context LLM training. It reaches 97.3% of FlashAttention-2 efficiency at 64K sequence length with numerical precision within the allowed range. - We provide extensive empirical evidence and comprehensive ablations that `\name{}`{=latex} consistently and substantially outperforms the strong open-source baseline, OLMo2, across large-scale corpora at multiple model scales, validating each design choice and establishing `\name{}`{=latex} as a reliable foundation for depth scaling in LLMs. Mixture-of-Depths Attention =========================== Preliminary ----------- Most modern large language models are built on the Transformer architecture [@vaswani2017attention], where self-attention is the primary token-mixing operator. Given a sequence of $T$ tokens $X=(x_1, x_2, \ldots, x_T)\in\mathbb{R}^{T\times D}$ (with hidden dimension $D$), self-attention first projects tokens into queries ($Q$), keys ($K$), and values ($V$) via trainable matrices $W_Q\in\mathbb{R}^{D\times (H_qd)}$ and $W_K, W_V\in\mathbb{R}^{D\times (H_kd)}$. Under grouped query attention (GQA) [@ainslie2023gqa], $H_q=G H_k$, $H_k=H_v$, and $D=H_qd$: $$Q = X W_Q,\ \ \ \ \ \ \ \ \ \ K = X W_K,\ \ \ \ \ \ \ \ \ V = X W_V,$$ where $Q\in\mathbb{R}^{T\times (H_qd)}$ and $K,V\in\mathbb{R}^{T\times (H_kd)}$. The attention operator computes pairwise similarity between queries and keys, applies a softmax to obtain per-head attention weights $A_h\in \mathbb{R}^{T\times T}$, and returns a weighted sum of values: $$\mathrm{Attention}(Q, K, V) = \mathrm{Concat}_{h=1}^{H_q}\left(\mathrm{softmax}\left(\frac{Q_h K_{\phi(h)}^T}{\sqrt{d}} + \mathcal{M}\right)V_{\phi(h)}\right)$$ where $Q_h\in\mathbb{R}^{T\times d}$, $K_j,V_j\in\mathbb{R}^{T\times d}$, and $\phi(h)=\lceil h/G\rceil$ maps each query head to its shared key-value head. Here, $\mathcal{M}\in \mathbb{R}^{T\times T}$ is an additive attention mask. For causal attention, $\mathcal{M}_{ij}=0$ if $j\leq i$ and $\mathcal{M}_{ij}=-\infty$ otherwise. For full attention, $\mathcal{M}$ is all zeros. ![ Conceptual comparison of mechanisms that utilize the depth stream. (a) **Depth Residual** [@he2016resnet] is the standard residual connection along depth: it reads the current representation and writes back by addition. (b) **Depth Dense** [@huang2017densenet; @pagliardini2024denseformer] 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** (`\name`{=latex}), 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. ](figures/dsa_variants_v2.0.png){#fig:dsa_variants width="100%"} Stacking Transformers Along the Depth Stream -------------------------------------------- Deep neural networks have enabled breakthroughs across domains, especially after the introduction of residual connections [@he2016resnet]. Scaling studies [@kaplan2020scalinglaw; @hoffmann2022chinchilla; @hestness2017deepscaling] further show that increasing depth can substantially improve performance [@simonyan2014vdcn; @szegedy2015godeep]. This motivates a natural question: *Is the residual connection the optimal mechanism for propagating information through depth stream?* Along the depth stream, we can view a Transformer block as a three-step procedure: **read, operate, and write**. We use this lens to describe different mechanisms for stacking Transformer blocks. For clarity, the first two mechanisms (Depth Residual [@he2016resnet], Depth Dense [@huang2017densenet; @pagliardini2024denseformer]) are reference designs used to define the depth-stream design space. We introduce Depth Attention as an intermediate formulation and conceptual bridge. Our major technical contribution in this section starts from **Mixture-of-Depths Attention** (`\name{}`{=latex}), which unifies sequence and depth retrieval in one unified softmax operator. ```{=latex} \myPara{Depth Residual} ``` In depth residual connections [@he2016resnet; @srivastava2015highway], the \`\`read" step is identity and the \`\`write" step is add. The \`\`operate" step is the token-mixing operator, i.e., attention, or the feed-forward network (FFN), denoted by $\mathcal{F}(\cdot)$. As shown in Fig. `\ref{fig:dsa_variants}`{=latex}(a), the structure of depth residual can be formulated as follows: $$X_{l}=X_0+\sum^{l-1}_{i=1}\mathcal{F}(X_i, \mathcal{W}_{i}),$$ where $\mathcal{W}_{i}$ is the set of trainable weight matrices for the $i$-th layer. This formulation alleviates vanishing gradients and enables training deep networks. However, the depth stream is continuously *compressed* into a fixed-size tensor $X_l\in \mathbb{R}^{T\times D}$ via repeated superposition, which dilutes salient features and leads to signal degradation. ```{=latex} \myPara{Depth Dense} ``` To mitigate signal degradation, depth-dense methods [@huang2017densenet; @pagliardini2024denseformer] connect all layers along the depth stream. At the \`\`read" step, they form the input to layer $l$ by linearly projecting the set of preceding representations $\{X_i\in\mathbb{R}^{T\times D}\}_{i=0}^{l-1}$ back to shape $T\times D$. At the \`\`write" step, the layer output is concatenated with the historical set along depth. As shown in Fig. `\ref{fig:dsa_variants}`{=latex}(b), the structure of depth dense can be formulated as follows: $$\{X_{i}\}_{i=0}^{l}=\{X_0, \mathcal{F}(\{X_0\}, \mathcal{W}_{1}), \mathcal{F}(\{X_0, X_1\}, \mathcal{W}_{2}), \cdots ,\mathcal{F}(\{{X_i}\}_{i=0}^{l-1}, \mathcal{W}_{l})\},$$ where $\mathcal{W}_{i}$ is the set of trainable weight matrices for the $i$-th layer. Depth-dense connections propagate information through depth *losslessly*, because concatenation does not compress the historical set. However, they incur high cost and enforce a fixed connectivity pattern: the computation grows as $O(TL^2D^2)$ in dominant terms, which is prohibitive for large models. ```{=latex} \myPara{Depth Attention} ``` To reduce cost while retaining adaptive connectivity, we propose *depth attention* that reads historical depth information using attention in a data-dependent way, as illustrated in Fig. `\ref{fig:dsa_variants}`{=latex}(c). At the \`\`read" step, in the GQA-group view ($H_kd=D/G$), we denote one query-group representation by $Q_{l-1}\in\mathbb{R}^{T\times \frac{D}{G}}$ and the corresponding historical key-value sets by $\{K_i\in\mathbb{R}^{T\times \frac{D}{G}}\}_{i=0}^{l-1}$ and $\{V_i\in\mathbb{R}^{T\times \frac{D}{G}}\}_{i=0}^{l-1}$. The resulting input $X_{l}^{\text{in}}$ is then fed into the \`\`operate" step: $$X_{l}^{\text{in}} = \mathrm{Attention}(Q_{l-1}, \{{K_i}\}_{i=0}^{l-1}, \{{V_i}\}_{i=0}^{l-1}),$$ where attention is performed along the *depth* dimension: for token $t$, the query $Q_{l-1,t}$ attends only to the depth keys and values $\{K_{i,t}, V_{i,t}\}_{i=0}^{l-1}$ from the *same* token position across layers. After the \`\`operate" step, the current layer output $X_{l}^{\text{out}}$ is fed to the \`\`write" step, which produces new query/key/value projections: $$Q_{l} = X_{l}^{\text{out}} W_{Q, l}^\text{W},\ \ \ \ \ \ \ \ \ \ K_l = X_{l}^{\text{out}} W_{K, l}^\text{W},\ \ \ \ \ \ \ \ \ V_l = X_{l}^{\text{out}} W_{V, l}^\text{W},$$ where $W_{Q, l}^\text{W}, W_{K, l}^\text{W}, W_{V, l}^\text{W}\in\mathbb{R}^{D\times \frac{D}{G}}$ are trainable matrices for the layer-$l$ \`\`write" operation, and $Q_l,K_l,V_l\in\mathbb{R}^{T\times \frac{D}{G}}$ denote per-group projections. We concatenate $K_l$ and $V_l$ along depth for future reads, while $Q_l$ is passed forward to the next layer. Compared with depth-dense connections, depth attention reads historical information adaptively with much lower cost. Its computation scales as $O(TL^2D)$, which is a factor of $\frac{1}{D}$ smaller than depth dense. ```{=latex} \myPara{Mixture-of-Depths Attention} ``` Building upon the Depth Attention, we now propose *mixture-of-depths attention* (`\name{}`{=latex}). `\name{}`{=latex} adds depth-level information to standard sequence-level attention and fuses these operations into a single operator. As illustrated in Fig. `\ref{fig:dsa_pipeline}`{=latex} and Fig. `\ref{fig:dsa_variants}`{=latex}(d), `\name{}`{=latex} reads the current hidden state $X_{l-1}$ and the historical depth KV stream $\{(K_i, V_i)\}_{i=0}^{l-1}$. During the \`\`operate" step, we apply `\name{}`{=latex} to enable each token to attend to both the sequence-level keys and values and its own historical depth-wise keys and values, with all attention scores normalized jointly under a single softmax function. The implementation detail of `\name{}`{=latex} is presented in Alg. `\ref{alg:moda_forward}`{=latex}. At the \`\`write" step, for the attention layer, we append the current layer's key-value pair to the depth stream so that subsequent layers can access them. For the FFN layer, we obtain its corresponding key-value pair via a light-weight KV projection. Overall, `\name{}`{=latex} provides an efficient, data-dependent mechanism for exploiting depth history with substantially lower overhead than dense cross-layer connectivity. Furthermore, aggregating the sequence and depth information in one softmax operation provides a uniform representation space. ```{=latex} \footnotesize ``` ```{=latex} \setlength{\tabcolsep}{4pt} ``` ::: {#tab:complexity} Methods Depth Dense Depth Attention Mixture-of-Depths Attention -------------------------------------------- --------------------- --------------------- ----------------------------- Is data-dependent? `\ding{56}`{=latex} `\ding{52}`{=latex} `\ding{52}`{=latex} Is unified softmax? `\ding{56}`{=latex} `\ding{56}`{=latex} `\ding{52}`{=latex} Parameters [$O(L^2D^2)$]{style="color: blue"} [$O(LD^2)$]{style="color: blue"} [$O(L\frac{D^2}{G})$]{style="color: blue"} Decoding Cache [$O(LD)$]{style="color: blue"} [$O(L\frac{D}{G})$]{style="color: blue"} [$O(L\frac{D}{G})$]{style="color: blue"} Prefilling Cache [$O(TLD)$]{style="color: blue"} [$O(TL\frac{D}{G})$]{style="color: blue"} [$O(TL\frac{D}{G})$]{style="color: blue"} Decoding FLOPs [$O(L^2D^2)$]{style="color: blue"} [$O(L^2D)$]{style="color: blue"} [$O(L^2D)$]{style="color: blue"} Prefilling FLOPs [$O(TL^2D^2)$]{style="color: blue"} [$O(TL^2D)$]{style="color: blue"} [$O(TL^2D)$]{style="color: blue"} : 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) [@ainslie2023gqa], $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. ::: ```{=latex} \myPara{Complexity analysis} ``` Complexity analysis is critical for modern LLM design, we also present the detailed complexity analysis among depth-aware designs, e.g., depth dense, depth attention, and `\name{}`{=latex}. Table `\ref{tab:complexity}`{=latex} reports complete complexity and dominant asymptotic terms, where $T$ is sequence length, $D$ is model width, $L$ is the number of layers, head dimension $d$, and $G$ is the GQA group size. Notably, $H_q = G H_k$. From Table `\ref{tab:complexity}`{=latex}, Depth Dense is dominated by quadratic depth growth. Its parameter term is $O(L^2D^2)$, decoding cache is $O(LD)$, and both decoding and prefilling FLOPs contain quadratic-depth and quadratic-width terms, i.e., $O(L^2D^2)$ and $O(TL^2D^2)$. The proposed Depth Attention is a data-dependent method, which removes the dominant quadratic-width projection accumulation across depth, reducing parameters to $O(LD^2)$. It also lowers cache to $O(LD/G)$ and compute to $O(L^2D)$ and $O(TL^2D)$ for decoding and prefilling, respectively. Compared with Depth Attention, `\name{}`{=latex} keeps the same favorable FLOPs order and cache order, but further reduces parameter complexity from $O(LD^2)$ to $O(LD^2/G)$. The key reason is that `\name{}`{=latex} reuses the query projection from sequence attention, so no extra depth-query projection is introduced. Especially in GQA settings, only grouped depth key/value projections are needed. This makes `\name{}`{=latex} the most parameter-efficient option in Table `\ref{tab:complexity}`{=latex}, while preserving linear-in-width compute behavior and low-cache scaling. Overall, Table `\ref{tab:complexity}`{=latex} shows that `\name{}`{=latex} keeps the data-dependent behavior of attention while avoiding the dominant quadratic-depth parameter growth overhead of dense cross-layer connections. `\name{}`{=latex} aggregates sequence and depth information with a unified softmax operator, which provides better representation and efficiency in practice, especially in regimes with large $L$ and long $T$. Hardware-aware efficient `\name`{=latex} {#sec:hardware_efficient} ======================================== Naïvely PyTorch-implemented [@paszke2019pytorch] `\name{}`{=latex} requires non-contiguous reads of historical depth states, which degrades GPU utilization. We develop a hardware-aware implementation that reorganizes depth-stream tensors to enable contiguous memory access and fused computation. ```{=latex} \begin{algorithm}[t]\caption{\name{}: Hardware-aware Forward Pass} \label{alg:moda_forward} \begin{algorithmic}[1] \STATE \textbf{Input:} $\mathbf{Q}\in\mathbb{R}^{T_q\times (H_kd)}$, $\mathbf{K},\mathbf{V}\in\mathbb{R}^{T_{kv}\times (H_kd)}$, $\mathbf{K}^{\text{depth}},\mathbf{V}^{\text{depth}}\in\mathbb{R}^{(T_{kv}L)\times (H_kd)}$, group number $G$ \STATE \textbf{Output:} $\mathbf{O}\in\mathbb{R}^{T_q\times (H_kd)}$% \STATE Partition $\mathbf{Q},\mathbf{K},\mathbf{V}$ into hardware-friendly query/key/value blocks \STATE Ensure each query block aligns with $G$ for correct base-time mapping \FOR{each query block index $b_q$} \STATE Load $\mathbf{Q}_{[b_q]}$ from HBM to SRAM (on chip) \STATE Initialize on-chip states: $m\leftarrow-\infty,\ acc\leftarrow 0,\ o\leftarrow 0$ \STATE For each query row index $i_q$ in block $b_q$, compute base-time: $t_{base}(i_q)=\lfloor i_q/G\rfloor$ \STATE Let $t_{base}^{start}=\min_{i_q\in b_q} t_{base}(i_q)$ and $t_{base}^{last}=\max_{i_q\in b_q} t_{base}(i_q)$ \STATE Define $t_{base}^{end}=t_{base}^{last}+1$ as the exclusive upper bound \FOR{sequence key block $b_s$ with $b_s