Suppressing Final Layer Hidden State Jumps in Transformer Pretraining

TL;DR

The jump-suppressing regularizer (JREG) is proposed which penalizes this jump during pre-training, thereby encouraging more balanced capability usage across the middle layers, and 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.

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.

Figures (15)

• Figure 1: Many recent open-weight language models exhibit minimal hidden state displacement in their middle layers, but exhibit a pronounced “jump” at their final layer. Suppressing this jump during pre-training could improve overall performance by fostering a more balanced use of capabilities across the middle layers.
• Figure 2: Layer-wise hidden state displacement $\Psi_\ell$ for next-word prediction on 100 samples from the LAMBADA dataset. Across all model architectures, the displacement at the final layer tends to be larger than that of the middle layers.
• Figure 3: Analysis of checkpoint-wise hidden state displacement using fine-grained Pythia and OLMo pre-training checkpoints, revealing that as training progresses, the final layer exhibits large "jump" displacements.
• Figure 4: Hidden state displacement for the 170M Llama-based model pre-trained on 100B tokens with two different update steps. (light blue): 20K steps, (dark blue) 200K steps. Jump rates were $\zeta_{L}=1.93$ for 20K steps and $\zeta_{L}=7.24$ for 200K steps. Despite being trained on the same number of tokens, models with more update steps exhibit larger jumps in hidden state displacement.
• Figure 5: Weight coefficients $w_{\ell}$ for different values of the hyperparameter $\alpha$ in \ref{['eq:L_cos']}. When $\alpha=0.0$, the weight coefficients are uniform, when $\alpha=0.1$, they are almost linear, and for larger $\alpha$, they increase exponentially.

Theorems & Definitions (2)

• Definition 1: Displacement between the input and output hidden state vectors
• Definition 2: Jump rate