Self-Improving In-Context Learning

2026-05-21

AI Summary

This work proposes a general test-time prompt calibration method to enhance in-context learning (ICL) by improving task comprehension and generalization without requiring fine-tuning, additional data, or predefined labels. The approach optimizes continuous embeddings of a fixed few-shot prompt by constructing a bounded, self-supervised confidence proxy based on the model’s output log-probabilities and efficiently adjusts them via zeroth-order optimization. Applicable to both classification and free-form generation tasks, the method consistently matches or surpasses existing baselines across diverse ICL settings, significantly outperforming approaches designed solely for classification. Moreover, improvements in the proposed proxy metric exhibit strong correlation with gains in downstream task accuracy.

Abstract

We propose to improve in-context learning (ICL) by optimizing the continuous embeddings of a fixed few-shot prompt at test time. The key observation is that the log-probabilities a model assigns to its demonstrated outputs$\unicode{x2013}$available from a single forward pass without generating any tokens$\unicode{x2013}$provide a meaningful signal for how well the model has inferred the task from its demonstrations. We formalize this signal as a bounded, self-supervised confidence proxy and maximize it via zeroth-order optimization over the prompt embeddings, yielding a test-time calibration procedure. The approach requires no finetuning, no token generation, no predefined label set, and no external data, making it equally applicable to both classification and free-form generation tasks. Across a comprehensive suite of ICL tasks, the proposed calibration consistently matches or improves upon the base model and outperforms classification-specific baselines on most tasks. The statistically significant correlation between proxy improvement and downstream accuracy gain confirms that the proposed proxy encodes a reliable optimization signal for in-context learning.

Problem

• in-context learning
• prompt optimization
• self-supervised signal
• zero-shot calibration
• language models

Innovation

• in-context learning
• prompt embedding optimization
• self-supervised confidence proxy
• zeroth-order optimization
• test-time calibration