Enhancing In-Context Learning via Implicit Demonstration Augmentation

Brown et al. 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.

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. 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 upon this approach. This study introduces an algorithm for implicitly augmenting demonstrations within the realm of ICL.

Considering a PLM $\mathcal{G}$, this study focuses on the following task: given a query input text $\boldsymbol{x}$ and a candidate answer set $\mathcal{Y}\!=\!\{y_{1},y_{2},\cdots,y_{|\mathcal{Y}|}\}$, we aim to predict the answer $\hat{y}$ based on $m$ demonstration examples $\mathcal{C}\!=\!\{{c}_{1},c_{2},\cdots,c_{m}\}$, where each $c_{i}$ represents a training example $(\boldsymbol{x}_{i},y_{i})$ after template formulation and $m$ denotes the quantity of demonstration examples for each test sample. Formally, give a model $\mathcal{G}$, we first compute the probability of each answer $y_{j}$:

$$P_{\mathcal{G}}\left(y_{j}\mid\mathcal{C},\boldsymbol{x}\right).$$

(1)

Subsequently, the ultimate prediction $\hat{y}$, characterized by the highest probability is chosen from the candidate answer set $\mathcal{Y}$:

$$\hat{y}=\arg\max_{y_{j}\in\mathcal{Y}}P_{\mathcal{G}}\left(y_{j}\mid\mathcal{C},\boldsymbol{x}\right).$$

(2)

To simplify, the contextual input is denoted as $\boldsymbol{\tilde{x}}\!=\![\mathcal{C},\boldsymbol{x}]$ in the subsequent text. Then, the probability of answer $y_{j}$, represented as $P_{\mathcal{G}}(y_{j}|\boldsymbol{\tilde{x}})$, is computed using the Softmax function¹¹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.:

$$P_{\mathcal{G}}(y_{j}|\boldsymbol{\tilde{x}}):=P_{\mathcal{G}}(y_{j}|\boldsymbol{h}_{\boldsymbol{\tilde{x}}})=\frac{e^{\boldsymbol{w}_{y_{j}}^{T}\boldsymbol{h}_{\tilde{\boldsymbol{x}}}+b_{y_{j}}}}{\sum_{k}e^{\boldsymbol{w}_{k}^{T}\boldsymbol{h}_{\tilde{\boldsymbol{x}}}+b_{k}}},$$

(3)

where $\boldsymbol{h}_{\tilde{\boldsymbol{x}}}=\mathcal{G}(\tilde{\boldsymbol{x}})$ signifies the hidden state of the last block at the final position for $\tilde{\boldsymbol{x}}$. $\boldsymbol{w}_{k}$ and $b_{k}$ are the weight vector and bias corresponding to the final fully connected layer for the $k$-th token.

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 $\boldsymbol{h}_{\tilde{\boldsymbol{x}}}$ lies within the subspace spanned by $\boldsymbol{h}_{\mathcal{C}}$ and $\boldsymbol{h}_{{\boldsymbol{x}}}$: $\boldsymbol{h}_{\tilde{\boldsymbol{x}}}=\alpha\boldsymbol{h}_{\mathcal{C}}+\beta\boldsymbol{h}_{\boldsymbol{x}}$, where $\boldsymbol{h}_{\mathcal{C}}$ and $\boldsymbol{h}_{{\boldsymbol{x}}}$ represent the components of $\boldsymbol{h}_{\tilde{\boldsymbol{x}}}$ 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 $\alpha\boldsymbol{h}_{\mathcal{C}}$ and $\beta\boldsymbol{h}_{\boldsymbol{x}}$ as $\boldsymbol{h}_{\mathcal{C}}$ and $\boldsymbol{h}_{\boldsymbol{x}}$.

To augment $\boldsymbol{h}_{\mathcal{C}}$, we randomly sample vectors from the deep feature space of demonstrations. In particular, vectors are drawn from a multivariate normal distribution $\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\Sigma})$, where $\boldsymbol{\mu}$ and $\boldsymbol{\Sigma}$ denote the feature mean and covariance matrix. These statistical properties are estimated from the deep features of the demonstration set $\mathcal{D}$, which includes demonstration examples linked to all queries. The feature mean $\boldsymbol{\mu}$ is computed as

$$\boldsymbol{\mu}=\frac{1}{|\mathcal{D}|}\sum\nolimits_{i=1}^{|\mathcal{D}|}\boldsymbol{h}_{i},$$

(4)

where $\boldsymbol{h}_{i}=\mathcal{G}(c_{i})$ represents the hidden state of the last block at the final position for the $i$-th demonstration example $c_{i}$ in $\mathcal{D}$, and $|\mathcal{D}|$ denotes the size of $\mathcal{D}$. The covariance matrix $\boldsymbol{\Sigma}$ is computed as

$$\boldsymbol{\Sigma}=\frac{1}{|\mathcal{D}|}{\sum\nolimits_{i=1}^{|\mathcal{D}|}(\boldsymbol{h}_{i}-\boldsymbol{\mu})^{T}(\boldsymbol{h}_{i}-\boldsymbol{\mu})}.$$

(5)

Subsequently, ${\boldsymbol{h}}_{\mathcal{C}}$ is shifted in the extracted semantic vectors, resulting in augmented features, $\tilde{\boldsymbol{h}}_{\mathcal{C}}$, which follows

$$\tilde{\boldsymbol{h}}_{\mathcal{C}}\sim\mathcal{N}\left({\boldsymbol{h}}_{\mathcal{C}}+\lambda\boldsymbol{\mu},\lambda\boldsymbol{\Sigma}\right),$$

(6)

where $\lambda$ refers to a positive coefficient controlling the strength of semantic augmentation. In real-world applications, it can be directly assigned a value of 0.5. Sensitivity tests for $\lambda$ are discussed in Section 5.4.

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

$$\hat{y}=\arg\min_{y_{j}\in\mathcal{Y}}P_{\mathcal{G}}\left(y_{j}\mid\boldsymbol{h}_{\tilde{\boldsymbol{x}}}\right)^{-1}.$$

(7)

Assume that each $\boldsymbol{h}_{\mathcal{C}}$ is augmented for $\mathcal{M}$ times, resulting in an augmented demonstration feature set $\{\tilde{\boldsymbol{h}}_{\mathcal{C}}^{1},\cdots,\tilde{\boldsymbol{h}}_{\mathcal{C}}^{\mathcal{M}}\}$ with size $\mathcal{M}$. Here, $\tilde{\boldsymbol{h}}_{\mathcal{C}}^{i}$ represents the $i$-th augmented feature for $\boldsymbol{h}_{\mathcal{C}}$. Then, the final prediction for the query $\boldsymbol{x}$ depends on all augmented features of $\boldsymbol{h}_{\mathcal{C}}$ and can be expressed as

$${P}_{y_{j}}^{\mathcal{M}}(\boldsymbol{\tilde{x}})=\frac{1}{\mathcal{M}}{\sum\nolimits_{i=1}^{\mathcal{M}}P_{\mathcal{G}}(y_{j}|\tilde{\boldsymbol{h}}_{\mathcal{C}}^{i},\boldsymbol{h}_{\boldsymbol{x}})^{-1}},$$

(8)

$$\hat{y}={\arg\min_{y_{j}\in\mathcal{Y}}{P}_{y_{j}}^{\mathcal{M}}(\boldsymbol{\tilde{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 $\mathcal{M}\rightarrow\infty$, on the basis of the aforementioned decomposition of $\boldsymbol{h}_{\tilde{\boldsymbol{x}}}$, the expected prediction for answer $y_{j}$ (denoted as $P^{\infty}_{y_{j}}$) within the augmented feature set can be expressed as follows:

$\displaystyle P_{y_{j}}^{\infty}(\boldsymbol{\tilde{x}})\!=\!\mathbb{E}_{\boldsymbol{\tilde{h}}_{\mathcal{C}}}[{\sum_{k}e^{\Delta\boldsymbol{w}_{k,y_{j}}^{T}(\boldsymbol{\tilde{h}}_{\mathcal{C}}+\boldsymbol{h}_{\boldsymbol{x}})+\Delta b_{k,y_{j}}}}],$

(10)

where $\Delta\boldsymbol{w}_{k,y_{j}}=\boldsymbol{w}_{k}-\boldsymbol{w}_{y_{j}}$ and $\Delta b_{k,y_{j}}=b_{k}-b_{y_{j}}$.

However, accurately calculating $P^{\infty}_{y_{j}}$ is challenging. Alternatively, we proceed to derive a surrogate calculation for it. Applying the linearity of expectation, Eq. (10) can be expressed as:

$\displaystyle P_{y_{j}}^{\infty}(\boldsymbol{\tilde{x}})\!=\!\sum_{k}\mathrm{E}_{\tilde{\boldsymbol{h}}_{\mathcal{C}}}[e^{\Delta\boldsymbol{w}_{k,y_{j}}^{T}(\tilde{\boldsymbol{h}}_{\mathcal{C}}+\boldsymbol{h}_{\boldsymbol{x}})+\Delta b_{k,y_{j}}}].$

(11)

Given that $\boldsymbol{\tilde{h}}_{\mathcal{C}}$ is a Gaussian random variable conforming to $\mathcal{N}\left({\boldsymbol{h}}_{\mathcal{C}}+\lambda\boldsymbol{\mu},\lambda\boldsymbol{\Sigma}\right)$, we know that $\Delta\boldsymbol{w}_{k,y_{j}}^{T}\tilde{\boldsymbol{h}}_{\mathcal{C}}$ follows the multivariate normal distribution: $\mathcal{N}(\Delta\boldsymbol{w}_{k,y_{j}}^{T}\left({\boldsymbol{h}}_{\mathcal{C}}+\lambda\boldsymbol{\mu}\right),\lambda\Delta\boldsymbol{w}_{k,y_{j}}^{T}\boldsymbol{\Sigma}\Delta\boldsymbol{w}_{k,y_{j}})$. Then, utilizing the moment-generating function

$$\mathbb{E}[e^{{t}{X}}]=e^{{t}{\mu}+\frac{1}{2}{t}^{2}\sigma^{2}},X\sim\mathcal{N}(\mu,\sigma^{2}),$$

(12)

Eq. (11) can be derived as

$\displaystyle P_{y_{j}}^{\infty}(\boldsymbol{\tilde{x}})\!=\!\sum_{k}M_{k,y_{j}}N_{k,y_{j}}e^{\Delta\boldsymbol{w}_{k,y_{j}}^{T}(\boldsymbol{h}_{\mathcal{C}}+\boldsymbol{h}_{\boldsymbol{x}})+\Delta b_{k,y_{j}}},$

(13)

where $M_{k,y_{j}}=\exp(\lambda{\Delta\boldsymbol{w}_{k,y_{j}}^{T}\boldsymbol{\mu}})$ and $N_{k,y_{j}}=\exp({{\frac{\lambda}{2}}\Delta\boldsymbol{w}_{k,y_{j}}^{T}\boldsymbol{\Sigma}\Delta\boldsymbol{w}_{k,y_{j}}})$.

Subsequently, our newly proposed prediction function, referred to as IDA-Softmax, is defined as

$\begin{aligned} P_{y_{j}}^{\text{IDA}}(\boldsymbol{\tilde{x}}):=\sum_{k}M_{k,y_{j}}N_{k,y_{j}}e^{\Delta\boldsymbol{w}_{k,y_{j}}^{T}\boldsymbol{h}_{\tilde{\boldsymbol{x}}}+\Delta b_{k,y_{j}}}.\end{aligned}$

(14)

Consequently, instead of conducting the augmentation process explicitly, we can directly employ IDA-Softmax, $P_{y_{j}}^{\text{IDA}}$, for prediction. IDA-Softmax essentially utilizes two modulating factors associated with statistical properties derived from $\mathcal{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. Menon et al. (2021), which adjusts predictions by considering the class proportions within $\mathcal{D}$. Thus, the prediction for answer $y_{j}$ is computed as

$$\tilde{P}_{y_{j}}^{\text{IDA}}(\tilde{\boldsymbol{x}})=P_{y_{j}}^{\text{IDA}}(\tilde{\boldsymbol{x}})+\tau\log\pi_{y_{j}},$$

(15)

where $\tau$ is a positive hyperparameter, and $\pi_{y_{j}}$ demotes the proportion of answer $y_{j}$ in $\mathcal{D}$. In practical applications, the value of $\tau$ can be fixed at 1. This approach compensates for predictions of minor classes. When different answers are uniformly distributed, $\tau\log\pi_{y_{j}}$ exerts an equal influence on all answer types. Consequently, the final prediction is given by

$$\hat{y}={\arg\min_{y_{j}\in\mathcal{Y}}\tilde{P}_{y_{j}}^{\text{IDA}}(\boldsymbol{\tilde{x}})}.$$

(16)

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.