Technical Summary: Training Prompt Matters: State-Adaptive Optimization for Robust Fine-Tuning

1. Problem Statement

While prompt engineering is widely recognized as critical for maximizing Large Language Model (LLM) performance during inference, its role during the training phase remains underexplored. Prevailing fine-tuning paradigms typically treat training prompts as static, arbitrary inputs, operating under the assumption that semantically equivalent instructions yield identical learning outcomes.

This paper challenges that assumption, revealing that while paraphrased prompts may produce comparable in-task performance, they induce drastically different cross-task impacts. Specifically, the choice of training prompt significantly influences catastrophic forgetting of previously learned tasks and generalization to unseen tasks. The core problem addressed is the lack of a systematic approach to identify and utilize "superior" training prompts that align with the model's evolving state to mitigate interference and enhance robustness during continual learning.

2. Methodology: State-Adaptive Prompt Optimization (SAPO)

The authors propose State-Adaptive Prompt Optimization (SAPO), a lightweight, plug-and-play training strategy that shifts task formulation from a static input to a dynamic, state-adaptive variable. The methodology consists of three primary steps executed prior to learning each new task:

1. Prompt Expansion: A paraphrasing model generates a pool of semantically equivalent candidate prompts (e.g., 20 candidates) based on the original task instruction.
2. State-Adaptive Alignment Evaluation: The method evaluates the alignment between each candidate prompt and the model's current state. This is achieved by computing the pre-update loss (causal language modeling loss) on a small subset of the training data for each candidate.
3. Optimized Formulation Integration: The prompt yielding the lowest pre-update loss is selected for the subsequent fine-tuning phase. This same prompt is used for both training and evaluation of that specific task.

Mechanism: The paper posits that low-loss prompts indicate that the necessary task-specific logic is already well-aligned with the model's intrinsic capabilities. By selecting these prompts, SAPO minimizes the need for the model to internalize extensive, conflicting task-specific adaptations, thereby reducing inter-task gradient conflicts and facilitating the acquisition of generalizable knowledge.

3. Key Contributions

A. Systematic Study of Training Prompt Impact

The paper provides the first systematic investigation into the role of training prompts in LLM fine-tuning. Through extensive probing experiments, it demonstrates that:

• In-task Stability vs. Cross-task Sensitivity: Semantically equivalent prompts yield negligible differences in performance on the current task but induce drastic variability in forgetting (on prior tasks) and generalization (on unseen tasks).
• Positive Correlation: Variations in prompt choice are not random; there is a consistent positive correlation where prompts that mitigate forgetting also tend to facilitate generalization.

B. Existence and Identification of Superior Prompts

The study establishes the existence of "superior" training prompts that consistently improve cross-task performance. Crucially, it identifies pre-update loss as a robust, reliable proxy for selecting these prompts. The authors demonstrate a strong positive correlation between lower pre-learning loss and higher post-training performance on non-training tasks across diverse models (Llama-2, Qwen3) and task sequences.

C. The SAPO Framework

SAPO is introduced as a method to dynamically align task instructions with the model's state. Unlike existing continual learning methods that focus on parameter regularization, data replay, or modularization, SAPO operates orthogonally by optimizing the input formulation itself. It transforms fixed, state-agnostic training processes into state-adaptive ones without requiring complex architectural changes.

4. Experimental Results

Comprehensive evaluations were conducted on diverse benchmarks, including SuperNI (covering classification, generation, and mixed sequences) and TRACE (incorporating complex tasks like math reasoning and code generation). The experiments utilized models ranging from 7B to 14B parameters (Llama-2, Qwen3).

• Performance Gains: SAPO consistently improved the performance of various state-of-the-art continual learning baselines (including LoraInc, EWC, O-LoRA, and InsCL) across all metrics: Average Performance (AP), Backward Transfer (BWT), and Forward Transfer (FWT).
• Robustness: The method proved effective across different task sequence types (homogeneous generation, homogeneous classification, and mixed) and model families.
• Ablation Studies:
  • Prompt Optimization vs. Pessimization: Selecting the prompt with the highest loss (SAPP) resulted in performance declines, confirming that gains are driven by the specific efficacy of selecting low-loss prompts, not merely avoiding poor human prompts.
  • Paraphraser Robustness: SAPO remained effective regardless of the paraphrasing model used (Gemini, GPT, Qwen), indicating the framework does not rely on a specific generator's capabilities.
  • Mechanism Analysis: Gradient angle analysis revealed that low-loss prompts reduce the cosine similarity angle between gradients of the current task and prior/future tasks, indicating reduced optimization conflict.

5. Significance and Claims

The paper claims that training prompt formulation is a critical yet underexplored dimension in LLM fine-tuning. Its significance lies in several areas:

• Redefining Fine-Tuning Dynamics: It highlights that fine-tuning primarily modulates input activation pathways rather than creating new capabilities. Strategic prompt variation allows for the identification of pathways with minimal conflicts, managing interference during fine-tuning.
• Practical Recipe for Robustness: SAPO offers a practical, low-overhead strategy (adding only ~10-12% time overhead in empirical tests) to significantly mitigate catastrophic forgetting and improve zero-shot generalization.
• Safety and Alignment: The authors suggest that by preserving pre-existing safety guardrails and ethical constraints through reduced interference, SAPO contributes to the development of more trustworthy AI systems.
• Future Directions: The work opens new avenues for state-aware data formulation, suggesting that the principles of adaptive prompt optimization could extend beyond robust fine-tuning to wider training scenarios.

The paper concludes that the choice of training prompt is not merely a surface-level adjustment but a pivotal lever for orchestrating a model's global capability landscape during fine-tuning.