# Suppressing Final Layer Hidden State Jumps in Transformer Pretraining

arXiv: https://arxiv.org/abs/2601.18302 · HTML: https://mlsanity.com/paper/2601.18302/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 52 references that are in the map.
- First posted on arXiv 26 Jan 2026.

## Builds on (most cited first, 30 of 40)

- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [WizardLM: Empowering large pre-trained language models to follow complex instructions](https://mlsanity.com/paper/2304.12244/) (2023)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [RACE: Large-scale ReAding Comprehension Dataset From Examinations](https://mlsanity.com/paper/1704.04683/) (2017)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Eliciting Latent Predictions from Transformers with the Tuned Lens](https://mlsanity.com/paper/2303.08112/) (2023)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [How Far Can Camels Go? Exploring the State of Instruction Tuning on Open Resources](https://mlsanity.com/paper/2306.04751/) (2023)
- [ShortGPT: Layers in Large Language Models are More Redundant Than You Expect](https://mlsanity.com/paper/2403.03853/) (2024)
- [SWAG: A Large-Scale Adversarial Dataset for Grounded Commonsense Inference](https://mlsanity.com/paper/1808.05326/) (2018)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding](https://mlsanity.com/paper/2404.16710/) (2024)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [FastBERT: a Self-distilling BERT with Adaptive Inference Time](https://mlsanity.com/paper/2004.02178/) (2020)
- [The Remarkable Robustness of LLMs: Stages of Inference?](https://mlsanity.com/paper/2406.19384/) (2024)
- [On the Effect of Dropping Layers of Pre-trained Transformer Models](https://mlsanity.com/paper/2004.03844/) (2020)
- [Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping](https://mlsanity.com/paper/2010.13369/) (2020)
- [Organize the Web: Constructing Domains Enhances Pre-Training Data Curation](https://mlsanity.com/paper/2502.10341/) (2025)

## Built on by (most cited first, 3 of 3)

- [LoopUS: Recasting Pretrained LLMs into Looped Latent Refinement Models](https://mlsanity.com/paper/2605.11011/) (2026)
- [Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers](https://mlsanity.com/paper/2605.07772/) (2026)
- [Before It Fades: Reinforcing Temporal Representations at Inference Time in VideoLLMs](https://mlsanity.com/paper/2610.01595/) (2026)

## Abstract (arXiv)

This paper discusses the internal behavior of Transformer language models. Many recent pre-trained models have been reported to exhibit only slight changes in the angular distance between the input and output hidden state vectors in the middle Transformer layers, despite a disproportionately large ``jump'' in the angular distance occurring in or around the final Transformer layer. To characterize this, we first introduce a quantitative metric for the jump strength around the final layer, and then demonstrate its prevalence across many open-weight models, as well as its amplification throughout pre-training. Assuming such jumps indicate an undesirable property, we propose the jump-suppressing regularizer (JREG) which penalizes this jump during pre-training, thereby encouraging more balanced capability usage across the middle layers. Empirical evaluations of three model sizes of Llama-based models, trained with the proposed JREG method, reveal improved task performance compared to the baseline without altering the model architecture.