GradES: Significantly Faster Training in Transformers with Gradient-Based Early Stopping

TL;DR

GradES tackles the high cost of validation-based early stopping in transformer fine-tuning by monitoring per-weight-matrix gradient-change magnitudes and freezing converged components, thereby removing the need for costly validation inferences. It formalizes a gradient-based criterion $G_{\mathbf{W}}^{(l)}(t) = \|\nabla \mathbf{W}^{(l)}_t - \nabla \mathbf{W}^{(l)}_{t-1}\|_1$ and employs a grace period and gradient-flow preservation to ensure stable optimization, with an extension to LoRA using the combined gradient of the low-rank factors. Empirically, GradES yields 1.57–7.22× training-time speedups and improves language benchmark accuracy by about 1.2% and multimodal benchmark accuracy by about 3.88%, while maintaining or increasing FLOP efficiency across 0.6B–14B parameter models and both FP and LoRA fine-tuning. The approach is compatible with common optimizers and complementary efficiency techniques, offering practical gains for resource-constrained deployment and suggesting broader applicability to pretraining and other architectures in future work.

Abstract

Early stopping monitors global validation loss and halts all parameter updates simultaneously, which is computationally costly for large transformers due to the extended time required for validation inference. We propose \textit{GradES}, a novel gradient-based early stopping approach that operates within transformer components (attention projections and Feed-Forward layer matrices). We found that different components converge at varying rates during fine-tuning for both language and vision-language models. \textit{GradES} tracks the magnitude of gradient changes in backpropagation for these matrices during training. When a projection matrix's magnitude of gradient changes fall below a convergence threshold $τ$, we exclude that projection matrix from further updates individually, eliminating costly validation passes while allowing slow converging matrices to continue learning. \textit{GradES} speeds up training time by 1.57--7.22$\times$ while simultaneously enhancing generalization through early prevention of overfitting, resulting in 1.2\% higher average accuracy in language tasks and 3.88\% on multimodal benchmarks.

Figures (4)

• Figure 1: Element-wise $L_1$ norms for the gradient matrix of components in layer 7 for Qwen3-0.6B qwen2024qwen25. Each step consists of processing one training batch through the complete forward pass, loss computation, backpropagation, and parameter update cycle. The seven tracked matrices comprise attention projections ($\mathbf{W}_q^{(7)}$, $\mathbf{W}_k^{(7)}$, $\mathbf{W}_v^{(7)}$, $\mathbf{W}_o^{(7)}$) and MLP components ($\mathbf{W}_{\text{gate}}^{(7)}$, $\mathbf{W}_{\text{up}}^{(7)}$, $\mathbf{W}_{\text{down}}^{(7)}$). MLP projections exhibit 2 to 3$\times$ higher gradient magnitudes than attention projections throughout training, with $\mathbf{W}_{\text{up}}^{(7)}$ and $\mathbf{W}_{\text{down}}^{(7)}$ maintaining the largest magnitude of gradient changes. The red dotted line indicates our convergence threshold $\tau = 1.183 \times 10^{-3}$.
• Figure 2: GradES architecture. Color indicate gradient magnitude changes for each weight matrix. Components with low gradient changes (below threshold $\tau$) are stopped (lock icons) while others continue training.
• Figure 3: Cumulative frozen components during training across model scales. Fraction of weight matrices frozen over time for five different LLMs.
• Figure 4: Comparison of gradient norm evolution patterns across different model architectures during fine-tuning.

Theorems & Definitions (5)

• Theorem 1
• proof
• Theorem 2: Convergence of GradES
• proof
• Corollary 3