Mechanistic Fine-tuning for In-context Learning
This work addresses the efficiency problem for researchers and practitioners using large language models by offering a more cost-effective fine-tuning approach, though it is incremental as it builds on existing mechanistic insights.
The paper tackles the high computational cost of fine-tuning language models for in-context learning by proposing Attention Behavior Fine-Tuning (ABFT), which focuses on adjusting attention scores rather than final outputs, resulting in improved performance, robustness, unbiasedness, and efficiency with only about 0.01% data cost compared to previous methods.
In-context Learning (ICL) utilizes structured demonstration-query inputs to induce few-shot learning on Language Models (LMs), which are not originally pre-trained on ICL-style data. To bridge the gap between ICL and pre-training, some approaches fine-tune LMs on large ICL-style datasets by an end-to-end paradigm with massive computational costs. To reduce such costs, in this paper, we propose Attention Behavior Fine-Tuning (ABFT), utilizing the previous findings on the inner mechanism of ICL, building training objectives on the attention scores instead of the final outputs, to force the attention scores to focus on the correct label tokens presented in the context and mitigate attention scores from the wrong label tokens. Our experiments on 9 modern LMs and 8 datasets empirically find that ABFT outperforms in performance, robustness, unbiasedness, and efficiency, with only around 0.01% data cost compared to the previous methods. Moreover, our subsequent analysis finds that the end-to-end training objective contains the ABFT objective, suggesting the implicit bias of ICL-style data to the emergence of induction heads. Our work demonstrates the possibility of controlling specific module sequences within LMs to improve their behavior, opening up the future application of mechanistic interpretability.