Steering Recurrent Reasoners at Inference Time with Readout Feedback
Organizations: Graduate School of Engineering, The University of Tokyo
Abstract
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 and a small gain in one. The four positive pairs achieve 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.
Figures & tables
| Model | Hierarchy | Latent Space | Initial Value | Input Injection |
|---|---|---|---|---|
| AKOrN | single-level | Random, unlearned | Every step | |
| ItrSA++ | two-level | Random, unlearned | Every steps | |
| TRM | two-level | Fixed, learned | Every step |
Appendix figures & tables31 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | |||||||
|---|---|---|---|---|---|---|---|
| Sudoku | 500 | 30 | 4 | ||||
| Maze | 100 | 30 | 2 |
| Task Model | best val acc | metric | ||||
|---|---|---|---|---|---|---|
| AKOrN Sudoku | 1.949 | 0.281 | 16 | 1.552 | 100.0 | board accuracy |
| AKOrN Maze | 0.394 | 0.195 | 128 | 0.141 | 75.0 | shortest path accuracy |
| ItrSA++ Sudoku | 0.028 | 0.012 | 64 | 0.461 | 70.6 | board accuracy |
| ItrSA++ Maze | 0.165 | 0.327 | 104 | 0.028 | 85.0 | shortest path accuracy |
| TRM Sudoku | 0.829 | 0.182 | 88 | 1.462 | 74.8 | board accuracy |
| TRM Maze | 0.167 | 0.216 | 88 | 0.939 | 91.0 | shortest path accuracy |
| Model-task pair | ||||||||
|---|---|---|---|---|---|---|---|---|
| AKOrN Sudoku | 256 | 256 | 16 | 500 | 20,000 | |||
| AKOrN Maze | 256 | 256 | 128 | 100 | 1,000 | |||
| ItrSA++ Sudoku | 256 | 256 | 64 | 500 | 20,000 | |||
| ItrSA++ Maze | 256 | 256 | 104 | 100 | 1,000 | |||
| TRM Sudoku | 16 | 288 | 88 | 500 | 20,000 | |||
| TRM Maze | 16 | 192 | 88 | 100 | 1,000 |
| Model | ( ) | ( ) | ( ) | |||||||
| Base | RoFB | Base | RoFB | Base | RoFB | |||||
| AKOrN | 1 | 63.3 | 59.8 | 74.4 | 73.9 | 79.1 | 81.7 | |||
| 2 | 70.6 | 69.7 | 79.2 | 81.8 | 82.4 | 87.5 | ||||
| 4 | 76.3 | 78.3 | 82.9 | 87.9 | 85.1 | 91.5 | ||||
| 8 | 81.2 | 85.1 | 85.9 | 91.8 | 87.5 | 94.2 | ||||
| 16 | 85.3 | 90.4 | 88.7 | 94.8 | 89.6 | 96.2 | ||||
| Model | ( ) | ( ) | ( ) | |||||||
| Base | RoFB | Base | RoFB | Base | RoFB | |||||
| AKOrN | 1 | 69.8 | 69.8 | 72.2 | 72.2 | 72.6 | 75.8 | |||
| 2 | 71.6 | 71.6 | 73.4 | 73.4 | 73.7 | 76.5 | ||||
| 4 | 73.1 | 73.1 | 74.0 | 74.0 | 74.1 | 77.2 | ||||
| 8 | 73.8 | 73.8 | 74.0 | 74.0 | 73.9 | 76.9 | ||||
| 16 | 73.8 | 73.8 | 74.3 | 74.3 | 74.3 | 77.2 | ||||
| Model on Task | Baseline | RoFB | 95% CI of | min – max | Seeds with | ||
| AKOrN on Sudoku | 64 | 76.5 0.2 | 78.2 0.1 | 0.23 | – | – | 8/8 |
| 128 | 82.9 0.1 | 87.5 0.1 | 0.17 | – | – | 8/8 | |
| 256 | 85.1 0.1 | 91.5 0.1 | 0.19 | – | 8/8 | ||
| AKOrN on Maze | 64 | 72.4 0.5 | 72.4 0.5 | 0.00 | – | – | 0/8 |
| 128 | 73.8 0.5 | 73.8 0.5 | 0.00 | – | – | 0/8 | |
| 256 | 74.2 0.4 | 77.0 0.3 | 0.25 | – | 8/8 |
| Parameter | Five validation values relative to the selected value | Test values |
|---|---|---|
| Feedback strength | ||
| Start time | ||
| Entropy threshold | ||
| Gate temperature |
| Geometry measure | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku: all instances | |||
| Change in | |||
| Change in silhouette | |||
| Maze: all instances | |||
| Change in | |||
| Change in silhouette | |||
| Geometry measure | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku: incorrect under both conditions | |||
| Number of instances | 271 | 566 | 477 |
| Change in | |||
| Change in silhouette | |||
| Maze: incorrect under both conditions | |||
| Number of instances | 209 | 140 | 72 |
| Intervention | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku | |||
| Removal magnitude-matched RoFB | |||
| Magnitude-restored removal RoFB | |||
| Maze | |||
| Removal magnitude-matched RoFB | |||
| Magnitude-restored removal RoFB | |||
| Alternative | Method-specific parameter searched jointly with |
|---|---|
| Margin repulsion | Distance margin |
| Probability-weighted contrastive loss | Contrastive temperature ; separate from the fixed gate temperature |
| Within-class variance plus centroid margin | Centroid-distance margin ; both loss coefficients remain one |
| K-means centroid attraction | Cluster count on Sudoku or on Maze |
| Intervention | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku | |||
| Original RoFB ( ) | |||
| Attractive ( ) | |||
| Mixed ( ) | |||
| Const ( ) | |||
| Maze | |||
| Intervention | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| (a) Candidate-selected accuracy (%) | |||
| Sudoku | |||
| Original RoFB | 93.2 | 77.0 | 71.0 |
| Warmup | 93.8 | 78.0 | 72.0 |
| Entropy gate | 89.0 † | 77.0 | 69.4 |
| Both | 82.0 | 76.2 | 63.0 |
| Condition / difference | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku | |||
| No-force reference | |||
| RoFB | |||
| RoFB reference | |||
| Maze | |||
| No-force reference | |||
| Intervention | AKOrN | ItrSA++ | TRM |
|---|---|---|---|
| Sudoku | |||
| Fixed-scale entropy-gradient intervention | |||
| Magnitude-matched entropy-gradient intervention | |||
| Online RoFB | |||
| Maze | |||
| Fixed-scale entropy-gradient intervention | |||
| Architecture | Natural frequencies | Gain for each training seed | Mean SD | Positive seeds |
|---|---|---|---|---|
| Modified | Off | 2/3 | ||
| Modified | On | 3/3 | ||
| Original (reproduced) | Off | 0/3 | ||
| Original (reproduced) | On | 3/3 |
| pass@1 | pass@5 | pass@100 | ||||
|---|---|---|---|---|---|---|
| Aug. | Baseline | RoFB | Baseline | RoFB | Baseline | RoFB |
| 1 | 35.333 | 35.500 | 37.833 | 37.500 | 37.833 | 37.500 |
| 3 | 38.833 | 39.000 | 43.500 | 43.167 | 43.500 | 43.167 |
| 7 | 39.333 | 39.000 | 46.667 | 46.667 | 47.333 | 47.000 |
| 20 | 41.000 | 40.667 | 49.500 | 49.500 | 51.333 | 51.333 |
| 100 | 41.667 | 41.667 | 50.167 | 50.167 | 55.000 | 55.000 |
| Sudoku Extreme | Maze Hard | |
| Architecture | ||
| oscillator dim | 4 | 4 |
| embed width | 512 | 512 |
| # blocks | 1 | 1 |
| # heads | 8 | 8 |
| (step size) | 1.0 | 1.0 |
| Sudoku Extreme | Maze Hard | |
| Architecture | ||
| embed_dim | 384 | 384 |
| num_heads | 12 | 12 |
| num_layers | 1 | 1 |
| num_iter ( ) | 16 | 32 |
| num_rep_attn ( ) | 4 | 4 |
| Sudoku-Extreme | Maze-30x30-Hard | |
| Architecture | ||
| 3 | 3 | |
| 6 | 4 | |
| 0 (shared with ) | 0 (shared with ) | |
| 2 | 2 | |
| hidden_size | 512 | 512 |