cs.AIMar 22, 2026

A transformer architecture alteration to incentivise externalised reasoning

Authors: Elizabeth Pavlova, Mariia Koroliuk, Karthik Viswanathan, Cameron Tice, Edward James Young, Puria Radmard

Organizations: MARS · Geodesic Research, equal contribution · University of Amsterdam

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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 ±\pm 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\lambda=1.5 and β=0.25\beta=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.
ModelAvg. CoherenceEarly Exit RateAvg. Layer Savings
KL factor 0.251.159%22.4%
KL factor 0.55.549%8.1%
KL factor 1.07.547%6.4%
KL factor 2.08.443%5.3%
KL factor 4.08.537%4.6%
Base Model8.90%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.
More in this topic

Efficient Language Model Inference

Efficient language model inference reduces the compute, memory, latency, and energy needed to generate useful outputs.

See all 317 papers

Explore similar work

CardsList

  1. Turbo Connection: Reasoning as Information Flow from Higher to Lower Layers

    Feb 20, 2026Mohan Tang, Sidi LuTransformerLarge Reasoning Models

    UCLA.

  2. Shorthand for Thought: Compressing LLM Reasoning via Entropy-Guided Supertokens

    Apr 29, 2026Zhenyu Zhao, Sander Land, Daniel M. Bikel +1LLM CompressionLLM Interpretability

    Writer, Inc.

  3. T^2MLR: Transformer with Temporal Middle-Layer Recurrence

    Jul 16, 2026Ziyang Cai, Xingyu Zhu, Yihe Dong +2Recurrent TransformersLLM Reasoning

    Princeton Language and Intelligence, Department of Computer Science Princeton University