Mechanistic Fine-tuning for In-context Learning

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Cho, Hakaze, Luo, Peng, Kato, Mariko, Kaenbyou, Rin, Inoue, Naoya
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866909811700924416
author Cho, Hakaze
Luo, Peng
Kato, Mariko
Kaenbyou, Rin
Inoue, Naoya
author_facet Cho, Hakaze
Luo, Peng
Kato, Mariko
Kaenbyou, Rin
Inoue, Naoya
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_14233
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Mechanistic Fine-tuning for In-context Learning
Cho, Hakaze
Luo, Peng
Kato, Mariko
Kaenbyou, Rin
Inoue, Naoya
Computation and Language
Artificial Intelligence
Machine Learning
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.
title Mechanistic Fine-tuning for In-context Learning
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2505.14233