English

Suppressing Final Layer Hidden State Jumps in Transformer Pretraining

Computation and Language 2026-01-27 v1

Abstract

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.

Keywords

Cite

@article{arxiv.2601.18302,
  title  = {Suppressing Final Layer Hidden State Jumps in Transformer Pretraining},
  author = {Keigo Shibata and Kazuki Yano and Ryosuke Takahashi and Jaesung Lee and Wataru Ikeda and Jun Suzuki},
  journal= {arXiv preprint arXiv:2601.18302},
  year   = {2026}
}

Comments

Accepted to the Findings of EACL 2026