Take Off the Training Wheels! Progressive In-Context Learning for Effective Alignment (2024)

Abstract: Recent studies have explored the working mechanisms of In-Context Learning (ICL). However, they mainly focus on classification and simple generation tasks, limiting their broader application to more complex generation tasks in practice. To address this gap, we investigate the impact of demonstrations on token representations within the practical alignment tasks. We find that the transformer embeds the task function learned from demonstrations into the separator token representation, which plays an important role in the generation of prior response tokens. Once the prior response tokens are determined, the demonstrations become redundant. Motivated by this finding, we propose an efficient Progressive In-Context Alignment (PICA) method consisting of two stages. In the first few-shot stage, the model generates several prior response tokens via standard ICL while concurrently extracting the ICL vector that stores the task function from the separator token representation. In the following zero-shot stage, this ICL vector guides the model to generate responses without further demonstrations. Extensive experiments demonstrate that our PICA not only surpasses vanilla ICL but also achieves comparable performance to other alignment tuning methods. The proposed training-free method reduces the time cost (e.g., 5.45×) with improved alignment performance (e.g., 6.57+). Consequently, our work highlights the application of ICL for alignment and calls for a deeper understanding of ICL for complex generations. The code will be available at https://github.com/HITsz-TMG/PICA.

## introduction in-context learning (icl) has attracted growing attention alongside the scaling of large language models (llms) @xcite . by conditioning on a handful of input-label pairs as examples, corresponding author. llms achieve notable improvements and produce impressive few-shot performance across a range of downstream tasks @xcite . after that, numerous studies have explored the working mechanism of icl and propose several effective methods to enhance icl @xcite @xcite . however, these works mainly focus on classification tasks and simple generation tasks, which limits the exploration of these methods in more complex generation tasks, such as aligning llms with human preferences. as a complex and practical task, alignment typically requires training the model, such as supervise fine-tuning (sft) @xcite and reinforcement learning from human feedback (rlhf) @xcite . a recent work @xcite proposed urial, a simple method using in-context examples to align several powerful base llms and achieves notable instruction-following performance. the success of urial demonstrates the feasibility of in-context alignment and encourages us to explore and optimize icl in the alignment task. in this paper, we investigate the impact of demonstrations during in-context alignment. we visualize the token distribution kl-divergence of instructions and responses in zero-shot and few-shot settings (figure 1 ). to reduce context noise, we set up two few-shot settings with different demonstrations as control groups and have the following observations through comparative experiments: @xcite the model likely stores the task function learned from the demonstration in the separator token representation. @xcite demonstrations play a crucial role in prior response generation but are redundant in posterior response generation. these observations highlight the influence of demonstrations on token representation in icl for alignment tasks, indicating that demonstrations are not always indispensable during the entire response generation stage. motivated by these findings, we propose a progressive in-context alignment (pica) method to enhance both the efficiency and effectiveness of regular icl. specifically, our approach involves a two-stage progressive generation strategy: the few-shot stage and the zero-shot stage. during the few-shot stage, the model generates prior part of the response using the standard icl settings. after generating a specific number of tokens, we transition the model into the zero-shot stage, eliminating the need for further demonstrations to generate the remaining part of the response. to capitalize on the task-related information embedded in the separator tokens, we introduce an icl vector guidance method. inspired by the work of task vector in icl @xcite @xcite , we extract the icl vector from the hidden states of specific transformer layers. this vector is then used to steer the model during the zero-shot stage by intervening in the forward pass. pica minimize the need for demonstrations while improving output quality, thereby reducing the computational cost associated with demonstrations and enhancing overall performance. extensive experiments show that pica outperforms regular icl in both of efficiency and effectiveness. as a trainingfree method, it is also comparable to other alignment methods (i.e., sft and rlhf). for example, on average, our pica boosts the performance of mistral-7b to reach 90% of the performance of gpt-4-0613. these results support our observations and show the effectiveness of our method in various aspects of alignment. additionally, we conduct ablation studies to investigate the robustness and generalizability of our method. our contributions are summarized as follows: in contrast to training-based methods, lin et al. in-context learning working mechanism. recent studies have explored the working mechanisms within icl. several works try to theoretically demonstrate a strong similarity between the attention patterns in icl and the process of gradient descent @xcite . from a more practical perspective, another line of research suggests that the icl may function by learning a mapping function from demonstrations, which it then applies to input queries to make predictions @xcite @xcite . @xcite extract an icl task vector from the hidden states and utilize it for intervention during zero-shot inference. todd et al. (2023) extract a function vector from attention activations using the causal mediation method, which is subsequently added to the hidden states of certain transformer layers during inference. li et al. (2024) derive a state vector from attention activations and propose several optimization strategies. unlike these works, we focus on using comparative experiments to explore the impact of demonstrations on token representation, and leverage these findings to enhance the efficiency of icl. ## motivation in this section, we aim to shed light on the working mechanisms of in-context learning by investigating the following question: what is the impact of demonstration on token representation in incontext alignment? to explore this, we design a comparative experiment to highlight how token representations differ between zero-shot and few-shot settings. we use token probability distributions as a proxy for token representations and utilize kl-divergence to measure the shifts in these distributions. by visualizing and quantifying the shifts in token probability distributions caused by demonstrations, we can understand the role of demonstrations in aligning the model and provide further optimization for in-context alignment. regarding the experimental setup, we randomly selected 100 data instances of similar length from ultra-chat @xcite , a commonly used dataset for alignment tuning, as our experimental dataset. for the input prompt, we use a straightforward design by adding several tokens at the end of the query to serve as separator tokens, explicitly distinguishing between the query and the response. we present the visualization results based on the llama2-7b model in the figure 1 , while the results for other models are provided in appendix c. we break the token distribution of the whole instance into the input and output parts. a straightforward reason is that the input token distribution shift represents differences in understanding the instruction, while the output token distribution shift represents the ability to respond. by observing and analyzing the visualization, we have two hypotheses: (1) the icl alignment task function might be encoded into the separator token representation. (2) the quality of response is highly reliant on the quality of prior response tokens. input token distribution. by comparing the input token probability distributions between zeroshot and few-shot settings, a significant shift is observed in both the prior tokens of the query and the separator tokens. the kl-divergence decreases as the number of query tokens increases. by comparing the experimental group and the control group, we find that the shift in the query distribution also occurs in the control group. however, this shift in the separator tokens is not consistent across different demonstration settings, suggesting distinct underlying causes for these shifts. we attribute the shift in the query's prior token distribution to a "context shift", and we attribute the shift in the separator tokens distribution to a "task shift". given that llms are trained to predict the next token based on the provided context, altering the context directly impacts the token distribution, which we refer to as the "context shift". however, as the number of query tokens increases, the decision space gradually aligns for both zero-shot and few-shot settings, leading to higher consistency in query token prediction and thus a reduced kl-divergence. on the contrary, the trend observed in the query distribution is not mirrored in the separator token distribution. in the control group, the separator token representations remain highly similar. we attribute the large kl-divergence observed in the separator token distribution of the experimental group to the differing tasks, indicating that separator tokens likely encode task-specific information during icl. we reasonably speculate that the primary impact of demonstration on instruction understanding is reflected in the encoding of separator tokens, where the alignment task function learned through icl is stored. this hypothesis aligns with prior work @xcite , yet our findings contribute additional evidence supporting this perspective. ## method observations from §3 reveal that demonstrations are not always indispensable during the entire response generation stage. to minimize the need for demonstrations while preserving alignment per- generated output stage 1: few-shot demonstration query separator token stage 2: zero-shot replace discard figure 2: overview of pica, which include few-shot stage and zero-shot stage. the gray block denotes the hidden state and orange block denotes the separator token hidden state that forms the icl vector. the blue block denotes the generated answer token from few-shot stage. therefore, in the zero-shot stage, the model completes the response based on the existing prior response tokens: in-context learning vector guidance. we observe that transformers show task-specific encoding behavior with the separator token. recent works @xcite have similar observations, demonstrating that functions learned by icl can be represented through compressed vectors derived from transformers and can perform simple generation tasks in zero-shot settings. building on this, we propose the icl vector guidance to assist the model in generating highquality responses during the zero-shot stage. unlike these previous works that intervene single hidden state of the last separator token, we intervene in the initial l layer of all separator tokens. our preliminary experiments found that this method is more effective for the alignment task, where the output is much longer than that of the simple generation tasks focused on in previous works. specifically, during the forward pass in the fewshot generation, we extract the separator token hidden state h few i from the first l layers, which we combine and refer to as the icl vector. subsequently, in the zero-shot stage, we intervene in the separator token representation by replacing the hidden state with the extracted hidden state from the few-shot stage: ## analysis 6.1 layer selection we delve into the impact of layer selection on the extraction of the icl vector. we evaluate the performance based on the win rate compared to gpt3-text-davinci-003 on the alpaca-eval datasets, as shown in figure 3 . our results reveal a dual-phase trend: initially, increasing the number of layers improves performance, but this improvement stops or slightly declines in the later layers. this indicates that the icl function is dynamically stored within the separator token representation. in the initial layers, transformers primarily focus on learning and encapsulating the icl function within the hidden state, where additional layers enhance the richness of the functional information in the icl vector. in contrast, the later layers prioritize applying this learned information for prediction tasks. here, additional layers tend to introduce noise, causing a slight drop in performance. this also suggests that /odpde /odpde 0lvwudoey :lq5dwh 9dqlood,&/ 3,&$ figure 5: the mean and standard error of icl and pica performance with five demonstration across three models. ## conclusion in this paper, we investigate and analyze the impact of demonstrations on token representation in in-context alignment through comparative experiments. based on our observations and analyses, we introduce a novel progressive in-context alignment method that significantly reduces the need for demonstrations while preserving alignment performance. extensive experiments indicate that pica outperforms tuning-free baselines in both effectiveness and efficiency, achieving performance that is better or comparable to sft or rlhf. our experiments and analyses provide in-depth insights for future research on icl in alignment. in the future, we aim to further explore the mechanisms and optimizations of icl in more complex tasks.