Divide, Reweight, and Conquer: A Logit Arithmetic Approach for In-Context Learning

## Abstract

Abstract In-Context Learning (ICL) emerges as a key feature for Large Language Models (LLMs), allowing them to adapt to new tasks by leveraging task-specific examples without updating model parameters. However, ICL faces challenges with increasing numbers of examples due to performance degradation and quadratic computational costs.
In this paper, we propose L ogit A rithmetic R eweighting A pproach (LARA), a novel framework that enhances ICL by using logit-based ensembling of multiple demonstrations.
Our approach divides long input demonstrations into parallelizable shorter inputs to significantly reduce memory requirements, and then effectively aggregate the information by reweighting logits of each group via a non-gradient optimization approach.
We further introduce Binary LARA (B-LARA), a variant that constrains weights to binary values to simplify the search space and reduces memory usage by filtering out less informative demonstration groups.
Experiments on BBH and MMLU demonstrate that LARA and B-LARA outperform all baseline methods in both accuracy and memory efficiency.
We also conduct extensive analysis to show that LARA generalizes well to scenarios of varying numbers of examples from limited to many-shot demonstrations. 1 1 1 Our code is available at https://github.com/Chengsong-Huang/LARA .

## 1 Introduction

In-Context Learning (ICL) (Brown et al., ) is one of the emergent abilities of Large Language Models (LLMs) as they are scaled to billions of parameters (Wei et al., ).
ICL enables LLMs to adapt to new tasks
by utilizing task-specific examples within the input context (Dong et al., ), and does not require any updates to or access to model parameters.
While ICL has achieved impressive performance across various domains, it
encounters significant challenges when dealing with an increasing number of examples.
Longer context window size often leads to performance degradation (Xiong et al., ).
This is due to the low density of useful information within longer prompts, and the reduced sensitivity to positional information, both of which diminish the capability of the model to effectively capture and utilize key content.
Additionally, the quadratic growth of computational cost with the input length makes it particularly expensive for large-scale models.

Previous works primarily focus on two directions to address these challenges. The first direction is input compression, which aims to shorten the input length (Jiang et al., ; Pan et al., ; Xu et al., ; Wingate et al., ) or selectively retrieve relevant portions of demonstrations to be included in the prompt (an Luo et al., ). However, these methods risk losing critical information,
which may negatively impact model performance.
The second direction involves aggregating hidden states within LLMs to simulate the effect of in-context demonstrations (Hao et al., ; Li et al., ; Hendel et al., ). These methods, however, are not applicable to closed-source models like GPT-4, as they require direct access to the model internal weights. Additionally, they contradict the core advantage of in-context learning, which is the ability to operate without modifications to hidden states or model parameters.

In this study, we propose a novel framework, Logit Arithmetic Reweighting Approach (LARA), which aims to combine the strengths of both input compression and hidden state approaches. Our method first divides demonstrations into subgroups to allow LLMs to focus on shorter inputs and reduce computational requirements.
We then design a weighted sum aggregation approach to combine the output logits from the language model given each subgroup of examples.
This ensures that the relevant information from each subgroup could potentially be captured by the language model.
One key innovation in LARA is that we use a non-gradient approach to optimize the weights of logits for each subgroup. We employ the Covariance Matrix Adaptive Evolution Strategy (CMA-ES) (Hansen & Ostermeier, ) to efficiently explore the weight vector space via resampling based on best-performing candidates.
This allows us to optimize the contribution of each subgroup without any gradient updates.
We further develop Binary-LARA (B-LARA) by constraining the weight values to $\{0,1\}$,
which can be interpreted as a process of subgroup selection. This not only reduces the computational cost but more importantly, leads to better performance due to the simplified search space for the binary weight vector.

Figure: Figure 1: Illustration of the differences between few-shot in-context learning and LARA (ours) during inference. Unlike few-shot in-context learning, which concatenates all demonstrations as a prefix to the input, our method splits the in-context examples into different groups. The next token is then generated based on a weighted average of logits, with weights precomputed using the framework described in Sec. 3.3.
Refer to caption: x1.png

Our experiments on BBH and MMLU benchmarks show that both LARA and B-LARA consistently outperform direct in-context learning and simple retrieval-based demonstration selection across various models,
with the additional benefit of lower GPU memory usage.
Further analysis reveals that the method excels in both low-resource scenarios with few examples and settings with abundant demonstrations, consistently delivering superior performance. Moreover, our ablation study highlights the critical role of the reweighting steps, although even logit averaging alone outperforms standard in-context learning.

To summarize, our main contributions are as follows:

- •
To the best of our knowledge, we are the first to propose ensembling information through logit arithmetic from different ICL demonstrations. We introduce LARA, a non-gradient optimization framework that reweights the information of different demonstration groups to improve ICL performance.
- •
We conduct extensive experiments on Llama3.1-8B (Dubey et al., ), Mistral-7B (Jiang et al., ), and Gemma-7B (Mesnard et al., ) on BBH Srivastava et al. () and MMLU Hendrycks et al. (), and show that LARA outperforms all baseline methods across all three models.
- •
Our comprehensive analysis reveals the broad applicability and efficiency of LARA and B-LARA. We demonstrate that our methods consistently outperform baselines across a wide range of example quantities, from fewer than 5 to more than 200. We also demonstrate the applicability of our methods to black-box LLMs.

## 2 Prelimaries

#### In-Context Learning.

Traditional In-Context Learning leverages $N$ labeled examples in the input prompt, represented as $\mathcal{D}_{\text{train}}=\{(\bm{x_{i}},\bm{y_{i}})\}_{i=1}^{N}$ to provide hints for language model generation. Each pair $(\bm{x_{i}},\bm{y_{i}})$ is converted into a semantically meaningful demonstration $d_{i}=\tau(\bm{x_{i}},\bm{y_{i}})$ using a predefined template $\tau$. These demonstrations are then concatenated to form a comprehensive context $\mathcal{C}=d_{1}\oplus d_{2}\oplus\cdots\oplus d_{N}$, with appropriate separators (e.g., newlines or special tokens) between each demonstration. For each test input $\bm{x_{\text{test}}}$, the language model receives the concatenated prompt $\mathcal{C}\oplus\bm{x_{\text{test}}}$ to generate a response.

#### Logit-based Generation.

We consider decoding approaches for language generation, where the language model receives an input prompt $\mathcal{C}\oplus\bm{x_{\text{test}}}$ and produces coherent and logical responses. The term “logit” refers to the raw, unnormalized scores output by the model before they are converted into probabilities by a softmax function. These logits are generated by passing the input sequence through the LLM. Formally, given the logit $\bm{z}$, the probability of the next token $x_{t}$ given the previous tokens $x_{1:t-1}$ is computed using the softmax function:

$$ $P(x_{t}\mid x_{1:t-1})=\frac{\exp(\bm{z}_{x_{t}})}{\sum_{x^{\prime}\in V}\exp( \bm{z}_{x^{\prime}})}$ (1) $$

where $\bm{z}_{x_{t}}$ is the logit corresponding to the token $x_{t}$, and $V$ is the vocabulary set.

## 3 Methodology

Figure: Figure 2: Illustration of the LARA framework. The input demonstration set $\mathcal{D}_{\text{train}}$ is divided into subsets $\mathcal{S}_{1},\mathcal{S}_{2},\dots,\mathcal{S}_{k}$, which are further split into two groups: one for candidate examples and the other for validation examples. For each token, logits are generated using Logit-Arithmetic Decoding, which aggregates the output logits from all subsets. After generating all tokens, the cross-entropy loss is computed based on the weighted-average logits and the ground truth from the validation subset. The subset weights are then resampled and adjusted to minimize the loss. This process of token generation, loss calculation, and weight resampling is repeated iteratively. After optimizing the weights for the first group of candidate examples, the roles of the candidate and validation examples are swapped.
Refer to caption: x2.png

In this section, we provide an overview of LARA. Figure 2 illustrates the overall framework of our approach. Unlike directly concatenating $\mathcal{D}_{\text{train}}$ into a single sequence, we first divide the $N$ examples into subgroups, which are used as inputs to the LLM. The output logits from these subgroups are then aggregated, and we assign weights to each subgroup using a non-gradient search algorithm. During inference, the precomputed weights are used to combine the logits from each group.

In Sec. 3.1, we explain the partition strategy to divide examples into subgroups. Then we introduce how the outputs are aggregated across different subgroups in Sec. 3.2, and the reweighting strategy for optimal combination in Sec. 3.3. Furthermore, we show in Sec. 3.4 that imposing a hard constraint for our reweighting strategy could further reduce memory usage and computational resources. Finally, we discuss in Sec. 3.5 the inference efficiency brought by our proposed approach.

### 3.1 Partition Strategy

Given $N$-shot in-context examples, we first split $\mathcal{D}_{\text{train}}$ into $k$ disjoint subsets each containing $L$ in-context examples, such that $\mathcal{D}_{\text{train}}=\mathcal{S}_{1}\cup\mathcal{S}_{2}\cup\ldots\cup
\mathcal{S}_{k}$ with $|\mathcal{S}_{i}|=L$ for all $i\in\{1,\ldots,k\}$. When inputting a subgroup $\mathcal{S}_{i}$ to an LLM, we concatenate all of its elements to get $\mathcal{C}_{i}=d_{(i-1)L+1}\oplus d_{(i-1)L+2}\oplus\cdots\oplus d_{iL}$, and the complete input for the $i$-th subgroup to LLM is $\mathcal{C}_{i}\oplus\bm{x}_{\text{test}}$.
We assume that $N$ is divisible by $k$ in our experiments, so that $L=N/k$. In practice, in cases where $N$ is not divisible by $k$, we could truncate the last subset and only retain $L(k-1)$ examples.

### 3.2 Logit-Arithmetic Decoding

Previous studies (Li et al., ; Liu et al., ; Dekoninck et al., ) have utilized logit offsets to control the outputs of large language models for better generation quality or instruction following.
Inspired by these work, we propose a novel method that combines information from multiple in-context demonstrations through logit-arithmetic decoding. Specifically, our approach focuses on aggregating the logits produced by the language model outputs for various contextual inputs.
With the input query $\bm{x}_{\text{test}}$ and the example subset being $\mathcal{S}_{i}$, we can compute the logit outputs of the language model, denoted as $f_{\theta}(\mathcal{S}_{i},\bm{x}_{\text{test}})=\log p(y\mid\mathcal{S}_{i},
\bm{x}_{\text{test}})$.
We then combine these logits using a weighted sum to get the generation probability over the output token:

$$ $p(y\mid\bm{x}_{\text{test}},\bm{w})=\text{softmax}\left(\sum_{i=1}^{k}w_{i} \cdot f_{\theta}(\mathcal{S}_{i},\bm{x}_{\text{test}})\right)$ (2) $$

where $k$ is the number of example subsets, and $w_{i}$ are weights that indicate the importance of the contribution of each subset, with $\sum_{i=1}^{k}w_{i}=1$. As a baseline approach, we could set uniform weighting, where $w_{i}=1/k$. However, this may not be optimal for all tasks, as the quality and relevance of different subgroups may vary. In the following section, we introduce a reweighting strategy to optimize these weights to enhance model performance.

### 3.3 Reweighting Logits by Non-Gradient Optimization

To further enhance the model performance, we employ non-gradient optimization methods to optimize the weights $w_{i}$ based on the loss calculated from $p(y\mid\bm{x_{\text{val}}})$. Given the combined probability $p(y\mid\bm{x_{\text{val}}})$, our objective is to minimize a cross-entropy loss function $\mathcal{L}(\bm{w})$ over the predicted probabilities and the ground truth. Specifically, we utilize the following cross-entropy loss function for the generation model:

$$ $\mathcal{L}(\bm{w})=-\sum_{(\bm{x_{\text{val}}},\bm{y_{\text{val}}})\in \mathcal{D}}\sum_{t=1}^{T}\log p(y_{t}\mid\bm{x_{\text{val}}},\bm{w})$ $$

where $D$ represents the validation dataset, $T$ is the length of the sequence, $y_{t}$ is the true word at time step $t$, $\bm{x_{\text{val}}}$ is the input sequence, $\bm{w}$ denotes the weight vector, and $p(y_{t}\mid\bm{x_{\text{val}}},\bm{w})$ represents the predicted probability of the true word $y_{t}$ at time step $t$, given the input sequence $\bm{x_{\text{val}}}$ and the weight vector $\bm{w}$.

To avoid introducing additional labeled data, we employ a cross-validation strategy. We partition the demonstration set $\mathcal{S}$ into two subsets: $\mathcal{S}_{A}=\mathcal{S}_{1}\cup\mathcal{S}_{2}\cup\ldots\cup S_{\lfloor k/
2\rfloor}$ and $\mathcal{S}_{B}=\mathcal{S}_{\lfloor k/2\rfloor+1}\cup\mathcal{S}_{\lfloor k/2
\rfloor+2}\cup\ldots\cup\mathcal{S}_{k}$. When optimizing weights for $\mathcal{S}_{i}\in\mathcal{S}_{A}$,
we use $\mathcal{S}_{B}$ as the validation set, and vice versa.

We choose non-gradient optimization methods over gradient-based alternatives due to two key factors: (1) The loss function $\mathcal{L}(\bm{w})$ is non-differentiable, since updating the weight vector
$\bm{w}$ affects the logits of subsequent tokens, leading to possibly different decoding results of subsequent tokens. (2) The dimensionality of the weight vector $\bm{w}$ is relatively low, specifically equalled to the number of groups $k$.

In our empirical experiments, we refer to Liu et al. () and employ the Covariance Matrix Adaptive Evolution Strategy (CMA-ES) (Hansen & Ostermeier, ).
CMA-ES is a stochastic, derivative-free optimization algorithm.
During each iteration, CMA-ES samples a set of candidates in the space of the weight vector $\bm{w}$ from a multivariate normal distribution, evaluates $\mathcal{L}(\bm{w})$ for each candidate, and then updates the mean and covariance matrix of the distribution based on the best-performing candidates.
This allows for an efficient exploration over the weight space.

### 3.4 Binary Constraints for LARA

We further propose a variant of LARA, named as B-LARA, by imposing a hard constraint on the weight vector $\bm{w}$ to binary values $\{0,1\}$.
This binary constraint offers two key advantages: first, it simplifies the search space and potentially leads to faster convergence; second, it allows for direct elimination of demonstration groups with zero weight, thereby improving inference efficiency.
Intuitively, the binary optimization of $\bm{w}$ can be seen as a form of subset selection to identify the most relevant demonstrations in $\mathcal{D}_{\text{train}}$ benefitting model performance on specific tasks.

To solve this binary optimization problem, we employ the simplest evolution strategy (1+1)-ES (Rechenberg, ).
It involves a simple cycle: a single parent produces one offspring per generation through mutation—adding a small, random change. If this offspring performs as well or better than the parent based on a predefined fitness criterion, it becomes the new parent for the next generation. Otherwise, the original parent remains. The overall sampling procedure is shown in Algorithm 1.

Figure: Algorithm 1 B-LARA Optimization Algorithm with Updated Index

The simplicity of this method in repeated mutation and selection makes it particularly suitable for our binary optimization scenario.

### 3.5 Computational Complexity

We analyze the computational complexity of LARA and B-LARA compared to standard ICL.
During inference, the self-attention mechanism in Transformer models is the primary bottleneck for GPU memory requirement, with the memory complexity being $O(n^{2})$, where $n$ is the input sequence length. This quadratic scaling is due to the pairwise interactions between tokens in the attention matrix.

By splitting the input sequence into $k$ groups, each of length around $\frac{n}{k}$, LARA and B-LARA can leverage parallel computing resources more effectively.
The complexity for LARA becomes $O(\frac{n}{k}^{2}*k)$= $O(\frac{n^{2}}{k})$.
B-LARA further reduces computational complexity by selecting only a subset of groups. If $m$ out of $k$ subgroups are assigned non-zero weights, then the complexity of B-LARA becomes $O(\frac{mn^{2}}{k^{2}})$.
We show the empirical GPU memory usage in Sec. 5.2.

## 4 Experiments

In this section, we provide details of our main experiments. We first give an overview of the experimental setup and implementation details in Sec. 4.1, and then present our findings along with the results in Sec. 4.2.

### 4.1 Experimental Setup

#### Datasets and Evaluation.

We evaluate our methods using two well-established benchmarks: Big-Bench Hard (BBH) (Srivastava et al., ) and Massive Multitask Language Understanding (MMLU) (Hendrycks et al., ). BBH tests models on challenging reasoning tasks across domains including arithmetic reasoning, commonsense reasoning, and linguistics. MMLU measures generalization across 57 diverse subjects, covering both humanities and STEM fields, offering a comprehensive evaluation of knowledge and problem-solving abilities of LLMs. For both benchmarks, we use exact match (EM) as our evaluation criterion, which requires model predictions to perfectly match the correct answers. We report the accuracy scores in our experiment results. The details about dataset analysis and prompts can be found in Appendix A.

#### Models.

Our proposed LARA for in-context learning is applicable to any LLM.
To demonstrate its generality, we evaluate it on three open-source, decoder-only models: Llama3.1-8B (Dubey et al., ), Mistral-7B (Jiang et al., ), and Gemma-7B (Mesnard et al., ). Llama-3.1-8B is known for strong performance across various NLP tasks, Mistral-7B is optimized for efficiency and is balanced between computational cost and accuracy. Gemma-7B focuses on advanced reasoning and language comprehension. These models represent diverse architectures and training strategies, allowing us to test the adaptability of our methods.
By using open-source models in evaluation, we ensure the reproducibility of our proposed method and validate its broad applicability across state-of-the-art model architectures.

#### Hyperparameter Setting.

In our main experiment, we use $\mathcal{D}_{\text{train}}$ consisting of $N=32$ in-context examples for our methods. For each task, $\mathcal{D}_{\text{train}}$ is split into subsets of size $L\in\{2,4,8\}$, and for each $L$ we perform up to $J=20$ iterations for weight optimization. We compare the minimum validation loss across different settings of $L$ to determine the optimal configuration for the final inference phase.
The baseline methods also use the same $\mathcal{D}_{\text{train}}$ as input.
For our method and all baselines, we set the temperature to 0 to enforce greedy decoding. Our experiments are conducted on a single A100 80GB GPU.

#### Compared Methods.

We introduce several primary baseline methods: Direct In-Context Learning (ICL), KNN-Augmented In-ConText Example Selection Liu et al. () (KATE), Rationale-Augmented Ensembles (RAE) (Wang et al., ) and In-context Vector (ICV) (Liu et al., ) as the representative of parameter access methods.
We use the same 32 in-context examples as inputs to all baseline methods as our proposed method.
For Direct ICL, all 32 examples are concatenated with the prompt. For KATE, we apply the Top-K selection from Liu et al. () that uses a smaller model(^2^22[https://huggingface.co/sentence-transformers/all-distilroberta-v1](https://huggingface.co/sentence-transformers/all-distilroberta-v1)) to retrieve the most similar input-output pairs from $\mathcal{D}_{\text{train}}$ as in-context demonstrations. We evaluate KATE with 2, 4, and 8 demonstrations as baselines. For RAE, we divide the examples into different groups and use each group as in-context examples to generate separate results. The final output is determined by applying majority voting across these individual group-based results. In ICV, we follow the original paper to set $\lambda=0.1$ and average the ICV given by all 32 examples.
We report results with group sizes of 2, 4, and 8 to ensure the same memory usage as our method.

### 4.2 Main Results

**Table 1: Accuracy of all methods on BBH and MMLU. The results shown are the average performance across datasets within each benchmark. Please refer to appendix B.2 for breakdown results of each dataset. The subscript of KATE indicates the number of selected ICL demonstrations as input to LLMs.**
|  | BBHaverage | MMLUaverage |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Llama3.1-8B | Gemma-7B | Mistral-7B | Llama3.1-8B | Gemma-7B | Mistral-7B |
| ICL | 45.64 | 37.08 | 42.91 | 65.63 | 61.44 | 62.84 |
| KATE2 | 43.60 | 37.07 | 43.16 | 66.62 | 56.28 | 63.99 |
| KATE4 | 44.03 | 38.83 | 43.16 | 66.75 | 55.78 | 63.48 |
| KATE8 | 44.47 | 37.03 | 42.96 | 67.19 | 54.13 | 63.93 |
| IRE2 | 44.59 | 40.24 | 43.95 | 66.88 | 65.18 | 62.99 |
| IRE4 | 45.23 | 40.44 | 44.49 | 66.40 | 65.01 | 62.99 |
| IRE8 | 44.06 | 39.85 | 44.07 | 67.09 | 64.80 | 63.61 |
| ICV | 45.93 | 42.16 | 44.50 | 66.97 | 64.99 | 64.02 |
| LARA | 47.46 | 41.77 | 44.77 | 66.54 | 64.36 | 63.93 |
| B-LARA | 47.69 | 42.75 | 45.03 | 67.80 | 65.56 | 64.12 |

Results from Table 1 demonstrate the effectiveness of our proposed methods, LARA, and B-LARA, across BBH and MMLU benchmarks. B-LARA consistently outperforms most of baseline methods across three model architectures. Notably, B-LARA achieves the highest accuracy and improves over direct ICL by $2.05$, $5.67$, and $2.12$ points on BBH dataset across three models respectively.
Moreover, our methods and can consistently outperform retrieval or simple ensemble baselines like KATE and IRE, indicating that our method is more effective in combining information from multiple demonstration subgroups.
Compared to the ICV baseline, which has the advantage of access to model parameters, our methods still achieve better performance without access to the hidden state, which further demonstrates the efficacy of our methods in aggregating information without direct access to model internal parameters.

An interesting finding is that B-LARA performs better than LARA despite a more constrained search space for the weight vector.
We believe this is because we only use 20 iterations for weight optimization, and
the binary constraint brings more benefits by introducing a simplified optimization landscape and providing a regularization effect to prevent overfitting.

## 5 Analysis

In this section, we present a comprehensive analysis of our proposed method LARA under various conditions.

### 5.1 How does the Reweighting Step Affect Model Performance?

**Table 2: Average performance of Llama3.1-8B our methods without reweighting. For the ablation “w/o reweight”, the subscript means the size $L$ of each group of demonstrations. The results for other models are shown in Appendix B.1**
| Method | BBH | MMLU |
| --- | --- | --- |
| LARA | 47.46 | 66.54 |
| B-LARA | 47.69 | 67.80 |
| w/o reweight2 | 43.33 | 67.23 |
| w/o reweight4 | 44.50 | 67.58 |
| w/o reweight8 | 43.02 | 67.43 |

We conduct an ablation study to assess the effectiveness of the reweighting step, denoted as “w/o reweight” which simply averages over the output logits of the LLM across different demonstration groups.

In our ablation study, removing the reweighting step used in LARA also demonstrated its value by outperforming traditional baseline methods. For instance, it achieved a notable $67.58$ with Llama3.1-8B in the MMLU benchmark, which is better than directly ICL ($65.63$). This performance highlights that logit-arithmetic can successfully combine the information in different groups of demonstrations.

The results further emphasize the importance of the reweighting step in LARA. LARA outperforms the non-reweight version in most settings. This underscores the reweighting process as critical for enhancing model accuracy. The worse performance of non-reweight offers clear evidence of how significant reweighting is to optimizing the model’s contextual handling.

### 5.2 How Does LARA Enhance Memory Efficiency?

Figure: Figure 3: GPU Memory usage of LARA in gigabytes on a single A100 80GB GPU with different input sequence lengths and number of subgroups. Note that when the number of subgroups equals to 1, the setting is the same as ICL. The sequence length is denoted in thousands of tokens. We set the batch size equal to 4. Data points indicating Out-Of-Memory (OOM) are omitted.
Refer to caption: x3.png

We empirically evaluate the computational efficiency of LARA by measuring GPU memory usage with different input sequence lengths and subgroup configurations. We set the number of groups $k$ with 1,2,4,8. Specifically, when $k$ is set as 1, LARA will degrade to ICL.

Results in Figure 3 demonstrate that LARA is more memory-efficient compared to standard ICL, especially when handling long sequences. Standard ICL results in Out-of-Memory (OOM) errors when the input length exceeds 10k tokens on a Mistral-7B model with a batch size of 4 on an A100 80GB GPU. In contrast, our method handles input lengths over 25k tokens with 4 and 8 subgroups, demonstrating that LARA efficiently utilizes larger amounts of training data.

### 5.3 Can LARA Perform Well with More Examples?

**Table 3: Accuracy of methods on GoEmotion and TacRED. The subscript of IRE means the number of groups in IRE.**
|  | GoEmotion | TacRED |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Llama3.1-8B | Gemma-7B | Mistral-7B | Llama3.1-8B | Gemma-7B | Mistral-7B |
| ICL | 18.60 | 15.60 | 17.80 | 38.20 | 43.80 | 55.40 |
| IRE2 | 22.20 | 22.20 | 21.60 | 43.80 | 45.40 | 55.40 |
| IRE4 | 21.00 | 22.40 | 21.40 | 45.60 | 45.00 | 52.40 |
| IRE8 | 21.20 | 19.00 | 20.40 | 36.20 | 39.00 | 49.20 |
| ICV | 18.80 | 20.80 | 18.40 | 44.40 | 46.80 | 54.40 |
| LARA | 21.00 | 20.80 | 19.20 | 48.60 | 47.40 | 54.20 |
| B-LARA | 24.00 | 22.80 | 23.80 | 48.60 | 49.00 | 59.00 |

We investigate the performance of LARA with an increased number of demonstrations, leveraging the LongICLBench (Li et al., ), a benchmark tailored for addressing challenges in long in-context learning. For our experiments, we select two datasets: GoEmotion and TacRED. Following the LongICLBench setup, we employ multiple rounds of examples, where each round includes several examples, each labeled with a distinct class.
To align with the input limit constraints of ICL, we sampled 8 rounds (224 examples) of examples for GoEmotions and 4 rounds (164 examples) for TacRED. For LARA and B-LARA, we choose 4, 8, and 16 as the potential candidate number of groups.
We report the accuracy of different methods on these datasets in Table 3.

The experimental results clearly highlight the advantages of LARA, which demonstrates consistent improvements over baseline methods across both GoEmotion and TacRED datasets, showcasing its effectiveness in diverse tasks. Notably, the B-LARA variant further amplifies this performance, outperforming all competing approaches on both datasets and across various models. This suggests that B-LARA can work well in many shot settings.

### 5.4 Can LARA Perform Well with Limited In-Context Examples?

In previous experiments, we primarily explore the many-shot in-context learning (ICL) setting. In this subsection, we focus on a more constrained scenario, where only a limited number of in-context examples are available.
This analysis aims to understand the relationship between the number of demonstrations and the performance of LARA compared to baseline methods with limited examples.

Figure: Figure 4: Accuracy of LARA on BBH using different numbers of examples. B-LARA uses different settings due to differences in example usage during training and inference. We use two lines to highlight this difference. The accuracy means the average accuracy on BBH dataset.
Refer to caption: x4.png

We set the number of examples $N$ within $\{2,4,8,16\}$ and compare our proposed method with ICL on the BBH dataset with Mistral-7B. Figure 4 demonstrates that both LARA and B-LARA consistently outperform the baseline ICL, and the performance gap increases with the number of examples used. Note that we do not plot the performance of LARA and B-LARA under $N=2$. This is because LARA and B-LARA are simplified to our non-reweighting ablation when the size of each subgroup becomes $1$ and no reweighting is required. We also show the performance of performance without reweighing here. We set the number of group $k$ as 2 in this experiment. While there is a significant gap between the non-reweight version and B-LARA, the non-reweight version still demonstrates effectiveness compared to ICL.

Since B-LARA has a weight constraint of $\{0,1\}$, subgroups with zero-weights are pruned during inference for efficiency.
As shown in Figure 4, the real number of examples used by B-LARA in inference is substantially lower than other methods.
In the 32-shot setting, only 45% of subgroups of B-LARA are assigned non-zero weights,
reducing more than half of the computational load without compromising performance.
Additionally, as the total number of examples increases, the proportion of examples used in inference decreases, indicating that B-LARA is particularly suitable for resource-constrained environments.

### 5.5 Is LARA Applicable to Black-Box LLMs?

**Table 4: Average performance of various methods of GPT-4o-mini on the BBH benchmark.**
| ICL | LARA | B-LARA |
| --- | --- | --- |
| 53.17 | 56.06 | 57.41 |

One advantage of our method is that it could also be applied to LLM APIs, since it only uses output logits for example reweighting or selection.
In these scenarios, techniques such as in-context vector or task vector, which often rely on internal state visibility, cannot be applied.

We evaluate our method with GPT-4o-mini (^3^33gpt-4o-mini-2024-07-18) on BBH dataset.
The results in Table 4 demonstrate that LARA and B-LARA outperform ICL. We note that the OpenAI API only provides top $20$ logits for each output token, while our methods are still able to achieve competitive results.
This indicates that our method generalizes well to black-box LLMs,
and can be applied to situations where internal weights of models are restricted and only output logits are available.

## 6 Related Work

### 6.1 Long In-Context Learning

Recent studies on long-context learning problems in LLMs can be categorized into two main strategies: enhancing the impact of in-context examples and compressing input sequences.
Structured prompting leverages rescaled attention mechanisms to effectively integrate grouped examples (Hao et al., ). Methods such as task vectors (Hendel et al., ) and function vectors (Todd et al., ) further refine this strategy by generating vectors that assess the contribution of each example based on the offset of hidden state, which improves model adaptability.
Liu et al. () generate task-specific vectors that steer model behavior in latent space based on the in-context examples.
Regarding input compression, methods like prompt pruning (Jiang et al., ; Pan et al., ) and additional summarization models (Xu et al., ; Gilbert et al., ) directly shorten inputs while maintaining essential content. Soft prompt-based compression (Wingate et al., ; Mu et al., ) intends to generate a soft-prompt that includes most of the information.

### 6.2 Logit Arithmetic

Several works have employed logit arithmetic across various domains and downstream tasks. Contrastive decoding  (Li et al., ) improves performance by utilizing the difference in logits from models of different sizes. Proxy tuning (Liu et al., ) enhances a larger model’s capabilities by adding the logit differences of a smaller model, recorded before and after training, to simulate training effects. In model arithmetic  (Dekoninck et al., ), logits adjusted with various prompts steer the generation processes of large language models. Huang et al. () propose using logit subtraction to facilitate the selective forgetting of knowledge in LLMs. Additionally, logit arithmetic has been leveraged to enhance the safety of generated outputs (Xu et al., ).

### 6.3 Non-gradient Optimization of LLMs

Due to the high memory requirements associated with gradient-based optimization methods, recent research has shifted towards non-gradient techniques for neural network optimization. Zhang et al. (); Malladi et al. () propose training large language models (LLMs) using non-gradient methods to mitigate these memory constraints. These approaches have also been applied in federated learning, exploring their effectiveness in distributed settings (Xu et al., ). Additionally, a gradient-free method has been used to optimize manifold neural networks (Zhang et al., ). Similarly, LoraHub (Huang et al., ) utilizes non-gradient techniques to dynamically reweight different LoRA modules, enhancing adaptation to new downstream tasks. Guo et al. () also introduces non-gradient methods to prompt engineering to search for better prompts.

## 7 Conclusion

We propose LARA, a novel framework that enhances in-context learning by ensembling logits from multiple demonstrations, improving performance without requiring parameter updates. Our method reduces computational complexity while achieving better accuracy. Additionally, Binary LARA further optimizes efficiency by selectively removing less informative demonstrations. Experiments on BBH and MMLU benchmarks show that both LARA and B-LARA outperform traditional ICL methods in terms of efficiency and performance.
Future research directions include extending our study to combine logits from different sources beyond just in-context learning (ICL) examples—such as different models or varying instructions—and building a distributed inference system based on LARA.

## Reproducibility Statement

We provide detailed information on the dataset we used in Appendix A. The codes for our main experiment can be found in [https://github.com/Chengsong-Huang/LARA](https://github.com/Chengsong-Huang/LARA).

## Acknowledgement

We would like to thank Xuezhi Wang (Google DeepMind) for her insightful suggestions and constructive feedback, which significantly improved the quality of this work. We thank Changze Lv (Fudan University) and our labmates Yuyi Yang and Jixuan Leng for their proofreading.

## Appendix A Dataset Details

### A.1 Prompts for Inference

**Table 5: Prompt examples for each dataset in One-shot learning.**
| Dataset | Prompt |
| --- | --- |
| BBH | Question: {question} |
|  | Answer: {answer} |
|  | Question: {question} |
|  | Answer: |
| MMLU | The following are multiple choice questions (with answers) about {subject}. |
|  | Question: {question} Answer: {answer} |
|  | Question: {question} Answer: |
| GoEmotion | Given a comment, please predict the emotion category of this comment. The predict answer must come from the demonstration examples with the exact format. |
|  | The examples are as follows: |
|  | comment: {question} |
|  | emotion category: {answer} |
|  | comment: {question} |
|  | emotion category: |
| TacRED | Given a sentence and a pair of subject and object entities within the sentence, please predict the relation between the given entities. |
|  | You can only select from the following words: {potential relation} |
|  | sentence: {question} |
|  | the relation between the two entities is: {answer} |
|  | sentence: {question} |
|  | the relation between the two entities is: |

### A.2 Dataset Statistics

**Table 6: Dataset Statistics.**
| Dataset | #Tokens/Shot | Description |
| --- | --- | --- |
| BBH | 55 | A collection of challenging tasks from the BIG-Bench Hard benchmark. |
| MMLU | 65 | Multiple-choice questions across various subjects. |
| GoEmotion | 28 | Annotated Reddit comments for emotion classification. |
| TacRED | 80 | A dataset for relation extraction tasks. |

## Appendix B Full Results

### B.1 Full Ablation Study

**Table 7: Ablation Study Results**
|  | BBHaverage | MMLUaverage |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Llama3.1-8B | Gemma-7B | Mistral-7B | Llama3.1-8B | Gemma-7B | Mistral-7B |
| ICL | 45.64 | 37.08 | 42.91 | 65.63 | 61.44 | 62.84 |
| KATE2 | 43.60 | 37.07 | 43.16 | 66.62 | 56.28 | 63.99 |
| KATE4 | 44.03 | 38.83 | 43.16 | 66.75 | 55.78 | 63.48 |
| KATE8 | 44.47 | 37.03 | 42.96 | 67.19 | 54.13 | 63.93 |
| IRE2 | 44.59 | 40.24 | 43.95 | 66.88 | 65.18 | 62.99 |
| IRE4 | 45.23 | 40.44 | 44.49 | 66.40 | 65.01 | 62.99 |
| IRE8 | 44.06 | 39.85 | 44.07 | 67.09 | 64.80 | 63.61 |
| ICV | 45.93 | 42.16 | 44.50 | 66.97 | 64.99 | 64.02 |
| LARA | 47.46 | 41.77 | 44.77 | 66.54 | 64.36 | 63.93 |
| B-LARA | 47.69 | 42.75 | 45.03 | 67.80 | 65.56 | 64.12 |
| w/o reweight2 | 43.33 | 43.56 | 42.83 | 67.23 | 65.61 | 62.95 |
| w/o reweight4 | 44.50 | 41.98 | 44.78 | 67.58 | 65.87 | 63.32 |
| w/o reweight8 | 43.02 | 39.35 | 44.84 | 67.43 | 65.04 | 63.55 |

### B.2 Full Main Results

Here we will show the full results of our three models in BBH and MMLU benchmark. The methods include LARA, B-LARA, KATE, ICL, LAG(logit-average-generation which is the ablation study in our paper, together with IRE and ICV.

Figure: Table 8: Performance scores across tasks in BBH (Mistral-7B)

Figure: Table 9: Performance scores across tasks in MMLU (Mistral-7B) Part 1

Figure: Table 10: Performance scores across tasks in MMLU (Mistral-7B) Part 2

Figure: Table 11: Performance scores across tasks in MMLU (Mistral-7B) Part 2

Figure: Table 12: Performance scores across tasks in BBH (Gemma-7B)

Figure: Table 13: Performance scores across tasks in BBH (Gemma-7B) Part 1

Figure: Table 14: Performance scores across tasks in BBH (Gemma-7B) Part 2

Figure: Table 15: Performance scores across tasks in MMLU (Mistral-7B) Part 3

Figure: Table 16: Performance scores across tasks in BBH (Llama3.1-8B)

Figure: Table 17: Performance scores across tasks in MMLU (Llama3.1-8B) Part 1

Figure: Table 18: Performance scores across tasks in MMLU (Llama3.1-8B) Part 2

Figure: Table 19: Performance scores across tasks in MMLU (llama3.1-8B) Part 3

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. For the experiments in Section 5.5, please provide additional details. For instance, does decoding each token require re-feeding the entire prefix to the API? If so, this would result in significant costs.
2. Where are the results for different candidate numbers of groups in Table 1 and 3?
soundness: 2
strengths: 1. The proposed logits re-ensemble provides partial evidence supporting the effectiveness of grouped ICL, serving as an additional rationale for future research on extending this approach toward long ICL.
2. The proposed Binary LARA can further enhance performance while reducing costs.
summary: This paper proposes LARA, a new framework that enhances in-context learning (ICL) by ensembling logits from multiple demonstrations, achieving accuracy gains without requiring any parameter updates. The authors also introduce Binary LARA (B-LARA), which further refines efficiency by selectively excluding less informative demonstrations. Experiments on BBH and MMLU benchmarks validate that both LARA and B-LARA outperform standard ICL.
weaknesses: 1. The proposed method does not effectively address or mitigate the disadvantages of the related work introduced in the introduction. For instance, while OpenAI's API provides some logits information, LARA fundamentally remains a white-box approach. Additionally, Binary LARA also carries the risk of losing critical information. As a result, the motivation for the method is unclear, and there is a lack of in-depth analysis regarding why LARA is effective.
2. The selected baseline does not effectively demonstrate the validity of the proposed method. The comparison baseline does not specifically address the issue of large-scale demonstrations. Additionally, other relevant grouping-based methods, such as [1] and [2], were not included in the comparison. Additionally, the main text uses RAE, while the table uses IRE.
3. LARA requires iterative computation of optimal weights for the model and dataset before inference, which introduces latency. When there are few samples to infer, or even just a single sample—common in real-world interactions with LLMs—the proportion of additional costs relative to the inference becomes particularly high.

  [1] Structured Prompting: Scaling In-Context Learning to 1,000 Examples.

  [2] Focused Large Language Models are Stable Many-Shot Learners.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: As mentioned in the above weaknesses.
soundness: 3
strengths: 1. The issues of performance degradation and quadratic computational costs in many-shot in-context learning (ICL) are currently hot research topics in the AI community. The approach presented in this paper, which involves Logit Arithmetic concatenation after grouping inputs, is both novel and intriguing. The experiments effectively validate the method's improvements on these problems.

2. The method is supported by a solid theoretical foundation. The design of validation and cross-entropy loss can be viewed as a form of non-gradient fine-tuning for special tasks in the ICL context. B-LARA represents a further optimization of the LARA method's weights. The paper provides detailed formulas and algorithmic derivations to substantiate these claims.
summary: This paper presents a method to address the issues of performance degradation and quadratic computational costs in the context of many-shot in-context learning (ICL). By grouping demonstrations for parallel input and employing Logit Arithmetic for reweighting and concatenation, the authors cleverly leverage the effective information from demonstrations. This approach alleviates the resource consumption associated with large-shot inputs and enhances the performance of ICL to some extent.
weaknesses: 1. The use of Logit Arithmetic reweighting and concatenation to leverage effective information from demonstrations is indeed a good approach. However, I am curious about how the authors address the issue of output prefix inconsistency during the decoding process after grouping different demonstrations as inputs. For example, if the first token of the output differs under different shot inputs, how can the subsequent tokens be rigorously and effectively subjected to Logit Arithmetic reweighting? This concern arises because most of the ground truth in the datasets used in this paper consists of a single token.

2. In the "Related Work" section, should the authors discuss more relevant methods for addressing many-shot ICL, such as FocusICL、STRUCTICL? Additionally, would it be beneficial to include these methods as baselines for comparison?

3. ICL is a broadly applicable and effective paradigm; however, the benchmarks used in the experiments lack diversity. For instance, I am curious about the performance of the proposed method on mathematical reasoning problems, such as those found in MATH or GSM8K.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: In Figure 4, when number of examples used is 16, LARA without weight re-weighting seems performs better than the one with reweighting, could you provide more insights. And how could you verify the effectiveness of the re-weighting strategy if so.
soundness: 3
strengths: 1. LARA enhances ICL with a non-gradient optimization prompting strategy, and it shows performance gain in different settings.
2. The division strategy can handle the situation of longer context naturally.
3. The paper is clearly write and easy to follow.
summary: The paper proposed a new approach LARA to enhance in-context learning in LLM. The devised framework divides the demonstration examples into several subgroups, and ensemble the logit results from different groups based on Covariance Matrix Adaptive Evolution Strategy. They also proposes a binary version of weight selection, which is equivalent to subgroup pruning, and only selects more informative groups as the contributing demonstration. The evaluation experiments show the effectiveness of the proposed methods compared to original ICL and other variations.
weaknesses: 1. In weight optimization algorithm for B-LARA, it requires about 20 iterations to obtain the final weights, which may results in inefficiency. Although finally only a subset of groups is involved, the intermediate steps still consume a lot. And there is no ablation study on the impact of the iteration number.
2. The method is limited to the case that the vocabulary space of the target answer is already known.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: My main concerns have been listed above. I look forward to the authors' response and am willing to reopen and adjust the score upward.
soundness: 2
strengths: 1. Clear Presentation: The methodology is explained thoroughly, making LARA and B-LARA easy to understand.
2. Practical Relevance: LARA’s memory and computational efficiency make it well-suited for real-world applications, including black-box models.
3. Comprehensive Experiments: Extensive experiments across multiple benchmarks and ablation studies validate the effectiveness of the proposed approach.
summary: The paper presents the Logit Arithmetic Reweighting Approach (LARA) for enhancing in-context learning (ICL) in large language models (LLMs). LARA improves ICL performance by dividing long prompts into shorter, parallelizable inputs and using logit-based ensembling to optimize memory usage and accuracy. This method also includes a Binary LARA (B-LARA) variant, which constrains the weight of each subgroup to binary values, reducing memory needs and focusing on the most informative inputs. The experiments on the BBH and MMLU benchmarks demonstrate that LARA and B-LARA outperform baseline methods in both efficiency and accuracy across different models.
weaknesses: 1. While the empirical results are strong, the paper lacks a deeper theoretical analysis explaining why the logit reweighting improves performance, which could strengthen the foundational understanding of the approach.
2. The proposed method achieves good results on the MMLU and BBH datasets; however, it appears applicable only to multiple-choice questions or those with single-token answers, as it relies on the exact next-token logits for prediction, suggesting limited generalization capability.

Meta-review: This paper introduces the Logit Arithmetic Reweighting Approach (LARA), a novel framework that enhances ICL for LLMs by using logit-based ensembling of multiple demonstrations while reducing memory requirements through parallelization. Additionally, it presents Binary LARA (B-LARA), which simplifies the search space by constraining weights to binary values, and shows that both LARA and B-LARA outperform baseline methods in accuracy and memory efficiency in experiments on BBH and MMLU. Despite its contributions, the reviewers' feedbacks suggest that this paper still has to be improved on: the inadequacy of baseline comparisons, the latency introduced by weight optimization, potential output prefix inconsistencies during decoding, the lack of diverse benchmarks, inefficiencies in the weight optimization algorithm, and the limitation of requiring a known vocabulary space for target answers.

Decision: Reject