Invariant-Measure Reasoners: Stable Representations for Latent Reasoning

Yuto Inui; Takuya Konishi; Yoshinobu Kawahara

cs.LGOct 7, 2026

Invariant-Measure Reasoners: Stable Representations for Latent Reasoning

Authors: Yuto Inui, Takuya Konishi, Yoshinobu Kawahara

Organizations: The University of Osaka · RIKEN

Abstract

Latent reasoning models repeatedly update a latent state using the same recurrent block. As the recurrent depth increases, the sequence of latent states may converge to a compact subset of the state space without necessarily converging to a fixed point. Existing models typically predict by applying a prediction head to a single latent state. However, the latent state can continue to change even after many updates, potentially making predictions unstable across recurrent depths. To address this instability, we introduce invariant-measure reasoners (ImR), a framework that uses an invariant measure as a stable representation. This measure describes the long-run distribution of latent states on the compact subset and is invariant under updates by the recurrent block. ImR predicts from the expectation of the prediction head's output under this measure. We use ImR in two ways: fine-tuning only the prediction head of existing models and training models from scratch. Both approaches reduce prediction instability and improve accuracy in many settings on maze and Sudoku tasks. In some settings, models trained with ImR exhibit non-fixed-point behavior more frequently than existing models yet achieve high accuracy even with such behavior, unlike existing models. These results suggest that ImR can leverage otherwise destabilizing dynamics for latent reasoning.

Figures & tables

Figure 1: (a) An example of prediction instability in equilibrium reasoners (EqR) ( Huang et al., 2026 ) on a maze problem with the latent trajectory projected onto the first two principal components. Trajectory color indicates recurrent depth, and pale green and red regions approximately indicate correct and incorrect predictions, respectively. (b) Schematic of ImR. For an input ${\bm{x}}$ , repeated applications of the recurrent block $f_{{\bm{x}}}$ to the latent state ${\bm{s}}_{t}$ generate a latent trajectory. ImR predicts from the expectation $\mathbb{E}_{{\bm{s}}\sim\mu^{\star}}[g({\bm{s}})]$ under an invariant measure $\mu^{\star}$ .

Figure 2: Accuracy across recurrent depth on Maze-Unique (a) and Sudoku-Extreme (b), with enlarged views (c) over the final 25 recurrent depths. For both fine-tuned models and ImR, curves show means and shaded regions show standard deviation across 3 training seeds.

Sudoku-Extreme : dimension of ${\bm{s}}^{(H)}$ : $49.7\mathrm{K}$ ; parameters $5.03\mathrm{M}$
		Fixed			Non-fixed
Model	Non-fixed rate	Stable correct	Stable incorrect	Unstable	Stable correct	Stable incorrect	Unstable
EqR †	$4.1$	$100.0$	$0.0$	$0.0$	$3.6$	$0.0$	$96.4$
ImR (FT)	$4.1\pm 0.0$	$99.2\pm 0.0$	$0.0\pm 0.0$	$0.8\pm 0.0$	$3.4\pm 0.0$	$79.7\pm 0.1$	$\bm{16.9\pm 0.1}$
ImR	$1.2\pm 0.2$	$99.7\pm 0.1$	$0.0\pm 0.0$	$0.2\pm 0.1$	$\bm{8.0\pm 12.2}$	$44.5\pm 5.8$	$47.5\pm 8.0$
Sudoku-Extreme : dimension of ${\bm{s}}^{(H)}$ : $41.5\mathrm{K}$ ; parameters $4.86\mathrm{M}$

Table 1: Trajectory groups and prediction stability. Entries with

\pm

report mean and standard deviation across 3 training seeds. EqR † reports results for both EqR and EqR (FT), which match on all reported metrics. For each dataset and state dimension of

{\bm{s}}^{(H)}

, bold marks the highest stable correct proportion and the lowest unstable proportion in the non-fixed group.

Degree of the start cell: 2
Model	Stable correct (%)	Stable incorrect (%)	Unstable (%)
DT	$76.6$	$17.3$	$6.1$
DT (FT)	$81.5\pm 0.1$	$13.9\pm 0.1$	$4.5\pm 0.1$
ImR (FT)	$\bm{84.5\pm 0.0}$	$15.1\pm 0.1$	$\bm{0.4\pm 0.1}$
Degree of the start cell: 3
DT	$25.3$	$57.9$	$16.8$

Table 2: Prediction stability on OOD mazes over depths

4{,}900

–

5{,}000

. Fine-tuned models report mean

\pm

standard deviation across 3 training seeds. For each degree of the start cell, bold marks the highest stable correct proportion and the lowest unstable proportion.

Figure 3: Accuracy across recurrent depth on OOD mazes whose start cell has degree 2 (a) or 3 (b), with enlarged views (c) over the final 20 recurrent depths. For both fine-tuned models, curves show means and shaded regions show standard deviation across 3 training seeds.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Train	Validation	Test
Maze-Unique	$1{,}000$	$1{,}000$	$1{,}000$
Sudoku-Extreme	$1{,}001{,}000$	$10{,}000$	$422{,}786$
Maze	$40{,}000$	$10{,}000$	$1{,}000$

Appendix

Table 3: Dataset sizes (top) and counts of generated test mazes by the degree of the start cell (bottom). The degree counts are pooled across all maze sizes. The Sudoku-Extreme training count includes augmentation.

Figure 4: Examples of ID and OOD mazes used for DT. The first column shows a $9\times 9$ ID validation maze. The remaining columns show $19\times 19$ OOD mazes whose start cells have degrees 1 through 4. The top row contains the inputs, and the bottom row marks the target solution path in blue, the start cell in green, and the goal cell in red.

Dataset	Position mixing	Puzzle prefix	$f_{{\bm{x}}}$	$g$	Total
Maze-Unique	Self-attention	With	$263{,}040$	$768$	$264{,}066$
		Without	$262{,}912$	$768$	$263{,}938$
Sudoku-Extreme	MLP	With	$5{,}022{,}720$	$5{,}632$	$5{,}029{,}378$
		Without	$4{,}848{,}640$	$5{,}632$	$4{,}855{,}298$
Maze	Convolution	Without	$740{,}736$	$39{,}312$	$783{,}504$

Appendix

Table 4: Model architectures and the number of parameters. For EqR, parameters used to embed the input are included under

f_{{\bm{x}}}

. The totals also include the halting head in EqR and the neural network used to initialize the latent state in DT.

Architecture	Dataset	Batch size	Optimizer updates
EqR	Maze-Unique	768	$1{,}000$
EqR	Sudoku-Extreme	768	$1{,}000$
DT	Maze	50	$1{,}600$

Appendix

Table 5: Hyperparameters for fine-tuning the prediction head.

Model	Dataset	Puzzle prefix	$L$	$K$	Batch size	Optimizer updates
ImR	Maze-Unique	With / Without	36	4	768	$150{,}000$
ImR	Sudoku-Extreme	With / Without	36	4	768	$50{,}000$
EqR	Maze-Unique	Without	36	4	768	$150{,}000$
EqR	Sudoku-Extreme	Without	3	1	768	$50{,}000$

Appendix

Table 6: Hyperparameters selected for training from scratch with the architecture of EqR. Here,

L

is the number of recurrent updates per inner loop, and gradients through the recurrent block are retained for the final

K

updates.

Figure 5: The EqR trajectory in Figure (a) and its actual predictions. Left: the trajectory over depths 0–512 and prediction regions on the reconstructed PCA plane. Trajectory color indicates recurrent depth; the circle marks the initial state. Right: whole-maze correctness from the original head applied to each actual latent state. Pale green and pale red indicate correct and incorrect predictions, respectively.

Model	Puzzle prefix	Accuracy (%)
Maze-Unique , $T=512$
EqR	With	$89.4$
EqR (FT)	With	$89.4\pm 0.0$
ImR (FT)	With	$91.6\pm 0.1$
ImR	With	$98.8\pm 0.9$
EqR	Without	$82.9\pm 20.9$

Appendix

Table 7: Accuracy at the final recurrent depth. Entries with

\pm

report the mean and standard deviation across 3 training seeds. The remaining entries use a single released model.

Figure 6: PCA projections of the part of the latent state used by the prediction head for ImR models trained without the puzzle prefix. Each row shows an input and the corresponding latent trajectory projected onto PC1–PC2, PC3–PC4, and PC5–PC6. The first two rows show Maze-Unique examples, and the last two show Sudoku-Extreme examples. Axis labels give the percentage of variance explained by each principal component. Colors indicate recurrent depths 897–1024; circles and stars mark the first and last states in this interval, respectively. All trajectories were generated with noise set to zero.

Figure 7: Accuracy across recurrent depth for each maze size. The first two columns show sizes

9\times 9

through

99\times 99

, with 100 test mazes per size, combining all degrees of the start cell. The right column shows enlarged views over depths

4{,}981

–

5{,}000

. Curves for DT (FT) and ImR (FT) report means across 3 training seeds.

Explore similar work

Jun 28, 2026cs.LG

Invariant Reasoning Directions in Latent Trajectories of Language Models

Latent reasoning models perform multi-step inference directly in hidden-state space, yet the structure of these latent reasoning trajectories remains poorly understood. We show that contrastive refinement signals between stronger and weaker reasoning trajectories exhibit a highly concentrated low-rank structure, while unconstrained latent updates remain sensitive to paraphrases, checkpoint choice, and trajectory perturbations. These observations suggest that latent reasoning trajectories contain stable invariant directions mixed with unstable instance-specific variation. We introduce \textbf{Trajectory-Invariant Latent Refinement (TILR)}, a training-free intervention framework for identifying and manipulating stable reasoning directions in latent space. TILR first learns a low-rank invariant subspace from contrastive trajectory differences across inputs, then constrains latent interventions to this subspace while suppressing poorly aligned updates through an adaptive alignment gate. Across six reasoning benchmarks, we find that a small number of latent directions explain most variation between strong and weak reasoning trajectories. Interventions on these directions causally improve reasoning consistency and reduce trajectory instability under paraphrases and perturbations. TILR improves answer consistency under paraphrase by ~10% and reduces latent trajectory variance by up to

50\%

while preserving reasoning accuracy. These results support a geometric view of latent reasoning in which transferable reasoning behavior emerges from stable low-dimensional structure within hidden-state trajectories.

Arun Vignesh Malarkkan, Manan Roy Choudhury, Utkarsh Byahut +3

Arizona State University, Tempe, Arizona.

Reasoning Consistency in Language Models Implicit Reasoning in Language Models Causal Interventions in Language Models Language Model Robustness Representation Geometry in Language Models Latent Reasoning

Ask PDF

Jun 16, 2026cs.LG

Learning to Refine Hidden States for Reliable LLM Reasoning

Large language models show strong reasoning ability, but their internal reasoning process can remain unstable in complex multi-step settings, where early hidden-state errors may propagate to incorrect predictions. We propose ReLAR, a reinforcement-guided latent refinement framework that iteratively updates hidden representations before decoding. ReLAR maintains a compact latent reasoning state and uses learned depth and action controllers to adaptively determine both the number and direction of refinement steps. The controllers are trained with a policy gradient objective based on step-wise likelihood improvement, enabling efficient input-dependent reasoning without explicit chain-of-thought generation. Experiments on medical, mathematical, multi-hop reasoning, and open-ended generation benchmarks show that ReLAR improves accuracy, generation quality, and reasoning stability with substantially lower inference overhead than explicit reasoning baselines.

Chia-Hsuan Hsu, Jui-Ming Yao

RL for Language Model Reasoning Efficient Language Model Reasoning LLM Reliability Implicit Reasoning in Language Models Adaptive Computation Latent Reasoning

Ask PDF

Sep 3, 2026cs.LG

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

Repeating a small block of middle layers increases a language model's effective inference depth without adding parameters or generating extra tokens, and recent work shows that this latent recurrence improves reasoning. However, two design choices limit these gains. Each iteration sees only the previous output and cannot directly access earlier computations. Moreover, a fixed loop count wastes depth on easy inputs while leaving hard ones with too little computation. We introduce RecurTrace, which addresses both limitations using the loop's own trajectory. Specifically, Loop Memory Attention lets each looped layer attend to its own states from previous iterations along the loop-time axis, so the model can revisit earlier computations instead of relying on the latest state alone. A halting head then reads the loop state and predicts whether to continue, with supervision from an oracle that identifies when additional depth still reduces loss. In a controlled MathQA comparison on the same looped backbone, RecurTrace achieves 56.9% accuracy with an average of 2.0 loops, exceeding the best fixed loop depth by 2.2 points at matched compute. By comparison, ACT and PonderNet collapse to one loop, and CALM reaches only 54.1% with 5.6 loops, while the stronger LoopUS-Conf and TaH-Mismatch baselines reach 55.3% at 3.2 loops and 55.7% at 2.1 loops. Finally, RecurTrace improves generation accuracy over same-budget fine-tuned baselines at 0.6B, 1.7B, 4B, and 8B, with the gain growing with model size from 0.6 to 3.4 points.

Yuxiang Wang, Kunyu Feng, Yingda Shen +3

The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2

Memory-Augmented Language Models Recurrent Transformers LLM Inference Acceleration Adaptive Computation

Ask PDF