A Theoretical Framework for LLM Fine-tuning Using Early Stopping for Non-random Initialization
Zexuan Sun, Garvesh Raskutti
TL;DR
This work develops a theoretical framework for fine-tuning pretrained LLMs by integrating early stopping with attention-NTK analyses, extending NTK theory to non-random initializations and proving convergence guarantees governed by the eigenvalue decay of the NTK. It formalizes a linearization of transformer-attention networks around pretrained weights, analyzes mixed initialization settings, and explains task arithmetic via task vectors within the RKHS/NTK lens. The authors connect the optimal stopping rule to kernel ridge regression, derive convergence rates that depend on the NTK spectrum, and validate predictions with experiments on GPT-Neo-1.3B across multiple tasks, including multitask settings. The framework offers principled guidance for early stopping in LLM fine-tuning and contributes to understanding how task similarity, eigenvalue decay, and weight-space edits interact in practice, with implications for efficient, robust downstream adaptation.
Abstract
In the era of large language models (LLMs), fine-tuning pretrained models has become ubiquitous. Yet the theoretical underpinning remains an open question. A central question is why only a few epochs of fine-tuning are typically sufficient to achieve strong performance on many different tasks. In this work, we approach this question by developing a statistical framework, combining rigorous early stopping theory with the attention-based Neural Tangent Kernel (NTK) for LLMs, offering new theoretical insights on fine-tuning practices. Specifically, we formally extend classical NTK theory [Jacot et al., 2018] to non-random (i.e., pretrained) initializations and provide a convergence guarantee for attention-based fine-tuning. One key insight provided by the theory is that the convergence rate with respect to sample size is closely linked to the eigenvalue decay rate of the empirical kernel matrix induced by the NTK. We also demonstrate how the framework can be used to explain task vectors for multiple tasks in LLMs. Finally, experiments with modern language models on real-world datasets provide empirical evidence supporting our theoretical insights.