PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Awni Altabaa, Siyu Chen, John Lafferty, Zhuoran Yang

arXiv:2510.14095v1cs.LGcs.AI

TL;DR

Robust OOD generalization remains a challenge for Transformers, and chain-of-thought training can yield brittle algorithms. This paper studies recursive latent-space reasoning with discrete latent-space supervision on modular arithmetic graphs. The authors report near-perfect performance across the OOD splits examined and improved generalization with discrete latent anchoring.

  • Problem

    Chain-of-thought methods can yield brittle algorithms that fail to generalize, while robust OOD generalization remains a challenge for Transformers.

    Source: §1 · p. 4, §4.1 · p. 9

  • Method

    The approach uses recursive computation with discrete latent-space supervision, incorporating discretization alongside mechanisms 1 and 2.

    Source: §4.2 · pp. 12, 14

  • Results

    The full Discrete Latent Space Supervision ⟲1 model achieves near-perfect performance across all examined OOD splits.

    Source: §4.2 · p. 13

  • Takeaways & Limitations

    Discrete latent anchoring allows computational depth to scale without accumulating noise, and the authors report significantly improved OOD generalization.

    Source: §4.2 · p. 14

    The reported comparison concerns standard Transformer and chain-of-thought baselines, whose OOD generalization is limited in this setting.

    Source: §3.2 · p. 8, §4.2 · p. 12

Abstract

from arXiv · show

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.

1. Introduction

The introduction frames robust out-of-distribution (OOD) length generalization as a persistent challenge for Transformers and asks which architectural biases support algorithmic generalization. It investigates four mechanisms on modular arithmetic graphs and reports robust generalization alongside mechanistic interpretability analysis.

1. Introduction

  • Transformers remain limited in OOD length generalization, especially when test problems are larger or more complex than training instances.The introduction notes that chain-of-thought techniques help to some degree, but generalization far beyond training remains an obstacle.Source: §1 · p. 3
  • The study investigates architectural mechanisms for OOD algorithmic generalization using modular arithmetic on computational graphs as a controlled, compositional testbed.Graph size and depth parameterize complexity, and the task captures core features of mathematical reasoning benchmarks such as GSM8K.Source: §1 · p. 3
  • The proposed approach combines input-adaptive recurrence, algorithmic supervision, discrete-bottleneck anchored latents, and explicit error correction.Together, these mechanisms are presented as an approach to native, scalable latent-space reasoning in Transformers.Source: §1 · p. 4
  • The method achieves perfect generalization on mathematical reasoning inputs several times larger than those seen during training.The authors also report a mechanistic interpretability analysis of the learned circuits and how they support robust OOD generalization.Source: §1 · p. 4

2. Related Work

The related work connects this study to recurrence, adaptive computation, discrete latent states, algorithmic supervision, and mechanistic interpretability. The paper uses these threads to motivate recursive computation and depth-invariant representations for scaling beyond training instances.

2. Related Work

  • The authors investigate scaling recurrent iterations with input size at inference to solve more complex instances beyond the training distribution.They situate this approach within adaptive-computation work that varies computation according to input demands.Source: §2 · p. 5
  • The paper treats recurrence as an inductive bias for discovering scalable recursive algorithms, and discrete latent states as anchors for a common, depth-invariant representation space.The stated aim of the anchoring is to scale computation beyond training while avoiding representational shift across depth.Source: §2 · p. 5
  • The paper uses algorithmic supervision not only as a baseline but also to supervise the proposed model’s internal states.The related-work discussion presents chain-of-thought training as explicit alignment to a problem-solving procedure.Source: §2 · p. 6
  • The interpretability analysis identifies a two-stage copying mechanism: retrieve parent-variable names, then retrieve their computed values from other latent states.The paper relates this mechanism to induction-head-style copying and reports observing a modular-addition algorithm in its models.Source: §2 · p. 6

3. Problem Setup

The paper defines a modular-arithmetic task on computational graphs and evaluates whether Transformers can generalize to larger graphs. It finds that chain-of-thought training improves on end-to-end training, but its OOD performance deteriorates beyond moderately larger graphs.

3.1. Task Description: Modular Arithmetic on Computational Graphs

  • The task represents computations as a directed acyclic graph in which nodes are variables and edges encode dependencies; leaves receive values and non-leaves apply modular operations to parent values.The experiments use addition, multiplication, and subtraction with prime modular base p = 23.Source: §3.1 · p. 6
  • The model must compute every node value correctly for an instance to count as correct.Inputs are generated from randomly assigned graph structure, variable names, leaf values, and operations.Source: §3.1 · p. 7
  • OOD evaluation trains on graphs with N ≤ 32 nodes and tests on varying sizes up to N = 128, a fourfold increase.The task requires processing larger inputs and scaling computation time at test time.Source: §3.1 · p. 7

3.2. Limitations of Standard Transformers with CoT Training

  • Chain-of-thought models solve larger graphs than end-to-end models, but their OOD performance deteriorates beyond moderately larger graphs.The best CoT models generalize from training graphs with N ≤ 32 to about N ≈ 40, after which performance rapidly declines.Source: §3.2 · p. 9

4. Reasoning in Latent Space with Algorithmic Supervision

This section examines recurrent latent-space reasoning and its role in OOD generalization, focusing on input-adaptive recurrence, discrete latent anchoring, error correction, and test-time iteration. It presents the full architecture’s performance across the examined OOD splits and discusses support for generalization to larger problems.

4.1. Mechanisms for Effective OOD Generalization.

  • The proposed mechanisms support iterative algorithms in latent processing, addressing the brittleness of token-based autoregressive CoT as reasoning sequences grow longer.The architecture combines recurrent Transformer blocks, algorithmic supervision, latent discretization, and self-correction.Source: §4.1 · p. 9
  • Input-adaptive recurrence uses repeated Transformer blocks to compute one graph layer per iteration, scaling iteration count linearly with graph depth.Each step can process the entire context in parallel rather than materializing a token-by-token computational trace.Source: §4.1 · p. 10
  • Latent state algorithmic supervision predicts node values at each recurrent step and trains the model to compute all nodes of depth t or less by iteration t.This directly aligns latent states with the target layer-by-layer algorithm, unlike CoT’s token-space supervision.Source: §4.1 · pp. 10–11
  • Discretization projects continuous hidden states into a shared structured symbolic space after each iteration, anchoring representations and supporting stable computation beyond training depths.The anchoring constrains the model to learn a depth-invariant computational process.Source: §4.1 · p. 11
  • Self-correction trains the model to detect and repair incorrect intermediate values before continuing computation, addressing error propagation across steps.Training randomly corrupts selected value components in discrete latent states at recurrent iterations.Source: §4.1 · p. 12

4.2. Experimental Results & Discussion

  • End-to-end models fail to learn the task beyond small graph sizes, while Chain-of-Thought reaches near-perfect in-distribution performance but generalizes only to a limited degree OOD.The evaluated variants progressively add recurrence and algorithmic supervision, discretization, and finally error correction.Source: §4.2 · p. 12
  • The full Discrete Latent Space Supervision ⟲ model, incorporating all proposed mechanisms, achieves near-perfect performance across all examined OOD splits.Performance improves as more ingredients are incorporated.Source: §4.2 · p. 13, Figure 4 · p. 13
  • Recurrence enforces a depth-invariant solution description at every step, enabling a recursive algorithm to solve much larger and more complex instances.The computational description remains the same across solution steps, allowing scaling to depths beyond those seen during training.Source: §4.2 · pp. 13–14

The Importance of Anchored Discrete Representations

  • Continuous latent supervision outperforms Chain-of-Thought but degrades on larger inputs as recurrent computation accumulates noise and representations drift from the learned manifold.Larger inputs require increasing recurrent depth and computation time, exacerbating noise in continuous representations.Source: §4.2 · p. 14
  • Discrete latent supervision discretizes and re-embeds states with a shared embedder, anchoring them in a common semantically consistent space and significantly improving OOD generalization.This anchoring allows computational depth to increase without accumulating noise.Source: §4.2 · p. 14

Error-Correction Leads to Greater Robustness in Scaling

  • Training with randomly corrupted latent states teaches the model to correct errors in the next recurrent iteration, yielding near-perfect OOD generalization but requiring more recurrent-block layers.Effective correction requires greater computational depth per step to identify and correct prior errors before computing the current step.Source: §4.2 · p. 14

Robust Test-time Scaling

  • Increasing recurrent iterations enables solving incrementally larger and harder instances, with the architectural mechanisms supporting robust test-time scaling beyond training.The model scales computation by increasing iterations of its latent reasoning process.Source: §4.2 · p. 14

Details, Extensions & Further Ablations

  • Hyperparameter choice, especially positional encoding and model depth, is critical; reported results use the best searched model per method and best results across explored Chain-of-Thought schemes.Hyperparameter-search details are provided in the appendix.Source: §4.2 · p. 14

5. Mechanistic Interpretability

The section tests hypotheses about each model block with controlled causal interventions, then identifies a variable-copying and retrieval circuit plus a modular-addition MLP that the authors link to robust OOD generalization.

5. Mechanistic Interpretability

  • Controlled causal interventions test hypotheses about first-layer attention, second-layer attention, and the final MLP’s functions.The analysis asks what algorithm the trained model implements and why it generalizes to OOD data.Source: §5 · p. 15

Induction Head & Modular Addition Mechanism

  • First-layer attention copies variable-name embeddings to the RHS, preparing them for value retrieval in the second layer.Head groups attend to different variable positions, and their projections select variable embeddings.Source: Induction Head & Modular Addition… · pp. 15–16, Figure 7 · p. 17
  • Second-layer attention uses copied variable names to retrieve their values from earlier equations through a variable-dependent induction-head mechanism.Heads grouped by retrieved variable find its first occurrence, where its value was computed, and copy that value embedding to the current RHS.Source: Induction Head & Modular Addition… · p. 16
  • The final MLP computes the three variable values’ sum modulo 23 by amplifying diagonal Fourier frequencies that encode the sum.These components represent sinusoidal functions of x + y + z, whose periodicity handles the modulo operation.Source: Induction Head & Modular Addition… · pp. 16–17
  • The authors link robust OOD generalization to a length-invariant induction-head algorithm and mechanisms that support operation on longer contexts despite limited-size training graphs.They identify input-adaptive recurrence, intermediate supervision, and discretization as contributing mechanisms.Source: Induction Head & Modular Addition… · pp. 15, 17

6. Conclusion

The conclusion reports near-perfect OOD performance from recurrent latent-space reasoning mechanisms and identifies broader task settings as future work. The appendix adds details on end-to-end and chain-of-thought baselines and their supervision formats.

6. Conclusion

  • The integrated mechanisms yielded near-perfect OOD performance through reasoning in latent representations, while the interpretability analysis illuminated how components support systematic generalization.The authors identify synthetic mathematical reasoning as a clear testbed and call for exploration in more diverse, less-structured, and multi-task settings.Source: §6 · pp. 17–18

A.1. End-to-End Baselines

  • End-to-end baselines include feedforward models with independently learned layer parameters and recurrent models that repeatedly apply a Transformer block.The appendix describes a hyperparameter search over recurrence, block depth, hidden dimension, and positional encoding.Source: §A.1 · p. 24

A.2. Chain-of-Thought Baselines

  • Chain-of-thought baselines are causal Transformer language models trained by next-token prediction on sequences containing step-by-step solutions, then evaluated by greedy generation.The appendix provides further experimental details and results for these and end-to-end baselines.Source: A. Experimental Details on… · p. 24, §A.2 · p. 27

A.2.1. Chain-of-Thought Trajectories

  • The two CoT trajectory formats differ in detail: one recalls and computes variable values, while the other also recalls each variable’s defining equation.The equation-first format is longer and provides richer supervision.Source: §A.2.1 · p. 27

A.2.2. Experimental Details & Additional Results

CoT baselines show limited OOD generalization, while the appendix details factorized latent inputs, iterative supervision, discretization, and training for error correction.

A.2.2. Experimental Details & Additional Results

  • NoPE performed best among the evaluated CoT positional encodings, generalizing well to 40 nodes after training on N ≤32, while other methods did not generalize beyond training and none generalized robustly beyond 40 nodes.Some models recalled equation structure correctly but computed values unreliably, suggesting arithmetic computation was a common error source.Source: §A.2.2 · p. 29, Figure 10 · p. 30

B.1. Latent State Embedding Structure

  • Inputs are tokenized into syntax, variable, operation, and value factors, with variable tokens’ initially empty value factors filled as the model processes the input.The model embeds each factor separately and adds the embeddings.Source: §B.1 · p. 33

B.2. Latent State Supervision

  • Latent supervision trains recurrent iterations to compute node values one graph layer deeper, applying losses to variables whose depth is at most the current iteration.The overall training loss sums the individual factor losses at each iteration.Source: §B.2 · pp. 33–34

B.3. Discretization of Intermediate States

  • With a discrete bottleneck, linear readouts decode latent embeddings into discrete factor predictions, which are re-embedded as input to the next iteration.This discretization enables teacher-forcing across recurrent iterations and more efficient training.Source: §B.3 · p. 34

B.4. Self-Correction Mechanism

  • The self-correction method randomly corrupts value components of discrete latent states during training, requiring the model to detect and correct incorrect values before proceeding.The motivation is that an error can affect all subsequent computation in a reasoning path.Source: §B.4 · pp. 34–35

B.5. Experiment Details & Additional Results

The experiments compare discrete latent supervision with and without re-correction across input sizes, finding that re-correction calls for a deeper recurrent block.

B.5. Experiment Details & Additional Results

  • DeBERTa relative positional encoding performed best by far, and two recurrent-block layers sufficed for Discrete Latent Space Supervision without re-correction.The best-performing model for that method was DeBERTa-L2H16D256.Source: §B.5 · p. 35
  • The re-correction method required a deeper recurrent block; two layers showed no significant improvement for it.The reported best-performing re-correction configuration used four layers.Source: §B.5 · p. 35
  • With a 4-layer recurrent block, re-correction enabled near-perfect OOD performance; the reported best re-correction model was DeBERTa-L4H16D384.Figure 13 reports average “% Fully Solved” across 8–128 nodes after training on at most 32 nodes.Source: §B.5 · p. 35, Figure 13 · p. 36

C. Details of Mechanistic Interpretability Analysis

The interpretability analysis uses controlled probe equations to examine a trained discrete-latent model and traces the computation at the right-hand-side position.

C. Details of Mechanistic Interpretability Analysis

  • The analysis uses a DeBERTa-L2H16D256 model trained with Discrete Latent Space Supervision, selected from same-configuration seeds for its 99.98% OOD “% Fully Solved” score.Training used modular-23 addition with maximum graph size 32; interpretation tests used graph size 128.Source: C. Details of Mechanistic… · p. 38

C.1. Technique Overview

  • For a probe equation, the analysis separates identifying dependent variables, retrieving their values, and computing their modular sum; it assigns these operations to first-layer attention, second-layer attention, and the final MLP layer.The illustrated computation is at the equation’s RHS position.Source: C. Details of Mechanistic… · p. 38, §C.1 · p. 39

C.1.1. Interpretation Technique for Attention Layers

Controlled probes and weight-matrix analysis reveal specialized attention-head groups: the first layer routes variable identities, the second retrieves their values, and the MLP computes modular sums in the frequency domain.

C.1.1. Interpretation Technique for Attention Layers

  • Controlled probe equations and linear-algebra analysis reveal attention-head groups specialized for distinct subtasks in the arithmetic computation.The probes vary inputs to identify attention patterns, then the weight matrices are analyzed to decode information flow.Source: §C.1.1 · pp. 39–40

C.2. First Layer Attention: Variable Copying

  • First-layer groups identify variable names, while second-layer groups retrieve the corresponding values for the computation.The second-layer analysis identifies heads retrieving values for each of the three variable positions.Source: §C.1.1 · p. 40

C.1.2. Interpretation Technique for the Second MLP Layer

  • The second-layer MLP computes modular addition through sinusoidal basis functions, amplifying diagonal frequencies (a, a, a) that encode x + y + z mod 23.The analysis applies a three-dimensional DFT to representations across input values and finds the diagonal-frequency pattern becomes dominant.Source: §C.1.2 · p. 41

C.2.1. Group Structure in the First Layer Attention

Controlled variation of variable names and relative-variance measurements reveal disjoint first-layer head groups that respond to specific variable positions at the RHS.

C.2.1. Group Structure in the First Layer Attention

  • The experiment varies one variable name across 128 valid probe sequences while holding the other variables fixed, then measures each head’s relative variance.Higher relative variance indicates greater sensitivity to the varied position.Source: §C.2.1 · pp. 42–43
  • The resulting grouping indicates that first-layer attention copies information from LHS variables to the RHS position.The analysis focuses on the RHS query position, where output sensitivity identifies which variable each head attends to.Source: §C.2.1 · p. 43

C.2.2. First Layer Attention Copying the Variable Identity

The first attention layer routes variable identities to the RHS with little change from its MLP; the second layer routes variable values there for subsequent modular arithmetic.

C.2.2. First Layer Attention Copying the Variable Identity

  • First-layer attention head groups amplify variable embeddings more than other factored embeddings, supporting their interpretation as copying variable identities to the RHS.The comparison is based on norm-amplification histograms across embedding types and head groups.Source: §C.2.2 · p. 45

C.2.3. First Layer MLP Residual Stream Does Not Change the Residual Stream Significantly

  • The first-layer MLP produces relatively small changes to the residual stream at the RHS position.The reported relative L2 error is relatively small across 256 samples.Source: §C.2.3 · p. 45

C.2.4. Conclusion

  • The first block’s attention copies variable-factor embeddings from the LHS variables to the RHS, while its MLP changes the residual stream little.The attention-group analysis finds substantially greater norm amplification for the variable factor than for other embedding types.Source: §C.2.4 · p. 45, §C.2.2 · p. 45, §C.2.3 · p. 45

C.3. Second Layer Attention: Value Copying

  • Second-layer head groups retrieve the values of var0, var1, and var2, with each group specialized for one variable position.The groups are {0, 8, 15}, {5, 10}, and {2, 3, 4, 7, 9}, respectively.Source: §C.3 · p. 47

C.3.1. Conclusion

  • Second-layer attention writes the three variables’ value-factor embeddings into the residual stream at the RHS position.The grouped OV matrices amplify the value factor more than the other embedding types.Source: §C.3.1 · p. 47, §C.3 · p. 47

C.4. Second Layer MLP: Module Addition in the Frequency Domain

  • The second-layer MLP analysis begins by characterizing the structure of its input representations.Those inputs are the values copied by second-layer attention.Source: §C.4 · p. 47

C.4.1. Understanding the Input Structure to the Second Layer MLP

The second-layer MLP receives transformed value embeddings with a structured geometry and computes modular sums through diagonal frequency components; the error analysis identifies feedforward calculation as the major error source.

C.4.1. Understanding the Input Structure to the Second Layer MLP

  • Across variables, transformed value embeddings are near-orthogonal; within each variable, the 23 actual-value embeddings have an approximately circulant similarity structure.The special empty and N/A embeddings are nearly orthogonal to the 23 actual-value embeddings and to each other.Source: Figure 21 · p. 49, §C.4.1 · p. 49

C.4.2. Module Addition in the Frequency Domain.

  • As the MLP processes signals, the (0, 0, 0) bias diminishes while the norm of diagonal (a, a, a) frequencies increases, encoding x + y + z mod 23.The frequency-domain analysis examines representations at multiple positions in the network.Source: §C.4.2 · p. 50

C.4.3. Conclusion

  • The second-layer MLP computes modular addition through sinusoidal basis functions, amplifying diagonal frequencies that encode the sum and suppressing the bias term.The diagonal components have the form (a, a, a), for a in {1, . . . , 22}.Source: §C.4.3 · p. 50

C.5. Error Analysis

  • The feedforward layer’s calculation is the major source of errors among the three analyzed sources.The accounting distinguishes wrong variable-position copying, incorrect value copying, and feedforward miscalculation.Source: §C.5 · p. 51

C.5.1. Identifying Different Sources of Errors

The analysis assigns errors to first-layer variable attention, second-layer value copying, or feedforward calculation, and identifies feedforward calculation as the major source.

C.5.1. Identifying Different Sources of Errors

  • 40 errors among 23k examples were reported, with feedforward calculation identified as the major source.Table 2 attributes test-set errors by source.Source: Table 2 · p. 51, §C.5.1 · p. 52
  • First-layer attention errors are designated-variable probabilities below 0.9; observed probabilities cluster near 0 or 1, making failures easy to identify.The probability averages attention from a head group at the RHS position to its assigned variable position.Source: §C.5.1 · p. 52
  • Second-layer copy errors are cosine similarities below 0.9 between head-group outputs and target value embeddings; correct copies are near 1, unlike incorrect copies.The threshold separates the two patterns clearly and is not critical to the analysis.Source: §C.5.1 · p. 52
  • Feedforward calculation errors occur when both attention checks pass but the model still predicts incorrectly after the second-layer MLP.The analysis counts an error at this stage only after the first-layer attention and second-layer copy checks succeed.Source: §C.5.1 · p. 52

C.5.2. Additional Error Analysis

The error analysis finds distinguishable attention and copying error patterns, but reports no self-correction at equation positions after iterations reach equation depth.

C.5.2. Additional Error Analysis

  • For equations with depth at least the iteration count, ⟨var2⟩ attention is less concentrated on its assigned variable and has more errors than the other head groups.Incorrect predictions are associated with group-wise attention probabilities dropping to nearly 0.Source: §C.5.2 · p. 53
  • Second-layer copy outputs usually have cosine similarity near 1 to target value embeddings, while incorrect predictions can show similarity near 0.This contrast makes second-layer copy errors readily identifiable.Source: §C.5.2 · p. 53
  • Beyond equation depth, first-layer attention does not focus on the correct variable, and the authors report no further computation or self-correction at equation positions.The first two rows of Figure 23 analyze equations with depth at least the iteration count; the bottom row covers equations beyond that depth.Source: §C.5.2 · p. 53, Figure 23 · p. 58
  • The authors propose training on imperfect or corrupted data so extra iterations can teach recovery from errors made in earlier iterations.They note that perfect training data gives the model no need to compute beyond equation depth.Source: §C.5.2 · p. 53

Related papers

Loading 2510.14095v1…