Training Prompt Matters: State-Adaptive Optimization for Robust Fine-Tuning
Abstract
While prompt engineering is instrumental in maximizing the capabilities of Large Language Models (LLMs) during inference, the role of prompts during training remains critically underexplored. Prevailing fine-tuning paradigms typically treat training prompts as mere surface forms, assuming that semantically equivalent instructions yield identical learning outcomes. However, we reveal that this equivalence is deceptive: while paraphrased prompts often lead to comparable in-task performance, they induce drastically different cross-task impacts regarding catastrophic forgetting and generalization. Crucially, these impacts are positively correlated across tasks, indicating the existence of superior prompts that consistently yield better performance. Furthermore, we discover that these superior prompts can be robustly identified by task loss prior to learning. Leveraging these insights, we introduce State-Adaptive Prompt Optimization (SAPO), a lightweight yet effective training strategy that shifts task formulation from a static input to a dynamic, state-adaptive variable. Comprehensive experiments on diverse benchmarks confirm its effectiveness, which significantly mitigates forgetting while improving generalization, achieving substantial performance gains over state-of-the-art methods. These results provide insights into how training prompts shape learning dynamics and offer a practical recipe for robust fine-tuning. Our code is available at https://github.com/Eric8932/SAPO.
Lay Summary
When training AI models, developers typically assume that differently worded instructions with the same meaning will produce identical learning outcomes. However, we discover that paraphrased instructions drastically alter whether an AI forgets older skills or generalizes to new ones. To solve this, we develop State-Adaptive Prompt Optimization (SAPO). Before learning a new task, SAPO tests several wording variations and selects the instruction that best aligns with the AI's current knowledge. This simple adjustment significantly prevents the AI from forgetting previous capabilities while boosting its performance on unseen tasks, providing a practical recipe for building reliable and versatile AI systems.