Steering Recurrent Reasoners at Inference Time with Readout Feedback

Shunsuke Kamiya; Masanori Koyama; Seongcheol Jeong; Fumiya Uchiyama; Kenji Kubo; Kohei Hayashi; Masahiro Suzuki; Yutaka Matsuo

Source-linked AI summary

Steering Recurrent Reasoners at Inference Time with Readout Feedback

Shunsuke Kamiya, Masanori Koyama, Seongcheol Jeong, Fumiya Uchiyama, Kenji Kubo, Kohei Hayashi, Masahiro Suzuki, Yutaka Matsuo

arXiv:2608.24136v1Aug 25, 2026cs.LG

TL;DR

Existing inference-time methods for recurrent reasoners mainly run longer trajectories or sample more of them, without using information revealed within each trajectory. The paper introduces RoFB, which steers frozen latent dynamics using readout probabilities, and finds gains in four of six model-task pairs at comparable or lower computational cost. The evidence is limited to two puzzle benchmarks and three compact recurrent architectures.

Problem
Existing inference-time strategies for recurrent reasoners passively lengthen trajectories or sample multiple trajectories, motivating active steering of frozen latent dynamics.
Method
RoFB converts intermediate readout probabilities into a coupling signal injected into recurrent latent dynamics during inference without changing model weights.
Results
RoFB improves performance in four of six model-task pairs, reaching operating points unattainable by longer evaluation or increased voting in the positive pairs.
Takeaways & Limitations
Readout-based steering can complement longer rollouts and multi-trajectory voting for controlling frozen recurrent dynamics at test time.
Takeaways & Limitations
Evidence is confined to Sudoku Extreme and Maze Hard with three small-scale recurrent architectures and does not cover language models, where direct transfer is non-trivial.

Abstract

from arXiv · show

Recurrent models, which repeatedly update latent states with shared computation blocks, have emerged as powerful architectures for solving complex reasoning tasks. Existing inference-time methods scale computation by running more steps or sampling more trajectories, but ignore information revealed within each trajectory. Here we show that recurrent models can be improved at inference time by using their own readout probabilities to steer latent dynamics without retraining. We introduce Readout Feedback (RoFB), a test-time intervention that converts intermediate predictions into token-wise pairwise coupling forces injected into the latent dynamics. Across three recurrent models (AKOrN, ItrSA++, TRM) on Sudoku and Maze, RoFB yields clear gains in four of six model-task pairs, achieving performance unattainable by merely running more steps or selecting from multiple trajectories, at comparable or lower computational cost. These results suggest that closed-loop steering of latent dynamics can serve as a complementary inference-time control mechanism for recurrent reasoning models.

1 Introduction

Inference-time computation can improve recurrent reasoning, but existing approaches passively run longer or sample more trajectories. RoFB instead uses intermediate readout probabilities to steer latent dynamics without changing model weights.

1 Introduction: Recurrent reasoners expose inference-time computation through repeated latent-state updates with shared computation blocks.Their computation depth is controlled by the number of recurrent updates.
1 Introduction: Longer rollouts can saturate or remain trapped in unsuccessful dynamical regimes, while multi-trajectory voting increases compute roughly linearly with candidate count.
1 Introduction: RoFB actively steers each latent trajectory by injecting readout-probability-based coupling into the dynamics while leaving all model weights unchanged.The coupling encourages token clustering according to distances between tokens’ readout probabilities.
1 Introduction: RoFB is evaluated on Sudoku and Maze using AKOrN, ItrSA++, and TRM, with gains complementary to confidence-based voting when voting applies.The evaluation also examines inference steps, trajectory count, and normalized inference compute.

2 Preliminaries

The considered recurrent models embed input tokens, repeatedly update latent states with shared modules, and produce token-wise readout probabilities. AKOrN, ItrSA++, and TRM differ in state geometry, coupling, hierarchy, and readout design.

2 Preliminaries: Recurrent models comprise an input embedding map φ, recurrent module R, and readout head ψ.The recurrent module repeatedly updates latent states, while the readout produces per-token outputs.
2 Preliminaries: After T recurrent updates, the readout head produces per-token probability vectors through a softmax over C output classes.
2 Preliminaries: The study selects AKOrN, ItrSA++, and TRM for strong reasoning performance and diverse architectural designs, including hierarchical structures and latent-state geometry.Table 1 summarizes their latent geometry, coupling mechanisms, and readout strategies.
2.2 AKOrN: AKOrN places latent tokens on a product of spheres and uses Kuramoto-like dynamics with learnable all-to-all coupling from self-attention.Its updates include tangent-space projection and token-wise normalization.
2.3 ItrSA++: ItrSA++ uses a two-level hierarchy in which low-level z captures fast-changing details and high-level y captures slower-changing global structure.The high-level process is updated every L steps, and RMS normalization identifies its latent state space with a sphere.
2.4 Tiny Recursive Model (TRM): TRM uses a two-level hidden-state hierarchy, with a transformer-based low-level process for local details and a high-level process updated every L steps for global structure.The high-level state is read out after a specified number of recursions.

3 Readout Feedback (RoFB)

RoFB is motivated by phase-dependent token dynamics: successful trajectories form class-based latent clusters, while unsuccessful ones remain less structured. It uses readout probabilities to inject gated, pairwise coupling into recurrent updates while preserving manifold constraints and original forward information.

Inter-token Clusterization: Successful trajectories pass through wandering, “Aha,” and solution phases, with entropy dropping sharply when predictions become correct.For AKOrN on Sudoku, the “Aha” phase typically occurs around t ≈120, followed by a correct and confident solution phase.
Inter-token Clusterization: The “Aha” phase occurs before t = 256 for solvable trajectories but is absent within t = 256 for unsolved trajectories.Its timing depends on the initial latent value, consistent with trajectories escaping local minima before converging to a global minimum.
Inter-token Clusterization: Solved puzzles produce class-based token clusters: same-class tokens align near cosine similarity 1.0, whereas different-class tokens are less similar.Unsolved samples lack this structured clustering and retain broadly high, less target-dependent similarities.
RoFB: RoFB converts distances between readout probability vectors into pairwise feedback, encouraging alignment for similar distributions and discouraging alignment for dissimilar ones.The experiments use the repulsive instantiation h(d) = −d, so the implemented method primarily discourages collapse among dissimilar-readout tokens.
RoFB: The feedback term is added to recurrent dynamics through a gate and scaling factor, then projected onto the manifold’s tangent space and back onto the manifold.This design retains the original forward information while keeping latent states on the manifold.
RoFB: For Sudoku and Maze, RoFB applies the update only to blank-cell tokens and can activate it during the wandering phase using an entropy-based confidence gate.The gate may use a minimum step, entropy threshold, and sigmoid temperature; λ, tmin, α, and τ are tunable hyperparameters.

4 Related Work

Related work studies recurrent computation, internal dynamics, inference-time interventions, and output-space refinement. RoFB differs by using intermediate readout probabilities from a frozen recurrent model to steer latent dynamics at inference time rather than training or selecting outputs.

Iterative Reasoning Models: Iterative reasoning models reuse shared computation blocks to solve complex tasks, including algorithm learning and synthesis.This line includes Universal Transformer, Deep Equilibrium Models, and recurrent networks that iterate computation.
Internal Dynamics of Recurrent Reasoners: The internal mechanisms of compact distilled recurrent reasoners such as HRM and TRM remain less understood than those of recurrent LLMs.RoFB is designed to leverage dynamical characteristics of these models for inference steering.
Test-time Intervention: Inference-time methods commonly prompt, sample, allocate compute, or select among candidate trajectories rather than intervene directly in latent computation.C-voting is identified as a recurrent-reasoner-specific neighbor that selects among multiple latent trajectories.
Output-space Refinement: Output-space refinement methods exploit confidence, neighborhood structure, prediction similarity, or consistency to improve predictions or representations.Examples include dense CRFs, CRF-RNN, NRC, AaD, and deep clustering methods.
Output-space Refinement: RoFB shares the use of relational structure in output distributions but applies it to intermediate readouts of a frozen recurrent model during inference, without optimizing a training or adaptation loss.This distinguishes RoFB from source-free adaptation, clustering, and other output-space refinement approaches.

5 Experiments

RoFB was evaluated on Sudoku and Maze using three recurrent models under compute-matched inference settings. It improved performance in four of six model-task pairs, sometimes reaching baseline-unattainable accuracy at substantially lower inference cost.

5.1 Performance: Performance generally improved with increasing Kvote and Teval for every model-task pair except TRM on Maze.TRM on Maze plateaued, with RoFB showing only a 0.1% decrease as inference steps increased.
5.1 Performance: RoFB improved performance in four of six model-task pairs at the same or lower computational cost.The gains occurred for AKOrN on both tasks, ItrSA++ on Maze, and TRM on Sudoku; the remaining two pairs were largely unchanged.
5.1 Performance: 83.6% shortest-path accuracy from ItrSA++ on Maze exceeded the 81.7% baseline at 64 votes and 256 steps with roughly 100× less inference compute.RoFB used Kvote=1 and Teval=128, while the baseline used Kvote=64 and Teval=256.
5.1 Performance: 91.5% board accuracy for AKOrN on Sudoku at Kvote=4 and Teval=256 was 6.4% above the baseline at the same operating point and exceeded every baseline result at Kvote=64.This operating point was not attained even by the baseline with the larger voting budget.
5.1 Performance: RoFB increased TRM Sudoku board accuracy from 68.4% to 74.2% (+5.8%) at Nblock=16, despite voting being unavailable for TRM.TRM uses a fixed initial latent value rather than sampling initial values, so confidence-based voting is not applicable.
5.2 Mechanistic Analysis of Coupling: Flipping the feedback sign caused large deterioration in the four positive model-task pairs but little to no change in the other two.The result supports the importance of pushing latent states along the clustering direction toward the correct cluster.

6 Limitations and Discussion

The evidence is limited to two puzzle benchmarks and three small-scale recurrent architectures, leaving transfer to language models unresolved. Two model-task pairs also show no meaningful RoFB improvement, and TRM on Maze has little headroom because its baseline is already highly accurate.

6 Limitations and Discussion: The evidence covers only Sudoku Extreme, Maze Hard, and three small-scale recurrent reasoning architectures, not language models.Direct transfer to language models is non-trivial, and large vocabularies may dilute the inter-token coupling RoFB relies on.
6 Limitations and Discussion: Two of six model-task pairs, ItrSA++ on Sudoku and TRM on Maze, show no meaningful improvement under RoFB.Neither RoFB nor its sign-flipped counterpart changes accuracy by more than 1% in these pairs, and no a priori diagnostic currently predicts when RoFB will help.
6 Limitations and Discussion: TRM on Maze begins at 92.0% shortest-path and 98.8% valid-path accuracy, leaving little room for validation-based hyperparameter search to distinguish RoFB configurations.The resulting flat validation objective may contribute to the null result independently of whether the latent dynamics has a useful coupling direction.

7 Conclusions

RoFB improves inference performance in four of six evaluated model-task pairs without retraining and at lower computational cost, while suggesting readout-based steering as a complementary inference-time strategy.

RoFB improves inference performance in four of six evaluated model-task pairs without requiring retraining.
The reported gains provide evidence that readout-based steering can complement longer rollouts and multi-trajectory voting.
The evidence is limited to puzzle-style benchmarks and compact recurrent reasoners.

A Experimental setup

The experiments evaluate three iterative reasoners on shared Sudoku Extreme and Maze Hard datasets, using model-specific training protocols and a four-parameter RoFB sweep.

Three iterative reasoners—AKOrN, ItrSA++, and TRM—are evaluated on Sudoku Extreme and Maze Hard.The datasets use shared training and test splits across models.
Sudoku Extreme uses a 9×9 grid, while Maze Hard uses a 30×30 grid with 1,000 training and 1,000 test examples.
All models are trained from scratch with AdamW, truncated backpropagation through time, and exponential moving-average parameters.
Each model-task pair receives a sweep over λ, α, tmin, and τ, with the coupling function fixed to h(d) = −d.
The sweep uses Optuna TPE with Successive Halving and 30 trials per model-task pair.

C Computational cost of RoFB

RoFB adds a per-step feedback block and a one-time hyperparameter-sweep cost to inference, yet its Pareto improvement persists when total FLOPs and wall-clock costs are considered.

RoFB’s two overhead sources are a one-time hyperparameter sweep and a per-step block inserted into inference.
The normalized operating-point cost combines per-inference forward cost with a one-time-per-cell sweep cost.
The main figure excludes sweep cost because it is incurred once per cell and is below 2% for the four AKOrN/ItrSA++ cells.
The sweep cost ranges from 0.38 to 1.65 baseline test evaluations across the four AKOrN/ItrSA++ cells.
The per-step RoFB block adds a batched (B, N, N) × (B, N, D) matrix multiplication, with overhead ranging from 1.2% to 19.5%.
At the canonical horizon, RoFB forward overhead ranges from 1.008 on TRM × Sudoku to 1.098 on AKOrN × Maze.
Including sweep cost shifts RoFB operating points right by a constant per-cell offset while preserving Pareto improvement across all six cells.
Wall-clock comparisons must account for fp32 execution on AKOrN and ItrSA++ versus bf16 execution on TRM.

E Detailed Training Setups

The detailed setup tables provide the full architectures and training hyperparameters for AKOrN, ItrSA++, and TRM.

Table 6 lists AKOrN’s full architecture and training hyperparameters.
Table 7 lists ItrSA++’s full architecture and training hyperparameters, using the affine-free canonical baseline.
Table 8 lists TRM’s full architecture and training hyperparameters.

F Compute Resources

Training and evaluation used NVIDIA GH200 GPUs, with configurations ranging from single-GPU runs to 8-GPU nodes. Approximately 100 GPU-days supported training across six cells, plus about 10 GPU-days for RoFB hyperparameter sweeps.

Approximately 100 GPU-days were used for training across all six cells.
RoFB hyperparameter sweeps used approximately 10 GPU-days on the same hardware, amortized over the K-curve test evaluation suite.
ItrSA++ cells and TRM × Sudoku used 8-GPU nodes, TRM × Maze used a 4-GPU node, and AKOrN cells used a single GPU.
All training and evaluation runs were performed on NVIDIA GH200 GPUs.