Source-linked AI summary
Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves
Jonas Knupp, Jan Hendrik Metzen, Jeremias Bohn, Georg Groh, Kristian Kersting
TL;DR
The paper addresses missing fair resource-matched evaluations and architectural bottlenecks that restrict scalable depth-recurrent latent reasoning. It introduces Dreamer, which combines sequence, depth, and sparse expert attention, and reports consistent gains over tightly matched state-of-the-art Transformers in language reasoning benchmarks. The authors also identify open questions around dynamic depth and generalization beyond trained depths.
Problem
Prior depth-recurrent work lacks combined FLOP-, parameter-, and memory-matched baselines, while constant hidden sizes restrict many-step latent reasoning.
Method
Dreamer is a modular depth-recurrent framework combining sequence attention, depth attention, and sparse expert attention to address scaling bottlenecks.
Results
Dreamer consistently outperforms tightly FLOP-, parameter-, and memory-matched state-of-the-art Transformers on natural language reasoning benchmarks.
Takeaways & Limitations
Attention mixtures provide a modular view of information access across sequence, depth, and experts for understanding and controlling depth-recurrent model behavior.
Takeaways & Limitations
Reliable generalization of depth-recurrent models beyond the trained depth regime remains underexplored, and dynamic depth is left for future work.
Abstract
from arXiv · showhide
Depth-recurrence facilitates latent reasoning by sharing parameters across depths. However, prior work lacks combined FLOP-, parameter-, and memory-matched baselines, underutilizes depth-recurrence due to partially fixed layer stacks, and ignores the bottleneck of constant hidden-sizes that restricts many-step latent reasoning. To address this, we introduce a modular framework of depth-recurrent attention mixtures (Dreamer), combining sequence attention, depth attention, and sparse expert attention. It alleviates the hidden-size bottleneck through attention along depth, decouples scaling dimensions, and allows depth-recurrent models to scale efficiently and effectively. Across language reasoning benchmarks, our models require 2 to 8x fewer training tokens for the same accuracy as FLOP-, parameter-, and memory-matched SOTA, and outperform ca. 2x larger SOTA models with the same training tokens. We further present insights into knowledge usage across depths, e.g., showing 2 to 11x larger expert selection diversity than SOTA MoEs.
1. Introduction
The paper targets inefficiencies in discrete chain-of-thought and bottlenecks that limit scalable depth-recurrent latent reasoning. It introduces Dreamer, a modular attention-mixture framework, and evaluates it with tightly matched comparisons.
- Discrete chain-of-thought limits expressivity while producing long, expensive sequences for training and generation.
- Depth recurrence reuses parameters across depths, enabling open-ended latent reasoning without increasing model parameter count.
- Constant hidden-state size creates a bottleneck for complex latent reasoning, while dense depth-recurrent scaling is computationally intractable.
- Dreamer combines sequence attention, depth attention, and sparse expert attention into a modular depth-recurrent framework.
- The experiments use tightly FLOP-, parameter-, and memory-matched comparisons against state-of-the-art mixture-of-experts models and report consistent, large reasoning gains.
- The analysis examines knowledge allocation and reuse across depths using high-resolution expert configurations.
2. Related Work
Prior latent-reasoning and depth-recurrence approaches provide useful alternatives but leave gaps in fair scaling comparisons, sequence-parallel training, or attention over depth. Dreamer frames sequence, depth, and expert interactions as coordinated attention dimensions.
- Earlier depth-recurrent work often lacked FLOP-matched evaluations, used dense cores, or omitted sparse baselines for fair comparison.
- Dreamer contrasts with latent-token approaches that increase sequence length or do not support naive sequence-parallel training and knowledge reuse across depths.
- Prior skip-connection and hidden-state aggregation methods do not implement proper dot-product attention along depth.
3. Methods
Dreamer repeatedly applies a layer with attention over sequence positions, previous depths, and sparse MLP experts. Depth attention reuses Transformer mechanisms along depth, while expert attention supplies scalable sparse routing.
- Dreamer applies a layer combining attention over sequence, depth, and experts repeatedly to build latent recurrent reasoning.
- Standard self-attention aggregates values using scaled dot-products between queries and keys, with causal softmax masking by default.
- Depth attention applies attention over previous depths of the same token, treating depth as the sequence dimension and alleviating the depth hidden-size bottleneck.
- Depth attention can reuse optimized attention implementations by treating sequence as batch and depth as sequence, with a key-value cache.
- Depth-position RoPE and per-token cache overwriting support depth-aware attention while keeping its memory overhead constant with respect to sequence length.
- Expert attention turns MLP experts into an attention variant whose keys are learnable weights and whose scores retain low-rank, depth-dependent routing characteristics.
- Sparse expert attention executes only a subset of experts, while routing and balancing strategies address sparsity-related nondifferentiability and load-balancing challenges.
- The balancing bias is updated using expert-usage counts relative to the median, with update rate λ and zero statistical-average bias.
3.6. Attention Mixture
The attention mixture combines depth, sequence, and expert attention into a modular framework, with sequential or parallel execution options. Parallelization improves throughput but can increase benchmark error rates.
- Execution variants: The three attention components can be executed sequentially or parallelized to improve throughput.The paper describes partially and fully parallel variants in addition to the sequential formulation.
- Attention mixture design: At each depth, depth attention consumes previous hidden states, while sequence and expert attention operate on the latest depth’s states.In practice, only the latest hidden state is propagated, alongside cached keys and values for depth and sequence attention.
- Execution variants: 15% speedup was observed for each additional parallelization at ca. 1B parameters, with a cost of up to 10–20% higher benchmark error rates.The reported tradeoff comes from small-scale tests and may differ for larger models.
- Attention mixture design: Attention mixtures unify depth attention, sequence attention, and expert attention across their respective dimensions.The framework is designed as a modular foundation that can incorporate additional attention dimensions.
3.7. Depth Recurrence (DR)
The depth-recurrent design addresses scaling and optimization challenges through residual normalization, sparse expert attention, and a shared routing baseline. Model comparisons use coordinated FLOP and parameter matching to maintain similar resource requirements.
- Depth-recurrent architecture: Residual-stream RMSNorm addresses incompatible residual scales across depths and supports learning depth-generalized experts.The depth-recurrent update is written as x_l+1 = Norm(y_l).
- Depth-recurrent architecture: Sparse MoEs replace reused attention projections in depth-recurrent attention modules while tying routing scores across the projection and aggregation MoEs.This preserves a fully depth-recurrent architecture while reducing routing overhead and stabilizing training.
- Expert routing: A shared always-active expert supplies a baseline output scaled by the selected expert’s score, with stopped gradient flow.The baseline output helps learn useful and balanced routing.
- Expert routing: The shared expert can be folded into routable experts before inference, eliminating its compute and memory overhead.Only the slight MoE routing overhead remains relative to standard attention and is included in FLOP matching.
- Resource matching: FLOP and parameter matching uses alternating binary searches over expert MLP sizes and expert counts to align models with a baseline.The matching procedure repeats FLOP adjustment after changing the number of experts.
- Resource matching: Table 1 reports closely aligned FLOPs per token and total memory across matched models, measured over 1024 generated tokens.The matching is intended to make resource comparisons meaningful.
4. Experiments
The experiments compare classical layering, depth recurrence, and depth recurrence with depth attention under matched scaling conditions. Results show strong reasoning gains from recurrent designs, with depth attention becoming especially valuable as model depth increases, while attention analyses reveal structured routing across depths and experts.
- Experimental setup: The study evaluates LA, DR, and DR+DA at depths 16 and 32, using sparse expert attention and matched computational resources.Models at the same depth are adjusted for FLOP and parameter matching, with memory matching also reported in the evaluation setup.
- Reasoning and language modeling: DR and DR+DA strongly outperform the classically layered baseline across scales and tests.DR+DA at depth 16 outperforms depth-32 LA on all reasoning benchmarks, corresponding to nearly 2× lower parameter, FLOP, and memory usage.
- Reasoning and language modeling: At depth 32, DR+DA strongly outperforms both LA and DR, and its gains over LA increase with depth.At depth 16, DR outperforms DR+DA on half the reasoning benchmarks because its smaller MLP is used for FLOP matching.
- Reasoning and language modeling: For depths ≤16, pure DR may be preferable in some cases, whereas depth attention is especially effective in deeper models.The authors suggest deeper models face a stronger hidden-size bottleneck because more information accumulates across depths, while noting that rigorous scaling laws remain future work.
- Depth attention: Depth-attention maps show nontrivial recall patterns: early depths often revisit initial depths, middle depths are recalled later, and middle-to-high depths inform the final output.These patterns are presented as evidence that depth attention extends expressivity beyond uniform skip-connections and supports targeted routing between depths.
- Expert attention: DR uses 2–11× more experts per depth than LA, while global expert balancing remains strong with a Gini coefficient of 0.075.Lower depths use up to approximately 5× more experts than higher depths, and approximately 50% of experts are widely depth-generalized.
5. Discussion and Future Work
The framework's depth attention is computationally worthwhile despite memory-movement overhead, while depth-recurrent scaling remains bounded by efficiency and generalization questions. The proposed attention-mixture view also opens modular extensions, although their practical benefits on sparsity-averse hardware remain unclear.
- Discussion: Depth attention’s compute and memory costs remain worthwhile in FLOP- and memory-matched comparisons.Its key-value cache can be overwritten after each token, keeping memory constant with respect to sequence length.
- Discussion: A single depth-attention head keeps memory-movement overhead small while consistently achieving accuracy gains.The authors note that depth 16 with depth recurrence and depth attention can outperform latent attention with depth 32.
- Future Work: Future depth-attention variants could mitigate memory movement or reduce memory consumption as depth increases further.Suggested directions include sliding windows, dilation, linear attention, and modern RNNs or SSMs along depth.
- Future Work: Reliable generalization beyond the trained depth regime remains a critical, underexplored problem for depth recurrence.The paper also identifies dynamic depth and alternatives to or scaling strategies for RoPE as open directions.
- Attention Mixture: Attention mixtures unify sequence, depth, and expert information access and provide a modular basis for adding further attention dimensions.The paper discusses extensions such as parameter attention, learned latents, text embeddings, asynchronous components, and conditional execution.
- Attention Mixture: The practical benefit of additional attention dimensions remains unclear under the constraints of sparsity-averse GPUs and TPUs.
6. Conclusion
Dreamer addresses depth recurrence’s layer-size and hidden-size bottlenecks through sparse expert attention and depth attention. Its unified attention-mixture architecture consistently outperforms tightly matched Transformer baselines and improves scaling behavior.
- Conclusion: Sparse expert attention alleviates the layer-size bottleneck, allowing depth recurrence to scale efficiently.
- Conclusion: Depth attention alleviates the hidden-size bottleneck, allowing depth recurrence to scale effectively.
- Conclusion: Dreamer consistently outperforms tightly FLOP-, parameter-, and memory-matched SOTA Transformers on natural language reasoning benchmarks.
- Conclusion: Depth attention provides a qualitative extension to expressivity that particularly improves scaling behavior.
- Conclusion: Attention mixtures offer a modular framework that models major information access through sequence, depth, and expert attention.
8. Broader Impact Statement
The paper presents its framework as improving reasoning efficiency and model capability while reducing the computational resources needed for complex reasoning tasks. It frames these efficiency gains as potentially useful for deploying advanced AI more sustainably.
- Broader Impact Statement: The framework is described as providing 2–8× better data efficiency and approximately 2× better parameter- and FLOP-efficiency.
- Broader Impact Statement: The resulting efficiency could support complex reasoning with significantly fewer computational resources and facilitate deployment of advanced AI.
A. Experiment Details
The experiment-details passages identify tables for hyperparameters and training datasets. These tables provide the paper’s configuration and data references for interpreting its experiments.
- Experiment Details: Table 3 reports the model hyperparameters used in the experiments.
- Experiment Details: Table 4 reports the training datasets used in the experiments.
B. Further Analysis
Figure 7 complements Figure 5 by showing P(expert|depth) with the same expert ordering, revealing depth-specialization patterns across experts.
- Figure 7 shows P(expert|depth), complementing Figure 5’s P(depth|expert) while preserving the same expert-ordering method.
- Higher depths tend to use more depth-specialized experts.
- The second-to-last depth is an exception, using some depth-generalized experts.