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

TL;DR

This paper introduces Dreamer, a modular depth-recurrent attention framework that improves latent reasoning efficiency, reduces training data needs, and enhances model diversity across depths in language reasoning tasks.

Contribution

It proposes a novel depth-recurrent attention mixture framework that addresses hidden-size bottlenecks and scales efficiently for reasoning models.

Findings

Models require 2 to 8x fewer training tokens for the same accuracy.

Outperforms larger SOTA models with the same training tokens.

Shows increased expert selection diversity across depths.

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.