Beam Search as Test-Time Self-Distillation via Counterfactual Contexts

Su Ee Tan; Xiaotong Ji; Rasul Tutunov; Haitham Bou-Ammar; Matthieu Zimmer

cs.LGSep 29, 2026

Beam Search as Test-Time Self-Distillation via Counterfactual Contexts

Authors: Su Ee Tan, Xiaotong Ji, Rasul Tutunov, Haitham Bou-Ammar, Matthieu Zimmer

Organizations: Huawei Noah’s Ark Lab · UCL Centre for AI

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.

Figures & tables

Model	Method	MATH500	HumanEval	GPQA
Qwen3.5-0.8B	Base $\tau=1.0$	0.198	0.085	0.020
	Low Temperature $\tau=0.25$	0.342	0.225	0.101
	Beam Search	0.398	0.207	0.091
	Power Sampling $\tau=0.25$	0.430	0.287	0.131
	Power Sampling $\tau=1.0$	0.300	0.262	0.060
	CBS $1/\alpha=0.25$	0.480	0.238	0.106

Table 1: Performance comparison of CBS across 4 models and 3 benchmarks.

Figure 1: Inference efficiency of (a) Qwen2.5-7B and (b) DeepSeek-Math-7B-Instruct on the MATH500 benchmark.

Figure 2: Pass@

k

on the MATH500 dataset between Power sampling and CBS (Qwen3.5-2B).

Context	Content	MATH500
No context	-	0.398
Contrastive	Positive: ‘‘This is an example for a response with excellent reasoning:’’ Negative: ‘‘This is an example for a response with wrong reasoning:’’	0.509
Neutral	Positive: ‘‘This is an example for a response to the question assigned to control group one:’’ Negative: ‘‘This is an example for a response to the question assigned to control group two:’’	0.380

Table 2: Context ablation for Qwen3.5-0.8B on MATH500.

Method	Qwen3.5-0.8B	Qwen2.5-7B	Qwen3.5-2B	DeepSeek-Math-7B-Instruct
Beam search ( $N=24$ )	0.391	0.600	0.676	0.452
CBS $1/\alpha=0.25$	0.480	0.610	0.672	0.462
CBS $1/\alpha=0.7$	0.436	0.640	0.694	0.458

Table 3: Beam-population budget ablation on MATH500.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Prompt Template
MATH500	System: You are a helpful AI Assistant that provides well-reasoned and detailed responses. You first think about the reasoning process as an internal monologue and then provide the user with the boxed answer. Respond in the following format: <think> … </think> <answer> \boxed{…} </answer>. User: {problem}
HumanEval	System: Please reason step by step internally. Then output ONLY the final Python code that completes the task within triple backticks. Do not include explanations, markdown, or \boxed{}. User: Complete the following Python function: {prompt}
GPQA	System: You are a helpful AI Assistant. Please reason step by step, and put your final answer within \boxed{}. User: Answer the following multiple choice question. The last line of your response should be of the following format: ‘\boxed{$LETTER}’ (without quotes) where LETTER is one of ABCD (ex. ‘\boxed{A}’). Think step by step before answering. {question} A) {choice_A} B) {choice_B} C) {choice_C} D) {choice_D}

Appendix

Table 4: Task Prompt Templates

Role	Serialized Qwen prompt content
System	<\|im_start\|>system You are a helpful AI Assistant that provides well-reasoned and detailed responses. You first think about the reasoning process as an internal monologue and then provide the user with the boxed answer. Respond in the following format: <think> … </think> <answer> \boxed{…} </answer>. <\|im_end\|>
User query and positive context	<\|im_start\|>user Solve for $x$ : $2^{x+1}=32$ . This is an example for a response with excellent reasoning: <\|im_end\|>
Assistant prefix	<\|im_start\|>assistant

Appendix

Table 5: Example serialization of a MATH500-style query with the positive contrastive context under the Qwen chat template.

Method	Hyperparameters
Standard sampling	$N=1$ , $W=1$ , $K=32$ , $\tau=1.0$
Low-temperature sampling	$N=1$ , $W=1$ , $K=32$ , $\tau=0.25$
Beam search	$N=16$ , $W=4$ , $K=32$ , $\tau=1.0$ , $1/\alpha=0$
Contrastive beam search (low $\alpha$ )	$N=16$ , $W=4$ , $K=32$ , $\tau=1.0$ , $1/\alpha=0.25$
Contrastive beam search (high $\alpha$ )	$N=16$ , $W=4$ , $K=32$ , $\tau=1.0$ , $1/\alpha=0.7$
Power Sampling	$K_{t}=16$ , $M_{t}=16$ , $B=32$ (equivalent to $K$ ), $\tau=0.25$ (low) and $\tau=1.0$ (high), $\alpha=4$

Appendix

Table 6: Main hyperparameter settings for each method.

Figure 4: Pass@

k

performance

k\in\{1,2,3,4\}

on MATH500 dataset for Power sampling (PS), standard beam search (BS) and contrastive beam search (CBS) using Deepseek-Math-7B-Instruct.

Panel of Figure 5 — Figure 5: Best-of-16 accuracy for ten contrastive context pairs on MATH500, HumanEval and GPQA (Qwen2.5-7B), at the two deployed steering strengths. Dashed line: no-context baseline.

Content	Contrastive Context Pair
Reasoning	Positive: ‘‘This is an example for a response with excellent reasoning:’’ Negative: ‘‘This is an example for a response with wrong reasoning:’’
Completeness	Positive: ‘‘This is an example for a response with a complete treatment of the problem:’’ Negative: ‘‘This is an example for a response with an incomplete treatment of the problem:’’
Step verification	Positive: ‘‘This is an example for a response that carefully verifies each step of its reasoning:’’ Negative: ‘‘This is an example for a response that skips verification and rushes to a conclusion:’’
Reliability	Positive: ‘‘This is an example for a response that is reliable and trustworthy:’’ Negative: ‘‘This is an example for a response that is unreliable and error-prone:’’
Logical validity	Positive: ‘‘This is an example for a response with logically sound arguments:’’ Negative: ‘‘This is an example for a response with logical fallacies:’’
Reviewer judgment	Positive: ‘‘This is an example for a response that a careful reviewer would approve without hesitation:’’ Negative: ‘‘This is an example for a response that a careful reviewer would reject immediately:’’

Appendix

Table 7: Examples of 10 contrastive context pairs. Every pair shares same initial framing with “This is an example for a response …”. Each positive context is paired with its counterfactual negative.

Explore similar work

Activation-Conditioned Self-Distillation
Sep 29, 2026Zhexi Lu, Subhajit Chaudhury, Tejaswini Pedapati +2Unsupervised On-Policy Self-Distillation Self-Distillation Framework
Rensselaer Polytechnic Institute · IBM Research
Localizing Credit at the Divergence: Path-Conditioned Self-Distillation for LLM Reasoning
Jun 14, 2026Yu Li, Shu Hong, Tian LanUnsupervised On-Policy Self-Distillation Online-Policy Distillation
Department of Electrical and Computer Engineering, George Washington University
Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information
May 12, 2026Guobin Shen, Xiang Cheng, Chenxiao Zhao +4Offline Reinforcement Learning Mutual Information
1Xiaohongshu Inc. · Institute of Automation, Chinese Academy of Sciences

Sep 29, 2026, cs.LGJ/K move · Enter open · S save

Activation-Conditioned Self-Distillation

Zhexi Lu, Subhajit Chaudhury, Tejaswini Pedapati, Keerthiram Murugesan+1

Rensselaer Polytechnic Institute · IBM Research

Unsupervised On-Policy Self-Distillation Self-Distillation Framework

Open paper AskPDF arXiv

On-policy self-distillation uses a model as its own teacher to provide dense supervision for reasoning, often through reference-solution conditioning. Providing privileged information does not by itself ensure effective token-level supervision throughout long responses. We introduce Activation-Conditioned Self-Distillation (ACSD), which extracts a steering vector by contrasting activations of self-generated trajectories that reach verified correct answers within a generation budget with those of all remaining trajectories. A frozen copy of the base model applies this vector at each prediction position, and the student learns from its next-token distributions on student-generated prefixes. Outcome verification is used for direction construction and calibration; distillation requires neither problem-specific reference text nor teacher parameter updates. The distilled student is used alone at inference. On each of five models, ACSD achieves the highest mean accuracy over four mathematical benchmarks among the evaluated methods. On DeepSeek-R1-0528-Qwen3-8B, mean mathematical accuracy reaches 71.9% and LiveCodeBench v6 pass@12 reaches 70.9%, compared with 69.0% and 66.3% for the reference-conditioned OPSD baseline. Contrasts among correct trajectories also support distillation, and extracted directions can be reused across mathematical training datasets. On fixed student trajectories, ACSD maintains more stable late-position logit-update magnitudes than OPSD.

Abstract

Explore similar work

Activation-Conditioned Self-Distillation

Localizing Credit at the Divergence: Path-Conditioned Self-Distillation for LLM Reasoning

Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information