In-Context Learning State Vector with Inner and Momentum Optimization

TL;DR

The paper investigates how in-context learning in transformer models can be understood and improved via a compact state vector extracted from early attention activations. It introduces inner optimization, momentum optimization, and divide-and-conquer aggregation to refine and scale this state vector for test-time adaptation, achieving substantial gains on zero-shot and few-shot tasks across multiple models and datasets. By linking the state vector to a dual form of gradient descent and demonstrating improved robustness and efficiency, the work offers both practical ICL enhancements and a deeper mechanistic perspective on how demonstrations shape predictions. The findings suggest a principled path for scalable, interpretable ICL interventions and motivate further exploration on larger models and theoretical foundations.

Abstract

Large Language Models (LLMs) have exhibited an impressive ability to perform In-Context Learning (ICL) from only a few examples. Recent works have indicated that the functions learned by ICL can be represented through compressed vectors derived from the transformer. However, the working mechanisms and optimization of these vectors are yet to be thoroughly explored. In this paper, we address this gap by presenting a comprehensive analysis of these compressed vectors, drawing parallels to the parameters trained with gradient descent, and introduce the concept of state vector. Inspired by the works on model soup and momentum-based gradient descent, we propose inner and momentum optimization methods that are applied to refine the state vector progressively as test-time adaptation. Moreover, we simulate state vector aggregation in the multiple example setting, where demonstrations comprising numerous examples are usually too lengthy for regular ICL, and further propose a divide-and-conquer aggregation method to address this challenge. We conduct extensive experiments using Llama-2 and GPT-J in both zero-shot setting and few-shot setting. The experimental results show that our optimization method effectively enhances the state vector and achieves the state-of-the-art performance on diverse tasks. Code is available at https://github.com/HITsz-TMG/ICL-State-Vector

Figures (10)

• Figure 1: The overall framework of the proposed state vector. The state vectors are extracted from the output activations of attention heads. These state vectors are progressively optimized by inner optimization and momentum optimization, or be aggregated through a divide-and-conquer (D&C) aggregation. Finally, the processed state vector is utilized to intervene the inference forward pass.
• Figure 2: Performance of aggregation across number of examples. Avg. denotes the average aggregation baseline and D&C. denotes the divide-and-conquer aggregation. The X axis represents the number of examples, and the Y axis represents the accuracy.
• Figure 3: Performance comparison of gradient optimization algorithms. The method means the optimization algorithm applied to the $\texttt{opt}(\cdot)$ in Eqn. \ref{['grad_eqn']}.
• Figure 4: Average zero-shot performance across six datasets for each choice of the intermediate layer $L$. The solid line means the average value, while the shaded area indicates the standard deviation.
• Figure 5: The 2D PCA visualization of the state vector in the Antonym ,English-French and Product-Company task, where each color represents the state vector corresponding to examples occupying specific positions in the demonstration and the outlier is the first order.