Technical Summary: Self-Improving In-Context Learning

Problem Statement

In-context learning (ICL) allows large language models (LLMs) to perform new tasks by conditioning on a few-shot prompt without updating model parameters. However, this capability is notoriously fragile; performance is highly sensitive to the choice, formatting, and ordering of demonstrations. Existing test-time improvement methods generally fall into three categories: demonstration selection, demonstration ordering, and output calibration. These approaches share structural limitations: they often operate on discrete prompt-level decisions, require a predefined finite label set (restricting them to classification tasks), or rely on external data and model internals. Consequently, there is a lack of general-purpose, test-time calibration methods that can improve ICL for both classification and free-form generation tasks without requiring finetuning or external resources.

Methodology

The proposed method, Self-Improving In-Context Learning, addresses these limitations by optimizing the continuous embeddings of a fixed few-shot prompt at test time. The core strategy involves treating the prompt as a black box and iteratively adjusting its embedding matrix to maximize a self-supervised confidence proxy derived entirely from the model's own predictions.

1. The Optimization Signal (Confidence Proxy)

The method relies on the observation that the log-probabilities a model assigns to its demonstrated outputs (obtainable via a single teacher-forced forward pass without token generation) provide a meaningful signal of task inference quality. This signal is formalized as a bounded, scalar proxy $f(X)$, where $X$ is the prompt embedding matrix. The proxy is a weighted linear combination of three components:

• Absolute Predictive Confidence ($\bar{C}$): A length-normalized average of the geometric mean of token probabilities across all demonstrated output spans. This ensures comparability across labels of different lengths.
• Pooled Robustness ($R$): A measure of tail fragility. Instead of relying solely on the mean, this component calculates a fixed quantile (e.g., 10th percentile) of all true-token probabilities across all demonstrations. This penalizes cases where a label appears confident on average but contains low-probability outlier tokens.
• Information Gain ($G$): A measure of progressive improvement. It tracks whether the model's confidence increases as it processes later demonstrations, reflecting the sequential nature of autoregressive task inference.

The final proxy is defined as $f(X) = \alpha\bar{C}(X) + \beta R(X) + \gamma G(X)$, with weights summing to 1.

2. Zeroth-Order Optimization

Since the proxy is non-smooth and backpropagation through the model is prohibitively expensive at test time, the method employs zeroth-order optimization (specifically, Gaussian smoothing) to estimate the gradient of the proxy with respect to the input embeddings.

• Process: The algorithm estimates the gradient by evaluating the proxy on the current embeddings and on embeddings perturbed by random noise sampled from a Gaussian distribution.
• Constraints: To prevent embeddings from drifting into out-of-distribution regions, the method employs:
  • Gradient Clipping: To bound large updates.
  • Cosine Similarity Constraint: To ensure updated embeddings remain directionally close to the original.
  • Initial Proxy Gate: Optimization is skipped if the initial proxy value is too low, as the gradient signal is exponentially flat near zero.
• Scope: Perturbations are restricted to the demonstration region of the prompt; query tokens are not perturbed, as their log-probabilities are not part of the proxy calculation under causal attention.

3. Test-Time Calibration

The procedure is entirely plug-and-play. It requires no access to model parameters, no external data, no predefined label set, and no token generation during the optimization phase. It operates solely on the input embeddings and the model's log-probabilities.

Key Contributions

• Embedding-Space Intervention: The paper proposes a fundamentally different strategy from discrete prompt engineering by intervening directly on the continuous prompt embedding matrix, reshaping the output distribution while keeping the human-readable input intact.
• Task-Agnostic Proxy: The proposed confidence proxy is applicable to both classification and free-form generation tasks, overcoming the structural limitation of existing baselines that require finite label sets.
• Self-Supervised Signal: The method derives its optimization signal solely from the model's own predictions on the provided demonstrations, requiring no ground-truth labels for the query or external oracle.
• Zeroth-Order Efficiency: By using zeroth-order optimization, the method avoids the computational cost of backpropagation and the memory overhead of storing gradients, making it feasible for test-time application on large models.

Experimental Results

The method was evaluated on ICLEval, a benchmark designed to isolate ICL ability from pretraining knowledge by using hash strings and rule-learning tasks. The evaluation covered 12 tasks (exact copying and rule learning) across three models: Llama 3.1-8B, Qwen3-4B, and Gemma 2-2B.

• Performance: The proposed method consistently matched or improved upon the base model across all tasks and models. It never caused degradation.
  • Llama 3.1-8B: Mean accuracy improved from 0.67 to 0.71 (+6.0%).
  • Qwen3-4B: Mean accuracy improved from 0.65 to 0.72 (+10.8%).
  • Gemma 2-2B: Mean accuracy improved from 0.45 to 0.48 (+6.7%).
• Statistical Significance: One-sided McNemar tests confirmed statistically significant improvements for all three models ($p < 0.05$).
• Baseline Comparison: On the classification subset, the method outperformed or matched classification-specific baselines (Contextual Calibration, DC-PMI, DEmO) on most tasks, despite operating without a predefined label space.
• Proxy-Performance Correlation: A statistically significant positive correlation was found between the improvement in the proxy value and the downstream accuracy gain across tasks, confirming that the proxy encodes a reliable optimization signal.

Significance and Claims

The paper claims that the proposed method demonstrates that demonstration log-probabilities constitute a reliable optimization surface for in-context learning. The ascent direction of this surface aligns with downstream performance, suggesting that the model's internal confidence on its own examples is a sufficient signal for self-improvement.

The authors emphasize that this approach is orthogonal to existing families of ICL improvement (selection, ordering, calibration) and can be composed with them. The method's significance lies in its ability to improve ICL reliability without:

1. Finetuning model parameters.
2. Generating tokens during the calibration phase.
3. Requiring a predefined label set (enabling application to free-form generation).
4. Accessing external data or model internals beyond the embedding layer.

The work suggests that test-time embedding calibration is a viable, general-purpose mechanism for enhancing the robustness of in-context learning across diverse task types and model scales.