PaperlayerView exact revision on arXiv

Source-linked AI summary

Task-Specific Skill Localization in Fine-tuned Language Models

Abhishek Panigrahi, Nikunj Saunshi, Haoyu Zhao, Sanjeev Arora

arXiv:2302.06600v2cs.CLcs.LG

TL;DR

Where task-specific skills reside in fine-tuned language models has received limited study. This paper uses model grafting to locate sparse parameter regions, reporting that tiny regions retain task performance and that grafted models show calibration, OOD, and multi-task findings.

  • Problem

    Where newly learned task-specific skills reside inside fine-tuned language models has received limited study.

    Source: §1 · p. 1

  • Method

    The paper learns a sparse mask that grafts selected fine-tuned parameter values onto the pretrained model while leaving other parameters unchanged.

    Source: §1 · p. 1, §2.1 · p. 2

  • Results

    Grafted models retain task performance with tiny regions, improve calibration and often OOD performance over vanilla fine-tuning, and show somewhat disjoint task regions whose overlap seems to track task similarity.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    Preliminary experiments suggest grafting can benefit continual learning, with memory for retaining skills proportional to the number of tasks times region sparsity rather than model size.

    Source: §5.2 · p. 9

    One experiment considers only contradiction and entailment labels.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

Pre-trained language models can be fine-tuned to solve diverse NLP tasks, including in few-shot settings. Thus fine-tuning allows the model to quickly pick up task-specific ``skills,'' but there has been limited study of where these newly-learnt skills reside inside the massive model. This paper introduces the term skill localization for this problem and proposes a solution. Given the downstream task and a model fine-tuned on that task, a simple optimization is used to identify a very small subset of parameters ($\sim0.01$% of model parameters) responsible for ($>95$%) of the model's performance, in the sense that grafting the fine-tuned values for just this tiny subset onto the pre-trained model gives performance almost as well as the fine-tuned model. While reminiscent of recent works on parameter-efficient fine-tuning, the novel aspects here are that: (i) No further re-training is needed on the subset (unlike, say, with lottery tickets). (ii) Notable improvements are seen over vanilla fine-tuning with respect to calibration of predictions in-distribution ($40$-$90$% error reduction) as well as the quality of predictions out-of-distribution (OOD). In models trained on multiple tasks, a stronger notion of skill localization is observed, where the sparse regions corresponding to different tasks are almost disjoint, and their overlap (when it happens) is a proxy for task similarity. Experiments suggest that localization via grafting can assist certain forms of continual learning.

1. Introduction

Fine-tuning quickly gives language models task-specific skills, but where those skills reside in the model is an open question. The paper localizes them through sparse grafts and reports performance, calibration, OOD, and multi-task findings.

Method

  • A tiny grafting region transfers fine-tuned parameter values onto the pretrained model while leaving all other parameters unchanged.The paper describes this as a compact way to retain a fine-tuned model's task performance.Source: §1 · p. 1, Figure 1 · p. 2

Results

  • 0.01% of RoBERTa parameters recovers 95% accuracy of fine-tuned GLUE models without further retraining.For GPT-2, regions containing 0.05% of parameters recover 95% accuracy.Source: §1 · p. 1
  • Grafted models have better-calibrated outputs than vanilla fine-tuning and often much better OOD performance.The paper reports that retraining a sparse region does not provide the same OOD benefits.Source: §1 · p. 2
  • In multi-task fine-tuning, task regions are somewhat disjoint, and their overlap seems to correlate with intuitive task similarity.Grafting regions for a subset of tasks also works for those and related tasks, but not others.Source: §1 · p. 2

2. Skill Localization through Model Grafting

The section frames skill localization as finding task-relevant parameters in a fine-tuned model, then defines grafting and an optimization procedure to identify a sparse region. It distinguishes this approach from sparse fine-tuning and lottery-ticket pruning.

2. Skill Localization through Model Grafting

  • Learned regions outperform movement-ranked and random regions at low sparsity levels.The paper notes that identifying useful regions from parameter movement alone is not concise or closely tied to the task.Source: §2 · p. 2, Figure 2 · p. 2

2.1. Model Grafting

  • A binary mask selects a region whose fine-tuned parameter values are grafted onto the pretrained model, with other parameters left at pretrained values.The region should retain fine-tuning performance; the mask-weighted parameter change represents the acquired skill.Source: §2.1 · p. 2

2.2. Differences from Sparsity-based Fine-Tuning

  • Grafting uses pretrained values outside the region and requires no retraining, unlike lottery-ticket pruning; bias-only grafts do not localize skills well.The paper also reports that BitFit's parameter-efficient fine-tuning mechanism differs from standard fine-tuning and does not provide the same calibration benefits.Source: §2.2 · pp. 2–3, Figure 4 · p. 3

2.3. Optimization Procedure to Learn Grafted Models

  • The optimization reparameterizes the mask as a sigmoid over a candidate mask and takes a few gradient steps to make minimal changes while reducing task loss.The candidate mask can select parameters with the largest fine-tuning movement; learned regions outperform it at lower sparsity.Source: §2.3 · p. 3
  • Without retraining, lottery-ticket pruning that sets other parameters to zero performs poorly even at high density.The figure caption reports that pruning fails at any tested sparsity level; the text specifies failure even when 90% of parameters remain.Source: Figure 4 · p. 3

3. Evaluating Grafting for Skill Localization

The section evaluates whether tiny grafted parameter regions retain fine-tuned task performance across training setups. It compares optimized regions with simpler parameter-selection and retraining alternatives, and tests standard fine-tuning and AdamW.

3. Evaluating Grafting for Skill Localization

  • The experiments use prompt-based fine-tuning of RoBERTa-base, usually with SGD, on 13 tasks, with reported results generally from the 4096-shot setting.The evaluation also includes 64-shot experiments and compares performance across five randomly sampled datasets.Source: §3 · p. 3

3.1. Sparse Grafting Retain Skills

  • At sparsity of at most 0.01%, sparse grafting recovers at least 95% of fine-tuning performance across the reported datasets.The grafted models also agree with fine-tuned test labels 93% of the time for single-sentence tasks and 86% for two-sentence tasks.Source: §3.1 · pp. 3–4
  • The optimized grafting regions outperform top-movement parameter selection, especially at lower sparsity levels.The comparison uses regions selected by optimization versus the top-s most-changed parameters.Source: §3.1 · p. 4
  • Retraining a 0.01%-sparse graft region from pretrained initialization performs well, often better than grafting, and forms a trainable subnetwork.The reported region is substantially sparser than BitFit’s 0.1% bias parameters.Source: §3.1 · p. 4, Table 1 · p. 4
  • Without retraining, bias-only grafts and lottery-ticket-style regions fail to capture skill, even when the latter retain 90% of parameters.The paper contrasts these approaches with its 0.01%-sparse grafting regions.Source: §3.1 · p. 4

3.2. Other Fine-Tuning Paradigms

  • Sparse localization also works with standard fine-tuning using a linear head on the [CLS] representation.The experiments test a training setup different from prompt-based fine-tuning.Source: §3.2 · p. 4
  • AdamW fine-tuning does not yield high-performing sparse grafts, but an ℓ1 penalty of 0.001 on parameter movement recovers them.The comparison is reported for prompt-based fine-tuning.Source: §3.2 · p. 4, Figure 5 · p. 4

4. Calibration, OOD and Generalization

The section examines calibration, out-of-distribution performance, and generalization for sparse grafts. Grafted models retain most fine-tuned accuracy, improve calibration, can outperform fine-tuning under larger distribution shifts, and have a more tractable generalization analysis.

4. Calibration, OOD and Generalization

  • Grafted models retrieve more than 95% of fine-tuned accuracy while improving calibration; at 0.01% sparsity, single-sentence tasks lose 0.7% accuracy and improve calibration error by 5%.For two-sentence tasks, the reported accuracy drop is 3.4% and calibration-error improvement is 9.6%.Source: Table 2 · p. 5

4.1. Skill Localization improves Calibration

  • Sparse grafting reduces ECE by 40–90% when sparsity causes less than a 5% accuracy reduction.The paper attributes the calibration difference to vanilla fine-tuning’s overconfidence on wrong predictions, which grafted models avoid.Source: §4.1 · p. 5

4.2. Out-of-Distribution Generalization

  • For small distribution shifts, fine-tuned models are already robust and grafting provides little or no OOD advantage.The examples are SST-2 to IMDb and MNLI to SNLI.Source: §4.2 · p. 6
  • For larger shifts, grafted models match or outperform fine-tuning on OOD accuracy, with reported gains around 5% for MPQA and QNLI transfers.The examples include MPQA to SST-2 and Yelp, and QNLI to MNLI and SNLI.Source: §4.2 · p. 6
  • Grafting and WiSE-FT each perform better for different ID–OOD pairs, while WiSE-FT applied to grafted models achieves a competitive trade-off.The passage also reports that grafting may slightly reduce in-distribution performance while matching or outperforming fine-tuning out of distribution.Source: §4.2 · p. 6
  • At comparable sparsity, grafting has better calibration and OOD performance than the evaluated parameter-efficient methods.The paper notes that sparsity alone does not account for grafting’s calibration and OOD performance.Source: §4.2 · p. 6

4.3. Understanding Generalization for Fine-tuning

  • For datasets with a few thousand datapoints, retraining grafted parameters empirically keeps the train–test loss gap from exceeding 1%.The paper presents the small grafting region as making generalization analysis more tractable and explores it further using classical generalization theory.Source: §4.3 · p. 6

5. Multi-Task and Continual Learning

The section tests skill localization in jointly trained and sequentially trained models. In multi-task models, sparse regions overlap little except for similar tasks and their unions can retain gains for selected tasks; preliminary continual-learning experiments preserve earlier-task performance.

5.1. Multi-Task Learning

  • The multi-task experiment fine-tunes RoBERTa-base with SGD on eight 4096-shot datasets across four task classes, and achieves test accuracy comparable to task-specific fine-tuning.The regions are learned for each task by optimizing its loss from the multi-task model parameters.Source: §5.1 · p. 6
  • In multi-task models, task-specific graft regions overlap little overall, with somewhat greater overlap for similar pairs such as SST-2/CR and SNLI/MNLI.The authors describe this overlap as evidence of skill sharing across similar tasks.Source: §5.1 · p. 7
  • A task’s grafted region generally helps that task and a few related tasks, with transfer sometimes asymmetric.The passage reports transfer for SST-2/CR and SNLI/MNLI, and that the MRPC region helps QQP but not vice versa; QNLI’s region helps QQP paraphrasing.Source: §5.1 · p. 7
  • Unions of task regions retain around 70% of accuracy gains for selected tasks and related tasks, but not for others.Allowing a few gradient steps to purify the union improves accuracy gains on relevant tasks by around 80% without much affecting gains for other tasks.Source: §5.1 · pp. 7–8

5.2. Forget-free Continual Learning

  • Preliminary continual-learning results on QNLI, AG News, then SST-2 retain QNLI performance with grafting, unlike naive fine-tuning’s 20% QNLI accuracy drop.The procedure freezes prior task skills and updates nonintersecting regions for new tasks; inference uses the union of regions, requiring pretrained-parameter resets for each evaluation.Source: Table 4 · p. 8, §5.2 · pp. 8–9

6. Related Work

The related-work section distinguishes skill localization through post-hoc grafting from work on internal knowledge representations, parameter-efficient fine-tuning, lottery tickets, and multi-task learning.

6. Related Work

  • Prior work reports sparse feedforward activations in large pre-trained models, fact-correlated knowledge neurons in BERT, latent knowledge in language-model representations, knowledge localization in pre-trained models’ feed-forward layers, and skill neurons highly predictive of downstream tasks.Source: §6 · p. 9
  • Unlike parameter-efficient fine-tuning, grafting retains a small subset of existing updates post hoc and is framed as a tool for understanding fine-tuning, not training efficiency.The section describes PEFT approaches by Houlsby et al. (2019), Gordon et al. (2020), and BitFit (Ben Zaken et al., 2022).Source: §6 · p. 9
  • Prior multi-task methods often add task-specific parameters, whereas this work analyzes skill localization within a jointly trained model after training.The section cites Misra et al. (2016), Long et al. (2017), Lu et al. (2017), Collobert & Weston (2008), Liu et al. (2016), Gupta et al. (2016), and Liu et al. (2019a).Source: §6 · p. 9

7. Conclusions and future directions

The paper presents sparse grafting as an initial route to localizing fine-tuning skills and sees task-region unions as promising for small multi-task subsets.

7. Conclusions and future directions

  • For small task sets, unions of task-specific regions show promise for models that perform well on any selected subset, avoiding naive training for all 2^m subsets.The paper identifies finding such models for arbitrary task subsets as an open problem.Source: §7 · p. 9

A.1. Architecture Details on Fine-tuning

  • For standard fine-tuning, only the pre-trained model’s embeddings are fixed; prompt-based fine-tuning also fixes the language-model head.For prompt tuning, fixing embeddings improved performance with SGD; fixing the head matched learning it for RoBERTa and reduced region sizes.Source: §A.1 · p. 14

A.2. Hyperparameter settings

  • SGD uses batch sizes 2, 4, or 8, learning rates 10^-2, 5 × 10^-3, or 10^-3, and weight decay 10^-4.Training runs for 16 × n × k steps, capped at 16 × n × 512 for 4096-shot and larger experiments.Source: §A.2 · p. 14

A.3. Calibration Error

The section defines expected calibration error and explains grafting as a way to improve calibration when few-shot data make hold-out calibration difficult. It also describes where graft parameters tend to occur.

A.3. Calibration Error

  • Expected calibration error (ECE) sums each confidence bin’s absolute confidence–accuracy gap, weighted by the fraction of samples in that bin.Predictions are grouped into M equal-width confidence intervals; each bin’s accuracy and average confidence determine its gap.Source: §A.3 · p. 14
  • Grafting provides a direct way to improve calibration with minimal performance drop, where few-shot settings make hold-out calibration difficult.Learning a logit calibration function with ϵ ECE requires at least Ω(1/ϵ^2) samples, while few-shot settings may have fewer than 100 training examples.Source: §A.3 · p. 14

B.1. Distribution of graft parameters

  • Graft parameters concentrate mostly in middle layers for most tasks, except AG News, and chiefly in attention Value, feed-forward first-layer, and LayerNorm parameters.LayerNorm graft parameters are distributed uniformly across layers; Value and feed-forward first-layer distributions vary by task.Source: §B.1 · p. 14, Figure 9 · p. 14

B.2. Grafts with a larger training data

In larger-data settings, sparse grafts retain most fine-tuned performance and improve calibration; the reported graft size does not grow proportionally with training examples.

B.2. Grafts with a larger training data

  • 0.03% grafting regions retrieve >95% of fine-tuning accuracy and are better calibrated than the original model.Table 5 reports a 0.3% accuracy drop and 6% calibration-error improvement for single-sentence tasks, and a 2.8% accuracy drop and 4.0% calibration-error improvement for two-sentence tasks.Source: Table 5 · p. 15
  • The full-shot experiments allow up to 50k training examples per classification class and combine MR, CR, MPQA, and SST-2 into one Sentiment dataset.Training data are split randomly 95%–5% for training and validation; mask training uses 10× as many epochs as in the prior settings.Source: §B.2 · pp. 14–15
  • With a 10× increase in training examples, the necessary grafting region does not increase by 10×; the authors suggest task-specific intrinsic dimensionality for fine-tuning.The reported grafted models achieve >95% of final-model performance and an absolute 4.5% improvement in calibration error.Source: §B.2 · p. 15

B.3. OOD and calibration comparisons with parameter-efficient methods

The section compares grafting with parameter-efficient methods, examines overlap with FISH masks and GPT-2 grafts, and reports additional multi-task findings. Its comparisons indicate that sparsity alone does not explain grafting’s calibration and OOD performance.

B.3. OOD and calibration comparisons with parameter-efficient methods

  • At similar sparsity, grafting improves calibration error by 4.6% with 1.7% lower accuracy at 0.05–0.1% sparsity, and improves OOD performance by 7.9% with 1.0% lower ID accuracy at 0.05–1% sparsity.The comparisons are against LoRA and two adapter methods in the 4096-shot setting.Source: Table 6 · p. 16, Table 7 · p. 16
  • Sparsity alone is not established as the cause of grafting’s better calibration and OOD performance; optimization and the grafting procedure may also contribute.The authors leave deeper exploration of these factors for future work.Source: §B.3 · p. 16

B.4. Overlap with FISH (Sung et al., 2021) mask

  • Grafting regions overlap with FISH masks by less than 10%; sparse FISH masks are not suitable grafting regions, though their grafted-model performance improves as sparsity decreases.The authors infer that full-model fine-tuning need not use the pre-trained model’s parameters with the largest Fisher information.Source: §B.4 · p. 17

B.5. Experiments with auto-regressive language modeling

  • GPT-2 exhibits skill localization with 0.05% of parameters, compared with 0.01% for RoBERTa, while GPT-2 performance is worse than a similarly sized RoBERTa model.The authors note that GPT-2’s denser regions and worse generalization are consistent with their sparsity–generalization connection.Source: §B.5 · p. 17

C.1. Training details for multi-task training.

  • Multi-task training downsamples each dataset to 4096-shot, uses SGD for all tasks with eight times as many iterations, and selects the model with the best average validation performance.Batches are drawn from a randomly selected dataset.Source: §C.1 · p. 17

C.2. Additional Multi-task Experiments

  • Increasing the base graft-region size increases overlap among task-specific grafts, while the cross-task effects reported earlier persist in diminished form.The additional experiments vary the initialization γbase for task-specific regions.Source: §C.2 · p. 17
  • Multi-task AdamW models show skill localization without explicit ℓ1 regularization, unlike the single-task models, which required it for sparse grafts.Source: §C.2 · p. 17

C.3. Details for Continual Learning

The continual-learning procedure allocates each arriving task a grafting region and combines regions for inference; experiments report QNLI retention with minimal effects on other tasks. The generalization discussion connects sparse grafts to small train-test gaps and bounds, while ablations show that localization results depend on optimization settings.

C.3. Details for Continual Learning

  • At inference, the method combines regions for tasks seen so far, while training each arriving task only on parameters outside earlier regions.Retention requires memorizing each grafting region and its task’s arrival sequence number.Source: §C.3 · p. 17
  • The continual method retains QNLI performance with minimal effects on other tasks after sequential training on QNLI, AG news, and SST-2.Each dataset used a 4096-shot training set, and evaluation followed training on SST-2.Source: §C.3 · p. 17
  • When the shared base region grows, task-specific grafts show increasing intersection; the figures compare row-task regions’ effects on column tasks.The captions vary the base-region size and optimizer; the reported increase in intersection is attributed to the common γbase.Source: §C.3 · p. 19, Figure 11 · p. 19, Figure 12 · p. 19

D. Discussions on Generalization Theory

  • In the multi-task model, AdamW shows better skill localization than task-specific models, while explicit ℓ1 regularization does not improve it.The ablations compare task-region overlap and relative accuracy gains across tasks.Source: D. Discussions on Generalization Theory · p. 20, Figure 13 · p. 20, Figure 14 · p. 20
  • The generalization discussion treats a sparse graft as a simpler surrogate class, in contrast to vanilla fine-tuning’s large class and vacuous uniform-convergence bounds.For QNLI in the 4096-shot setting, vanilla fine-tuning reaches 100% training and about 88% test accuracy, leaving a second bound term of at least 12%.Source: D. Discussions on Generalization Theory · pp. 20–21
  • Sparse graft models using about 5,000 parameters have an empirical train-test gap of about 1%, while retaining capacity to perform well on the task.The discussion attributes the small gap to sparse grafts’ inability to fit the training sets perfectly, including in the 4096-shot setting.Source: D. Discussions on Generalization Theory · p. 21

D.1. Generalization Bound for Graft Re-training and Parameter-Efficient FT

The section derives a generalization bound for re-trained graft models under a quantization assumption, separating variance from capacity-related bias. It also extends the bound to directly grafted fine-tuned models under an additional near-ERM assumption.

D.1. Generalization Bound for Graft Re-training and Parameter-Efficient FT

  • The bound assumes bounded loss and that quantizing parameters to q possible values preserves model performance within ε1.The stated quantization condition applies to every model and data point.Source: §D.1 · p. 21
  • The re-trained graft bound separates variance from bias: sparsity controls variance, while the graft region’s capacity contributes to bias.The region count Θn is empirically expected to be small when same-task regions overlap substantially; for fixed-parameter methods such as BitFit, Θn is 1.Source: §D.1 · pp. 21–22
  • The variance term is bounded by eO(s/n), and smaller graft regions are therefore more sample efficient; re-trained grafts empirically have about a 1% train-test gap.Here s is the number of parameters in the region and n is the training-set size.Source: §D.1 · p. 21
  • Using the full model makes the bias term zero but leaves a large variance term, which can be at least about 10% in practice.The discussion attributes this variance to full-network fine-tuning’s ability to fit the training set perfectly.Source: §D.1 · p. 22

D.2. Generalization Bound for Grafted Models

  • The re-training bound also applies to parameter-efficient fine-tuning, while a bound for directly grafted fine-tuned models additionally assumes they are nearly ERM on the training samples.Under that additional assumption, the paper states that the re-trained-model bound transfers to the directly grafted model.Source: §D.1 · p. 22, §D.2 · p. 22

D.3. Proof of Theorem D.2

The proof decomposes the re-trained graft model’s test-loss difference into quantization, training, and generalization terms, then bounds those terms under the stated assumption. A training-trajectory experiment finds that grafted-model performance tracks fine-tuning before saturating earlier.

D.3. Proof of Theorem D.2

  • The proof decomposes the test-loss difference into seven terms, bounding quantization-related terms by ε1 and the re-trained model’s training-loss comparison by zero.The remaining train-test terms are handled using standard generalization arguments.Source: §D.3 · p. 22
  • The proof’s variance bound counts quantized graft models across possible regions, using at most Θn · q^s models for the union bound.The text states the resulting probability guarantee under Assumption D.1.Source: §D.3 · pp. 22–23

D.4. Core skills are learned first.

  • Across SST-2, QNLI, SNLI, and MNLI runs, grafted-model accuracy tracks fine-tuning during training before saturating earlier.Grafting regions were computed at an intermediate checkpoint selected where train and test accuracy appeared to diverge, then applied at successive checkpoints.Source: §D.4 · p. 23, §D.3 · p. 23, Figure 15 · p. 23
Loading 2302.06600v2…