Improving Input-label Mapping with Demonstration Replay for In-context Learning (2023)

Abstract: In-context learning (ICL) is an emerging capability of large autoregressive language models where a few input-label demonstrations are appended to the input to enhance the model’s understanding of downstream NLP tasks, without directly adjusting the model parameters. The effectiveness of ICL can be attributed to the strong language modeling capabilities of large language models (LLMs), which enable them to learn the mapping between input and labels based on in-context demonstrations. Despite achieving promising results, the causal nature of language modeling in ICL restricts the attention to be backward only, i.e., a token only attends to its previous tokens, failing to capture the full input-label information and limiting the model’s performance. In this paper, we propose a novel ICL method called Repeated Demonstration with Sliding Causal Attention, (RdSca). Specifically, we duplicate later demonstrations and concatenate them to the front, allowing the model to ‘observe’ the later information even under the causal restriction. Besides, we introduce sliding causal attention, which customizes causal attention to avoid information leakage. Experimental results show that our method significantly improves the input-label mapping in ICL demonstrations. We also conduct an in-depth analysis of how to customize the causal attention without training, which has been an unexplored area in previous research.

## introduction large language models (llms) have become the backbone of various natural language processing * work done during an internship at meituan. despite the promising results demonstrated by existing icl models @xcite , their causal nature in language modeling restricts each token's attention solely to its preceding tokens. as a result, these models fail to capture the complete input-label information, thereby limiting their overall performance. specifically, the pre-training objective of the current autoregressive llms focuses on predicting future tokens based on past ones @xcite , implemented with causal attention. while this approach works well for modeling regular sequences, it becomes less effective when applied to icl tasks. the limitation of causal attention restricts icl demonstrations to having only left context, which hampers the model's ability to fully comprehend and exploit the input-label relationship. unlike tokens in a sentence that possess sequential dependencies, there is no inherent sequential relationship between the demonstrations in the icl input. therefore, it is desirable for these demonstrations to interact with one another comprehensively, rather than relying solely on later demonstrations attending to earlier ones, while the reverse is not possible. intuitively, if we can enable each demonstration to attend to all the others, we can potentially obtain a more sophisticated context for the icl query. however, achieving this on an llm pre-trained with the objective of causal language modeling is not straightforward. simply removing the causal restriction and allowing the model to have access to the right context during inference is not feasible, as it would result in a significant disparity between training and inference conditions. in this work, we focus on capturing full inputlabel mapping information from demonstrations. to achieve this target, we propose two techniques. the first is repeated demonstration, where we replicate later demonstrations and concatenate them to the front. this allows the model to 'observe' the later information even under the causal restriction. for example, if we consider four demonstrations represented by however, simply duplicating demonstrations brings about a new problem: we do not want to attend a demonstration twice, as this may cause the model to take shortcuts by learning to repeat the answer of its first encounter, rather than learning the input-label mapping. to address this, we propose the second technique, the sliding causal attention, which customizes the original causal attention by restricting the attention window so that each demonstration can only attend to all other demonstrations once. in the case of four demonstrations, attention windows are through experiments, we demonstrate that our proposed method (repeated demonstrations with sliding causal attention, rdsca) significantly enhances the ability to learn input-label mapping from icl demonstrations. our proposed sliding casual attention is the first attempt that customizes the causal attention in the inference stage without further training. we investigate a number of different designs for customizing causal attention and reach some principles for the success of icl. for example, we find that the @xmath0. it should always be available to attend to no matter where the attention window slides. besides, the size of the attention window determines the richness of the semantic context, thus it affects the performance greatly. our contributions are summarized as: ## conclusion in this study, we introduce rdsca, a novel icl framework that enhances the learning of the inputlabel mapping from demonstrations. we propose that the causal attention mechanism of decoderonly llms restricts the model from fully exploiting the input-label mapping from demonstrations. to address this, we suggest repeating demonstrations to allow each demonstration to have full context and customizing the vanilla causal attention mechanism to prevent information leakage. experimental results show that our method consistently improves the input-label mapping ability of icl on llms of different scales. furthermore, we delve deeper into causal attention customization and show how different attention settings affect icl performance. additionally, this work is the first to customize the causal attention of a pre-trained autoregressive llm without further training, which may pave the way for further research in this direction.