Unveiling the Flow of Input-label Mappings for In-context Learning
Large Language Model; In-Context Learning
I have read and agree with the venue's withdrawal policy on behalf of myself and my co-authors.
@misc{
sun2024unveiling,
title={Unveiling the Flow of Input-label Mappings for In-context Learning},
author={Chenghao Sun and Zhen Huang and Yonggang Zhang and Le Lu and Houqiang Li and Xinmei Tian and Xu Shen and Jieping Ye},
year={2024},
url={https://openreview.net/forum?id=JxNjaxzVwv}
}
Large language models (LLMs) excel at processing complex tasks by in-context learning (ICL). However, the internal mechanisms behind ICL remain mysteri- ous, e.g., where does the input-label mappings store inside LLMs, and which modules allow for generalizing these mappings to new question. In this work, we make a substantial step towards reverse-engineering the ICL: (1) Applying a linear model, i.e., Principle Component Analysis (PCA), on the hidden states, we find that LLMs distill the semantic mappings into the principal components (PC) at only a small number of layers. (2) Traditional method to identify the ability-related modules relies heavily on pairs of reference and counterfactual samples, which are designed to activate and not activate the behavior, respectively. However, because of the persistent nature of ICL ability, it is difficult to design counterfactual discrete texts which do not involve ICL. To address this, we introduce PC Patching, which engineers the representation with the identified semantic mappings, rather than text. By subtracting the PC from the original features, the input-label mappings are suppressed, resulting in the counterfactual continuous activations for ICL. The results of PC Patching unveil that there is a small fraction (5%) of attention heads that drive LLMs to process the input-label mappings for the final answer. These insights prompt us to investigate the potential benefits of selectively fine-tuning these essential heads to boost the LLMs’ ICL performance. We empirically find that such precise tuning can yield notable enhancements on unseen ICL tasks. The promising applications on other scenarios, i.e., trustworthiness, further validate the effectiveness of our method. Our work serves as an exploration into the ICL and pave the way to scaling ICL for more intricate tasks.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
~Chenghao_Sun1 ~Zhen_Huang4 ~Yonggang_Zhang1 ~Le_Lu3 ~Houqiang_Li1 ~Xinmei_Tian1 ~Xu_Shen1 ~Jieping_Ye4
Chenghao Sun Zhen Huang Yonggang Zhang Le Lu Houqiang Li Xinmei Tian Xu Shen Jieping Ye
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Large Language Model; In-Context Learning
I certify that there is no acknowledgement section in this submission for double blind review.
sun|unveiling_the_flow_of_inputlabel_mappings_for_incontext_learning
interpretability and explainable AI
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Unveiling the Flow of Input-label Mappings for In-context Learning
ICLR 2025 Conference Withdrawn Submission
ICLR.cc/2025/Conference/Withdrawn_Submission