Enhancing In-Context Learning via Implicit Demonstration Augmentation (2024)

Abstract: The emergence of in-context learning (ICL) enables large pre-trained language models (PLMs) to make predictions for unseen inputs without updating parameters. Despite its potential, ICL’s effectiveness heavily relies on the quality, quantity, and permutation of demonstrations, commonly leading to suboptimal and unstable performance. In this paper, we tackle this challenge for the first time from the perspective of demonstration augmentation. Specifically, we start with enriching representations of demonstrations by leveraging their deep feature distribution. We then theoretically reveal that when the number of augmented copies approaches infinity, the augmentation is approximately equal to a novel logit calibration mechanism integrated with specific statistical properties. This insight results in a simple yet highly efficient method that significantly improves the average and worst-case accuracy across diverse PLMs and tasks. Moreover, our method effectively reduces performance variance among varying demonstrations, permutations, and templates, and displays the capability to address imbalanced class distributions.

heading: 1 Introduction
text: Large pre-trained language models (PLMs) have showcased exceptional abilities in in-context learning (ICL) (Brown et al., 2020; Wang et al., 2023; Rubin et al., 2022), which assists the model in discerning the underlying patterns within demonstrations and make more accurate predictions (Chan et al., 2022; Wu et al., 2023). As a new paradigm, ICL offers compelling advantages, allowing for natural language interaction with PLMs (Wei et al., 2022; Yang et al., 2023), as well as reduced computational costs (Li et al., 2023a; Rubin et al., 2022).
While promising, ICL’s performance is highly dependent on provided demonstrations and templates (Liu et al., 2022; Zhang et al., 2022b;
*Corresponding authors.
Semantic augmentation
Sorensen et al., 2022), resulting in subpar and unstable performance. This promotes research aimed at improving the quality (Rubin et al., 2022; Li et al., 2023b), quantity (Li et al., 2023a; Choi et al., 2022), and permutations (Lu et al., 2022; Tang et al., 2023) of demonstrations. Other research avenues include prediction adjustment (Zhao et al., 2021; Han et al., 2023; Fei et al., 2023) and learning process design (e.g., channel models (Min et al., 2022a) and meta-training frameworks (Min et al., 2022b)). Despite ongoing efforts, ICL still struggles with efficiently and reliably capturing sufficient knowledge from context, leaving performance stability as a persistent bottleneck.
In this study, we propose enriching contextual knowledge for PLMs by augmenting demonstrations. We first attempt to enhance the representation of demonstrations by transforming them along semantic directions sampled from the deep feature space of demonstration examples, as depicted in Figure 1. This operation stems from the observation that the deep features in a network are usually linearized (Bengio et al., 2013; Cheung and Yeung, 2021; Cho, 2016), implying the existence of numerous semantic directions within the deep feature space, hence potentially enabling us to incorporate richer contextual knowledge without extending input length. From this novel perspective, we theoretically prove that when the number of augmented
2810
pieces approaches infinity, its effect approximately equals a logit adjustment operation. Specifically, we derive a refined Softmax function that integrates the statistical properties of demonstrations. Consequently, rather than explicitly executing the augmentation procedure, we can efficiently conduct implicit demonstration augmentation using the derived prediction function, obtaining an improved ICL method with theoretical guidance.
We conduct extensive experiments across seven PLMs and various classification tasks. The empirical results demonstrate that our approach remarkably enhances prediction accuracy and reduces performance variability across different demonstrations, permutations, and templates. Notably, our method is straightforward, effective, and generalizable, enabling seamless integration with other ICL methods to enhance their performance.
Our contributions can be summarized as follows: • We introduce Implicit Demonstration
Augmentation-based ICL (IDAICL), a pioneering work that incorporates demonstration augmentation into ICL. Instead of solely enhancing demonstration quality, quantity, or order, our method explores context augmentation within the deep feature space, offering a new perspective to enrich demonstrations bypassing input length limitations.
• We theoretically establish that as the number of augmented pieces approaches infinity, our augmentation strategy approximates a logitadjusted prediction function that integrates statistical properties derived from the input data distribution. Equipped with this function, IDAICL provides a straightforward yet theoryguided solution to enhance ICL.
• Extensive experiments conducted across diverse tasks and PLMs conclusively illustrate that IDAICL considerably improves average and worst-case accuracy compared to existing ICL methods. Moreover, it effectively enhances performance stability.

heading: 2 Background and Related Work

heading: 2.1 In-Context Learning
text: Brown et al. (2020) showcased the ICL capability of PLMs, wherein PLMs generate predictions solely based on a concatenation of training examples for few-shot learning without updating parameters. Subsequent studies (Holtzman et al., 2021;
Min et al., 2022a,b) have developed this approach, yielding promising outcomes across various tasks. Nevertheless, recent research has uncovered certain limitations. To begin with, the volume of input knowledge for each query is constrained by the maximum input length of PLMs (Hao et al., 2022), and the computational cost increases as the number of demonstrations grows (Li et al., 2023a), making it challenging to integrate significant knowledge from demonstrations to PLMs. Additionally, ICL’s performance is sensitive to the input of PLMs (Davison et al., 2019; Jiang et al., 2020), thus exhibiting high variance and poor worst-case accuracy (Perez et al., 2021; Lu et al., 2022).
Researchers have explored various techniques to address the biases and instability of ICL. These techniques encompass learning process design (Min et al., 2022a,b), demonstration retrieval (Rubin et al., 2022; Zhang et al., 2022b), prompt engineering (Sorensen et al., 2022; Lu et al., 2022), and prediction calibration (Zhao et al., 2021; Fei et al., 2023). However, these methods have yet to fully address the issue of severely limited knowledge transfer from demonstrations to large PLMs.

heading: 2.2 Data Augmentation
text: Data augmentation (Chen et al., 2023), which involves artificially creating training data through transformations, is a well-established research area in machine learning. Although data augmentation techniques have undergone extensive exploration in diverse machine learning domains (Maharana et al., 2022; Shorten and Khoshgoftaar, 2019), applying them to text data poses challenges due to the complexity of preserving labels during textual transformations (Kobayashi, 2018). Nonetheless, data augmentations in the latent space, such as adversarial training (Zhang et al., 2022a; Zhu et al., 2020; Cheng et al., 2020), interpolation (Chen et al., 2022b; Wu et al., 2022), and generative techniques (Li et al., 2022; Malandrakis et al., 2019), have demonstrated notable enhancements when applied alongside large PLMs.
Recently, Wang et al. (2019) introduced the concept of implicit data augmentation in the context of image classification. This approach involves transforming training data within the deep feature space and boils down to the optimization of a novel robust loss function. Subsequent studies (Chen et al., 2022c; Li et al., 2021; Zhou and Wu, 2023a) for image classification tasks have further improved
Demonstration examples
PLM
3h
Features
h
Augmentation Augmented Demonstration
upon this approach. This study introduces an algorithm for implicitly augmenting demonstrations within the realm of ICL.

heading: 3 Methodology

heading: 3.1 In-Context Learning with PLMs
text: Considering a PLM G, this study focuses on the following task: given a query input text x and a candidate answer set Y = {y1, y2, · · · , y|Y|}, we aim to predict the answer ŷ based on m demonstration examples C={c1, c2, · · · , cm}, where each ci represents a training example (xi, yi) after template formulation and m denotes the quantity of demonstration examples for each test sample. Formally, give a model G, we first compute the probability of each answer yj :
PG (yj | C,x) . (1)
Subsequently, the ultimate prediction ŷ, characterized by the highest probability is chosen from the candidate answer set Y:
ŷ = argmax yj∈Y
PG (yj | C,x) . (2)
To simplify, the contextual input is denoted as x̃=[C,x] in the subsequent text. Then, the probability of answer yj , represented as PG(yj |x̃), is computed using the Softmax function1:
PG(yj |x̃) := PG(yj |hx̃) = e wTyjhx̃+byj
∑ k e wTk hx̃+bk , (3)
where hx̃ = G(x̃) signifies the hidden state of the last block at the final position for x̃. wk and bk are the weight vector and bias corresponding to the final fully connected layer for the k-th token.
1We begin by examining situations in which the answer comprises a single token, and our subsequent analysis is equally applicable to scenarios involving multiple tokens.

heading: 3.2 Demonstration Augmentation
text: Recognizing the established efficacy of data augmentation in machine learning (Feng et al., 2021), this study investigates demonstration augmentation and suggests enhancing the deep features of demonstrations by transforming them along semantic directions sampled from the deep feature space of demonstration examples. This strategy is motivated by the intriguing observation that the deep features in networks are often linearized (Bengio et al., 2013; Chen et al., 2022a). Building on this observation, we hypothesize that hx̃ lies within the subspace spanned by hC and hx: hx̃ = αhC + βhx, where hC and hx represent the components of hx̃ linked respectively to the demonstrations and the query. The necessity of this assumption stems from intricate relationships among token representations and the exclusive augmentation of the component related to demonstrations. Notably, this decomposition is not necessary in practical applications. In the subsequent text, we directly refer to αhC and βhx as hC and hx.
To augment hC , we randomly sample vectors from the deep feature space of demonstrations. In particular, vectors are drawn from a multivariate normal distribution N (µ,Σ), where µ and Σ denote the feature mean and covariance matrix. These statistical properties are estimated from the deep features of the demonstration set D, which includes demonstration examples linked to all queries. The feature mean µ is computed as
µ = 1 |D| ∑|D| i=1 hi, (4)
where hi = G(ci) represents the hidden state of the last block at the final position for the i-th demonstration example ci in D, and |D| denotes the size
of D. The covariance matrix Σ is computed as
Σ = 1 |D| ∑|D| i=1 (hi − µ)T (hi − µ). (5)
Subsequently, hC is shifted in the extracted semantic vectors, resulting in augmented features, h̃C , which follows
h̃C ∼ N (hC + λµ, λΣ) , (6)
where λ refers to a positive coefficient controlling the strength of semantic augmentation. In realworld applications, it can be directly assigned a value of 0.5. Sensitivity tests for λ are discussed in Section 5.4.

heading: 3.3 Novel Prediction Function
text: Selecting the answer with the highest probability is equivalent to favoring the answer with the lowest inverse probability. Therefore, the prediction can be determined by
ŷ = arg min yj∈Y
PG (yj | hx̃)−1 . (7)
Assume that each hC is augmented for M times, resulting in an augmented demonstration feature set {h̃1C , · · · , h̃ M C } with size M. Here, h̃ i C represents the i-th augmented feature for hC . Then, the final prediction for the query x depends on all augmented features of hC and can be expressed as
PMyj (x̃) = 1 M ∑M i=1 PG(yj |h̃ i C ,hx) −1, (8)
ŷ = arg min yj∈Y
PMyj (x̃). (9)
Given that the performance of ICL benefits from an increased number of demonstration instances (Liu et al., 2022; Wu et al., 2023), we explore the scenario of augmenting an infinite number of times for the deep representation of demonstrations. Subsequently, an easily computable surrogate for the expected prediction can be derived, resulting in a highly efficient implementation. The whole pipeline of IDAICL is depicted in Figure 2.
As M → ∞, on the basis of the aforementioned decomposition of hx̃, the expected prediction for answer yj (denoted as P∞yj ) within the augmented feature set can be expressed as follows:
P∞yj (x̃)=Eh̃C [ ∑
k
e ∆wTk,yj (h̃C+hx)+∆bk,yj ], (10)
where ∆wk,yj = wk−wyj and ∆bk,yj = bk−byj .
However, accurately calculating P∞yj is challenging. Alternatively, we proceed to derive a surrogate calculation for it. Applying the linearity of expectation, Eq. (10) can be expressed as:
P∞yj (x̃)= ∑
k
Eh̃C [e ∆wTk,yj (h̃C+hx)+∆bk,yj ]. (11)
Given that h̃C is a Gaussian random variable conforming to N (hC + λµ, λΣ), we know that ∆wTk,yj h̃C follows the multivariate normal distribution: N (∆wTk,yj (hC + λµ) , λ∆w T k,yj
Σ∆wk,yj ). Then, utilizing the moment-generating function
E[etX ] = etµ+ 1 2 t2σ2 , X ∼ N (µ, σ2), (12)
Eq. (11) can be derived as P∞yj (x̃)= ∑
k
Mk,yjNk,yje ∆wTk,yj (hC+hx)+∆bk,yj ,
(13) where Mk,yj = exp(λ∆w T k,yj µ) and Nk,yj = exp(λ2∆w T k,yj
Σ∆wk,yj ). Subsequently, our newly proposed prediction function, referred to as IDA-Softmax, is defined as P IDAyj (x̃) := ∑
k
Mk,yjNk,yje ∆wTk,yj hx̃+∆bk,yj .
(14) Consequently, instead of conducting the augmentation process explicitly, we can directly employ IDA-Softmax, P IDAyj , for prediction. IDA-Softmax essentially utilizes two modulating factors associated with statistical properties derived from D to calibrate the sample logits. Previous studies (Min et al., 2022c; Chan et al., 2022) have underscored the pivotal role of knowledge about the input data distribution in predictions made by PLMs. Intuitively, PLMs can better capture the patterns and underlying structures within data, such as the spatial relationships between demonstrations and queries, ultimately enhancing their prediction performance.
Furthermore, to mitigate the imbalance among different answer types in demonstrations (Holtzman et al., 2021; Zhao et al., 2021), we adopt a post-hoc adjustment approach inspired by Menon et al. (2021), which adjusts predictions by considering the class proportions within D. Thus, the prediction for answer yj is computed as
P̃ IDAyj (x̃) = P IDA yj (x̃) + τ log πyj , (15)
where τ is a positive hyperparameter, and πyj demotes the proportion of answer yj in D. In practical applications, the value of τ can be fixed at 1.
This approach compensates for predictions of minor classes. When different answers are uniformly distributed, τ log πyj exerts an equal influence on all answer types. Consequently, the final prediction is given by
ŷ = arg min yj∈Y
P̃ IDAyj (x̃). (16)

heading: 4 Experimental Setup

heading: 4.1 Models and Datasets
text: We evaluated the performance of IDAICL across seven large PLMs, including GPT-2 (Radford et al., 2019) (with 0.1B, 0.3B, 0.8B, and 1.5B parameters), GPT-Neo (Black et al., 2021) (with 2.7B parameters), and LLaMA (Touvron et al., 2023) (with 13B and 33B parameters). Following previous research (Min et al., 2022a; Han et al., 2023; Lu et al., 2022), our evaluation encompasses ten text classification datasets. Among these, SST-2 (Socher et al., 2013), SST-5 (Socher et al., 2013), MR (Pang
and Lee, 2005), CR (Hu and Liu, 2004), and Amazon (McAuley and Leskovec, 2013) are five sentiment classification tasks. Subj (Pang and Lee, 2004), TREC (Voorhees and Tice, 2000), DBPedia (Lehmann et al., 2015), and AGNews (Zhang et al., 2015) cater to subjectivity, question, ontology, and news classification tasks, respectively. Additionally, CB (De Marneffe et al., 2019) is utilized for natural language inference. Among these datasets, SST-5, Amazon, TREC, and CB are characterized by imbalanced training data. Details of all datasets are provided in Section A of the Appendix.

heading: 4.2 Compared Baselines
text: Besides Vanilla ICL, we compared and integrated IDAICL with three popular ICL algorithms, focusing on learning process design and demonstration retrieval. These include MetaICL (Min et al., 2022b), Channel ICL (Min et al., 2022a), and Efficient Prompt Retrieval (EPR) (Rubin et al., 2022). Moreover, we compared IDAICL with other ad-
vanced prediction calibration methods: Contextual Calibration (ConCa) (Zhao et al., 2021), Prototypical Calibration (PROCA) (Han et al., 2023), and Domain-Context Calibration (D-ConCa) (Fei et al., 2023). Introductions to all compared methods and comprehensive experimental settings are presented in Sections B and C of the Appendix.

heading: 5 Experimental Results

heading: 5.1 Main Results
text: Table 1 displays the comparison results between IDAICL and four ICL baselines (Vanilla ICL, MetaICL, Channel ICL, and EPR) across GPT-2 models (with 0.8B and 1.5B parameters) and the GPT-Neo model. These results lead to three main findings. Firstly, IDAICL consistently exhibits high effectiveness across various model sizes and datasets, highlighting its strong generalization capacity, even under scenarios involving imbalanced training data. Compared to Vanilla ICL, IDAICL outperforms by an average of 17.7% and 18.4% across diverse datasets and m values for GPT-2 with 0.8B and 1.5B parameters, respectively. Secondly, in comparison to other ICL baselines like Channel ICL, MetaICL, and EPR, the integration of
IDAICL consistently delivers notable performance improvements, emphasizing the efficacy of enhancing demonstrations for refined predictions. The inclusion of IDAICL led to an average performance boost of 7.3% for MetaICL and 8.2% for Channel ICL. Lastly, IDAICL notably enhances worstcase accuracy and diminishes performance variance across different seeds, showcasing its ability to improve prediction stability. Additional results on LLaMA and smaller GPT-2 models are available in Tables 7 and 8 of the Appendix.

heading: 5.2 Comparison with Calibration Methods
text: We compared IDAICL with three advanced prediction calibration methods (ConCa, PROCA, and D-ConCa) across three PLMs: GPT-2, GPT-Neo, and LLaMA. Table 2 presents the comparison results for the LLaMA models, where IDAICL consistently achieves state-of-the-art performance, except for TREC using the LLaMA model with 33B parameters. These findings suggest that IDAICL which leverages statistical information derived from the input data distribution for prediction calibration, generally outperforms methods relying on estimated biases for correction. Further comparison results can be found in Table 9 of the Appendix.

heading: 5.3 Stability Analysis
text: Previous studies (Zhao et al., 2021; Sorensen et al., 2022; Min et al., 2022a; Zhang et al., 2022b) have highlighted the considerable variability in ICL’s performance. In this section, we verified that IDAICL can effectively enhance performance stability across diverse scenarios.
Varying numbers of demonstrations We have presented the results across different numbers of demonstrations in Table 1. For a clearer depiction, the outcomes regarding GPT-Neo are illustrated in Figure 3. As the number of demonstration examples (represented by m) increases, both Vanilla ICL and IDAICL exhibit improved performance, emphasizing the importance of comprehensive statistical properties of the input data for IDAICL’s effectiveness. Notably, IDAICL significantly enhances performance stability across various numbers of demonstrations and consistently outperforms Vanilla ICL. The performance improvement is particularly pronounced when m takes on smaller values, indicating the efficacy of IDAICL in enriching the available knowledge for PLMs.
Varying demonstrations To confirm that augmenting demonstrations can enhance the robustness of the ICL strategy across various demonstrations, we investigated three distinct demonstration selection settings. Setting I: Training samples most similar to the test sample are chosen. Setting II: Samples are randomly selected from the training data. Setting III: Training samples exhibiting the greatest dissimilarity from the test sample are selected. As shown in Figures 4(a) and (b), IDAICL significantly outperforms Vanilla ICL and demonstrates greater robustness across the three selection settings. Additionally, our discoveries suggest that selecting demonstrations that are more similar to the test samples leads to better performance than
exclusively selecting dissimilar ones, which aligns with the findings obtained by Wang et al. (2022).
Varying templates To assess the performance of IDAICL across various templates, we employed fifteen templates on the SST-2 dataset following those outlined by Zhao et al. (2021). The templates are elaborated in Table 10 of the Appendix. Figures 4(c) and (d) display the performance of Vanilla ICL and IDAICL across six templates. Some templates achieve higher average performance than others. Nevertheless, IDAICL consistently enhances both average and worst-case accuracy, simultaneously reducing performance variance across different templates. The complete results are available in Figure 7 of the Appendix.
Impact of imbalance in labels Figures 5(a) and (b) depict comparison results among Vanilla ICL, MetaICL, Channel ICL, and IDAICL across different degrees of imbalances. It is evident that the performance of Vanilla ICL is sensitive to class imbalance, while that of IDAICL and Channel ICL exhibit robustness to the imbalance. Moreover, notable performance improvements are observed with higher levels of imbalance. Additionally, Figures 5(c) and (d) illustrate the confusion matrices for CR and Subj datasets, with the proportion of one category (i.e., "Negative" and "Subjective") in demonstrations setting to 0.1 and 0.2. IDAICL significantly improves the accuracy of the underrepresented classes when compared to Vanilla ICL, thereby contributing to enhanced fairness among classes. In the subsequent section, we demonstrate that the strong performance of IDAICL in handling imbalanced label distributions stems from both the statistical properties and the class proportion term.

heading: 5.4 Sensitivity and Ablation Studies
text: We conducted ablation studies on IDAICL to investigate the influence of the two modulating factors and the class proportion term. The parameters λ and τ govern the augmentation strength and the impact of the class proportion term, respectively. In Figure 6(a), a significant performance drop is observed when predictions are not calibrated using statistical properties derived from the demonstrations. Additionally, optimal performance is achieved when λ equals 0.5.
Figure 6(b) showcases the accuracy of SST-2 and MR datasets with the negative class proportion in demonstrations setting to 0.1. Results indicate that solely leveraging statistical properties (i.e.,
τ equals 0) enhances performance under imbalanced demonstrations, with further improvements observed upon the inclusion of the class proportion term. Additionally, optimal performance is attained when τ equals 1. Consequently, we recommend setting λ to 0.5 and τ to 1 for practical applications. More results are presented in Appendix F.

heading: 5.5 Further Discussion
text: To further investigate the effect of statistical properties within demonstrations on model performance, we exclusively employed queries along with statistical information for inference, excluding the inclusion of demonstrations for each test sample. These statistics were estimated using deep features of all training samples. As shown in Table 3, IDAICL relying solely on statistical properties distinctly outperforms Vanilla ICL across scenarios with zero, one, and even four demonstrations. This emphasizes the crucial role of prior statistics obtained from training data in PLMs’ predictions. This phenomenon is understandable as statistical properties inherently encompass richer global information compared to individual demonstrations.

heading: 6 Conclusion
text: This study introduces IDAICL, a novel ICL approach designed to enhance demonstrations by utilizing semantic directions sampled from the deep feature distribution of demonstration examples. Our augmentation strategy enriches the knowledge available to PLMs without extending the context length. A new prediction function is then theoretically established considering the number of augmented pieces approaching infinity. This eliminates the need for explicit augmentation and allows for direct utilization of this derived function for
predictions. Our extensive experiments, spanning various tasks and PLMs, demonstrate that IDAICL significantly enhances both prediction accuracy and stability when compared to other ICL baselines.
Limitations
While IDAICL proves to be competitive in few-shot learning, there are limitations that open up avenues for future research. First, due to the necessity of accessing the parameters of the final fully connected layer in PLMs, IDAICL is exclusively suitable for open-source models. Future research is expected to develop alternative augmentation strategies tailored for black-box PLMs. Second, our evaluation of IDAICL focused on seven PLMs and ten text classification tasks. We defer further explorations involving other PLMs and non-classification tasks for future work. Additionally, IDAICL relies on a small set of demonstrations to estimate the feature mean and covariance matrix. If such a collection is unavailable or extremely scarce, IDAICL may need to be used in conjunction with demonstration generation methods.
Other avenues for future work involve exploring more effective augmentation distributions. This entails exploring finer-grained distributions, such as category-level or sample-level distributions, to emphasize the unique characteristics of individual categories or samples, and extending these distributions beyond the constraints of training data. Furthermore, given the effectiveness of data augmentation in model training, future research could explore the utilization of our derived prediction function in both the training and fine-tuning phases of large PLMs.

heading: A Details of Applied Datasets
text: Table 4 presents comprehensive statistics for all datasets utilized in this study. The information includes task descriptions, average sentence lengths, class counts, and details on class imbalance. Additionally, Table 5 provides sample instances and label names for each of the datasets.

heading: B Details of Compared Baselines
text: The compared methods are described as follows:
• Vanilla ICL: We use the PLMs as they are and implement ICL by conditioning it on a concatenation of m training examples, following the approach outlined by Brown et al. (2020).
• MetaICL: The fundamental concept underlying MetaICL is to utilize a multi-task learning framework across a diverse range of metatraining tasks (Min et al., 2022b).
• Channel ICL: It employs a noisy channel approach for language model prompting in few-shot text classification (Min et al., 2022a).
• EPR: It employs language models to autonomously label examples that are suitable as effective prompts and subsequently trains a prompt retriever based on this acquired signal (Rubin et al., 2022).
• ConCa: It assesses the model’s inclination towards specific answers by introducing a dummy test input that lacks content (Zhao et al., 2021).
• PROCA: The prediction of PROCA is calibrated based on the likelihood of prototypical clusters (Han et al., 2023).
• D-ConCa: It initially assesses the impacts of various label biases by employing randomly sampled words from the task corpus. During inference, it utilizes the estimated label bias to calibrate the model’s output probabilities (Fei et al., 2023).

heading: C More Details of Experimental Settings
text: The entire implementation is conducted utilizing PyTorch (Paszke et al., 2019) and Transformers (Wolf et al., 2020). We follow the parameter configurations and details specified in previous research (Min et al., 2022a). The number of demonstrations is primarily set to m = 12, but we also explore m values of {1, 4, 8, 12, 16} in the ablations, with the specific settings detailed in the respective sections. Demonstration examples for each test sample are randomly selected from the training data, unless specific methods employ a specially designed selection method, such as EPR (Rubin et al., 2022). The values of the feature mean and covariance matrix are estimated from the demonstration set containing demonstration examples corresponding to all test samples. We depart from the assumption made in previous studies, which presupposes an equal distribution of training examples across all classes (Gao et al., 2021; Logan IV et al., 2022), in order to facilitate a more realistic and demanding evaluation.
Each experiment is repeated under five different random seeds. The batch size is set to 32, and the sequence length is configured to 128 for datasets with shorter texts, including SST2 (Socher et al., 2013), SST-5 (Socher et al., 2013), MR (Pang and Lee, 2005), CR (Hu and Liu, 2004), and TREC (Voorhees and Tice, 2000). On the other hand, for datasets with longer input texts, including AGNews (Zhang et al., 2015), DBPedia (Lehmann et al., 2015), Subj (Pang and Lee, 2004), CB (De Marneffe et al., 2019), and Amazon (McAuley and Leskovec, 2013), a batch size of 16 and a sequence length of 256 are employed. Regarding the hyperparameters in IDAICL, the values of λ and τ are fixed at 0.5 and 1, respectively, except in sensitivity tests. The settings used for the compared methods adhere to those specified in the original papers (Min et al., 2022a,b; Rubin et al., 2022; Zhao et al., 2021; Han et al., 2023; Fei et al., 2023). Accuracy serves as the primary evaluation metric, alongside the provided values of Macro-F1 for the LLaMA model. For each task, a specific template is utilized for inference, as detailed in Table 6. Additionally, we also examine the impact of different templates on the performance of IDAICL following those outlined by Zhao et al. (2021), which include question-answer templates, conversation-style templates, prompts resembling web pages, and variations on label names,
as listed in Table 10.

heading: D More Comparison Results
text: The comparison results between Vanilla ICL and IDAICL on LLaMA models with 13B and 33B parameters across various datasets are presented in Table 7. Additionally, the corresponding results for GPT-2 models with 0.1B and 0.3B parameters are outlined in Table 8. It is evident that IDAICL
consistently outperforms Vanilla ICL across all datasets and different model sizes, highlighting the high generalization capability of IDAICL. Additionally, IDAICL showcases reduced performance variance and significantly enhances the worst-case performance. Based on the findings presented in Table 9, IDAICL generally outperforms other prediction calibration methods, demonstrating the significance of statistical properties derived from the input data distribution in the predictions of PLMs.

heading: E More Results for Varying Templates
text: The comparison results between Vanilla ICL and IDAICL under all fifteen prompt templates are presented in Figure 7, illustrating that IDAICL consistently enhances both average and worst-case accuracy across all templates. Furthermore, the performance variance of IDAICL among different templates is notably smaller when compared to Vanilla ICL, highlighting the robustness of IDAICL’s performance across diverse templates.

heading: F More Sensitivity and Ablation Studies
text: We performed sensitivity tests on two hyperparameters within IDAICL: λ and τ . These values govern the strength of implicit augmentation and the influence of the class proportion term, respectively. As depicted in Figure 8, optimal performance is achieved when λ≈0.5 and τ≈1 for both datasets. Furthermore, Figures 9(a) and (b) illustrate the average performance of ten datasets across different hyperparameter settings. Much like the earlier findings, the best average performance is identified at λ = 0.5 and τ = 1. Consequently, setting λ as 0.5 and τ as 1 is recommended for real applications. Furthermore, the performance remains sta-
ble within the ranges of λ ∈ {0.25, 0.5, 0.75} and τ ∈ {0.5, 1, 1.5}, indicating that adjustments can be made within these stable ranges.

heading: G More Results for Imbalanced Labels
text: The imbalanced label distribution in the training data has a significant impact on the classification performance of the model (Zhou and Wu, 2023b; Zhou et al., 2022). We depicted the confusion matrices for the SST-2 and MR datasets under two imbalance levels in Figures 9(c) and (d), in which the proportion of the negative class in demonstrations is set to 0.1 and 0.2. These results manifest that IDAICL significantly enhances the performance of the underrepresented classes in comparison to
Vanilla ICL, thus proving its capability to address the class imbalance in demonstrations.
H Varying Demonstration Permutations
Research has substantiated that the performance of ICL is sensitive to the permutation of demonstrations (Lu et al., 2022; Zhao et al., 2021). We assessed the performance of IDAICL under varying demonstration permutations. Specifically, we selected ten different sets of twelve training examples from the SST-2 datasets. For each set of examples, we shuffled the order ten times and calculated the accuracy for each permutation. The findings are depicted in Figure 10, indicating that IDAICL exhibits relatively stable performance across different demonstrations and permutations, while Vanilla ICL demonstrates high variance.

heading: Future Work
text: Limitations
While IDAICL proves to be competitive in few-shot learning, there are limitations that open up avenues for future research. First, due to the necessity of accessing the parameters of the final fully connected layer in PLMs, IDAICL is exclusively suitable for open-source models. Future research is expected to develop alternative augmentation strategies tailored for black-box PLMs. Second, our evaluation of IDAICL focused on seven PLMs and ten text classification tasks. We defer further explorations involving other PLMs and non-classification tasks for future work. Additionally, IDAICL relies on a small set of demonstrations to estimate the feature mean and covariance matrix. If such a collection is unavailable or extremely scarce, IDAICL may need to be used in conjunction with demonstration generation methods. Other avenues for future work involve exploring more effective augmentation distributions. This entails exploring finer-grained distributions, such as category-level or sample-level distributions, to emphasize the unique characteristics of individual categories or samples, and extending these distributions beyond the constraints of training data. Furthermore, given the effectiveness of data augmentation in model training, future research could explore the utilization of our derived prediction function in both the training and fine-tuning phases of large PLMs