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

arXiv: https://arxiv.org/abs/2609.37041 · HTML: https://mlsanity.com/paper/2609.37041/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in On-policy and self-distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 29 of its 34 references that are in the map.
- First posted on arXiv 29 Sep 2026.

## Placed in

- Knowledge distillation / [On-policy and self-distillation for LLMs](https://mlsanity.com/distill/on-policy-and-self-distillation-for-llms/)

## Builds on (most cited first, 29 of 29)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Contrastive Decoding: Open-ended Text Generation as Optimization](https://mlsanity.com/paper/2210.15097/) (2022)
- [DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models](https://mlsanity.com/paper/2309.03883/) (2023)
- [Reinforcement Learning via Self-Distillation](https://mlsanity.com/paper/2601.20802/) (2026)
- [Trusting Your Evidence: Hallucinate Less with Context-aware Decoding](https://mlsanity.com/paper/2305.14739/) (2023)
- [Self-Distillation Enables Continual Learning](https://mlsanity.com/paper/2601.19897/) (2026)
- [Deep Think with Confidence](https://mlsanity.com/paper/2508.15260/) (2025)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [RL with KL penalties is better viewed as Bayesian inference](https://mlsanity.com/paper/2205.11275/) (2022)
- [Confidence Improves Self-Consistency in LLMs](https://mlsanity.com/paper/2502.06233/) (2025)
- [ARGS: Alignment as Reward-Guided Search](https://mlsanity.com/paper/2402.01694/) (2024)
- [Reasoning with Sampling: Your Base Model is Smarter Than You Think](https://mlsanity.com/paper/2510.14901/) (2025)
- [TreeBoN: Enhancing Inference-Time Alignment with Speculative Tree-Search and Best-of-N Sampling](https://mlsanity.com/paper/2410.16033/) (2024)
- [Scalable Power Sampling: Unlocking Efficient, Training-Free Reasoning for LLMs via Distribution Sharpening](https://mlsanity.com/paper/2601.21590/) (2026)
- [Improving Factuality in Large Language Models via Decoding-Time Hallucinatory and Truthful Comparators](https://mlsanity.com/paper/2408.12325/) (2024)
- [On Almost Surely Safe Alignment of Large Language Models at Inference-Time](https://mlsanity.com/paper/2502.01208/) (2025)
- [Distilling Morphology-Conditioned Hypernetworks for Efficient Universal Morphology Control](https://mlsanity.com/paper/2402.06570/) (2024)
- [Explaining and Improving Contrastive Decoding by Extrapolating the Probabilities of a Huge and Hypothetical LM](https://mlsanity.com/paper/2411.01610/) (2024)
- [Thinking by Subtraction: Confidence-Driven Contrastive Decoding for LLM Reasoning](https://mlsanity.com/paper/2602.18232/) (2026)
- [The Model Knows, the Decoder Finds: Future Value Guided Particle Power Sampling](https://mlsanity.com/paper/2605.02427/) (2026)

## Abstract (arXiv)

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.