DETAIL: Task DEmonsTration Attribution for Interpretable In-context Learning

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 10:15:02

Abstract: In-context learning (ICL) allows transformer-based language models that are pre-trained on general text to quickly learn a specific task with a few "task demonstrations"  without updating their parameters, significantly boosting their flexibility and generality. ICL possesses many distinct characteristics from conventional machine learning, thereby requiring new approaches to interpret this learning paradigm. Taking the viewpoint of recent works showing that transformers learn in context by formulating an internal optimizer, we propose an influence function-based attribution technique, DETAIL, that addresses the specific characteristics of ICL. We empirically verify the effectiveness of our approach for demonstration attribution while being computationally efficient. Leveraging the results, we then show how DETAIL can help improve model performance in real-world scenarios through demonstration reordering and curation. Finally, we experimentally prove the wide applicability of DETAIL by showing our attribution scores obtained on white-box models are transferable to black-box models in improving model performance.

Main paper: 
 introduction: 
 The rapid development of transformer-based language models (Bommasani et al., 2022; Brown et al., 2020; Chowdhery et al., 2022) has inspired a new in-context learning (ICL) paradigm (Brown et al., 2020) , which allows a language model sufficiently pre-trained on general text to quickly adapt to specific tasks. This lightweight approach of customizing a general model for specific tasks is in contrast to fine-tuning (Hu et al., 2022; Xia et al., 2024) that necessitates both access to the model parameters and resourceintensive step of tuning these parameters for model adaptation. In ICL, a few task demonstrations are included in the input text (i.e., prompt) together with a query to help the language model better understand how to answer the query. It has been shown that including task demonstrations in the prompt can enhance the capability of language models to apply common sense and logical reasoning (Wei et al., 2023; Yao et al., 2023) and learn patterns from the supplied demonstrations (Brown et al., 2020) , significantly enhancing the flexibility and generality of language models. In the ICL paradigm, each demonstration can be viewed as a "training data point" for ICL. Analogous to how the performance of a conventional supervised machine learning (ML) model depends on the quality of training data, the performance of ICL depends on the quality of task demonstrations (Liu et al., 2023a) . A research question naturally arises: How to attribute and interpret ICL demonstrations that are helpful or harmful for model prediction?Though there are many prior works on interpreting and attributing model prediction for conventional ML models (Ghorbani & Zou, 2019; Koh & Liang, 2017; Ribeiro et al., 2016) , these methods are not readily applicable to ICL due to its unique characteristics. Firstly, many existing attribution techniques require either computing the gradients (Sundararajan et al., 2017) or multiple queries to the model (Cook, 1977) , both of which are slow and computationally expensive. In contrast, ICL is often applied in real-time to a large foundation model (Bommasani et al., 2022) that necessitates the attribution approaches for ICL to be fast and efficient. Secondly, ICL is known to be sensitive to ordering: The same set of demonstrations can result in significantly different model performance under different permutations (Liu et al., 2023b; Lu et al., 2022) . However, conventional methods do not explicitly consider the ordering of training examples. Thirdly, ICL demonstration is usually supplied as a sentence comprising a sequence of tokens, rendering conventional token-level attribution methods ineffective, as they do not capture contextual information of each ICL demonstration (Bahdanau et al., 2016; Sundararajan et al., 2017) . Lastly, ICL does not update model parameters, rendering conventional techniques that analyze model parameter change (Koh & Liang, 2017 ) not applicable. Moreover, the absence of the need to update model parameters also allows a good attribution result for ICL to be transferable across different language models.To address these challenges, we propose DETAIL, a novel technique that takes advantage of a classical attribution approach while tackling the unique characteristics of ICL. We adopt the perspective that transformers formulate an internal optimizer (Akyürek et al., 2023; Chen et al., 2024; Von Oswald et al., 2023; Xie et al., 2022) for ICL. Based on this internal optimizer, we design a method to understand the impact of each demonstration on the transformer's prediction. Notably, this approach allows us to leverage powerful existing analysis tools for transformer-based models in ICL, where otherwise the characteristics of ICL make applying these tools difficult. Specifically, we describe an intuitive (re-)formulation of the influence function (Koh & Liang, 2017) , a popular attribution method for conventional ML, on the internal optimizer and show that DETAIL addresses the challenges of computational cost, sensitivity to order, and attribution quality. Then, we empirically verify that our formulation can identify demonstrations helpful for model prediction and outlier demonstrations. Additionally, we apply our method to tasks with real-world implications including prompt reordering, noisy demonstration detection, and demonstration curation to show its effectiveness. We demonstrate that DETAIL achieves improved performance when applied to typical white-box large language models (LLMs). Furthermore, as many powerful LLMs are currently closed-source (thus black-box), we show that our DETAIL score obtained on a white-box LLM (e.g., Vicuna-7b (Zheng et al., 2023) ) exhibits transferable characteristics to the performance on a popular black-box model (ChatGPT). 
preliminaries: 
 In-context learning (ICL). ICL is a learning paradigm that provides a few task demonstrations with formatted input and output for a pre-trained transformer (e.g., an LLM) to learn the mapping function from inputs to outputs in context (i.e., via the forward pass) (Brown et al., 2020) . Formally, a transformer model f takes in a prompt p(S, x query ) comprising a formatted sequence of demonstrations S = (z 1 , z 2 , . . . , z n ) where z i = {x i , y i } from a specific downstream task along with a query input x query (i.e., prompt) to predict its label as ŷquery = f (p(S, x query )). An visual for an example prompt is provided in App. D. We wish to attribute the model prediction ŷquery to each z i ∈ set(S).ICL as implementing an internal optimizer. With the growing interest in the internal mechanism of transformers, previous works (Akyürek et al., 2023; Von Oswald et al., 2023; von Oswald et al., 2023; Zhang et al., 2024a) have theoretically shown that ICL can be treated as the transformer implementing an internal optimizer on the ICL demonstrations. Specifically, (Von Oswald et al., 2023; von Oswald et al., 2023; Zhang et al., 2024a) formulated the objective of ICL optimizer with (regularized) mean-squared error on a linear weight applied to the token itself in linear self-attentions (LSAs) (Zhang et al., 2024a) or a transformation of tokens (i.e., kernelized mean-squared error) if an extra multi-layered perception is attached before the LSA (Von Oswald et al., 2023, Proposition 2) in a recurrent transformer architecture. The transformer layers then function as performing gradient descent on the weight to minimize the objective (Von Oswald et al., 2023; Zhang et al., 2024a) .Influence function. Influence function (Koh & Liang, 2017) approximates the change of the loss of a test data point z test when up-weighting a training data point z i . Formally, the influence of z i on predicting z test is1 EQUATIONwhere L(z test , θ) refers to the loss (function) on a test point z test of the model parameterized by θ andH θ := 1/n n i=1 ∇ 2 θ L(z i , θ) is the Hessian. However, this definition cannot be directly applied to ICL since there is no model parameter change during ICL, unlike the learning settings in (Barshan et al., 2020; Koh & Liang, 2017) . We show how to adapt the formulation of Eq. ( 1) to ICL in our proposed method DETAIL next.Detailed discussion about prior works that attempted to understand ICL and to provide data attribution (e.g., influence function) is included in App. C. 
influence function on internal kernel regression: 
 Following the idea that transformers learn in context by implementing an internal kernelized least-square objective, we present our formulation of DETAIL by computing the influence function on a kernelized linear regression (Hainmueller & Hazlett, 2014) . Specifically, we build the regression w.r.t. the following kernelk(x, x ′ ) := m(x) ⊤ m(x ′ ) (2)where m(x) ∈ R 1×d refers to (the mapping of an ICL demonstration 2 to) an internal representation of x (e.g., hidden state of a transformer layer) with output dimension d.Let X := (x 1 , x 2 , • • • , x n ) and Y := (y 1 , y 2 , • • • , y n )be the vectors of inputs and outputs in S respectively. The equivalent kernel regression can be written as Ŷ := m(X)βwhere β ∈ R d×1 is the weight vector over the kernelized feature space. In practice, the dimension d of m is usually much larger than the number of demonstrations, causing severe over-parameterization. Such over-parameterization renders the influence values fragile (Basu et al., 2021) . As such, we follow (Basu et al., 2021) and adopt an ℓ 2 regularization on β controlled by a hyper-parameter λ, which forms a kernelized ridge regression (Murphy, 2012) with loss:EQUATIONTaking the 2nd derivative of Eq. ( 3), we obtain the hessian H β asEQUATIONAdopting a matrix multiplication form for the summation in Eq. ( 4), we write the influence of training data on the model parameters β as follows,EQUATIONwhere K := m(X) ⊤ m(X) ∈ R d×d is the Gram matrix and I reg ∈ R d refers to the influence of a particular demonstration (x, y) w.r.t. the kernel regression weights β. Then, combining Eqs. (1), ( 3) and ( 5), we can express the DETAIL score as the influence of a demonstration z on a query z test :EQUATIONwhere β has a closed-form expression (shown in Alg. 1 in App. B). While inverting matrices in Eq. ( 6) requires O(d 3 ) time, d is usually in the thousands: A typical LLM like Llama-2-7b (Touvron et al., 2023) has an embedding size d = 4096, allowing reasonable computation time of I (e.g., a few seconds). In practice, this computation is accelerated by the techniques already implemented in existing scientific computation libraries admitting sub-cubic complexity for matrix inversion.Computing self-influence. One important application of the influence function in ML is identifying outliers via selfinfluence (Barshan et al., 2020; Koh & Liang, 2017) . The conventional definition of I trivially admits computing selfinfluence simply by replacing z test in Eq. ( 1) with z i . While the same approach applies to DETAIL, there are two shortcomings: (i) As the embedding is sensitive to the position, placing the same demonstration at the end of the prompt (as a query) or in the middle (as a demonstration) results in different embeddings, leading to unreasonable influence score.(ii) For each demonstration, it needs one forward pass of the model to compute the self-influence, which can be costly when the ICL dataset size is large. Instead, we implement self-influence for DETAIL by reusing the demonstration's embedding. This way, we keep the two sides of Eq. (1) consistent and only require one forward pass of the model to compute the self-influence for all demonstrations.Further speedup via random matrix projection. While the current formulation in Eq. ( 6) is already computationally cheap, a relatively large embedding size (e.g. d = 4096 for Llama-2-7b) can become a bottleneck as inverting the matrix can be relatively slow. We apply an insight that for ICL, much of the information in the embedding m is redundant (we do not need a 4096-dimensional β to fit 20 demonstrations). Hence, we project m to a much lower dimensional space via a random matrix projection while preserving the necessary information, following the Johnson-Lindenstrauss lemma (Dasgupta & Gupta, 2003; Johnson & Lindenstrauss, 1984) , precisely represented as a projection matrix P ∈ R d×d ′ with each entry i.i.d. from N (0, 1/d ′ ). We provide a more detailed discussion in App. D. Empirically, we show that we can compress m to a much smaller dimension d ′ ≤ 1000, resulting in a 10× computation speedup on a typical 7B LLM on an NVIDIA L40 GPU (see Sec. 4.2).A visualization of our proposed method is in Fig. 1 and a pseudo-code implementation is in App. B. 
empirical evaluation: 
 We evaluate the effectiveness of DETAIL (i.e., Eq. ( 6)) on two metrics: computational time (via logged system time required) and effectiveness in attribution (via performance metrics for tasks). We start by visualizing how the DETAIL scores, particularly I(z test , z i ) (test influence, abbreviated as I test ) and I(z i , z i ) (self influence, abbreviated as I self ) following (Koh & Liang, 2017, Sections 5.1 & 5.4) , attribute demonstrations to a query first on a custom transformer and then on LLMs. Note that the hyper-parameter λ varies under different scenarios and we discuss some heuristics for setting λ in App. D.Enforcing ICL behavior. We consider tasks where transformers learn from the demonstrations and form an internal optimizer. To evaluate the effectiveness of our method, we enforce the ICL behavior by mapping the labels of the demonstrations to a token that carries no semantic meaning, This way, pre-trained transformers cannot leverage memorized knowledge to produce answers but have to learn the correct answer from the supplied demonstrations. Specifi- before and after the random projection since the projection is optional.cally, we map all labels to one of {A, B, C, D, E} depending on the number of classes. More details about the specific mapping are in each section.We use the MNIST dataset (Deng, 2012) to visualize how I test can be applied to attribute model prediction to each demonstration. We design a task where the transformer needs to learn a mapping of the digit image to a label letter in context. Specifically, for each image of 28 × 28 pixels, we flatten the pixels and concatenate them with a mapped label to form a 785-dimensional token vector. For simplicity of illustration, we only use images of digits in {0, 1}. For each ICL dataset, we assign to each distinct digit a letter in {A, B, C, D, E} randomly. We build a recurrent transformer based on the design in (Von Oswald et al., 2023) with 10 recurrent layers each consisting of 15 attention heads. We pre-train the transformer with random label mappings from randomly selected digits so that the transformer cannot memorize the mapping but has to infer from the demonstrations. We use the hidden state after the 1st layer as m to compute I test . A qualitative visualization of I test 's attribution and a quantitative plot showing how the test prediction varies by removing demonstrations with the highest (lowest) I test are in Fig. 2 and Fig. 3 respectively. Fig. 2 shows that removing tokens (represented by the image pixels and the corresponding label) with the largest I test makes the model make wrong predictions, whereas removing tokens with the lowest I test can retain the correct model predictions. Fig. 3 shows that removing tokens with the lowest I test results in a slower decrease in prediction accuracy than removing the highest I test , demonstrating that tokens with the highest I test 's are more helpful and vice versa.With the insight obtained in Sec. 4.1, we now apply DETAIL to full-fledged LLMs. We start with demonstration perturbation and noisy demonstration detection to demonstrate that DETAIL can be used to interpret the quality of a demonstration. Then, leveraging these results, we further show how we can apply the DETAIL scores to tasks more closely related to real-world scenarios.A distinction between LLMs and the custom transformer used above is that demonstrations in LLMs are usually a sequence of tokens, whereas demonstrations in the custom transformer are single tokens representing the actual numerical values of the problems (see Fig. 3 ). This distinction makes it difficult to find an appropriate internal representation for each demonstration (i.e., m). To overcome this challenge, we draw inspiration from prior works (Wang et al., 2023; Xie et al., 2022) which suggest that information flows from input words to output labels in ICL. As an implementation detail, we take the embedding of the last token before the label token (hereafter referred to as the target position) in the middle layer where most of the information has flown to the target token positions. We include ablation studies on using different layers' embeddings in App. E.6 and using different target positions in App. E.7.Setup. We consider (for white-box models) mainly a Vicuna-7b v1.3 (Zheng et al., 2023) and also a Llama-2-13b (Touvron et al., 2023) on some tasks using greedy decoding to show that DETAIL works on models with different training data and of varying sizes. While our theoretical backing stands for transformer-based models, we experiment DETAIL on Mamba-2.8b (Gu & Dao, 2023) , a state-space model architecture that has received increased attention recently in App. E.8. We primarily evaluate our method on AG News (4 classes) (Zhang et al., 2015) , SST-2 (2 classes) (Socher et al., 2013) , Rotten Tomatoes (2 classes) (Pang & Lee, 2005) , and Subj (2 classes) (Conneau & Kiela, 2018) datasets which all admit classification tasks. Due to space limits, some results are deferred to App. E.Demonstration perturbation. We show that DETAIL can explain LLM predictions by showing how perturbation (i.e., corrupting the labels of some demonstrations to an incorrect class or removing some demonstrations) with the high/low I test affects the model's predictive power. We randomly pick 20 ICL datasets each comprising 20 demonstrations and 1 query from AG News and find the average and standard error of the accuracy of predicting the query after perturbation using Vicuna-7b and Llama-2-13b, shown in Fig. 4 (results for other datasets deferred to App. E.1). It can be observed that perturbing demonstrations with low I test results in a slower drop (or even improvement) in accuracy and vice versa, similar to the trend observed in Fig. 3 , showcasing the applicability of DETAIL to LLMs. We also experiment on Falcon-7b (Almazrouei et al., 2023) and Llama-2-7b (Touvron et al., 2023) in App. E.1 where we perturb 10 demonstrations and observe a similar accuracy gap.Noisy demonstration detection. We utilize the DETAIL score to detect noisy demonstrations with corrupted labels. The experiment setup largely follows (Koh & Liang, 2017, Section 5.4 ). We randomly draw 100 ICL datasets each consisting of 20 demonstrations and 1 query. For each ICL dataset, we randomly corrupt the labels of 4 demonstrations (i.e., flipping the label to an incorrect class). The demonstrations are then ranked in descending order of their I self . The fraction of noisy demonstrations detected is plotted in the first 3 figures of Fig. 5 (result for other datasets deferred to App. E.2). We compare our method with the leave-one-out (LOO) score (Cook, 1977) where the difference in crossentropy loss of the model output is used as the utility. be observed that LOO performs close to random selection, whereas our method has a much higher identification rate w.r.t. the number of demonstrations checked. We also note that our method not only outperforms LOO in effectiveness but is also around 10× faster than LOO which requires multiple calls to LLM inference for each demonstration in the ICL dataset.Dimension reduction via random projection. We analyze the impact of random projection on the effectiveness of DETAIL. Intuitively, dimension reduction trades off the effectiveness of DETAIL with computational efficiency, specifically the O(d 3 ) cost for inverting K I . To understand the trade-off, we follow the same setup as the noisy demonstration detection experiment and compare the change in AUC ROC of detection and system wall time as the dimension d ′ of the projection matrix P decreases.The result for Subj is in the last figure of Fig. 5 (results for other datasets deferred to App. E.3), showing that wall time stays minimal (≈ 0.3s) for project dimensions up to 1000 generally before it exponentially increases. Effectiveness measured in terms of AUC reaches optimal with d ′ ≥ 1000.The results suggest a "sweet spot" -d ′ ≈ 1000 -for a low running time and high performance. stration detection and I test for demonstration perturbation, can be applied to real-world scenarios, achieving superior performance and speed.ICL order optimization. One distinctive trait of ICL compared to conventional ML is that the order of demonstrations affects the model's predictive performance (Liu et al., 2023b; Lu et al., 2022) . We show that I self helps reorder the demonstrations with improved model predictive performance. We first show, using a Vicuna-7b model, that moving demonstrations with lower quality to the front (or back) of the prompt tends to improve the test accuracy of the model. To see this, we corrupt the label of a random demonstration and allocate this corrupted demonstration to different positions of the ICL dataset (each with 20 demonstrations with permutations drawn from a Sobol sequence (Mitchell et al., 2022) to capture the average performance better). A general trend with decreasing-then-increasing accuracy can be observed in Fig. 6 : Allocating noisy demonstrations to the front (or the back) results in much higher test accuracy. Leveraging this insight, we utilize I self to reorder a random permutation of ICL demonstrations and show the reordered prompt improves the test accuracy. For each randomly ordered prompt, I self for each demonstration is computed (note that this computation only requires 1 pass of the LLM). Then, based on the trend observed in Fig. 6 , for Subj and Rotten Tomatoes datasets, the demonstrations are reordered by placing the two demonstrations with the largest I self in front followed by the rest in ascending order. For SST-2, the demonstrations are reordered in descending order of I self . We compare our DETAIL score with other metrics proposed for demonstration attribution/selection or can be directly extended to attributing demonstrations. We analyze both the attributability via a demonstration curation experiment and the computational cost via recording the system wall time for performing the attribution. We select representative conventional approaches from different paradigms, including integrated gradients (IG) (Sundararajan et al., 2017) and LIME (Ribeiro et al., 2016) (Attention (Bahdanau et al., 2016) in App. E.5). As these methods are originally designed for token-level attribution, we use the sum of the scores of all tokens in each demonstration as the attribution Task DEmonsTration Attribution for Interpretable In-context Learning score. We also compare recent efforts on demonstration selection (Chang & Jia, 2023; Nguyen & Wong, 2023; S. et al., 2024) . We select an ICL dataset of 20 demonstrations, compute the attribution scores on a validation set of 20 demonstrations, and record the accuracy after removing 10 demonstrations in place on 120 test queries. The results are tabulated in Table 2 . For hyper-parameter, we choose M = 100 and k = 1 for (Nguyen & Wong, 2023) , 5 iterations of LiSSA update (Agarwal et al., 2017) for (S. et al., 2024) , and M = 10, K = 4 for datamodel (Chang & Jia, 2023) . We do not perform batch inference for a fair comparison of computational time as some approaches do not admit batch inference. We also use a projection of d ′ = 1000 to compute I test . It can be observed that DETAIL outperforms all other attribution methods in test accuracy. Our computation is efficient (with wall time of 4 ∼ 10s), achieving over 5× speedup compared to other methods except (S. et al., 2024) which achieves a comparable computational time to ours but a lower accuracy. Notably, IG and LIME perform close to random removal, which is likely because these methods are designed for token-level attribution and generalize poorly to demonstration-level attribution. We evaluate the transferability of DETAIL on GPT-3.5, 5 a popular black-box model. We experiment with both the demonstration reordering and demonstration curation tasks where we compute the DETAIL scores on a Vicuna-7b model (white-box) and then test the performance on GPT.Our method produces promising results on both tasks as shown in Table 3 and Table 4 respectively. Notably, curating demonstrations with our method achieves a 17.9% average improvement in accuracy compared to random cu- 5 We use gpt-3.5-turbo-1106. See https://platform. openai.com/docs/models/gpt-3-5-turbo.rating on the demonstration curation task (Table 4 ). We also note an over 2% improvement in accuracy for reordering task if we corrupt 3 demonstrations (Table 3 ). With no corrupted demonstration, reordering with our approach does not improve performance on GPT-3.5, which we attribute to the stronger inference power of GPT-3.5, resulting in less variance w.r.t. demonstration orders, consistent with the findings in (Lu et al., 2022) .  
conclusion, limitation, and future work: 
 We tackle the problem of attributing demonstrations in ICL for transformers. Based on the well-known influence function commonly used for attributing conventional ML, we propose DETAIL, an innovative adoption of the influence function to ICL through the lens of treating the transformer as implementing an internal kernelized ridge regression.Combined with a dimension reduction technique using random projection, DETAIL can be computed in real-time with an impressive performance on various real-world related tasks such as demonstration order optimization and demonstration curation. One limitation of our approach is the need to access the internal state of the transformer, which we mitigate by additionally showing that DETAIL scores are transferable to black-box models. As a first step toward attributing demonstrations w.r.t. a transformer's internal optimizer, we hope this work serves as a building block for future research to develop attribution techniques for more generalized in-context learning settings such as chain-of-thought (Wei et al., 2023) . 
 
 Appendix: Hardware. All our experiments about 7B white-box models are conducted on a single L40 GPU. All experiments involving 13B white-box models are conducted on a single H100 GPU. Total budget for GPT-3.5 API calls to conduct the transferability experiments in Sec. 4.5 is estimated to be around US$500 (at the rate of US$0.5/1mln tokens for input and US$1.50/1mln tokens for output).Software. All our experiments are conducted using Python3.10 on a Ubuntu 22.04.4 LTS distribution. We use Jax (Bradbury et al., 2018) for the experiments in Sec. 4.1 and use PyTorch 2.1.0 (Ansel et al., 2024) for the experiments in Sec. 4.2 and Sec. 4.3. We adopt the implementations of the LLMs (transformer-based and SSM-based) provided in the Huggingface's "transformers" (Wolf et al., 2020) system throughout this work. The NLP datasets are also obtained from Huggingface's "datasets" API (Lhoest et al., 2021) . The precise repository references and other dependencies can be found in the code provided in the supplemental materials.We provide an implementation of computing our DETAIL score for LLM in Alg. 1. Line 6 shows the (optional) random projection where a random matrix P is multiplied by the embeddings. In line 8, β has a closed-form expression as the solution to a regularized ridge regression. In line 10 and line 14, we select the embeddings of the target positions. Then, depending on whether we use self-influence, we use different embeddings for computing ∇ β L as shown in lines 15-19. Line 20 computes the inverse hessian H -1 β = (K I + λI) -1 before finally calculating I i in line 21.Algorithm I i ← 0 for i ∈ {1, 2, • • • , t -1} 3: h 1 , h 2 , • • • , h L ← M (x [1:n] ) 4: p demo ← p [1:t-1]{Remove the last target which is the test query} 5: y demo ← one hot(y [1:t-1] ) {Remove the last label and convert to one-hot} 6: m demo ← h l [p demo ]P {Optional dimensionality reduction} 7:K β ← m demo m ⊤ demo {K β ∈ R (t-1)×(t-1) for speed-up as t ≪ d ′ } 8: β ← [(K β + λI) -1 m demo ] ⊤ y demo 9: p test ← p [t-1:t] 10: m test ← h l [p test ] 11: y test ← one hot(y [t-1:t] ) 12: K I ← m ⊤ demo m demo {K I ∈ R d ′ ×d ′ } 13: for i ∈ {1, 2, • • • , t -1} do 14: m i ← h l [p [i:i+1] ]P 15: if self influence then 16: ∇ β L ← m ⊤ i (m i β -y demo [i]) + λβ 17:else if test influence then 18:∇ β L ← m ⊤ test (m test β -y test ) + λβ 19:end if 20:I reg ← (K I + λI) -1 [m ⊤ i (m i β -y demo [i]) + λβ] {Eq.(5) with the constant dropped} 21: EQUATIONUnderstanding in-context learning. Prior works have attempted to understand the ICL capability of transformer-based language models (Ahn et al., 2023; Akyürek et al., 2023; Bai et al., 2023; Bhattamishra et al., 2024; Brown et al., 2020; Dai et al., 2023; Garg et al., 2023; Li et al., 2023; Olsson et al., 2022; Panwar et al., 2024; Von Oswald et al., 2023; Xie et al., 2022; Zhang et al., 2024a) . (Brown et al., 2020) empirically demonstrated that language models can act as few-shot learners for unseen NLP tasks. (Olsson et al., 2022) explained ICL by viewing attention heads as implementing simple algorithms for token sequence completion. (Panwar et al., 2024; Xie et al., 2022) studied the ICL capability of transformers by casting it as an implicit Bayesian inference. (Akyürek et al., 2023; Bhattamishra et al., 2024; Garg et al., 2023) further showed that transformers can learn (regularized) linear and discrete functions in context. (Bai et al., 2023) showed that transformers can adaptively implement appropriate algorithms for ICL. (Li et al., 2023) provided a statistical generalization bound for ICL. (Dai et al., 2023; Von Oswald et al., 2023; von Oswald et al., 2023; Zhang et al., 2024a) mathematically showed that transformers with specific parameters can perform gradient descent on parameters of an internal optimizer given the demonstrations. (Ahn et al., 2023; Zhang et al., 2024a) further proved that the parameters (of the internal optimizer) can converge during forward passing. Inspired by the theoretical grounding, which is the focus of these works, we design our novel attribution technique for task demonstrations by adopting a similar view (i.e. transformers learn in context by implementing an optimization algorithm internally).Data attribution. Past works have focused on explaining and attributing model performance to training data of conventional ML (Barshan et al., 2020; Cook, 1977; Ghorbani & Zou, 2019; Grosse et al., 2023; Koh & Liang, 2017; Ribeiro et al., 2016) . The rise of LLMs has inspired research efforts on attribution w.r.t. prompts with a focus on task demonstrations (Bohnet et al., 2022; Liu et al., 2023c; Machiraju et al., 2024; Sarti et al., 2023; Yue et al., 2023) , which is distinct from training data attribution since demonstrations are provided in context. Specifically, (Liu et al., 2023c) used human annotators to evaluate the verifiability of attributing model answers to a prompt. (Bohnet et al., 2022; Yue et al., 2023) relied on LLMs to evaluate attribution errors. These prior works are either computationally heavy (requiring additional queries of LLMs) or time-consuming (requiring human annotators). (Sarti et al., 2023) proposed an interpretability toolkit for sequence generation models using gradient-and perturbation-based methods. (Machiraju et al., 2024) , a contemporary work, proposed to use a decoder module on the token embeddings for per-token attribution but requires costly training to learn the decoder weights. Moreover, these methods do not specifically target demonstration attribution. Some prior techniques (Cook, 1977; Ghorbani & Zou, 2019 ) can be adapted for attributing ICL in LLMs but may be costly or ineffective. We empirically compare our method with those and the attribution methods consolidated in (Sarti et al., 2023) .Using influence in LLMs. Past works have attempted to apply the notion of influence in language models (Grosse et al., 2023; Kwon et al., 2024; Nguyen & Wong, 2023; S. et al., 2024; Xia et al., 2024; Zhang et al., 2024b) . (Nguyen & Wong, 2023) considered a simplification of the influence function for task demonstration curation. (Grosse et al., 2023; Kwon et al., 2024; Xia et al., 2024) applied influence to pre-training and fine-tuning data of LLMs. (Zhang et al., 2024b) used influence to select demonstration inputs for annotation. (S. et al., 2024) builds a classifier on the embeddings of demonstrations using a small LLM and computes influence w.r.t. the classifier for demonstration selection. In contrast, we demonstrate various use cases of our method including on-the-fly demonstration curation, reordering, and noisy demonstration detection. A contemporary work that shares technical similarity (S. et al., 2024) focuses on demonstration selection whereas we focus on attribution and (S. et al., 2024 ) is shown to be less effective than our method in Sec. 4.3. Additionally, compared to prior works leveraging influence to address specific problems, we apply influence function to provide a general attribution for demonstrations, with many applications that we empirically show.Potential societal impact. We propose an attribution technique for improving the interpretability of in-context learning. We believe our research has potential positive societal impacts in improving the safety of LLMs via filtering out corrupted/harmful demonstrations as demonstrated by our experiments as well as saving energy by curating the demonstration, hence reducing the cost of querying LLMs. We do not find any direct negative societal impact posed by our research contribution.Setting λ. Generally, there is no golden rule for the most appropriate λ that regularizes the ridge parameters β. Intuitively, a larger λ likely works better when the dimension of β is large since the model tends to be over-parameterized (i.e., in a LLM). Therefore, we set a relatively large λ = 1.0 for LLMs and a relatively small λ = 0.01 for our custom transformer. When detecting noisy demonstrations, we may not want to regularize β too much because we wish to retain the information captured by the eigenvalues of the hessian H which can be eroded with a larger λ. As such, for the noisy demonstration detection task, we set a very small λ = 10 -9 to retain most of the information captured by H while ensuring that it is invertible.Random projection matrix. We recall the Johnson-Lindenstrauss (JL) lemma (Dasgupta & Gupta, 2003; Johnson & Lindenstrauss, 1984) .Theorem D.1 (Johnson-Lindenstrauss Lemma). For any 0 < ϵ < 1 and any integer n, let d ′ be a positive integer such that d ′ ≥ 24 3ϵ 2 -2ϵ 3 log n , then for any set A of n points ∈ R d , there exists a mapping f : R d → R d ′ such that for all x i , x j ∈ A,(1 -ϵ)∥x i -x j ∥ 2 ≤ ∥f (x i ) -f (x j )∥ 2 ≤ (1 + ϵ)∥x i -x j ∥ 2 .A specific constructive proof is by settingA := 1 √ d ′ R where R i,j i.i.d.∼ N (0, 1). 6 In our work, we treat each embedding m as x and the projected embedding mP as f (x). The specific construction follows the abovementioned constructive proof definingP := 1 √ d ′ R ∼ N (0, 1 d ′ I) .Empirically, our threshold d ′ = 1000 corresponds to ϵ ⪅ 0.164, ensuring a good preservation of (Euclidean) distance between points.ICL prompt example. We include a visualization of a prompt for ICL below. Each input-output pair consists of a task demonstration. The query is appended at the end of the prompt with only the input and the output header.Input: tsai may be ploughing the same furrow once too often . Output: B Input: equilibrium the movie , as opposed to the manifesto , is really , really stupid . Output: B (More demonstrations...)Input: a friendly vacation for four old friends -two couples from college -turns ugly . . . then Output: A Input: he meets god and is given all the powers of god .Output:We include the full results for the demonstration perturbation task using a Vicuna-7b v1.3 model in Fig. 8 and using a Llama-2-13b model in Fig. 9 . A consistent trend can be observed across different datasets using both models. We compare (in addition to Vicuna-7b) a total of 3 LLMs: Llama-2-7b (Touvron et al., 2023) , Llama-2-13b (Touvron et al., 2023) and Falcon-7b (Almazrouei et al., 2023) . We reuse the same experimental setup as the demonstration label perturbation task and compare the accuracy by removing and corrupting 10 among 20 ICL data with high/low DETAIL scores computed in different models. The results are tabulated in Table 5 . A similar trend can be observed across these models where removing/corrupting demonstrations with high I test results in lower accuracy and vice versa. The results demonstrate that our method is robust against model pre-training/fine-tuning data as well as model size. We include the results on AG News, SST-2, Rotten Tomatoes, and Subj datasets using a Vicuna-7b model in Fig. 10 . Similar trends as in the main text is observed. A counterpart experiment using Llama-2-13b is in Fig. 11 , where a similar trend is observed. We additionally provide the result for the demonstration curation task with Attention (Bahdanau et al., 2016) and (S. et al., 2024) using 10 iterations of LiSSA update for computing the hessian vector product (Agarwal et al., 2017) and compare it with the case with 5 iterations. The result is shown in Table 7 . With 10 iterations, the wall time is much higher (over 20×) but accuracy is comparable to the case with 5 iterations. For Attention, the performance is comparable to random removal as shown in Table 2 . We experiment with the difference in the effectiveness DETAIL using the embeddings of different layers. We conduct experiments on demonstration removal, demonstration perturbation, and noisy label detection tasks. The results are shown in Fig. 14 . It can be observed that obtaining the DETAIL scores from the later layers of the model consistently produces desirable results. As a rule of thumb, for each demonstration, we generally want to take the embedding of its last few tokens because of the causal nature of inference and because information generally flows toward the end of the sequence (Wang et al., 2023) . We compare two possible choices of target position: the column position (immediately before the label) and the label position. We experiment on the demonstration removal task with these two choices of embeddings. The results are shown in Fig. 15 . Using embeddings of both positions achieves decent task performance as reflected by the clear distinction in accuracy between removing demonstrations with high/low DETAIL scores, demonstrating that our method is robust against the choice of token embeddings. In our experiments, we adopt the column position to isolate information about the label from the embedding. We consider experiments on a popular state-space model (SSM) architecture, a Mamba-2.8b model (Gu & Dao, 2023) , which consists of 64 layers with d = 2560. The tasks are described in the main text in detail. While we find that DETAIL can still successfully attribute Mamba on certain tasks and datasets, the performance is inconsistent. One hypothesis is that the largest currently available Mamba model (2.8B) is still significantly smaller than the 7B LLMs we conduct experiments on in the main text. A smaller model size reduces the inductive power of the model to formulate the "internal optimizer", leading to less interpretive DETAIL scores. We would also like to note that DETAIL is not designed to work on SSMs. Noisy demonstration detection. We further consider the noisy demonstration detection task as described in Sec. 4.2 on Mamba. Unfortunately, the performance is not consistent across datasets, as shown in Fig. 17 : detecting demonstrations with high I self performs close to random selection on SST-2 and Rotten Tomatoes datasets, although the inference speedup is still significant. We leave the analysis of these failure cases to future work. All experiments are repeated with 10 independent trials. Lines and shades represent the mean and standard error respectively. Demonstration curation. As DETAIL performs well using Mamba on the demonstration removal task, it is reasonable to hope that it works well on the demonstration curation task as well. As it turns out, DETAIL performs well on binary classification tasks as shown in Fig. 18 but performs poorly on AG News which is 4-way classification. We hypothesize that this is due to Mamba's worse inductive power to formulate an internal algorithm successfully. Following(Barshan et al.,2020) and the experiment implementation in(Koh & Liang, 2017), we drop the negative sign in our influence definition. The interpretation is that higher values imply a more positive impact.2 For LLMs, each demonstration may consist of more than 1 token. We discuss how to address this in Sec. 4.2.Note that a key difference between demonstration curation task and demonstration perturbation task is that for curation, the test dataset is unknown when computing the DETAIL scores.At the time of this writing, GPT-4 API costs $10/1mln tokens. See https://openai.com/pricing.Lecture notes. 


Decision: The paper uses influence function-based method for in-context learning. All reviewers agree that the submission can be accepted. The method is simple yet effective. The results are positive and can be transferred from white-box models to other black-box models.

Reviews: review: paper_summary: paper_summary: The paper proposes a novel attribution method for demonstrations in ICL. The proposed method takes a perspective that the transformers learn in context by formulating an internal optimizer. The influence function is approximated as an internal kernelized ridge regression, where the representations are taken from the intermediate layers of the transformers and white-box LLMs. The paper also demonstrate that the attribution obtained from the white-box LLM exhibits transferable characteristics to black-box models. The paper also showcases a few promising applications of DETAIL.
strength_weakness: strength_weakness: - It would have been interesting to see how well this method would identify adversarial attacks for LM.
- Some additional commentary on how DETAIL can help with LLM training would have been very helpful.
questions: questions: - Could you comment on whether DETAIL can be used to discover adversarial attacks of LLMs?
- Could you comment on how DETAIL can help with LLM training especially for instruction tuning?
limitations: limitations: Yes, the author addressed the limitations.
paper_summary: paper_summary: The paper introduces DETAIL, a novel influence-function based attribution technique to estimate the influence of each example in the demonstration sequence for the given target query for in-context learning (ICL). The authors empirically validate DETAIL’s effectiveness in stylized experiments and LLMs. Additionally, they demonstrate DETAIL’s applicability in real-world tasks like demonstration reordering and curation and the transferability of DETAIL’s attribution scores from white-box to black-box models.
strength_weakness: strength_weakness: 1. Strong Assumptions: The work is built on the assumption that transformers implement an internal optimizer. This assumption, while supported by some theoretical proof in stylized settings, may not universally hold. The proof further assumes a learn-to-learn setting for ICL, which differs from the definition provided in lines 88-105 in this paper. Additionally, Equation 3 assumes the loss function of ICL is a ridge regression function without robust theoretical guarantees.

2. Lack of Analysis for Equation 6: Given the strong assumptions underlying the derived Equation 6, its reliability may be questionable. 

3. Sensitivity to Demonstration Ordering: Recent work has demonstrated that ICL is not sensitive to ordering with more powerful demonstration retrievers and advanced language models, contradicted  with the conclusion in  line 35.

4. Limited Experimental Comparisons: The experimental validation of DETAIL in ICL demonstration curation could be strengthened. The authors should compare their method with other learning-free demonstration selection methods, such as those based on BERT embeddings.

5. Computational Cost: The computational cost of DETAIL is relatively high for demonstration curation or selection. Existing methods that utilize BERT embeddings are more computationally efficient.
questions: questions: 1. How do you obtain $y_{test}$ for $I_{test}$?
2. Have you explored initializing $m(x)$ and $y$ with BERT embeddings? This could be a potential area for improvement, leveraging pre-trained embeddings for better initialization.
limitations: limitations: Yes
paper_summary: paper_summary: This paper proposes a new method to estimate the influence of ICL examples to a query. This influence estimation can better help ICL learning, for example, reorder ICL examples and curating ICL examples.
strength_weakness: strength_weakness: - More comprehensive experiments: for applying DETAIL, the authors only showed a few tasks (all classification) and on one major public model (Vicuna-7b). To show the effectiveness of DETAIL, more comprehensive experiment results should be provided: e.g., Table 1 and 3 should use more tasks other than classification, and more models (can the authors add results on Llama-2-13b?).

- DETAIL for ICL order optimization: based on the description in Section 5.3, it seems like the authors re-ordered the examples not just based on I_{self}, but also based on the trend plot in Figure 5. Does test accuracy w.r.t. perturbed position alway need to be computed first before deciding on the orders? This seems really expensive and also the choice of putting "two demonstrations" in the front seems very arbitrary. How would one use such a method for a random new task?

- The need to use white-box models and poor transfer results: DETAIL requires white-box model access, and the transfer result in Table 3 is not significant on the more realistic setting (no corruption). Also, can the authors show the transfer results with more models, e.g., LLama-2 to GPT-3.5, and Vicuna-7b to GPT-4?
questions: questions: See above.
limitations: limitations: Discussed in Section 6.
paper_summary: paper_summary: This paper introduces DETAIL, a novel technique for attributing and interpreting in-context learning (ICL) demonstrations in transformer-based language models. The authors propose an adaptation of the influence function, typically used in conventional machine learning, to address the unique characteristics of ICL. DETAIL treats transformers as implementing an internal kernelized ridge regression, allowing for efficient and effective attribution of demonstrations. The method is evaluated on various tasks, including demonstration perturbation, noisy demonstration detection, and real-world applications such as demonstration reordering and curation. The authors demonstrate DETAIL's superiority over existing attribution methods in terms of both performance and computational efficiency.
strength_weakness: strength_weakness: 1. Section 5.1 presents an evaluation on a custom transformer using the MNIST dataset. While this provides an initial demonstration of DETAIL's capabilities, the paper doesn't clearly justify why this evaluation is necessary given the subsequent experiments on large language models. It's not immediately apparent how insights from this simplified setting transfer to more complex LLMs, potentially making this section feel disconnected from the main contributions of the paper.
2. Figure 3 shows that even the Llama-2-13b model achieves only 60-70% accuracy on the AG News dataset without any perturbation. This is significantly lower than the typical performance range of 85-96% reported in the literature for this dataset. Could the authors give more insight about the performance of in context learning on the AG News dataset?
questions: questions: see the weakness section
limitations: limitations: Yes, the authors adequately addressed the limitations

Comments: title: Thank you for acknowledging our response
Thank you for acknowledging our response and raising the score
Reply to Authors' Rebuttal
comment: Thank you very much for acknowledging our response and suggesting a new experiment setting with RAG for search retrieval. We will include your suggestions, along with our clarifications and additional experimental results, in the revision of our work.
Thank you very much for acknowledging that the additional experiments in our response are detailed and helpful and for raising the score. We will incorporate the experimental results and our clarifications in our revision.
Thanks for the authors' detailed responses to my questions. The additional experiments and discussions are helpful and provide more evidence supporting its usability. Hence, I would like to raise my score.
Thanks the authors for the rebuttal and adding the new results. Given all the results so far, the overall trend becomes a bit more clear: 
(1) the method gives more significant improvement under the corruption setting (which is less realistic, and given the authors' argument on similarity to search retrieval, it would be more convincing to provide such results, even with a simple RAG setup); 
(2) the methods does not yield significant gains when the baseline performance is already high (for GPT-4o it's ok, but for GPT-3.5 the gains still vary quite a bit, like there's no gain on "Subj" under the no corruption setting). It would be more convincing if the authors can show more consistent results across models/settings, and provide results on more challenging tasks where the baseline performance is not good enough, and see if DETAIL can still yield gains.

I will maintain my score given the limited effectiveness of the proposed method, but I won't object to accepting this paper.
Dear reviewers: as you are aware, the reviewer-author discussions phase ends on Aug 13. We request you to kindly make use of the remaining time to contribute productively to these discussions. If you have not read and/or responded to author rebuttal, please do it asap so that the authors get a chance to respond to you. If you have more questions to ask or want further clarification from the authors, please feel free to do it.