PaperlayerView exact revision on arXiv

Source-linked AI summary

Invariant-Measure Reasoners: Stable Representations for Latent Reasoning

Yuto Inui, Takuya Konishi, Yoshinobu Kawahara

arXiv:2610.10996v1cs.LG

TL;DR

Latent reasoning models can produce depth-dependent predictions when their latent states continue to change without converging to a fixed point. ImR predicts using the expectation of the prediction head under an invariant measure, and its use reduces instability and improves accuracy in many maze and Sudoku settings. The results support stable prediction even when latent states continue to evolve.

  • Problem

    Latent trajectories may approach a compact subset without converging to a fixed point, while predictions from a single latent state can vary across recurrent depths.

    Source: §3 · p. 3

  • Method

    ImR represents latent dynamics with an invariant measure and predicts from the expectation of the prediction head’s output under that measure.

    Source: §1 · p. 2

  • Results

    ImR reduces prediction instability and improves accuracy in many maze and Sudoku settings, with stable, correct predictions even when latent states continue to evolve.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The findings support stable prediction through convergence of empirical measures without requiring individual latent states to converge to a fixed point.

    Source: §7 · p. 10

    The experiments used fixed stopping depths; adaptive stopping rules balancing predictive accuracy and computation remain future work.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Latent states can keep changing after entering a bounded region, making predictions unstable across recurrent depths. ImR predicts from an invariant measure over latent states; its use reduces instability and improves accuracy in many maze and Sudoku settings.

1 INTRODUCTION

  • Latent trajectories may approach a compact region without reaching a fixed point, while final-state predictions can alternate between correct and incorrect.Increasing recurrent depth permits more test-time computation, but predictive performance can deteriorate beyond the depth used in training.Source: §1 · pp. 1–2
  • ImR predicts from the expectation of the prediction head’s output under an invariant measure, with fixed-point prediction as a special case.In practice, it approximates this prediction by averaging outputs over a finite latent trajectory.Source: §1 · p. 2
  • Both head-only fine-tuning and training from scratch with ImR reduce prediction instability and improve accuracy in many maze and Sudoku settings.In some settings, ImR-trained models also remain non-fixed-point while producing stable, high-accuracy predictions.Source: §1 · p. 2

2 PRELIMINARIES

The paper frames structured reasoning as mapping grid inputs to solutions and models repeated latent updates as a dynamical system. A prediction head maps latent states to logits, and the solution can depend on the stopping depth.

2 PRELIMINARIES

  • Structured reasoning tasks map input configurations to solutions represented over a finite set of grid cells.For mazes, the input specifies walls and start and goal cells, while the output encodes a path; Sudoku inputs are partially filled 9 × 9 grids and outputs are completed grids.Source: §2 · p. 2
  • A shared recurrent block updates a latent state representing the grid, defining a discrete-time dynamical system for a fixed input.The recurrent block can be a neural network with parameters shared across steps, and the initial state may be computed from the input or sampled from an input-independent distribution.Source: §2 · p. 3
  • Stopping the updates at depth T yields a finite latent trajectory, and a prediction head maps each latent state to logits over classes for each grid cell.The solution is obtained from the logits at the chosen stopping depth, so it can depend on T.Source: §2 · p. 3

3 INVARIANT-MEASURE REASONERS

ImR addresses prediction changes across recurrent depths by representing latent dynamics with an invariant measure and predicting from expected logits. The section develops this representation and prediction rule, then describes fine-tuning and from-scratch training procedures.

3.1 INVARIANT MEASURES OF LATENT DYNAMICS

  • When trajectory empirical measures converge weakly, their limit is invariant under recurrent updates and represents long-term behavior, including fixed points, periodic orbits, and more complex dynamics.The invariant measure may depend on the initial state; convergence to the same measure from different initial states is not required.Source: §3.1 · p. 4

3.2 PREDICTION FROM AN INVARIANT MEASURE

  • ImR predicts from the expectation of the prediction head under the limiting measure, approximated at finite depth by averaging logits along the trajectory.The invariant expected logits remain unchanged under recurrent updates even when the prediction head's output for an individual state changes.Source: §3.2 · p. 4
  • Under the stated assumptions, averaged logits converge, and with a unique maximum at each cell the solution becomes constant at sufficiently large depths.Finite-depth averages can still depend on stopping depth, but averaging can reduce that sensitivity; ImR retains the recurrent block and prediction head without adding trainable parameters.Source: §3.2 · p. 4
  • Discarding an initial transient can reduce its influence on finite-depth estimates, and prediction from only part of the latent state uses the corresponding invariant marginal.For fixed burn-in B, the modified empirical measure converges to the same invariant measure under the section's convergence assumption.Source: §3.2 · pp. 4–5

3.3 TRAINING WITH PREDICTION FROM AN INVARIANT MEASURE

  • Fine-tuning preserves a pretrained recurrent block and updates only its prediction head, whereas from-scratch training jointly optimizes both using averaged-logit loss.For fine-tuning, latent states are treated as constants so gradients pass only through the prediction head.Source: §3.3 · p. 5
  • From-scratch training uses deep supervision over inner loops, averages logits across each loop, and applies truncated backpropagation through time to limit gradient computation.Gradients pass through the recurrent block only for the final K of L updates; each loop forms a new empirical measure from states generated with the same parameters.Source: §3.3 · p. 5

4 RELATED WORK

ImR generalizes fixed-point representations to weakly convergent empirical measures and differs from learned halting approaches in how it weights latent states. It can also be combined with halting mechanisms.

Equilibrium models and stability

  • ImR extends fixed-point formulations to weak convergence of empirical measures, covering long-term dynamics beyond convergence to a single latent state.Its representation satisfies f#µ⋆ = µ⋆.Source: §4 · p. 5

Adaptive computation and halting

  • Unlike ACT and PonderNet, ImR weights outputs by latent-state visit frequencies rather than learned halting probabilities.ImR introduces no additional trainable parameters.Source: §4 · p. 6
  • ImR is complementary to halting mechanisms, which can determine when to stop collecting states for estimating the invariant measure.Source: §4 · p. 6

5 EXPERIMENTS

ImR improves accuracy and prediction stability in many evaluated settings, including non-fixed trajectories and OOD mazes, using either head-only fine-tuning or training from scratch. Experiments also show that the accuracy and stability effects depend on the task and training setup.

5 EXPERIMENTS

  • Experiments compare EqR and DT baselines with ImR fine-tuning; ImR predicts from the expectation of the prediction head over trajectory states, and stability is assessed per instance.Fine-tuning applies the task loss to the trajectory-averaged output, while baselines apply it to the final-state output.Source: §5 · p. 6

5.1 EQR ON MAZE-UNIQUE AND SUDOKU-EXTREME

  • With the puzzle prefix, scratch-trained ImR exceeds released EqR on both tasks at large depths, while ImR (FT) beats EqR (FT) only on Maze-Unique.For Sudoku-Extreme, EqR fine-tuning did not improve validation accuracy, so the original prediction head was retained.Source: Figure 2 · p. 7, §5.1 · p. 7
  • EqR evaluations cover Maze-Unique and Sudoku-Extreme, compare puzzle-prefix configurations, and classify trajectories by state residuals over the final 64 updates.The tasks require an exact predicted grid match for correctness.Source: §5.1 · p. 7
  • In non-fixed groups, ImR reduces instability; prefix-free scratch models average over 97% stable-correct predictions despite roughly 50% non-fixed rates on both tasks.Prefix-free models have higher mean non-fixed rates than those with the prefix, with substantial variation across seeds.Source: §5.1 · pp. 7–8, Table 1 · p. 8

5.2 DT ON OOD MAZES

  • DT fine-tuning and hyperparameter selection use ID 9×9 degree-1 mazes, while the 1,000-maze test set spans sizes from 9×9 to 99×99.All models solve degree-1 test mazes, including larger sizes than pretraining; the sole degree-4 test maze is unsolved by every evaluated model.Source: §5.2 · pp. 8–9
  • For degree-2/3 OOD mazes, ImR (FT) eventually surpasses released DT and finishes above DT (FT) in both groups.Its accuracy initially rises more slowly than released DT, with gains emerging at larger recurrent depths.Source: §5.2 · p. 8, Figure 3 · p. 9
  • Across both OOD degree groups, ImR (FT) reduces instability and raises stable-correct share versus DT (FT), although stable-incorrect share also rises, especially for degree 3.Fine-tuning and hyperparameter selection used ID mazes, and the recurrent block remained unchanged.Source: §5.2 · pp. 8–9, Table 2 · p. 9

6 DISCUSSION

Prefix-free ImR models use fewer latent state dimensions and show higher non-fixed rates, yet their non-fixed predictions are predominantly stable-correct. The experiments use fixed stopping depths, leaving adaptive stopping as future work.

6 DISCUSSION

  • Without the puzzle prefix, ImR models use fewer latent dimensions and have higher mean non-fixed rates than prefix-trained models.The passage suggests studying how latent trajectories and induced measures represent solutions across state dimensionalities to guide more efficient recurrent models.Source: §6 · p. 9
  • Despite those higher non-fixed rates, prefix-free models’ predictions in the non-fixed group are predominantly stable-correct.Source: §6 · p. 9
  • Experiments use fixed stopping depths; adaptive rules balancing predictive accuracy and computation remain future work.Source: §6 · p. 9

7 CONCLUSION

ImR stabilizes prediction by averaging the prediction head’s output under an invariant measure, allowing states to evolve without requiring fixed-point convergence. The appendix formalizes this view and notes that limiting measures can depend on initialization.

7 CONCLUSION

  • ImR’s invariant-measure predictions reduce instability and improve accuracy in many maze and Sudoku settings, including when latent states continue to evolve.The paper frames this as stable prediction without requiring individual states to converge to a fixed point.Source: §7 · p. 10

A.1 STOCHASTIC LATENT DYNAMICS

  • For stochastic recurrent dynamics, each update uses independently sampled Gaussian noise, and setting all noise samples to zero recovers the deterministic map.The input, model parameters, and noise distribution remain fixed throughout a trajectory.Source: §A.1 · p. 13

A.2 PUSHFORWARD AND KOOPMAN OPERATORS

  • An invariant measure is preserved by recurrent updates, so expectations of observables remain unchanged even when individual latent states change.The appendix describes state distributions through the pushforward operator and functions of states through its dual, the Koopman operator.Source: §A.2 · pp. 13–14

A AN OPERATOR-THEORETIC PERSPECTIVE ON LATENT REASONING

  • The operator-theoretic account identifies ImR’s prediction with the invariant component extracted from the prediction head, without requiring the head itself to be invariant.The appendix also relates pushforward and Koopman operators and distinguishes trajectory convergence from convergence across initializations.Source: A AN OPERATOR-THEORETIC PERSPECTIVE ON… · p. 12, §A.2 · p. 14

A.3 DEPENDENCE ON INITIALIZATION

  • The limiting invariant measure need not be independent of initialization, and some initial states may lead to subsets associated with incorrect predictions.Independence from initialization is described as a desirable property that may emerge through training, not an assumption.Source: §A.3 · p. 14

B EXPERIMENTAL AND TRAINING DETAILS

B.1 DATASETS AND VALIDATION

The appendix documents the EqR and DT datasets, validation splits, and generated DT maze tests spanning multiple sizes and start-cell degrees.

Datasets for EqR

  • Maze-Unique uses released 30 × 30 mazes with unique 100–140-move paths, while Sudoku-Extreme includes 1,001,000 augmented training examples from 1,000 base puzzles.Validation sets were constructed separately, and test sets were reserved for final evaluation.Source: §B.1 · p. 15

Datasets for DT

  • DT fine-tuning and validation use the same 50,000 9 × 9 ID mazes, split into 40,000 fine-tuning and 10,000 validation examples.The 80/20 split matches the original DT implementation.Source: §B.1 · p. 15
  • DT test mazes span sizes 9 × 9 through 99 × 99 and start-cell degrees 1–4, with 100 generated mazes at each size.The target is the unique path from start to goal; Figure 4 contrasts a 9 × 9 ID validation maze with 19 × 19 OOD examples.Source: §B.1 · p. 15, Figure 4 · p. 16
  • Generated mazes and paths are rendered as input images and target masks, with input dimensions H = W = 2m + 6 for an m × m maze.For example, a 9 × 9 maze produces a 24 × 24 input image.Source: §B.1 · p. 15

B.2 MODEL ARCHITECTURES AND THE NUMBER OF PARAMETERS

The experiments apply ImR to the existing EqR and DT architectures without changing their model architectures. EqR uses transformer or MLP-based recurrent blocks, while DT uses a convolutional recurrent network.

Model architectures and parameter counts

  • ImR uses the EqR and DT architectures unchanged, and Table 4 summarizes the architectures and parameter counts.The reported totals include input-embedding parameters under EqR’s recurrent block, EqR’s halting head, and DT’s latent-state initializer.Source: §B.2 · p. 16, Table 4 · p. 16

EqR

  • EqR uses a transformer for Maze-Unique and MLP-Mixer blocks for Sudoku-Extreme, with a shared recurrent block across internal updates and recurrent depths.The latent widths are 128 for Maze-Unique and 512 for Sudoku-Extreme; the puzzle prefix is optional.Source: §B.2 · pp. 16–17

DT

  • DT represents its latent state as a 128-channel spatial tensor and updates it using input concatenation, a convolution, and two residual blocks.One recurrent step comprises five convolutions; the shared convolutions preserve spatial resolution across evaluated maze sizes.Source: §B.2 · p. 17

Prediction heads

  • DT’s nonlinear prediction head maps each state to two pixel-class logits, so ImR applies the head separately to each state before averaging its outputs.By contrast, EqR’s linear prediction and halting heads can be applied after averaging the states in the empirical measure.Source: §B.2 · p. 17

B.3 TRAINING AND EVALUATION SETTINGS

The appendix specifies training and validation procedures for fine-tuning and training from scratch, then defines how trajectories and prediction stability are classified. Fine-tuning updates only the prediction head, while trajectory labels use changes over the final 64 recurrent updates.

B.3 TRAINING AND EVALUATION SETTINGS

  • Fine-tuning uses 384 recurrent updates; ImR averages prediction-head outputs over states 85–384, while the baseline predicts from state 384.Main evaluations include states from depth 1 through the stopping depth; fine-tuned heads are selected for ImR in all fine-tuning experiments.Source: §B.3 · p. 18
  • For validation, ImR predicts from the empirical-measure expectation over inner-loop states, whereas EqR predicts from the final state of that loop.Validation uses 16 outer steps, and evaluation selects the checkpoint with the highest validation accuracy.Source: §B.3 · p. 19
  • EqR fine-tuning did not improve Sudoku-Extreme validation accuracy, so the original prediction head was retained and no separate fine-tuned curve is shown.Fine-tuning changes only the prediction head; the other model parameters remain at their released values.Source: §B.3 · p. 19

B.4 CLASSIFICATION OF LATENT TRAJECTORIES

  • A trajectory is labeled fixed when relative changes remain small over all of the final 64 recurrent updates; otherwise it is non-fixed.This diagnostic does not guarantee convergence of the full latent state to a fixed point.Source: §B.4 · p. 20
  • Prediction-category proportions are computed separately within fixed and non-fixed groups over depths T − 63 through T, with group size as the denominator.For trained or fine-tuned configurations, the reported proportions are means and standard deviations across three seeds; released EqR models use one model per dataset.Source: §B.4 · p. 20

B.5 VISUALIZATION OF LATENT TRAJECTORIES AND PREDICTIONS

The appendix visualizes latent trajectories and reports final-depth accuracy. ImR trained from scratch outperforms EqR on both datasets, while fine-tuning results vary by task; trajectories continue changing across recurrent depths.

B.5 VISUALIZATION OF LATENT TRAJECTORIES AND PREDICTIONS

  • PCA projections of 513 states from a 512-update trajectory explain 77.0% of the variance in the first two components.The visualization excludes the puzzle prefix; PCA is applied only for display.Source: §B.5 · p. 20
  • The two-dimensional prediction boundary is approximate because omitted PCA components can change predictions at the same projected location.A point is classified as correct only when all 900 predicted maze cells match the target.Source: §B.5 · p. 21
  • At actual latent states, the original prediction head repeatedly alternates between correct and incorrect predictions, so the trajectory visits both regions beyond the PCA slice.Source: §B.5 · p. 21

C.1 ACCURACY AT THE FINAL RECURRENT DEPTH

  • For both prefix settings, scratch-trained ImR outperforms EqR on both datasets; fine-tuned ImR beats its baseline on Maze-Unique and both DT groups but trails on Sudoku-Extreme.Source: §C.1 · p. 21, Table 7 · p. 21

C.2 PROJECTIONS OF LATENT TRAJECTORIES

  • With noise set to zero, the projected Maze-Unique and Sudoku-Extreme trajectories continue changing across depths 897–1024.Source: §C.2 · p. 22, Figure 6 · p. 22

C.3 DT RESULTS BY MAZE SIZE

ImR (FT) shows reduced accuracy fluctuations at large recurrent depths across maze sizes.

C.3 DT RESULTS BY MAZE SIZE

  • ImR (FT) reduces accuracy fluctuations at large recurrent depths across maze sizes.Figure 7 reports means across three training seeds for DT (FT) and ImR (FT), with enlarged views of depths 4,981–5,000.Source: §C.3 · p. 22, Figure 7 · p. 23

Related papers

Loading 2610.10996v1…