PaperlayerView exact revision on arXiv

Source-linked AI summary

Two-stage LLM Fine-tuning with Less Specialization and More Generalization

Yihan Wang, Si Si, Daliang Li, Michal Lukasik, Felix Yu, Cho-Jui Hsieh, Inderjit S Dhillon, Sanjiv Kumar

arXiv:2211.00635v3cs.CLcs.LG

TL;DR

Fine-tuning can reduce an LLM’s general in-context learning by causing early format specialization. ProMoT first learns task format in removable soft prompts, then fine-tunes the model, preserving comparable supervised performance while reducing specialization and improving related-task generalization.

  • Problem

    Fine-tuning specialized tasks can harm pretrained LLMs’ general in-context learning, limiting their ability to handle unseen tasks without additional fine-tuning data.

    Source: §1 · p. 2

  • Method

    ProMoT uses two-stage fine-tuning: prompt tuning first absorbs task-specific format information, followed by model fine-tuning with the soft prompt attached.

    Source: §1 · p. 2

  • Results

    ProMoT reduces specialization while reaching similar supervised performance, and improves in-context generalization on semantically related tasks with different formats.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    ProMoT offers a way to build more general fine-tuned LLMs and can further improve generalization when combined with multi-task training.

    Source: §1 · p. 3, §6 · p. 9

    ProMoT cannot completely eliminate specialization and may still harm some unseen in-context learning tasks depending on the fine-tuning task.

    Source: §5.3 · p. 8

Abstract

from arXiv · show

Pretrained large language models (LLMs) are general purpose problem solvers applicable to a diverse set of tasks with prompts. They can be further improved towards a specific task by fine-tuning on a specialized dataset. However, fine-tuning usually makes the model narrowly specialized on this dataset with reduced general in-context learning performances, which is undesirable whenever the fine-tuned model needs to handle additional tasks where no fine-tuning data is available. In this work, we first demonstrate that fine-tuning on a single task indeed decreases LLMs' general in-context learning performance. We discover one important cause of such forgetting, format specialization, where the model overfits to the format of the fine-tuned task.We further show that format specialization happens at the very beginning of fine-tuning. To solve this problem, we propose Prompt Tuning with MOdel Tuning (ProMoT), a simple yet effective two-stage fine-tuning framework that reduces format specialization and improves generalization.ProMoT offloads task-specific format learning into additional and removable parameters by first doing prompt tuning and then fine-tuning the model itself with this soft prompt attached. With experiments on several fine-tuning tasks and 8 in-context evaluation tasks, we show that ProMoT achieves comparable performance on fine-tuned tasks to standard fine-tuning, but with much less loss of in-context learning performances across a board range of out-of-domain evaluation tasks. More importantly, ProMoT can even enhance generalization on in-context learning tasks that are semantically related to the fine-tuned task, e.g. ProMoT on En-Fr translation significantly improves performance on other language pairs, and ProMoT on NLI improves performance on summarization. Experiments also show that ProMoT can improve the generalization performance of multi-task training.

1 INTRODUCTION

Fine-tuning improves task performance but can damage general in-context learning through early format specialization. ProMoT addresses this by separating format learning into removable prompt parameters before tuning the model, reducing specialization while preserving or improving generalization.

1 INTRODUCTION

  • Fine-tuning on a single task reduces general in-context learning and can over-specialize models to a few downstream tasks.The paper motivates techniques that retain performance on supervised tasks while preserving unseen-task capabilities.Source: §1 · p. 2
  • Format specialization is an important cause of forgetting and emerges early, before the model fully learns task semantics.For example, binary classification can constrain outputs to “True” and “False,” reducing flexibility across other task formats.Source: §1 · p. 2
  • ProMoT first tunes removable soft prompts with the model frozen, then fine-tunes the main model while keeping those prompt parameters fixed.This design offloads task-specific format learning into additional parameters that can be removed when incoming tasks use different formats.Source: §1 · p. 2
  • ProMoT improves generalization on semantically related tasks with different formats, including summarization after NLI fine-tuning.The authors report better supervised performance than pretrained models and better general in-context learning than standard fine-tuning.Source: §1 · p. 2
  • ProMoT reduces specialization while reaching similar supervised performance across experiments on more than 10 NLP tasks.The framework can also combine with multi-task fine-tuning and 1-shot prompts to further improve generalization on unseen tasks.Source: §1 · p. 3

2 RELATED WORK

Prior work studies catastrophic forgetting, parameter-efficient tuning, and multi-task generalization, while ProMoT targets format differences across substantially different NLP tasks. It is presented as complementary to data-distribution changes and multi-task training.

2 RELATED WORK

  • ProMoT is orthogonal to methods that alter pretraining data distributions because it does not require changing the data distribution.This distinguishes the framework from approaches that ease tension between in-context learning and fine-tuning through distribution design.Source: §2 · p. 3
  • Earlier forgetting studies often examine tasks with similar formats, whereas this work considers contrasts such as classification versus long-form generation.The paper emphasizes that output format is especially important in these more divergent task pairs.Source: §2 · p. 3
  • Prompt tuning, adapters, and LoRA preserve more of the pretrained model but can underperform full fine-tuning and miss opportunities for broader improvement.ProMoT is positioned as a framework that retains parameter-efficient components while still fine-tuning the main model.Source: §2 · p. 3
  • Multi-task fine-tuning improves unseen-task performance, while ProMoT improves generalization for individual fine-tuning tasks and can be combined with multi-task training.Unlike approaches relying on engineered prompts, ProMoT learns soft prompts from data and uses them to absorb format information.Source: §2 · p. 3

3 FORMAT SPECIALIZATION IN FINE-TUNING CAUSES THE LOSS OF IN-CONTEXT LEARNING CAPABILITIES

Single-task fine-tuning reduces unseen-task in-context performance largely through rapid format specialization, which begins before semantic learning. Experiments show that models increasingly emit the fine-tuning task’s format even for unrelated prompted tasks.

3 FORMAT SPECIALIZATION IN FINE-TUNING CAUSES THE LOSS OF IN-CONTEXT LEARNING CAPABILITIES

  • The section establishes three findings: in-context abilities are lost during fine-tuning, format specialization contributes importantly to that loss, and format specialization occurs at the beginning.These findings motivate separating format learning from later semantic learning.Source: §3 · p. 3, §3.1 · p. 4, §3.2 · p. 4

3.1 LOSS OF IN-CONTEXT LEARNING CAPABILITIES DURING FINE-TUNING

  • The RTE experiment uses a 13B mT5 XXL model trained with PaLM’s default hyperparameters and input/output template.RTE requires predicting whether two sentences are entailed using “True” or “False.”Source: §3.1 · pp. 3–4
  • Fine-tuning improves RTE accuracy while drastically reducing performance on unseen 1-shot QA tasks.The experiment fine-tunes a 13B mT5 XXL model on RTE and evaluates TriviaQA and web_questions.Source: §3.1 · pp. 3–4, Figure 1 · p. 4

3.2 FORMAT SPECIALIZATION

  • Format specialization means overfitting to shared sequence characteristics such as language, lengths, styles, special tokens, punctuation, and output labels.Because single-task examples share these characteristics, their gradient signal favors the fine-tuning format even on semantically related tasks with different formats.Source: §3.2 · p. 4
  • RTE fine-tuning drives 90% of 1-shot TriviaQA outputs to “True” or “False” after 300 steps, and over 99% of 1-shot En-De outputs after 500 steps.This output collapse indicates specialization to the binary-label format rather than flexible response formatting.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 FORMAT LEARNING HAPPENS FIRST DURING STANDARD FINE-TUNING

  • Format learning dominates early because most fine-tuning examples share similar formats, producing a stronger gradient signal than nuanced semantic content.For RTE, format is membership in {True, False}, whereas semantics concern the input-label correlation.Source: §3.3 · p. 4
  • With randomized RTE labels isolating format, full and format-only gradients have cosine similarity approximately 1 at step 0 but about 0.21 after 400 steps.The alignment decline coincides with the True/False ratio approaching 100%, indicating that format is learned first.Source: §3.3 · p. 5, Figure 4 · p. 5

4 PROPOSED METHOD: PROMPT TUNING WITH MODEL TUNING (PROMOT)

ProMoT separates early format learning from later model learning using a trainable soft prompt followed by model fine-tuning with that prompt fixed. The framework is compatible with other tuning methods and removes the prompt at inference when the task format differs.

4 PROPOSED METHOD: PROMPT TUNING WITH MODEL TUNING (PROMOT)

  • ProMoT first learns task format in a soft prompt with the pretrained model frozen, then fine-tunes the model while keeping that prompt attached and fixed.This design aims to let model updates focus on semantic skills not already captured by the prompt.Source: §4 · pp. 5–6
  • Gradient alignment between full and format-only gradients is high at the start of standard fine-tuning but falls substantially by 400 steps, motivating ProMoT’s stage separation.The method uses this early format-learning pattern to allocate format information to separate parameters.Source: Figure 3 · p. 5
  • Stage 1 prepends a trainable continuous soft prompt to embedded inputs and optimizes it while the pretrained model remains frozen.The prompt is a small set of trainable embeddings, and its limited capacity is expected to restrict semantic learning.Source: §4 · pp. 5–6
  • Stage 2 freezes the learned soft prompt and fine-tunes the pretrained model with the prompt still prepended, encouraging learning beyond the stored format information.The prompt remains part of the input during model tuning rather than being discarded between stages.Source: §4 · pp. 5–6
  • ProMoT can combine with parameter-efficient methods such as LoRA, multi-task fine-tuning, and 1-shot prompting, although prompt tuning empirically absorbs early format information better than LoRA.The framework is presented as compatible with different tuning techniques in its respective stages.Source: §4 · p. 6
  • At inference, the learned prompt is used only when the task shares the fine-tuning format; otherwise it is removed and the original input is fed to the fine-tuned model.This makes the format-specific parameters removable for differently formatted tasks.Source: §4 · p. 6

5 EXPERIMENTS

Experiments compare ProMoT with standard and prompt tuning across fine-tuning tasks and unseen in-context tasks. ProMoT preserves supervised performance while reducing specialization and improving generalization, including in multitask settings.

5.2 SUPERVISED PERFORMANCE ON FINE-TUNING TASKS

  • ProMoT outperforms standard fine-tuning on 3 of 4 supervised target tasks and prompt tuning on all 4, without sacrificing fine-tuning performance.The comparison uses accuracy for RTE, SNLI, and OpenbookQA, and BLEU for WMT14 En-Fr.Source: §5.2 · p. 7, Table 2 · p. 7

5.3 GENERALIZATION WITH SINGLE TASK FINE-TUNING

  • Standard fine-tuning significantly reduces unseen-task in-context performance, reaching near-zero few-shot performance on 6 of 8 RTE evaluations.CB and WiC are exceptions because they share the RTE task format.Source: Table 3 · p. 7, §5.3 · p. 8
  • ProMoT reduces this loss and can improve semantically related tasks with different formats, including summarization after NLI fine-tuning and translation after English–French fine-tuning.The reported gains include summarization beyond the pretrained model and improved English-to-German and English-to-Romanian performance.Source: §5.3 · p. 8, Table 4 · p. 8

5.4 MORE GENERALIZATION WITH MULTITASK TRAINING

  • Multitask ProMoT significantly outperforms standard multitask fine-tuning on unseen-task generalization, with larger average improvements; adding 1-shot prompts boosts both methods further.The experiment mixes RTE and WMT14 En-Fr and evaluates normalized average performance.Source: §5.4 · p. 8, Table 5 · p. 9

5.5 ABLATION STUDY

  • Jointly tuning the soft prompt and model still causes specialization, while a fixed random prompt does not help; the two-stage process and learned prompt are therefore central.Natural-language prompts also perform far worse than ProMoT in the 1-shot comparison.Source: §5.5 · p. 9, Table 6 · p. 9

6 CONCLUSIONS AND LIMITATIONS

The paper identifies early format specialization as a cause of lost in-context ability and presents ProMoT as a way to preserve and extend generalization. It also notes missing theoretical guarantees, limited model scale, and possible misuse.

6 CONCLUSIONS AND LIMITATIONS

  • ProMoT uses trainable soft prompts to absorb task-specific formats before model fine-tuning, reducing specialization and enabling generalization across different tasks.The authors report effectiveness across diverse NLP tasks and describe the method as promising for fine-tuning with small datasets.Source: §6 · p. 9
  • The method has no theoretical guarantee for how much format specialization the first prompt-tuning stage can absorb, and experiments use models smaller than 15B parameters.The authors propose testing ProMoT on larger models.Source: §6 · p. 9

A BROADER IMPACTS

  • Improved general-purpose models may be misused to generate disinformation, motivating watermarking or AI-generated-text classifiers before release.Source: A BROADER IMPACTS · p. 14

B EXPERIMENT DETAILS

The appendix documents task templates, output processing, datasets, hyperparameters, training schedules, and implementation resources used to reproduce the experiments.

B.1 INPUT TEMPLATE USED IN EXPERIMENTS

  • The experiments use task-specific natural-language input templates and extract outputs between T5 sentinel tokens before trimming task-specific prefixes.Classification outputs are evaluated by checking whether the first output token is True or False.Source: §B.1 · p. 14, Table 7 · p. 14, §B.2 · p. 14

B.3 DATASET AND MODELS

  • The study reports dataset split statistics and uses mT5-based datasets and models available for research purposes.Source: §B.3 · p. 14, Table 8 · p. 14

B.4 HYPER-PARAMETERS

  • ProMoT tunes a 100-token soft prompt for 5,000 steps, then freezes it and fine-tunes the model for 1,000 steps, selecting checkpoints by target-task performance.Standard fine-tuning and prompt tuning use 1,000 and 5,000 steps respectively for comparison.Source: §B.4 · pp. 14–15
  • The joint prompt-and-model ablation fine-tunes both components for 1,000 steps with learning rate 0.001.Source: §B.4 · p. 15

B.5 HARDWARE AND IMPLEMENTATION

  • All experiments run on 64 parallel TPUs, with each training experiment completed within one day.Source: §B.5 · p. 15

C.1 ADDITIONAL RESULTS ON SINGLE TASK FINE-TUNING

Additional experiments on SNLI, OpenbookQA, and multitask evaluation support the main finding that ProMoT preserves supervised performance while improving unseen-task generalization.

C.1 ADDITIONAL RESULTS ON SINGLE TASK FINE-TUNING

  • The appendix compares prompt-tuning-plus-1-shot and ProMoT-plus-1-shot on fine-tuning tasks as complementary results to the main single-task experiments.Source: §C.1 · p. 15
  • Additional SNLI and OpenbookQA experiments are consistent with the main result: ProMoT maintains similar supervised performance while reducing forgetting and improving general in-context performance.Pretrained mT5 obtains zero accuracy on 0-shot OpenbookQA because it fails to output the required A–D labels.Source: §C.1 · p. 15, Table 10 · p. 15

C.2 4-SHOT EVALUATION RESULTS OF MULTI-TASK TRAINING

  • Additional 4-shot results compare pretrained, standard multitask, and multitask ProMoT models using normalized average performance.Source: §C.2 · p. 16, Table 11 · p. 16

C.3 TRAINING MORE STEPS: TRADE-OFF BETWEEN FINE-TUNING TARGET TASK AND IN-CONTEXT LEARNING ABILITIES

Longer training reveals a trade-off between fine-tuning-task performance and in-context learning, with ProMoT reducing forgetting at comparable target-task performance.

  • Figure 5 compares En-Fr BLEU with average classification accuracy during fine-tuning.Source: Figure 5 · p. 16, §C.3 · p. 16
  • Figure 6 compares En-Fr BLEU with average BLEU on other language pairs during fine-tuning.Source: Figure 6 · p. 16, §C.3 · p. 16
  • ProMoT checkpoints achieve higher in-context performance than standard fine-tuning at comparable fine-tuning-task performance.This indicates reduced forgetting in both classification and translation evaluations.Source: §C.3 · p. 16

C.4 ADDITIONAL EXPERIMENTS ON T5 XXL

Additional experiments test ProMoT across model sizes, architectures, first-stage parameterizations, and longer gradient-alignment trajectories, finding consistency with the main results and weaker performance for LoRA than prompt tuning.

C.4 ADDITIONAL EXPERIMENTS ON T5 XXL

  • On T5 XXL fine-tuned on RTE, the additional results are consistent with the main experiments.The evaluation covers the fine-tuned task and eight 1-shot tasks.Source: §C.4 · p. 17, Table 12 · p. 17

C.5 ADDITIONAL EXPERIMENTS ON MT5 XL

  • On mT5 XL fine-tuned on WMT14 En-Fr, the additional results are consistent with the main mT5 XXL experiments.The evaluation covers the fine-tuned task and eight 1-shot tasks.Source: §C.5 · p. 17, Table 13 · p. 17

C.6 ADDTIONAL EXPERIMENTS ON PALM 8B

  • On PaLM 8b fine-tuned on WMT14 En-Fr, ProMoT achieves similar supervised performance with less forgetting on general in-context tasks.The experiment uses prompt length 50 and learning rate 0.3 during prompt tuning.Source: §C.6 · p. 17, Table 14 · p. 18

C.7 USING LORA IN THE FIRST STAGE

  • Figure 7 measures cosine similarity between full and format gradients across MLP kernels and attention Query, Key, and Value components.Source: Figure 7 · p. 18
  • Prompt tuning in ProMoT’s first stage outperforms LoRA on both the supervised task and unseen 1-shot evaluations.The authors suggest soft prompts may align better with format descriptions in natural-language corpora.Source: §C.7 · pp. 17–18, Table 15 · p. 18

C.8 PLOTTING MORE STEPS FOR FIGURE 3

  • Gradient alignment drops significantly after 300 steps, matching the increase in the model’s true-to-false output ratio before that point.Source: §C.8 · p. 19

C.9 QUALITATIVE RESULTS ON FINE-TUNING WMT14 EN-FR TASK

Qualitative examples show that standard En-Fr fine-tuning transfers translation-specific output behavior to unrelated few-shot tasks, whereas ProMoT produces more task-appropriate outputs.

  • Standard En-Fr fine-tuning overfits format by producing French or repeating input-like translation behavior, while ProMoT alleviates this specialization.The comparison uses outputs from ground-truth targets, pretrained mT5, fine-tuned mT5, and ProMoT mT5.Source: §C.9 · p. 19
  • For a German summarization example, fine-tuning outputs a French translation, while ProMoT produces a German summary-like response.Source: §C.9 · p. 19
  • For a question-answering example, fine-tuning repeats the question-answer format, whereas ProMoT returns a concise answer related to the target.Source: §C.9 · p. 19
  • For an instruction-style task, fine-tuning produces French output, whereas ProMoT generates an English procedural response closer to the target format.Source: §C.9 · p. 19
Loading 2211.00635v3…