We propose a new architectural change, and post-training pipeline, for making LLMs more verbose reasoners by teaching a model to truncate forward passes early. We augment an existing transformer architecture with an early-exit mechanism at intermediate layers and train the model to exit at shallower layers when the next token can be predicted without deep computation. After a calibration stage, we incentivise the model to exit as early as possible while maintaining task performance using reinforcement learning. We provide preliminary results to this effect for small reasoning models, showing that they learn to adaptively reduce computations across tokens. We predict that, applied at the right scale, our approach can minimise the amount of excess computation that reasoning models have at their disposal to perform non-myopic planning using their internal activations, reserving this only for difficult-to-predict tokens.
Figures & tables
Figure 1: Overview of the early exit architecture. (A) Calibration via self-distillation: early-exit heads at intermediate layers sample exit decisions from learned probability distributions. Exit probabilities are calibrated against a teacher (the full-depth model) where KL divergence between intermediate and final layer logits determines targets, trained alongside token predictions (Appendix A ) (B) Early exit incentivisation via RL: the model generates sequences using the learned exit mechanism, with task reward from output tokens and early exit reward from skipped layers. Layers above the exit point are skipped, and the frozen residual stream representation is passed directly to the unembedding layer.
Figure 2: Left: RL training dynamics on Qwen3-4B for the Theory of Mind ( Sclar et al., 2024 ) task. Each panel shows individual runs (light lines), mean across runs (bold line with markers), and ± 1 SE error bars. The dashed line indicates the pre-SFT base model performance. Step 0 represents the SFT-calibrated model before RL. Accuracy improves while average compute decreases, demonstrating that the model learns to exit earlier without sacrificing performance. Coherence remains stable throughout. Evaluated on 60 theory of mind prompts per step. Right: Token-level visualisation of early exit behaviour. Each token is coloured by the layer at which computation was terminated. The model adaptively varies computational depth, using fewer layers for predictable tokens and the full depth for more complex ones. We used λ=1.5 and β=0.25 .
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Distribution of coherence scores across models with different KL factors. 3 3 3 Model trained on KL factor=0.25 is excluded from this visualisation due to its outlier behaviour (scores were heavily concentrated around 1, compressing the scale for other models) but is included in Table 1 . As the KL factor increases, coherence distributions shift toward higher scores, approaching base model performance.
Model
Avg. Coherence
Early Exit Rate
Avg. Layer Savings
KL factor 0.25
1.1
59%
22.4%
KL factor 0.5
5.5
49%
8.1%
KL factor 1.0
7.5
47%
6.4%
KL factor 2.0
8.4
43%
5.3%
KL factor 4.0
8.5
37%
4.6%
Base Model
8.9
0%
0.0%
Appendix
Table 1: Effect of KL factor on coherence, early exit rate, and average layer savings. Aggressive early exits (KL factor=0.25) achieve the highest computational savings but at a large cost to coherence. KL factor=1.0 balances reasonable computational savings with strong coherence.
Figure 4: Distribution of exit layers. The model’s learned exit distribution (orange bars) closely follows the target exit probability distribution from training data (grey bars), demonstrating effective learning of exit behaviour during Stage 1 calibration.
Complex problems, whether in math, logic, or planning, are solved by humans through a sequence of steps where the result of one step informs the next. In this work, we adopt the perspective that the reasoning power of Transformers is fundamentally limited by a fixed maximum number of steps along any latent path of computation. To address this, we introduce Turbo Connection (TurboConn), a novel architecture that overcomes the fixed-depth constraint by routing multiple residual connections from the higher-layer hidden states of each token t to the lower layers of token t+1. Fine-tuning pre-trained LLMs with our method not only yields accuracy gains of 0.9% to over 10% on benchmarks like GSM8K, Parity, and multi-step arithmetic, but also demonstrates that the density of these backward connections is critical; our dense interaction significantly outperforms "sparse" alternatives that only pass a single hidden state or vector. Notably, TurboConn can be integrated into pre-trained LLMs to overcome task-specific plateaus: while a fine-tuned Qwen-3-1.7B achieves only 53.78% on Parity, adding our architectural modification enables the model to reach 100% accuracy, all without the necessity to retrain the full model from scratch or sophisticated curriculum learning. Our results provide strong empirical evidence that the depth of the computational path is a key factor in reasoning ability, also offering a new mechanism to enhance LLMs without significantly affecting generation latency.
Reasoning in Large Language Models incurs significant inference-time compute, yet the token-level information structure of reasoning traces remains underexplored. We observe that reasoning tokens split into two functional types: low-entropy structural tokens (recurring phrases that scaffold the reasoning process) and higher-entropy organic tokens (problem-specific content that drives toward a solution). This asymmetry motivates a simple, model-agnostic compression pipeline: apply cross-word BPE merges on a model's own reasoning traces to derive supertokens that capture frequent structural patterns, then teach the model to adopt them via supervised fine-tuning. Across three model families and five mathematical reasoning benchmarks, our approach shortens reasoning traces by 8.1% on average with no statistically significant accuracy loss on any model--benchmark pair. Beyond compression, supertokens act as interpretable reasoning-move annotations (backtracking, verification, strategy shifts), exposing the model's high-level strategy at a glance. Analyzing transitions between structural categories reveals systematic differences between correct and incorrect traces: correct traces show productive recovery (backtracking followed by strategy shifts and verification), while incorrect traces are dominated by confusion cycles (repeated hedging and unresolved contradictions). These diagnostic signals suggest applications in reward shaping and early stopping for RL-based reasoning training.
Transformer reasoning is limited by autoregressive decoding, which repeat edly compresses rich hidden computation through token space and makes it difficult for intermediate reasoning states to persist across time. We in troduce Transformers with Temporal Middle-Layer Recurrence (T2MLR), a transformers-based latent reasoning architecture that fuses a cached middle layer representation from the previous token directly into an earlier layer of the current token position, enabling abstract intermediate computation to persist across decoding steps with little inference overhead. Across natural-language pretraining and multi-hop reasoning finetuning, T2MLR consistently outperforms data- and parameter-matched Transformer base lines. Moreover, applying recurrence to only a localized middle-layer block (as little as 20% of the network) often outperforms full-layer recurrence. Im portantly, T2MLR does not require pretraining from scratch: retrofitting the recurrent pathway into an existing pretrained 1.7B Transformer and briefly finetuning substantially improves math reasoning, lowering the barrier to practical adoption. These results suggest that effective latent reasoning in Transformers does not require looping over all layers as in previous works, but can instead emerge more strongly from targeted middle-layer recurrence.
Ziyang Cai, Xingyu Zhu, Yihe Dong +2
Princeton Language and Intelligence, Department of Computer Science Princeton University