A Kernel-Based View of Language Model Fine-Tuning
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, Sanjeev Arora
ICML 2023
cs.LG, cs.CL
2022-10-12
Recasts LM fine-tuning as kernel regression: derives an Adam-compatible NTK, shows prompt-based fine-tuning follows kernel dynamics on most of 14 NLP tasks, and explains LoRA via the Johnson-Lindenstrauss lemma.
Pre-training plus fine-tuning is something practitioners do every day, yet it has never had a proper theory. The most basic puzzle is statistical: a language model with a hundred million or more parameters, fine-tuned on a few dozen labeled examples, should overfit catastrophically by classical intuition, but it works. A cluster of everyday phenomena sit in the same blind spot. Adding a prompt like "It was [great/terrible]" shifts fine-tuning results noticeably; restricting updates to a low-rank subspace (LoRA and friends) somehow matches full fine-tuning.
All of these point at one question: how far does fine-tuning actually move the parameters? A wholesale rewrite, or a small correction? This Princeton work (ICML 2023) argues that in many cases fine-tuning behaves like kernel regression, with the model itself barely moving.
The Neural Tangent Kernel (NTK) was built to study gradient descent in infinitely wide, randomly initialized networks. The intuition: when a network is wide enough, the gradient at each parameter stays nearly fixed during training, so the whole process is captured by one kernel matrix K(ξ,ξ′)=⟨∇f(ξ),∇f(ξ′)⟩, and classification reduces to regression against that kernel. Two assumptions block this from applying to fine-tuning: it needs random initialization, whereas fine-tuning starts from learned, non-random weights; and it describes SGD, whereas real fine-tuning almost always uses Adam.
The authors close both gaps.
First, a kernel for Adam. Adam's full-trajectory dynamics resist a kernel description because each step depends on the entire gradient history. But early in training, Adam's moment estimates move in a small neighborhood and the update reduces to taking the sign of the gradient coordinate-wise, SignGD. The authors propose a SignGD kernel that replaces gradients with their signs, K^SignGD(ξ,ξ′)=⟨sign(∇f(ξ)),sign(∇f(ξ′))⟩, and prove it is the correct kernel analog. This matters because fine-tuning mostly lives in this early regime.
Second, why non-random initialization does not break kernel behavior. The key idea is a "natural task": with a suitable prompt, the downstream task is semantically a sub-case of the pre-training task (masked prediction). Then as width goes to infinity, the pre-trained model already nearly solves it, fine-tuning only nudges, and kernel behavior holds. The authors formalize this with Tensor Programs (Theorem 5.5). This also yields a clean rule for why prompting helps: it restates the task as fill-in-the-blank, making it "natural."
Experiments span 14 NLP tasks (sentiment, topic classification, entailment, paraphrase) on RoBERTa-base, with 5 k-shot splits per task at k=16 and k=64.
The eNTK (the kernel computed directly from pre-trained weights) solves 12 of the 14 tasks as well as fine-tuning; 8 tasks show full kernel behavior during fine-tuning (both linearization and fixed gradients). With a prompt, the kernel analog lands within 10 points of full fine-tuning on 10 tasks.
The prompt is decisive. Without it, standard fine-tuning sees kernel-vs-actual gaps up to 16 points; with it, the gap shrinks to about 3 points.
| Setting | Kernel vs. fine-tuning accuracy gap |
| Standard FT (no prompt), k=16/64 | up to 16 points |
| Prompt-based FT | 3 points |
| Prompt-based, kernel analog vs full FT (10/14 tasks) | <10 points |
| Adam-FT vs SGD-FT (with prompt) | <4 points |
A few tasks never show kernel behavior: TREC, MNLI, SNLI, QNLI, MPQA. The authors trace this to prompt design. For example, MNLI's neutral label "Maybe" wedged between premise and hypothesis produces ungrammatical sentences, so the task cannot be cast as natural fill-in-the-blank.
Section 7 explains LoRA along the way. LoRA writes the update as W+BA, the product of two thin matrices, which is a random projection of the full gradient into a low-rank subspace. By the Johnson-Lindenstrauss lemma, random projections preserve inner products, so every kernel entry is preserved; SGD-LoRA with large enough rank k shares the exact kernel and dynamics of full SGD fine-tuning. This is among the cleanest theoretical accounts of why low-rank fine-tuning does not lose accuracy.
This is not a stronger fine-tuning method; it is an understanding. Practitioners get three concrete things. First, prompting is not superstition: its job is to align the task with a form the model saw in pre-training, which guides prompt design. Second, LoRA's effectiveness has a theoretical guarantee, since large enough rank leaves optimization dynamics unchanged, which informs how big the rank should be. Third, the eNTK itself is a way to solve a downstream task without gradient updates, and in very-low-data, noisy-gradient settings it can be more stable than real fine-tuning.
The caveat is scope. These conclusions come from RoBERTa-scale masked models. Whether today's autoregressive models with tens of billions of parameters sit in the same kernel regime is not covered.
The authors draw explicit lines. Experiments cover only few-shot classification, a single masked language model, and a fixed set of prompts; scaling k or the model turns eNTK computation into a bottleneck. The SignGD approximation describes only early Adam, and it is unclear how far it extends. Concurrent work (Littwin & Yang 2023) suggests the Adam-to-SignGD reduction is precisely what gates whether kernel behavior appears, so the applicable window is narrower than it looks.
One more concern: the thresholds for declaring kernel behavior (linearization closing half the gap, gradient distance below 2.0) are hand-picked, and the paper admits the formal definition sets no numerical thresholds. The 8/14 figure is sensitive to these choices; different thresholds shift how strong the conclusion reads. The raw data is laid out in the tables for readers to check.