Beam Search as Test-Time Self-Distillation via Counterfactual Contexts
作者: Su Ee Tan, Xiaotong Ji, Rasul Tutunov, Haitham Bou-Ammar, Matthieu Zimmer
分类: cs.LG, cs.AI
发布日期: 2026-09-29
备注: Accepted at NeurIPS 2026 Workshop on Towards Test-Time Continual Learning Agents
💡 一句话要点
提出测试时自蒸馏方法以提升推理性能
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 自蒸馏 推理优化 反事实学习 自然语言处理 机器学习
📋 核心要点
- 现有的自蒸馏微调方法在推理阶段无法使用，限制了其应用范围。
- 本文提出通过反事实上下文在推理时进行自蒸馏，利用固定模板替代示范来生成奖励信号。
- 实验结果显示，该方法在多个任务上均优于标准采样和束搜索等基线，提升了模型性能。
📝 摘要（中文）
自蒸馏微调（SDFT）使语言模型能够自我指导，但其在推理时无法应用。本文提出了一种测试时自蒸馏的方法，通过反事实上下文提取信号，无需参数更新或训练数据。我们利用固定文本模板来替代示范，从而定义新的奖励信号。通过KL正则化策略，我们实现了基于Gibbs重加权的全局策略。实验表明，该方法在数学推理、代码生成和科学问答等任务上优于传统采样和束搜索方法，证明了自蒸馏原则在推理时的有效性。
🔬 方法详解
问题定义：本文旨在解决自蒸馏微调（SDFT）在推理阶段无法应用的问题。现有方法依赖于梯度更新和专家示范，限制了其在实际推理中的使用。
核心思路：我们提出了一种新的测试时自蒸馏方法，通过反事实上下文生成奖励信号，避免了参数更新和外部监督的需求。反事实上下文是固定的文本模板，能够模拟模型在不同推理质量下的表现。
技术框架：整体流程包括使用反事实上下文生成奖励信号，随后通过KL正则化策略优化模型的推理策略。该方法的核心在于全局重加权，而非逐token操作。
关键创新：最重要的创新在于利用反事实上下文替代传统的示范，从而在推理阶段实现自蒸馏。这一方法的全局重加权特性使得模型能够考虑未来的推理轨迹质量。
关键设计：我们设计了基于Gibbs重加权的策略，确保在推理过程中能够有效利用反事实上下文生成的奖励信号。该方法不需要额外的训练数据或模型参数更新，简化了推理过程。
🖼️ 关键图片
📊 实验亮点
实验结果表明，测试时自蒸馏方法在数学推理（MATH500）、代码生成（HumanEval）和科学问答（GPQA）任务上均优于标准采样、低温采样和束搜索等基线，平均提升幅度显著，验证了该方法的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的推理任务，如数学推理、代码生成和科学问答等。通过提升模型在推理阶段的表现，能够为实际应用提供更高效的解决方案，未来可能影响教育、编程辅助和科学研究等多个领域。
📄 摘要（原文）
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.