Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves

TL;DR

Depth-Recurrent Attention Mixtures (Dreamer) tackle two core bottlenecks in depth-recurrent transformers by introducing a modular, attention-based framework that unifies sequence attention, depth attention, and sparse expert attention. The approach decouples scaling dimensions and uses a single-layer depth-recurrent core with sparse MoEs to overcome the layer-size and hidden-size constraints, enabling efficient many-step latent reasoning. In tightly FLOP-, parameter-, and memory-matched experiments on language reasoning benchmarks, DR and DR+DA achieve substantial data efficiency, requiring roughly two to eight times fewer training tokens to reach the same accuracy, and outperform larger state-of-the-art models with the same token budgets. Analyses reveal richer knowledge usage across depths, with 2–11x higher expert selection diversity than MoEs in baselines, suggesting improved adaptability and compositional generalization in deeper reasoning tasks.

Abstract

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.

Figures (7)

• Figure 1: High-level illustration of our modular depth-recurrent attention mixture (Dreamer). This instance combines sequence attention (SA), depth attention (DA), and expert attention (EA) in a single depth-recurrent (DR) layer. It facilitates knowledge reuse, compositional generalization, and alleviates the hidden-size bottleneck, leading to better reasoning, data efficiency, and scaling behavior. The decoupled scaling dimensions (#params, #FLOPs, depth, and latent memory size) are individually adjustable.
• Figure 2: Unfolded high-level illustration of depth-recurrent attention mixture (Dreamer). It shows the three attention dimensions in this instantiation: sequence attention (horizontal), depth attention (vertical), and expert attention (z-axis). Filled boxes indicate active elements in current/latest token and depth, while unfilled boxes are inactive. Q/K/V = query/key/value.
• Figure 3: Math reasoning benchmark results (0-shot) during training, comparing the best accuracies of classical layering (LA), depth recurrence (DR), and DR with depth attention (DR+DA).
• Figure 4: Map of average DA scores, normalized and scaled per depth for visualization. Nontrivial patterns suggest extended expressivity beyond uniform skip-connections.
• Figure 5: Distribution of depths per expert. Experts are sorted by the number of depths in their top 90th percentile of sorted $P(\text{depth}|\text{expert})$, as a measure of depth-generalization. Within these groups, experts are sorted by sampled depth from their distributions. Observations: Higher depths use more depth-specialized experts. Ca. 50% of experts are widely depth-generalized.