ProMoT: Two-Stage LLM Fine-Tuning Framework

• ProMoT is a two-stage framework that separates prompt tuning from model tuning to prevent format specialization in LLMs.
• It preserves the inherent in-context learning abilities by confining format biases to a low-capacity prompt buffer.
• Empirical results show ProMoT boosts performance on tasks like RTE and translation while improving out-of-domain generalization.

Prompt Tuning with MOdel Tuning (ProMoT) is a two-stage fine-tuning framework designed to address format specialization in LLMs and improve their generalization to out-of-domain and in-context learning tasks. ProMoT achieves minimally specialized downstream performance without degrading or erasing the in-context learning capabilities intrinsic to pretrained LLMs. By explicitly separating format learning from semantic adaptation, ProMoT offers substantial empirical and practical benefits for both single-task and multi-task settings (Wang et al., 2022).

1. Motivation and Problem: Format Specialization in LLM Fine-Tuning

Standard fine-tuning of LLMs, such as mT5 or PaLM, on a single-task dataset rapidly erodes the model's in-context learning capabilities due to format specialization. When exposed to tasks with restricted output formats—e.g., Recognizing Textual Entailment (RTE) with binary labels “True” or “False”—the model disproportionately overfits to those output formats. Empirical evidence demonstrates that after a few hundred fine-tuning steps on RTE, the LLM attains approximately 90% RTE accuracy, but accuracy on one-shot TriviaQA, a general question answering task, collapses toward zero. Additionally, the fraction of TriviaQA outputs equal to “True” or “False” surpasses 90% in just 300 fine-tuning steps, demonstrating that the model has largely forgotten its prior generative skill set. Gradient-alignment studies show that the initial phases of fine-tuning are dominated by format learning, with gradients from the true task and from randomized-label variants (sharing only output format) being nearly colinear (Wang et al., 2022).

2. The ProMoT Framework: Two-Stage Separation

ProMoT comprises two sequential stages: prompt tuning and model tuning. The central premise is to offload task-specific format learning into a detachable, low-capacity soft prompt, thereby preserving the generalization capacity of the main model parameters.

Stage	Tunable Parameters	Objective
Prompt Tuning	Soft prompt $p$	Minimize $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$, freeze $\theta$
Model Tuning	Model $\theta$	Minimize $$L_\text{model}(\theta) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p^*;x])]$$, freeze $p^*$

• Soft Prompt Definition: $p \in \mathbb{R}^{p \times e}$, where $p$ is prompt length and $e$ is embedding dimension.
• Prompt Tuning: Model parameters $\theta$ are frozen; $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$0 is trained to absorb format-related patterns.
• Model Tuning: $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$1 is frozen; $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$2 is updated to learn semantics, conditioned on the prompt carrying format knowledge.

ProMoT's two-phase procedure is operationalized as follows:

$\theta$2 (Wang et al., 2022)

3. Theoretical Insights and Underlying Principles

ProMoT is motivated by the observation that format learning sharply dominates gradient dynamics in early standard fine-tuning. By explicitly channeling these gradients into the low-capacity prompt buffer $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$3, ProMoT prevents irreversible specialization in the high-capacity model parameters $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$4. Since $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$5 is explicitly limited in expressive power, it absorbs superficial features of the target task—such as label set, token structure, and output length—without confounding the semantic mapping capacity of $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$6. Subsequent model tuning (with $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$7 frozen) allows $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$8 to focus on the underlying semantics without reinforcing undesirable format behaviors. This mechanism enables $$L_\text{prompt}(p) = \mathbb{E}_{(x,y) \sim D} [-\log P_{f_\theta}(y|[p;x])]$$9 to preserve and enhance out-of-domain generalization, including in-context learning on tasks not encountered during fine-tuning (Wang et al., 2022).

4. Implementation and Practical Considerations

Empirical evaluations employ mT5 XXL models with soft prompt length $\theta$0, embedding size $\theta$1, and prompt initialization in $\theta$2. Prompt tuning is performed with Adafactor optimizer at $\theta$3, 5k steps, batch size 64. Model tuning utilizes Adafactor at $\theta$4, 1k steps, with dropout 0.1 and label smoothing 0.1. Inputs are truncated or padded for a maximum length of 1024 tokens, and prompts are prepended at the embedding layer. At inference, if the output format of the test task aligns with the fine-tuning dataset $\theta$5, $\theta$6 is attached; otherwise, it is removed and only $\theta$7 is used (Wang et al., 2022).

Recommended prompt length is 50–100 embeddings. Stage 1 should absorb format biases over approximately 5k steps, while stage 2 is kept shorter (∼1k steps) to minimize re-specialization. Robust learning rate defaults are $\theta$8 and $\theta$9.

5. Empirical Results and Applications

On RTE, standard fine-tuning achieves 92.06% accuracy; ProMoT attains 92.78%. For WMT14 En–Fr translation, ProMoT matches or slightly exceeds BLEU performance (41.80 vs. 41.30). In the domain of in-context generalization, dramatic improvements are observed. Averaged over eight unseen one-shot tasks (NLI, QA, translation, summarization), the pretrained mT5 attains a normalized average of 17.52. Fine-tuning on RTE degrades this to 15.43 (−2.09), while ProMoT recovers and improves to 20.10 (+2.58). On En–Fr, standard fine-tuning collapses to 9.15 (−8.37); ProMoT restores to 18.87 (+1.35). An extension, “ProMoT + 1-shot” (adding a single natural-language demonstration to model tuning), further increases generalization to as high as 22.3. These findings are robust across multiple seeds and model scales (mT5 XL, T5 XXL, PaLM 8B). In joint (multi-task) RTE+En–Fr training, multi-task ProMoT achieves a normalized average of 25.88 (+8.35 over pretrained), substantially exceeding standard multi-task fine-tuning (Wang et al., 2022).

6. Extensions, Limitations, and Practical Guidelines

ProMoT interfaces effectively with multi-task learning. In cross-task settings, ProMoT-trained models on binary NLI not only excel on NLI but also enhance summarization (e.g., XSum, WikiLingua) beyond the zero-shot pretrained baseline, suggesting cross-format and cross-semantic transfer effects. Recommended usage includes scenarios where both strong supervised performance and preservation of few-shot or zero-shot generalization are required. Limitations include: no formal guarantee on the proportion of specialization captured by $\theta$0; some added complexity in managing $\theta$1 during inference; and evaluation scale up to 13B parameters, though the approach is conceptually model-agnostic. A plausible implication is that ProMoT's separation principle could benefit other domains where task format learning compromises generalization (Wang et al., 2022).