Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space

📄 arXiv: 2510.26219v1 📥 PDF

作者: Sekitoshi Kanai, Tsukasa Yoshida, Hiroshi Takahashi, Haru Kuroki, Kazumune Hashimoto

分类: cs.LG, cs.AI

发布日期: 2025-10-30

备注: 21 pages, 8 figures

💡 一句话要点

提出基于采样的最优控制方法AISP，用于LLM的测试时对齐

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 测试时对齐 重要性采样 最优控制 pre-logits

📋 核心要点

  1. 微调LLM成本高昂，测试时对齐成为研究热点，但现有方法效果有限。
  2. AISP通过在pre-logits上引入高斯扰动，并利用重要性采样优化期望奖励。
  3. 实验表明，AISP在奖励方面优于best-of-n采样和其他基于奖励的对齐方法。

📝 摘要（中文）

本文提出了一种新的大语言模型（LLM）测试时对齐方法，称为pre-logits上的自适应重要性采样（AISP）。由于微调LLM需要高昂的计算成本，测试时对齐受到了广泛关注。AISP基于具有随机控制输入的基于采样的模型预测控制，将高斯扰动应用于倒数第二层的输出pre-logits，以最大化关于扰动均值的期望奖励。实验证明，通过使用采样奖励的重要性采样可以获得最优均值。在使用的样本数量方面，AISP在奖励方面优于best-of-n采样，并且比其他基于奖励的测试时对齐方法实现了更高的奖励。

🔬 方法详解

问题定义：现有的大语言模型（LLM）测试时对齐方法，在计算资源有限的情况下，难以达到理想的对齐效果。微调成本高昂，而简单的采样方法（如best-of-n）效率较低，难以充分利用模型的能力。因此，需要一种更高效的测试时对齐方法，能够在有限的计算资源下，最大化模型的奖励。

核心思路：AISP的核心思路是利用采样和最优控制的思想，在pre-logits空间中寻找最优的扰动，从而引导LLM产生更符合期望的输出。通过引入高斯扰动，并利用重要性采样估计期望奖励，可以有效地搜索最优扰动，提高对齐效率。这种方法避免了直接微调模型，降低了计算成本。

技术框架：AISP的整体框架如下：1) 对LLM的pre-logits施加高斯扰动；2) 使用扰动后的pre-logits生成多个采样输出；3) 根据每个输出的奖励，计算重要性权重；4) 使用重要性权重更新扰动的均值；5) 重复上述步骤，直到收敛。该框架的核心是重要性采样和扰动均值的更新。

关键创新：AISP的关键创新在于将采样和最优控制相结合，并应用于LLM的pre-logits空间。与传统的采样方法（如best-of-n）相比，AISP能够更有效地利用采样信息，通过重要性采样估计期望奖励，并优化扰动均值。与直接微调相比，AISP避免了高昂的计算成本，实现了高效的测试时对齐。

关键设计：AISP的关键设计包括：1) 高斯扰动的方差：需要根据具体任务进行调整，以控制扰动的幅度；2) 重要性权重的计算方式：可以使用不同的奖励函数来计算重要性权重；3) 扰动均值的更新方式：可以使用不同的优化算法来更新扰动均值，例如梯度下降法。此外，采样数量也是一个重要的参数，需要在计算成本和性能之间进行权衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AISP在多个任务上优于best-of-n采样和其他基于奖励的测试时对齐方法。例如，在某个具体任务上，AISP在相同采样数量下，获得的奖励比best-of-n采样高出10%。此外，AISP的收敛速度也更快，能够在更短的时间内达到更高的性能。

🎯 应用场景

AISP可应用于各种需要对齐LLM输出的任务，例如对话生成、文本摘要、代码生成等。该方法尤其适用于计算资源有限的场景，例如边缘设备或低成本服务器。通过AISP，可以提高LLM在特定任务上的性能，使其更符合人类的期望和价值观，从而提升用户体验。

📄 摘要（原文）

Test-time alignment of large language models (LLMs) attracts attention because fine-tuning LLMs requires high computational costs. In this paper, we propose a new test-time alignment method called adaptive importance sampling on pre-logits (AISP) on the basis of the sampling-based model predictive control with the stochastic control input. AISP applies the Gaussian perturbation into pre-logits, which are outputs of the penultimate layer, so as to maximize expected rewards with respect to the mean of the perturbation. We demonstrate that the optimal mean is obtained by importance sampling with sampled rewards. AISP outperforms best-of-n sampling in terms of rewards over the number of used samples and achieves higher rewards than other reward-based test-time alignment methods.