Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning

TL;DR

This work tackles the challenge of robust OOD generalization in Transformer-based reasoning by introducing a latent-space architectural suite that enables recursive, layer-by-layer computation. By solving a modular-arithmetic on computational-graphs task with input-adaptive recurrence, latent-state supervision, discretization, and self-correction, the approach learns a depth-invariant algorithm that generalizes to problem sizes well beyond training. Mechanistic interpretability reveals that the model internally learns variable-identification and value-copying steps, followed by frequency-domain modular addition, explaining its robust generalization. The findings suggest that native latent-space reasoning, anchored discretely and reinforced by self-correction, can outperform token-based CoT strategies for scalable, interpretable algorithmic generalization with potential impact on multi-step reasoning tasks.

Abstract

Systematic, compositional generalization beyond the training distribution remains a core challenge in machine learning -- and a critical bottleneck for the emergent reasoning abilities of modern language models. This work investigates out-of-distribution (OOD) generalization in Transformer networks using a GSM8K-style modular arithmetic on computational graphs task as a testbed. We introduce and explore a set of four architectural mechanisms aimed at enhancing OOD generalization: (i) input-adaptive recurrence; (ii) algorithmic supervision; (iii) anchored latent representations via a discrete bottleneck; and (iv) an explicit error-correction mechanism. Collectively, these mechanisms yield an architectural approach for native and scalable latent space reasoning in Transformer networks with robust algorithmic generalization capabilities. We complement these empirical results with a detailed mechanistic interpretability analysis that reveals how these mechanisms give rise to robust OOD generalization abilities.

Figures (46)

• Figure 1: Four key mechanisms enabling robust out-of-distribution generalization in transformer architectures.(a) Recurrence and input-adaptive computation allows models to dynamically allocate computational resources based on problem complexity. (b) Algorithmic supervision guides the learning process through structured intermediate representations. (c) Anchored discrete latent spaces provide stable reference points for compositional reasoning. (d) Error correction mechanisms enable iterative refinement of predictions through feedback loops. Together, these mechanisms enable transformers to develop recursive reasoning patterns that generalize beyond their training distribution.
• Figure 2: An illustration of an instance in modular arithmetic on computational graphs task. The goal is to compute the values of all nodes in the graph. For example, here $x_{23} = 20 + 2 = 22$ and $x_{55} = 6 \times 14 = 15$. Recall that we consider modular arithmetic with base $p = 23$.
• Figure 3: Overview of the proposed architecture for OOD generalization. It features a recurrent Transformer block, latent algorithmic supervision, and a discretization mechanism to anchor representations across iterations. Self-correction mechanism is not represented here.
• Figure 4: Out-of-Distribution generalization performance of different methods on the mathematical reasoning task.
• Figure 5: Effective out-of-distribution generalization via input-adaptive scaling of computation time. This depicts Discrete Latent Space Supervision $\circlearrowleft$