Extending Reasoning in Typed Decision Models via Looping
Authors: David Berghaus, Arber Zela, Behzad Shomali, Markus Frey, Mehdi Ali, Rafet Sifa, Volkan Cevher
Organizations: Lamarr Institute for Machine Learning and Artificial Intelligence, Bonn, Germany · Fraunhofer IAIS, Sankt Augustin, Germany · EPFL, Lausanne, Switzerland · University of Bonn, Bonn, Germany
Typed decision models choose among a fixed set of answer options in a single forward pass, so unlike generative language models they cannot spend more test-time compute on hard inputs by generating more tokens. Looping a block of layers offers this compute in latent space without adding parameters, but prior work found no gain from loops retrofitted into conventionally pretrained models. We show that such loops help when trained together with the decision model. Our method, RELOAD (retrofitted loops for adaptive decisions), fine-tunes a pretrained encoder or decoder with a damped loop over its middle block and a lightweight controller picks the number of passes adaptively for each input. Across six models from 0.4B to 4B parameters, looping improves both accuracy and calibration, mostly on problems that need longer chains of reasoning than any seen in training, where encoders gain up to 23 points. Neither matched training compute nor more data closes this gap, and separate layers for each pass come close only with nearly twice the parameters. The controller keeps most of the gain at about three quarters of the compute.
Figures & tables
Figure 1: Overview of RELOAD. (a) A middle block G of a pretrained model is applied K times with a damped update, and the hidden state at each answer option is scored. One model is trained on all budgets K∈{1,2,3,4} . (b) A controller picks the number of passes for each request from the prefix output h0 , before the loop runs. Bars show the share of StepGame test requests at each chosen K , by the number of statements that must be chained (Von, three seeds). (c) Top: StepGame test accuracy by chain length, for the same model. Bottom: mean number of passes per request. The controller matches the accuracy of four passes with 3.1 passes on average and saves most on short chains, where extra passes add little. Training problems need at most five statements.
In-distribution
Longer reasoning
Out-of-distribution
Model
Variant
Acc.
Brier
Acc.
Brier
Acc.
Brier
Layers
Single pass
92.8 ± 0.3
0.101 ± 0.004
70.6 ± 1.3
0.404 ± 0.025
50.9 ± 0.2
0.728 ± 0.005
28
RELOAD + controller
93.6 ± 0.2
0.088 ± 0.002
85.4 ± 1.0
0.203 ± 0.014
53.8 ± 1.0
0.682 ± 0.008
35/53/39
Von (395M)
RELOAD ( K=4 )
93.9 ± 0.2
0.085 ± 0.002
85.7 ± 0.4
0.198 ± 0.009
54.1 ± 0.7
0.675 ± 0.007
55
Single pass
93.0 ± 0.4
0.100 ± 0.003
72.3 ± 0.6
0.384 ± 0.008
47.9 ± 0.2
0.714 ± 0.015
28
RELOAD + controller
93.9 ± 0.2
0.087 ± 0.001
87.0 ± 1.7
0.183 ± 0.027
50.0 ± 0.6
0.696 ± 0.013
35/52/37
Table 1: Accuracy (%) and Brier score (lower is better), mean ± standard deviation over three seeds. In-distribution: held-out ProofWriter, StepGame with 1–5 statements and Kev decision. Longer reasoning: StepGame with 6–10 statements. Out-of-distribution: Kev transfer and eight unseen benchmarks, zero-shot. Qwen3 models use Kev’s encoding. RELOAD ( K=k ): one loop-trained checkpoint run with k passes; the single pass is trained separately. Every budget, including K=1 : Appendix Table 4 ; per-set results: Appendix A . Layers: layer applications per question; for the controller ( λ=0.005 , fitted on development data) per group (in-distribution / longer / out-of-distribution). Bold: best per model and column.
Figure 2: Reasoning horizon on generated StepGame stories whose statements all lie on the chain (three seeds; not used for training or model selection; training chains have at most five statements). (a) Von’s accuracy by chain length. Band: 95% bootstrap interval for K=4 . Chance: always answering the most frequent label. Dashed lines: horizon of the single pass and of K=4 , where a fitted sigmoid lies halfway between ceiling and chance (Appendix F ). (b) Horizon of the single pass (open circle) and of RELOAD at K=4 (filled circle) and K=8 (diamond, hidden for Qwen3 where it coincides with K=4 ). Gray bar: single pass to K=8 ; whiskers: standard deviation over seeds (single pass, K=4 ); right: gain of K=4 over the single pass.
Figure 3: Accuracy versus compute (Von, three seeds): mean test accuracy over ProofWriter, StepGame, Kev decision and Kev transfer against layer applications per question. Blue: RELOAD at fixed budgets (shaded: beyond the trained budgets). Green: RELOAD with the controller as λ varies, and with a controller that sees only input length and option count. Gray: the single pass, and the same model looped without loop training.
Figure 4: Test-time integrators for the trained loop of Von, without retraining (three seeds). K counts evaluations of G : Euler takes K steps of size 1/K ; the midpoint method and fourth-order Runge–Kutta (RK4) take two steps of two and four evaluations. Left: longer-reasoning accuracy, with error bars of one standard deviation over seeds; shaded: more passes than in training. Right: accuracy (%) and Brier score, mean ± standard deviation; in and out of distribution all Brier scores lie within 0.002 of each other. Last row: paired midpoint − Euler difference at K=4 with 95% interval. Both higher-order configurations were fixed on development data before testing.
Figure 5: Where the trained loop writes (Von, seed 1, StepGame, four passes). (a) Accuracy when each pass updates only one token group, h←h+41m⊙(G(h)−h) ; the block still reads all tokens, and the others keep their state from before the loop. Stories have a unique chain and at least one distractor; shared tokens (question, story, separators) do not attend to the options. (b) Accuracy change relative to the ordinary loop when attention edges (query to key) are removed in the block, or when four passes over the shared tokens precede four over the options (first three rows: all stories; last two: parsed stories). (c) Shapley values of the accuracy gain over token groups by chain length; per bin they sum, with a small share of the special tokens, to the gain of the full update over no update.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Variant
Params
ProofWriter
StepGame
Kev dec.
Kev transfer
Macro
Von
Single pass
395M
98.88
82.23
85.61
76.37
85.77
RELOAD ( K=4 )
395M
99.71
90.88
86.03
77.49
88.53
Depth up-scaling
726M
99.68
90.38
86.14
79.78
88.99
Ettin-400M
Single pass
395M
98.56
82.90
86.92
68.29
84.17
RELOAD ( K=4 )
395M
99.56
91.45
86.67
70.68
87.09
Depth up-scaling
726M
99.51
91.77
86.47
73.63
87.84
Appendix
Table 10: RELOAD versus depth up-scaling ( Kim et al., 2024 ) at equal compute ( K=4 ; held-out test accuracy, %, three seeds). Depth up-scaling gives passes 2–4 their own trained copies of the block.
Von
Accuracy (%)
Layers
RELOAD + controller, λ=0.005
88.16
37.1
RELOAD + controller, λ=0.002
88.35
40.0
Halting, τ=0.2
88.12
59.0
Halting, τ=0.1
88.37
62.0
Halting, c=0.95
88.19
57.4
Halting, c=0.99
88.42
67.6
Appendix
Table 13: Controller baselines (three seeds). Top: controller ( λ=0.005 ) minus a mix of the neighboring fixed budgets at the same mean number of passes, paired accuracy difference in points with 95% intervals. Bottom: Von, mean accuracy over ProofWriter, StepGame, Kev decision and Kev transfer, and layer applications per question, for the controller and for post-hoc halting rules (total variation τ between consecutive budgets, or top probability above c ).
3 swaps
6 swaps
8 swaps
Model
State
Answer
Probe
Answer
Probe
Answer
Probe
Single pass
100.0
95.4
76.5
40.1
44.2
27.8
RELOAD ( K=4 ), after pass 1
99.1
92.9
56.8
35.0
34.1
25.4
RELOAD ( K=4 ), after pass 2
100.0
99.3
80.7
58.2
46.2
34.0
RELOAD ( K=4 ), after pass 3
100.0
99.8
94.8
80.6
63.1
51.2
Von
RELOAD ( K=4 ), after pass 4
100.0
99.9
98.3
92.1
74.6
66.2
Appendix
Table 15: Object tracking with 3, 6 and 8 swaps (training covers up to 4; three seeds). Answer: accuracy of the answer read out after each pass of a four-pass run (for the single pass, its normal output). Probe: accuracy of a linear probe on the option-marker states at the output of the looped block that predicts the object’s final location (chance 25%).
K=1
2
3
4
6
α
RELOAD, answer (log-probability)
1.77
0.78
0.42
0.25
0.08
1.01
RELOAD, option-marker states
0.355
0.132
0.070
0.041
0.014
1.07
Looped only at inference, option-marker states
0.376
0.129
0.067
0.040
0.013
1.10
Appendix
Table 16: Distance from budget 8 at budget K and fitted exponent α (Von, mean over three seeds).
Model
Single pass
RELOAD ( K=4 )
RELOAD ( K=8 )
Von
18.9
4.1
2.9
Ettin-400M
17.0
3.2
2.0
Ettin-1B
16.1
2.9
0.9
Qwen3-0.6B
19.2
19.3
19.4
Qwen3-1.7B
23.4
22.9
22.6
Qwen3-4B
21.6
19.6
19.3
Appendix
Table 18: Accuracy lost (points) when five statements off the chain are added to the generated StepGame stories, mean over chain lengths 1–5 (three seeds).
Figure 6: Input injection for the four variants run on all five models, relative to each model’s plain loop at K=4 . Left: macro test accuracy; small points are seeds, large points their mean, and the band is the plain-loop seed spread ( ±0.65 ). Other panels: paired differences by group with 95% bootstrap intervals.
Model
∥h4−e∥/∥e∥
∥h4∥/∥e∥
cos(h4,e)
Ettin-400M
2.28
2.92
0.78
Qwen3-0.6B
2.42
3.03
0.71
Appendix
Table 24: Top: state drift of the plain loop at K=4 ( e is the prefix output; mean over seeds and all tokens). Bottom: size of the write relative to the state, and change in development accuracy at K=4 when the trained write is switched off at inference (rounds 1–3; range over models and variants).
Figure 7: Integration time of the trained loop (Von, seed 1, frozen weights). (a) Test accuracy along the four-pass trajectory ( Δt=1/4 ), continued past the trained end point T=1 (shaded); StepGame is split by story length. (b) Generated tracking problems with 5 to 12 swaps along the trajectory with step 1/16 ; diamonds mark the best tested time. (c) StepGame accuracy past T=1 with the ordinary readout and with every token rescaled to its norm at T=1 before the readout; numbers give the mean token norm relative to T=1 .
Figure 8: Shapley decomposition of the loop’s accuracy gain over token groups for every test set (Von, seed 1), with two passes (open bars) and four passes (filled bars, labeled). The value of a set of groups is the accuracy when only these groups are updated, so the values add up to the gain of the loop over no update, shown above each panel. Other: the remaining special tokens. Vertical scales differ between panels. ProofWriter uses 4,000 sampled test questions, and tracking 200 problems per depth.
Figure 9: Sentence attributions and the controller (Von, seed 1). (a) AUROC with which sentence-deletion Shapley values of the log-probability of the correct option separate relevant sentences from the others, for K passes and for the four-pass trajectory continued to T=2 . Relevant sentences are chain statements, or swaps whose removal alone changes the answer. (b) Mean number of passes chosen by the controller of Table 1 on StepGame stories with a unique chain, by story and chain length; empty cells cannot occur. (c) Shapley value per sentence of the controller’s predicted gain P^4−P^1 for chain statements and distractors, by chain length.
Figure 10: Controlled edits and answer sensitivity (Von, seed 1). (a) Change in StepGame test accuracy when an off-chain sentence is attached to a queried object or to an object inside the chain, relative to the same sentence with a letter that occurs nowhere else in the story. Solid: test-story sentences that mention a queried object; dashed: edited sentences that mention none; bars: 95% intervals. (b) Sensitivity of the margin between the correct and the strongest other option to a perturbation after pass k of four, in the metric of Figure 11 and relative to k=4 (medians and interquartile ranges over 60 StepGame stories and 20 tracking problems), with the reference (3/4)4−k for a block that ignores the perturbation.
Figure 11: Matched perturbations of the loop state (Von, seed 1, StepGame, four passes). (a) Accuracy when the same Gaussian perturbation, scaled by the per-feature standard deviation of the state before the loop, is added after pass k . (b) Norm of the resulting deviation of the final state relative to the inserted perturbation, with the reference (3/4)4−k for a block that ignores the perturbation.
Decision models answer typed questions about an input with one probability per option, computed in a single forward pass without generating text. Like the fast System 1 of dual-process theory, they answer at once, but their depth is fixed at training time, so they spend the same computation on easy and hard questions. Models that reason in generated text can think longer, but they pay for every extra step in decoded tokens. We present LEV, a decision model that thinks deeper by rerunning its layers in a loop instead of generating tokens. LEV builds on a language model pretrained as a loop, reads the options after every loop step and trains all steps jointly, so one checkpoint can answer at any depth. On three looped backbones, accuracy rises with depth on the same weights, mostly on hard questions, and a 2.6B LEV goes from 63.2% at depth one to 76.2% at depth four on JevBench. On-demand depth, set per question by a small stopper with a threshold fitted on unlabelled requests, saves 15% to 24% of the computation at a cost of under one accuracy point, level with or slightly above a fixed mix of depths. On a locked test set, the 1.4B and 3B models are 5.4 to 8.8 points more accurate than dense (non-looped) models and cascades of the same latency below 50 ms.
Dingjie Song, Zhiling Yan, Yuxuan Zhang +1
Lehigh University · University of British Columbia · Vector Institute
Typed decision models answer a declared question without generating text: a decision head returns a probability for each of the declared options in a single forward pass. A single pass is fast, intuitive System 1 thinking. We study what lies between one pass and generated reasoning: looping, in which the same layers are recursively applied several times before one typed readout. Each loop lets the model revise its hidden state before it commits to an answer, without generating a token; we call this System 1.5 thinking. We propose SanSi, which turns a pre-trained looped language model into a typed decision model. The option probabilities are read after every loop, and every loop is trained with a proper scoring rule, so that one model serves every budget from one loop to eight in a single run. On 10,027 test decisions from 59 sources, SanSi reaches 72.0% accuracy: 13.5 points above a non-looped model of the same shape trained with the same recipe, 5.3 points above a newer non-looped model of its size, and 1.8 points below one with three times the parameters. On two depth-controlled tasks, loops extend the solvable depth beyond the depths seen in training, where the larger single-pass model fails. Used as the judge for policy optimization with reinforcement learning, without gold answers, SanSi raises the share of the generator's correct answers from 37% to 52%.
Looped Language Models (LoopLMs) perform "latent reasoning" by recursively refining internal latent representations with shared weights, offering a more effective alternative to explicit verbal reasoning. Despite their effectiveness, we find that LoopLMs remain prone to loop instability: unstable refinement across iterations can produce localized uncertain "hard" tokens associated with reasoning errors. To address this, we propose LoopCD, loop-wise contrastive decoding that enhances the reasoning performance of LoopLMs by intervening on these tokens at inference time. Specifically, we exploit the internal dynamics of LoopLMs and contrast the logits from earlier iterations with logits from the last refined iteration to form the final sampling distribution. We find that this strategy is highly efficient, introducing only negligible inference overhead and requiring no additional training, while effectively improving reasoning performance by naturally refining reasoning-critical hard tokens. Extensive experiments show that our method improves the performance of recent representative LoopLMs across various reasoning tasks.
Byeongho Yu, Junhyuk So, Eunhyeok Park
Department of Computer Science and Engineering Pohang University of Science and Technology (POSTECH) · Graduate School of Artificial Intelligence Pohang University of Science and Technology (POSTECH)