Iterative Forward Tuning Boosts In-Context Learning in Language Models (2024)

Abstract: Despite the advancements in in-context learning (ICL) for large language models (LLMs), current research centers on specific prompt engineering, such as demonstration selection, with the expectation that a single iteration of demonstrations processing can generalize effectively to a given test sample. However, this perspective overlooks the potential benefits derived from multiple iterations involving demonstrations, a practice aligning more closely with the iterative decision-making process exhibited by humans, who often learn through analogy. In this study, we introduce a novel two-stage framework to boost ICL in LLMs. Specifically, our framework delineates the ICL process into two distinct stages: Deep-Thinking and test stages. The Deep-Thinking stage incorporates a unique attention mechanism, i.e., iterative enhanced attention, which enables multiple rounds of information accumulation. This mechanism operates by manipulating the Key-Value matrices without training, fostering enhanced understanding capabilities in LLMs by thinking demonstrations multiple times. We evaluated Deep-Thinking across a range of benchmarks and LLMs, showing its superior performance over vanilla ICL methods and its effectiveness in challenging tasks where demonstration selection is infeasible.

## introduction large language models (llms), e.g. openai gpts @xcite , llama @xcite and qwen @xcite , demonstrate the mysterious in-context learning (icl) ability, where llms make predictions directly by prepending demonstrations to the original input without updating model parameters. llms are expected to learn the patterns hidden in demonstrations and make predictions accordingly. as illustrated in figure 1 (a), an llm can correctly perform inference on an unseen task by conditioning on several demonstrations. the icl paradigm empowers llms to achieve impressive results in various downstream tasks with a few demonstrations, making language-model-as-a-service (lmaas) @xcite possible. since the performance of icl is sensitive to specific prompt settings, considerable efforts have been developed to improve the performance of icl by refining the prompt design from different perspectives, such as demonstration selection @xcite , instruction design @xcite , and intermediate chain-of-thought (cot) reasoning @xcite @xcite . these methods can facilitate llms to reduce inference variance and avoid poor worst-case accuracy to some extent by performing prompt engineering. the working mechanism of icl also draws a lot of attention. @xcite shed light on the connections between icl and explicit fine-tuning. specifically, icl computes meta-gradients via forward computation, while explicit fine-tuning obtains gradients by back-propagation. a dual form exists between attention and gradient descent-based optimization @xcite , directly connecting the test input to demonstrations. @xcite argue that label words in demonstrations act as anchors, enabling mapping from demonstrations to test input through information aggregation and label propagation. however, these studies assume that the models process demonstrations only once (i.e., perform a single forward computation), which is incoordinate with the human decision-making process by learning from analogy. humans usually learn from analogy via an iterative thinking process, such as analyzing demonstrations, reflecting on them, and forming abstract concepts. the models learned from demonstrations in inference time by "thinking for longer" or "thinking multiple times" @xcite . these findings inspire us to ask a question: can we boost the performance of icl by learning from demonstrations through several (iterative) forward inferences? in this paper, we propose a two-stage framework to boost the icl ability in llms. instead of simply concatenating demonstrations and test input together for inference, we decouple the icl process into a deep-thinking stage for demonstration training and a test stage, as illustrated in figure 1 (b) . in the deep-thinking stage, we introduce a new attention module that manipulates the updates of key-value matrices @xcite within the transformer's self-attention @xcite mechanism. this modification leverages key-value matrices as a bridge to change the information flow to accumulate and learn information over multiple forward iterations without any training. during the test stage, since the concepts contained in demonstrations are already stored in final key-value matrices, we only need to feed the test input into the model and utilize the key-value cache for inference. this deep-thinking strategy is motivated by humans' repeat logical thinking and reasoning process. llms are expected to extend their abilities to solve unseen, complex tasks by "thinking" demonstrations multiple times. to verify the effectiveness of the proposed deep-thinking, we initially conduct evaluations on conventional icl benchmarks across language models of various sizes. the experiments show that deep-thinking significantly outperforms vanilla icl in a variety of model sizes and tasks, surpassing previous state-of-the-art (sota) methods focused on selecting demonstrations. in addition, we introduce two more challenging benchmarks (i.e., mmlu @xcite and bbh (srivastava et al., 2023)) and conduct experiments on advanced llms, including llama2 @xcite and pythia @xcite . we argue that on these challenging benchmarks, demonstration selection becomes impractical due to the lack of a potential candidate pool. deep-thinking obtains a significant advantage over vanilla icl. ## preliminaries: in-context learning this paper focuses on in-context learning tasks. formally, given a nature language test input x test with a few (n -shot) input-output demonstrations c @xmath0, y i )} n @xmath1, the goal of in-context learning is to predict the label ŷ of x test from a predefined candidate label set @xmath2, y 2 , ..., y m } conditioned on n demonstrations. given an llm m (e.g., a gpt model), the prediction process can be formulated as follows: ## methodology in this paper, we propose a two-stage icl framework that improves performance through multiple forward iterations. as shown in figure 2 , we assign the demonstrations and test input to the deep-thinking and test stages, respectively, where key-value matrices serve as a bridge between the two stages. next, we describe these two stages in detail. ## related work in-context learning (icl) with llms has made a breakthrough and become mainstream in tackling various tasks @xcite @xcite . recently, great efforts have been made to improve the performance of icl from different perspectives, such as demonstrations selection @xcite , prompt template design @xcite , and intermediate chain-of-thought (cot) reasoning @xcite . for demonstration selection, liu et al. (2022) performed demonstration selection through a knnbased retriever, choosing the closest example to test input. @xcite proposed self-adaptive icl with a general select-and-rank framework for demonstration selection. in addition to example selection, @xcite investigated the sensitivity of icl to the permutation of demonstrations and proposed entropy metrics to determine their order. the above icl methods are usually restricted by the number of demonstrations. to mitigate such a challenge, @xcite attempted to scale icl by grouping demonstrations, which could increase the number of demonstrations to 1,000. the formatting function also plays a crucial role in icl, especially for tasks requiring complex reasoning steps, such as commonsense reasoning. wei et al. (2022b) @xcite @xcite attempt to uncover the underlying working mechanism of icl. in particular, @xcite showed that icl happened via bayesian inference, where certain concepts were implicitly predicted before the final prediction. subsequently, @xcite revealed that there are connections between icl and explicit fine-tuning and explained llms as meta-optimizers @xcite . unlike existing methods, to the best of our knowledge, we are the first to decouple icl into two stages and focus on how to deeply learn from fixed demonstrations rather than on demonstration selection or prompt engineering. this is advantageous for the world situation where provided samples are scarce, i.e., there is no large candidate set of demonstrations. ## conclusion in this paper, we introduce a novel two-stage framework aimed at enhancing the icl capabilities of llms by leveraging iterative forward inferences to learn from demonstrations. by decoupling the icl process into a dedicated deep-thinking stage for demonstration training and a subsequent test stage, we effectively mimic the decision-making process of humans by learning from analogy. this approach aligns with how humans engage in repeated logical thinking. the empirical evaluations across conventional icl benchmarks and more challenging datasets demonstrate that our deep-thinking strategy significantly outperforms previous icl approaches, particularly in scenarios where demonstration selection is impractical.