Take Off the Training Wheels Progressive In-Context Learning for Effective Alignment

TL;DR

The paper investigates how demonstrations shape in-context learning during alignment tasks and finds that the separator token encodes the task function learned from demonstrations. Building on this, it introduces Progressive In-Context Alignment (PICA), a two-stage generation method that uses an ICL vector extracted from the few-shot stage to guide zero-shot completion, reducing demonstration dependence. Empirical results across multiple models show PICA outperforms vanilla ICL in efficiency and often matches or comes close to training-based alignment methods like SFT and RLHF, achieving substantial speedups (e.g., around 5.45×) and approaching GPT-4 level performance on mid-size models. The work provides actionable insights into the ICL mechanism for complex generation tasks and presents a training-free path toward effective alignment, with broader implications for future research on ICL-driven alignment.

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.

Figures (12)

• Figure 1: The KL-divergence of token probability distributions on Llama2-7b. Experimental Group compares zero-shot and few-shot settings, while Control Group compares two few-shot settings with different demonstrations. We visualize the input and output separately and mark the prior query tokens and prior response tokens with purple circles.
• 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.
• Figure 3: Win rate comparing with GPT-3-text-davinci-003 on alpaca-eval for each choice of the intermediate layer $L$.
• Figure 4: Win rate comparing with GPT-3-text-davinci-003 on alpaca-eval for number of the prior token on three models. We normalize the result with vanilla ICL result.
• Figure 5: The mean and standard error of ICL and PICA performance with five demonstration across three models.