Mechanistic Fine-tuning for In-context Learning

TL;DR

This paper introduces Attention Behavior Fine-Tuning (ABFT), a cost-effective method to improve in-context learning in language models by adjusting attention scores, leading to better performance and interpretability.

Contribution

The paper proposes ABFT, a novel fine-tuning approach that modifies attention scores to enhance ICL, reducing computational costs and revealing insights into model mechanisms.

Findings

ABFT outperforms previous methods in accuracy, robustness, and efficiency.

ABFT requires only 0.01% of the data cost of prior approaches.

Analysis suggests ICL biases relate to the emergence of induction heads.

Abstract

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.