Language models improve answers by comparing good and bad prompts
Beam Search as Test-Time Self-Distillation via Counterfactual Contexts
Machine LearningArtificial Intelligence
Summary
Language models usually learn from training data with expert examples, but this paper finds a way to improve their answers during use without extra training. The authors show that by pretending the model is given either a great or poor example prompt and comparing the difference, they can guide the model to better answers. This method works during the answering process and uses a specialized search called beam search to find better responses. Tests on math, coding, and science problems show this approach improves performance over common decoding methods.
What this means in practice
- •For software developers: Generate higher-quality code completions by applying test-time self-distillation during generation.
- •For ai product engineers: Improve chatbot and QA system responses without retraining by using on-the-fly self-distillation with counterfactual prompts.
Authors
Su Ee Tan, Xiaotong Ji, Rasul Tutunov, Haitham Bou-Ammar, Matthieu Zimmer
Abstract
Self-Distillation Fine-Tuning (SDFT) enables a language model to act as its own teacher: by conditioning on a demonstration, the model produces an implicit reward via pointwise mutual information, which guides on-policy learning without external supervision. However, SDFT operates at training time: it requires gradient updates and access to expert demonstrations, making it inapplicable at inference. We propose test-time self-distillation, a decoding-time method that extracts a steering signal from the self-distillation framework without any parameter updates, reward models, or training data. Our key insight is that counterfactual contexts, i.e. fixed textual templates that hypothetically prime the model for excellent versus poor reasoning, can substitute for the demonstration. The log-odds ratio of a candidate answer under these two counterfactual conditions defines a new reward signal. We derive the optimal KL-regularized policy under this reward, which takes the form of a Gibbs reweighting of the base distribution. Crucially, this reweighting is global: it cannot be decomposed into independent per-token operations without ignoring future trajectory quality. We therefore approximate the target distribution via beam search. Experiments on mathematical reasoning (MATH500), code generation (HumanEval), and graduate-level science QA (GPQA) across multiple model scales show that test-time self-distillation improves over standard sampling, low temperature, beam search and power sampling baselines on average, demonstrating that the self-distillation principle can be operationalized at inference time.