Table of Contents
Fetching ...

Early-stopping for Transformer model training

Jing He, Hua Jiang, Cheng Li, Siqian Xin, Shuzhen Yang

TL;DR

This work tackles the challenge of costly Transformer training and overfitting by introducing a validation-free early-stopping criterion based on Random Matrix Theory. By focusing on the spectral tail of the first encoder self-attention Value matrix, en.0.s.a.V, and fitting it with a Power-Law, the authors identify a three-phase training dynamics and define a KS-based heavy-tailed metric with threshold $d^* = \frac{C}{\sqrt{n_{\text{tail}}}}$. The method triggers stopping at the epoch that maximizes $d^* - \tilde{d}$, achieving substantial training-time reductions (e.g., stopping around epoch 100 for several setups) with minimal loss-accuracy degradation compared to longer runs. This spectral-based, validation-set-free approach advances understanding of implicit regularization in Transformers and offers a practical pathway to more efficient LLM training, with potential generalization to other large models.

Abstract

This work, based on Random Matrix Theory (RMT), introduces a novel early-stopping strategy for Transformer training dynamics. Utilizing the Power Law (PL) fit to tansformer attention matrices as a probe, we demarcate training into three stages: structural exploration, heavy-tailed structure stabilization, and convergence saturation. Empirically, we observe that the spectral density of the shallow self-attention matrix $V$ consistently evolves into a heavy-tailed distribution. Crucially, we propose two consistent and validation-set-free criteria: a quantitative metric for heavy-tailed dynamics and a novel spectral signature indicative of convergence. The strong alignment between these criteria highlights the utility of RMT for monitoring and diagnosing the progression of Transformer model training.

Early-stopping for Transformer model training

TL;DR

This work tackles the challenge of costly Transformer training and overfitting by introducing a validation-free early-stopping criterion based on Random Matrix Theory. By focusing on the spectral tail of the first encoder self-attention Value matrix, en.0.s.a.V, and fitting it with a Power-Law, the authors identify a three-phase training dynamics and define a KS-based heavy-tailed metric with threshold . The method triggers stopping at the epoch that maximizes , achieving substantial training-time reductions (e.g., stopping around epoch 100 for several setups) with minimal loss-accuracy degradation compared to longer runs. This spectral-based, validation-set-free approach advances understanding of implicit regularization in Transformers and offers a practical pathway to more efficient LLM training, with potential generalization to other large models.

Abstract

This work, based on Random Matrix Theory (RMT), introduces a novel early-stopping strategy for Transformer training dynamics. Utilizing the Power Law (PL) fit to tansformer attention matrices as a probe, we demarcate training into three stages: structural exploration, heavy-tailed structure stabilization, and convergence saturation. Empirically, we observe that the spectral density of the shallow self-attention matrix consistently evolves into a heavy-tailed distribution. Crucially, we propose two consistent and validation-set-free criteria: a quantitative metric for heavy-tailed dynamics and a novel spectral signature indicative of convergence. The strong alignment between these criteria highlights the utility of RMT for monitoring and diagnosing the progression of Transformer model training.
Paper Structure (24 sections, 7 theorems, 44 equations, 14 figures, 5 tables, 1 algorithm)

This paper contains 24 sections, 7 theorems, 44 equations, 14 figures, 5 tables, 1 algorithm.

Key Result

Lemma 3.1

Suppose $\{X_i\}_{i=1}^n$ are generated independently from $p(x; \alpha, x_{\min}) = \frac{\alpha - 1}{x_{\min}}\left(\frac{x}{x_{\min}}\right)^{-\alpha}$, then the distance in Eq eq:KS formula satisfies (Here $O_p$ denotes the boundedness in probability.)

Figures (14)

  • Figure 1: Evolution of $\alpha$ and $x_{\text{min}}$ for en.0.s.a.v matrix of T2_En-Ch9.62M. The blue solid line (left Y-axis) represents the parameter $\alpha$, while the red dashed line (right Y-axis) corresponds to the PL lower bound parameter $x_{\text{min}}$.
  • Figure 2: PL epxonents density histograms of T1, T2, and T3 models
  • Figure 3: ESD and PL fits for the en.0.s.a.v matrix from the T3_En-Ch19.20M experiment at three representative training phases(1 st, 50th and 200th epochs).
  • Figure 4: Evolutionary Trajectory of the Early-Stopping Metric for T1_En-Ch 9.60M. (a) describes the variation of d and d* with training epochs for the three models, respectively; while (b) illustrates the changes in the difference between d* and d across training epochs for these three models, respectively.
  • Figure 5: P-value of the power-law fit for the 'en.0.s.a.v' matrix versus training epoch across different experimental setups.
  • ...and 9 more figures

Theorems & Definitions (8)

  • Lemma 3.1
  • Proposition 3.1
  • Corollary 3.1
  • Proposition A.1
  • Lemma A.1
  • Remark A.1
  • Lemma A.2
  • Proposition A.2: Convergence of $W^V$ Optimization