PaperlayerView exact revision on arXiv

Source-linked AI summary

A Kernel-Based View of Language Model Fine-Tuning

Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, Sanjeev Arora

arXiv:2210.05643v4cs.LGcs.CL

TL;DR

Pre-trained language model fine-tuning often succeeds with very little data, but its theoretical basis is poorly understood. This paper extends the kernel framework to Adam and analyzes when fine-tuning from pre-trained initialization can exhibit kernel behavior. Across 14 NLP tasks, meaningful prompts often induce such dynamics, and the kernel view is applied to subspace-based fine-tuning.

  • Problem

    The paper asks why fine-tuning very large pre-trained language models on a couple dozen examples does not overfit, despite limited theoretical understanding of the paradigm.

    Source: §1 · p. 1

  • Method

    The authors extend NTK analysis to Adam and use Tensor Programs to characterize conditions for kernel behavior from pre-trained initialization, including prompt-based fine-tuning.

    Source: §1 · p. 2

  • Results

    Across 14 diverse NLP tasks, meaningful prompts often induce kernel behavior, and kernel dynamics describe prompt-based fine-tuning on tasks the eNTK can solve.

    Source: §1 · p. 2, §8 · p. 9

  • Takeaways & Limitations

    The authors apply the kernel view to propose an explanation for the success of parameter-efficient subspace-based fine-tuning methods such as LoRA.

    Source: §1 · p. 2, §8 · p. 9

    Experiments cover few-shot classification with one masked language model and specific prompts, while the theory applies only to early-stage Adam training.

    Source: §4 · p. 4, §8 · p. 9

Abstract

from arXiv · show

It has become standard to solve NLP tasks by fine-tuning pre-trained language models (LMs), especially in low-data settings. There is minimal theoretical understanding of empirical success, e.g., why fine-tuning a model with $10^8$ or more parameters on a couple dozen training points does not result in overfitting. We investigate whether the Neural Tangent Kernel (NTK) - which originated as a model to study the gradient descent dynamics of infinitely wide networks with suitable random initialization - describes fine-tuning of pre-trained LMs. This study was inspired by the decent performance of NTK for computer vision tasks (Wei et al., 2022). We extend the NTK formalism to Adam and use Tensor Programs (Yang, 2020) to characterize conditions under which the NTK lens may describe fine-tuning updates to pre-trained language models. Extensive experiments on 14 NLP tasks validate our theory and show that formulating the downstream task as a masked word prediction problem through prompting often induces kernel-based dynamics during fine-tuning. Finally, we use this kernel view to propose an explanation for the success of parameter-efficient subspace-based fine-tuning methods.

1. Introduction

The paper asks why pre-trained language models can be fine-tuned on few examples without overfitting, and why prompting and low-rank updates perform well. It develops a kernel analysis of fine-tuning and tests when that view applies to NLP.

1. Introduction

  • The motivating gap is why fine-tuning very large models on a couple dozen examples does not overfit, and why outcomes depend on prompting or low-rank updates.Answering these questions requires understanding how parameter updates change across settings.Source: §1 · p. 1
  • Wei et al. (2022) found that inherited-initialization eNTK classification accuracy was within 6% absolute of fine-tuning on several image-recognition tasks.Their result suggests kernel classifiers may help explain fine-tuning sample efficiency.Source: §1 · p. 1
  • The paper extends NTK theory to Adam with a gradient-sign kernel formula proved correct for early-stage training.The formula is motivated by early-stage Adam dynamics.Source: §1 · p. 2
  • The analysis allows pre-trained, non-random initialization and proves a mechanism by which prompting can yield kernel behavior in complex architectures such as Transformers.It uses the Tensor Programs framework and characterizes conditions for kernel behavior.Source: §1 · p. 2
  • Across 14 diverse NLP tasks, meaningful prompts often induce kernel behavior during fine-tuning, and kernel dynamics describe prompt-based fine-tuning on tasks the eNTK can solve.The experiments also find that a meaningful prompt is crucial for good eNTK performance.Source: §1 · p. 2
  • The authors apply the kernel view to analyze why low-rank subspace fine-tuning methods such as LoRA can succeed.They present this as a possible aid to practical design and theoretical analysis of efficient variants.Source: §1 · p. 2

2. Related Work

Prior kernel analyses do not directly capture fine-tuning from pre-trained initialization with Adam. This paper situates its analysis among work on kernel training, transfer, prompting, and fine-tuning variants.

Kernel view of training

  • Earlier NTK work applies infinite-width kernels to network optimization and generalization, with reported performance on small datasets as well as underperformance on standard tasks.The passage cites Jacot et al. (2018), Arora et al. (2019b), and Arora et al. (2020).Source: §2 · p. 2
  • Earlier NTK analyses assume random initialization and SGD, whereas this work handles pre-trained initialization under a related-task assumption and derives kernels for early-stage Adam.The paper notes that Adam is important for Transformer optimization.Source: §2 · p. 2

Theory of self-supervised learning and transfer learning

  • Prior work studies transfer through linear probing, related-task fine-tuning, or gradient-based features; this paper characterizes kernel dynamics for pre-trained masked language models on language-understanding tasks.Chua et al. (2021) study kernel behavior for regularized fine-tuning when pre-training and downstream tasks are closely related.Source: §2 · p. 2
  • Model-agnostic analyses may not fully capture transfer between self-supervised pre-training and supervised fine-tuning, motivating an optimization- and architecture-aware kernel framework.Saunshi et al. (2021) analyze prompting for zero-shot autoregressive models but preclude an investigation of fine-tuning.Source: §2 · p. 2

Variants of fine-tuning methods

  • Standard fine-tuning updates all parameters, while later variants update subsets or restrict updates to low-dimensional subspaces.The passage cites Radford et al. (2018), Devlin et al. (2019), Lester et al. (2021), Ben Zaken et al. (2022), Li & Liang (2021), Hu et al. (2021), and Aghajanyan et al. (2021).Source: §2 · pp. 2–3
  • Prompt-based fine-tuning and parameter-efficient methods are established alternatives; the paper reports that good eNTK performance arises only for prompt-based fine-tuning and interprets subspace methods through kernels.Prompting work includes Schick & Schütze (2021) and Gao et al. (2021); LoRA is associated with Hu et al. (2021).Source: §2 · p. 3

3. Preliminaries

The preliminaries define the masked-language-model setup for downstream fine-tuning and specify when training dynamics can be represented by a fixed kernel. They connect the kernel analog to the optimizer and model output.

3.1. Pre-Training and Fine-Tuning Paradigm

  • Masked language models predict masked tokens independently; standard fine-tuning learns a new classifier on the [CLS] embedding, while prompt-based fine-tuning uses the pre-trained MLM to predict task-relevant masked words.For prompt-based fine-tuning, logits from task-relevant columns of the pre-trained word-embedding classifier serve as downstream-task surrogates.Source: §3.1 · p. 3
  • For standard fine-tuning experiments, the classifier head is initialized at the linear-probing solution before full fine-tuning.The authors choose this because kernel-behavior fine-tuning finds solutions close to initialization and hypothesize that its learned head is closer to linear probing than to random initialization.Source: §3.1 · p. 3

3.2. Kernel Behavior

  • Kernel behavior requires both a first-order linearization of network change and gradients that remain approximately fixed from their initialization values.The paper quantifies closeness to these approximations and formalizes them in Definition C.3.Source: §3.2 · p. 3
  • When training exhibits kernel behavior, the kernel analog expresses function changes using a fixed kernel; for SGD, this analog is the NTK inner product of initialization gradients.The eNTK construction depends on the model output and chosen optimizer; the paper considers SGD or Adam.Source: §3.2 · p. 4

4. Kernel Derivation for Adam

The paper derives a kernel analog for early-stage Adam by approximating its updates with SignGD. The resulting sign-based kernel is proved correct under kernel behavior, with the analysis bounded to early-stage dynamics.

4. Kernel Derivation for Adam

  • Early-stage Adam with a small learning rate reduces to coordinate-wise gradient-sign updates, called SignGD.The analysis assumes batch size 1 for simplicity.Source: §4 · p. 4
  • The paper reports that SignGD fine-tuning yields performance comparable to Adam fine-tuning.This empirical comparison is reported in Table 10.Source: §4 · p. 4
  • Under kernel behavior, the paper proves that SignGD training dynamics follow the derived kernel analog.The theorem states the result for a network trained with SignGD that exhibits kernel behavior.Source: §4 · p. 4
  • The SignGD kernel is the inner product of element-wise gradient signs at initialization, unlike the standard SGD NTK based on gradient inner products.The sign kernel reflects that early-stage Adam dynamics are agnostic to gradient scale.Source: §4 · p. 4
  • The analysis is limited to early-stage Adam; concurrent work suggests no kernel describes general late-stage Adam training when batch size is large.The authors note that the reduction of Adam to SignGD is crucial to observing kernel dynamics.Source: §4 · p. 4

5. Theory: Prompt-Based Fine-Tuning Can Exhibit Kernel Behavior

The section gives conditions under which prompt-based fine-tuning of a pre-trained network exhibits kernel behavior as width grows. Its account assumes a suitable relationship between pre-training and the downstream task, alongside stability and learnability conditions on the model.

5. Theory: Prompt-Based Fine-Tuning Can Exhibit Kernel Behavior

  • A pre-training scheme specifies the dataset, optimizer and hyperparameters, and model architecture; in a Transformer, width corresponds to embedding dimension.The model is the result of training that architecture on the dataset with the specified optimizer.Source: §5 · p. 5
  • Prompting casts the downstream task as fill-in-the-blank prediction, treated as a subcase of pre-training that an infinitely wide model can solve with a suitable prompt.The account assumes wider pre-trained networks become better at filling in masked tokens.Source: §5 · p. 5
  • Only well-suited prompts may make a task natural in the infinite-width limit; randomly initialized heads cannot meet the condition described.The task’s prompt is part of its dataset, and the stated condition requires the output derivative to vanish at the start of fine-tuning.Source: §5 · p. 5
  • A two-width test checks whether the output derivative decreases with width before any gradient updates, but this is necessary, not sufficient, for naturalness.The models use different widths, the same depth, and otherwise identical pre-training schemes.Source: §5 · p. 5
  • The model must be stable, non-trivial, and representable as a Tensor Program for the infinite-width analysis to apply.Stability means the output does not grow with width; non-triviality means fine-tuning can update it.Source: §5 · p. 5
  • Under the stated assumptions, prompt-based fine-tuning exhibits the Linearization and Fixed Features properties of kernel behavior.The theorem requires a natural downstream task and a stable, non-trivial model expressible as a Tensor Program.Source: §5 · p. 5
  • The theory covers autoregressive and masked language models, while the fine-tuning experiments are limited to masked language models.The paper says masked language models are known to perform better after fine-tuning.Source: §5 · p. 5

6. Experiments

Across 14 tasks, the eNTK matched prompt-based fine-tuning on 12, while 8 consistently induced kernel behavior. Prompt choice mattered, and additional measurements distinguished task solvability from whether fine-tuning followed kernel dynamics.

6. Experiments

  • 12 of 14 tasks were consistently solvable by the eNTK comparably to prompt-based fine-tuning, and 8 induced kernel behavior during fine-tuning.The results were measured across five subsampled datasets.Source: §6 · p. 6, Table 1 · p. 6

6.1. Kernel Performance on Downstream Tasks

  • Adam and SGD were within 4% absolute of each other in prompt-based fine-tuning.The difference between K(SignGD) and K(A-SignGD) was negligible on most tasks, although the asymmetric-problem solver could make K(A-SignGD) perform worse.Source: §6.1 · pp. 6–7
  • For 10 of 14 tasks, the kernel analog achieved accuracy within 10% of corresponding fine-tuning at both 16 and 64 shots.In standard fine-tuning, the K(SGD)–SGD-FT gap reached 16% absolute, versus 3% in the prompt-based setting; prompted kernels also consistently outperformed standard ones.Source: §6.1 · p. 6

6.2. Measuring Kernel Behavior

  • For every task the eNTK could solve, the linearized model recovered more than 50% of fine-tuning’s improvement over the pretrained model.The linearization property held for all tasks the eNTK could solve.Source: §6.2 · p. 7
  • Tasks the eNTK could solve had kernel distances below 2, indicating that the fixed-features property likely held.The study also found entailment anomalies: the eNTK could not solve MNLI and SNLI despite strong pretrained performance, while it could solve QNLI and RTE despite near-random pretrained performance; QNLI and RTE showed weaker linearization.Source: §6.2 · pp. 7–8

6.3. Tasks without Kernel Behavior

  • TREC, MNLI, SNLI, QNLI, and MPQA consistently did not induce kernel behavior; the authors relate this to prompts that do not format tasks as pretraining subcases.They point to insufficient or unnatural prompt signals for TREC and entailment tasks, and a mismatch between complete-sentence prompts and short MPQA examples.Source: §6.3 · p. 8
  • AG News was an exception: the eNTK solved it consistently, although Adam fine-tuning did not exhibit kernel behavior.The authors attribute this to a learning-rate difference: eNTK success suggested a very small rate, while the best-performing fine-tuning trajectory used a larger rate and more complex dynamics.Source: §6.3 · p. 8

7. Efficacy of Subspace-Based Fine-Tuning Methods

The kernel view gives a simple account of subspace-based fine-tuning: under kernel behavior, LoRA preserves the SGD fine-tuning kernel and its dynamics when its rank is sufficiently large. The explanation is conditional on the underlying fine-tuning exhibiting kernel behavior.

7. Efficacy of Subspace-Based Fine-Tuning Methods

  • Under kernel behavior, sufficiently large-rank SGD-LoRA leaves the SGD kernel and training dynamics unchanged.The paper also reports that prompt-based SGD-LoRA fine-tuning performs similarly to prompt-based SGD fine-tuning and to K(SGD).Source: §7 · pp. 8–9
  • LoRA’s gradient with respect to B is a random projection of the weight gradient, while its gradient with respect to A is zero at initialization.Applying Johnson–Lindenstrauss inner-product preservation across input pairs preserves kernel entries.Source: §7 · pp. 8–9
  • The result explains LoRA only when the base fine-tuning already exhibits kernel behavior, making prompt-based settings the meaningful case considered here.The theorem’s kernel invariance applies to both prompt-based and standard fine-tuning, but its explanation of LoRA’s success is conditional on kernel behavior.Source: §7 · pp. 8–9

8. Conclusion

The paper develops a kernel account of language-model fine-tuning, including an Adam kernel and conditions for prompt-based fine-tuning to exhibit kernel behavior. Experiments support that account in some settings, while the authors emphasize limits in scope and the dependence on task and prompt.

8. Conclusion

  • Experiments across 14 NLU tasks show that meaningful prompts often accompany kernel behavior, and kernel dynamics describe prompt-based fine-tuning on tasks the eNTK can solve.The conclusion also connects this kernel view to analysis of subspace-based fine-tuning.Source: §8 · p. 9
  • The experiments cover few-shot classification with one masked language model and specific prompts; the Adam theory applies only to early-stage training.The authors also note that extending experiments to additional settings and models is computationally costly because the eNTK is expensive to compute.Source: §8 · p. 9
  • Some tasks do not induce kernel behavior, so future theoretical analyses need to account for the downstream task and prompt choice.The paper suggests studying kernel spectra as a future route to understanding fine-tuning generalization.Source: §8 · p. 9

A. Experimental Details

The appendix describes the few-shot datasets and prompts, kernel computation, and procedures for solving the kernel models. It specifies dataset sampling and implementation choices as well as how pretrained-model outputs and multiclass labels enter kernel regression.

A.1. Datasets and Prompts

  • The experiments use 14 classification datasets with manual prompt templates and label words, sampling separate k-shot training and validation sets across five seeds.The set includes eight single-sentence and six sentence-pair tasks; demonstrations are not included in prompts.Source: §A.1 · p. 15, Table 3 · p. 15

A.2. Computing the Kernel

  • The eNTK is computed for 125M-parameter RoBERTa-base using backward-mode Jacobians and forward-mode Jacobian-vector products.K(SignGD) cannot use Jacobian-vector products and requires substantially more memory and runtime.Source: §A.2 · p. 15

A.3. Solving the Kernel

  • Kernel regression adds the pretrained model’s output to the kernel output and searches label scaling to use knowledge already present before fine-tuning.For multiclass tasks, each logit is treated as an independent output, scaling the dataset by the number of classes and allowing relationships between logits.Source: §A.3 · pp. 15–16
  • Symmetric-kernel models use the representer theorem and regularized logistic loss, while the asymmetric formulation uses distinct gradient and gradient-sign feature maps.The appendix also describes a combined discriminant function and notes that its stationary-point interpretation requires least-squares training.Source: §A.3 · p. 16
  • Fine-tuning and kernel hyperparameters are selected on few-shot validation data, with training-step counts adjusted to k and the number of classes.The appendix specifies optimizer grids and uses linear decay without warmup or weight decay.Source: §A.3 · pp. 16–17, Table 4 · p. 17

B. Additional Experimental Results

The appendix reports how additional tables and figures assess task solvability, prompt sensitivity, and kernel behavior. It also lays out the infinite-width assumptions and Tensor Programs framework used for the formal analysis.

B. Additional Experimental Results

  • Tables 5–7 provide numerical results for Figures 1–2 and measure how strongly the Fixed Features property holds, with smaller distances indicating stronger satisfaction.The appendix describes the tables’ scope but the supplied passages do not provide their cell values.Source: B. Additional Experimental Results · p. 17, Table 7 · p. 20

B.1. Solvable Task Experiments

  • The solvability check compares average loss-output derivatives across RoBERTa-base and RoBERTa-large, although the models differ in depth as well as width.The authors note that finding pretrained models with different widths and fixed depth is generally difficult.Source: §B.1 · p. 17

B.2. Robustness to Choice of Prompt

  • Prompt choice matters: on sentiment tasks, minimal null prompts produce a substantial gap between K(SGD) and fine-tuning, unlike the searched prompts.The authors hypothesize that this occurs because the task is no longer solvable in the infinite-width limit.Source: §B.2 · p. 17, Table 9 · p. 20
  • Figure 2 compares pretrained, linearized, and fine-tuned accuracies; recovery of fine-tuning performance by the linearized model indicates kernel behavior.It reports medians and ranges across five seeds and data splits at k = 64.Source: Figure 2 · p. 18

C. Kernel Behavior and the Parametrization

  • The formal discussion studies the infinite-width setting using Tensor Programs, modeling pretrained weights as a non-random initialization arising from prior training.The overview distinguishes kernel behavior, where gradients remain unchanged, from feature learning behavior.Source: C. Kernel Behavior and the… · p. 22

C.1. Preliminaries

  • Tensor Programs represent network computations and optimization through vector and scalar operations, enabling reasoning about their scales as width grows.Their Master Theorem gives deterministic limits for scalars as width tends to infinity.Source: §C.1 · p. 23

C.2. SignGD Kernel Derivation

This appendix derives conditions for kernel behavior under SignGD and prompt-based fine-tuning, using µP and a linear-network example to explain how feature-learning pre-training can give way to kernel-like fine-tuning. It also shows that parameter updates remain smaller-order than their initial counterparts and specifies the initialization and learning-rate conditions required for the LoRA application.

C.2. SignGD Kernel Derivation

  • Definition C.3 characterizes kernel behavior through linearization and fixed features; Theorem C.4 gives a conditional SignGD kernel result under Assumption C.2.The definition requires the network change to be approximated by its first-order Taylor expansion and gradients to remain approximately as before training.Source: §C.2 · p. 24

C.3. Prompt-based Fine-Tuning

  • For a downstream task Ξ natural for network f, fine-tuning exhibits kernel behavior under Assumption C.2; prompt-based fine-tuning exhibits kernel behavior when χ = o(1) before fine-tuning.Source: §C.3 · p. 26
  • The prompt fine-tuning theorem’s proof establishes bounded function updates and then kernel behavior when χt = O(χmax) = o(1).The result is shown first for an MLP with SGD and then extended to more general architectures and optimizers.Source: §C.3 · p. 26

C.4. µP for SGD and SignGD

  • For µP, SGD and SignGD have the same initialization and update scales, differing only in learning rates; a parameter rescaling leaves the function fixed, with distinct SGD and Adam learning-rate transformations.The lemma’s rescaling is A ← Aγ and B ← B/γ; C ← C/γ² for SGD and C ← C/γ for Adam, and its Adam conclusion extends to SignGD.Source: §C.4 · p. 28

C.5. Prompt-based Fine-Tuning: A Linear Example

  • In the linear example, µP pre-training exhibits feature learning, whereas fine-tuning’s o(1) χ makes its higher-order term vanish and its leading update kernel-like.The example identifies the higher-order term as ηWηUχ²ft−1(ξ), with χ mostly Θ(1) in pre-training and o(1) in fine-tuning.Source: §C.5 · pp. 29–30
  • At the tth step, the argument matches the second step because ΔU_t and ΔW_t remain smaller-order than ΔU_0 and ΔW_0.Source: §C.5 · p. 30

C.6. LoRA FT Exhibits Kernel Behavior

  • Theorem C.5 applies to any architecture, including LoRA, provided its initialization and learning rate satisfy Assumption C.2.For intermediate-layer LoRA, the proposed setup treats A and B as one of the W_j matrices under μP; it also generalizes to LoRA on U and V.Source: §C.6 · p. 30
  • For intermediate-layer LoRA, set k = αn with α a small constant, initialize A at Θ(1/√n), and use A/B learning rates Θ(1) for SGD or Θ(n^-1-d) for SignGD/Adam.This setup follows μP by treating A and B as one of the W_j matrices and generalizes to LoRA on U and V.Source: §C.6 · p. 30

D. Subspace-Based Fine-Tuning Methods

The section analyzes subspace-based fine-tuning through kernel dynamics, finding comparable performance for prompt-based SGD fine-tuning and SGD-LoRA fine-tuning. It also shows that sufficiently large random-projection dimensions preserve the IntrinsicDimension fine-tuning kernel.

D. Subspace-Based Fine-Tuning Methods

D.1. IntrinsicDimension FT

  • In the kernel view, IntrinsicDimension fine-tuning and LoRA fine-tuning induce similar optimization dynamics, though they were proposed for different purposes.IntrinsicDimension fine-tuning was proposed to measure downstream-task difficulty, while LoRA was proposed as an alternative fine-tuning method [Li et al., 2018; Aghajanyan et al., 2021].Source: §D.1 · p. 31
  • IntrinsicDimension fine-tuning preserves each kernel element relatively closely when the random-projection dimension k ≥ Θ(log N/ϵ^2).The theorem assumes a random projection and bounded entries of the SGD kernel.Source: §D.1 · p. 31

D.2. Proofs

The proofs relate subspace updates to kernel changes: LoRA’s kernel is close to the full-fine-tuning kernel under a dimension condition, and random projections preserve the relevant inner products. The LoRA argument uses its initialization and update structure, and extends to multiple applications of LoRA.

D.2. Proofs

  • If a fine-tuning method exhibits kernel behavior, the proof establishes that its LoRA version does too.The argument invokes the Tensor Programs framework.Source: §D.2 · p. 31
  • A random projection preserves inner products, allowing the IntrinsicDimension kernel difference to be bounded across all downstream-example pairs.The proof applies the projection lemma to pairs of example gradients and uses a union bound.Source: §D.2 · pp. 31–32
  • For LoRA, B is initialized to zero and W is fixed, so dA remains zero in the kernel setting while the update kernel is derived from the active factors.The construction applies to a fully connected layer with k ≪ n.Source: §D.2 · p. 32
  • With high probability, the LoRA and full-fine-tuning kernels differ by less than ϵ when k ≥ 20c^4 log N/ϵ^2; for L LoRA applications, the requirement becomes k ≥ Θ(Lc^4 log N/ϵ^2).The stated bound is first proved for a simple fully connected layer and then extended to L independent LoRA components.Source: §D.2 · p. 32
Loading 2210.05643v4…