2025-11-16T11:43:12.671286

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

Kanai, Yoshida, Takahashi et al.
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.
academic

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

基本信息

  • 论文ID: 2510.26219
  • 标题: Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space
  • 作者: Sekitoshi Kanai, Tsukasa Yoshida, Hiroshi Takahashi (NTT, Inc.), Haru Kuroki, Kazumune Hashimoto (The University of Osaka)
  • 分类: cs.LG cs.AI
  • 发表时间: 2025年10月30日 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2510.26219v1

摘要

大语言模型(LLM)的测试时对齐因避免高昂的微调成本而受到关注。本文提出了一种新的测试时对齐方法——预logit自适应重要性采样(AISP)，该方法基于带随机控制输入的采样模型预测控制。AISP对倒数第二层输出(pre-logits)施加高斯扰动，通过最大化扰动均值的期望奖励来实现对齐。论文证明了最优均值可通过对采样奖励进行重要性采样获得。AISP在样本使用效率上优于best-of-n采样，并在奖励值上超越其他基于奖励的测试时对齐方法。

研究背景与动机

要解决的问题

大语言模型的对齐是确保LLM安全和广泛应用的关键技术。传统的人类反馈强化学习(RLHF)方法需要微调LLM参数，带来巨大的计算成本。测试时对齐(test-time alignment)旨在不更新模型参数的情况下，使LLM生成符合人类偏好的响应。

问题的重要性

  1. 计算成本：微调大规模LLM需要大量GPU资源和训练时间
  2. 灵活性：测试时对齐允许在推理阶段动态调整模型行为
  3. 实用性：无需为每个特定任务重新训练模型

现有方法的局限性

  1. Best-of-N (BoN)采样：虽然简单有效，但不主动探索最优响应，样本效率较低
  2. RE-Control：需要训练值函数，需要大量数据集(如349,000个训练样本)和存储成本
  3. 传统最优控制：对于非线性、大规模的LLM系统不适用

研究动机

能否通过无需训练的方法来控制LLM探索最优响应？本文从控制理论角度出发，采用采样型模型预测控制(MPPI)技术，提出了一种无需训练的测试时对齐方法。

核心贡献

  1. 提出AISP方法：首次将采样型模型预测控制(MPPI)应用于LLM对齐，通过在pre-logit空间施加高斯扰动来实现无需训练的测试时对齐
  2. 理论贡献：
    • 证明了最优pre-logit分布可通过自由能(free energy)边界获得
    • 推导了基于自适应重要性采样的闭式解
    • 揭示了AISP与BoN的理论联系(在特定参数下AISP退化为BoN)
  3. 高斯假设的合理性分析：论证了pre-logit的高斯分布假设与神经网络softmax层的内在联系
  4. 性能提升：
    • 在样本效率上显著优于BoN(相同样本数下获得更高奖励)
    • 无需训练即超越RE-Control
    • 提出Batched AISP实现并行加速

方法详解

任务定义

给定输入提示 x=[x1,...,xTx]x = [x_1, ..., x_{T_x}]，LLM生成响应 y=[y1,...,yTy]y = [y_1, ..., y_{T_y}]。目标是在给定奖励模型 r(x,y)r(x,y) 的情况下，最大化期望奖励同时保持与基础LLM的KL散度约束：

min⁡UJ(x,U)=−EV∼QU,σ2[r(x,y(V))]+λDKL(QU,σ2∣P)\min_U J(x, U) = -\mathbb{E}_{V\sim Q_{U,\sigma^2}}[r(x,y(V))] + \lambda D_{KL}(Q_{U,\sigma^2}|P)

模型架构

1. 随机控制输入设计

与RE-Control使用确定性控制输入不同，AISP使用随机控制输入 vt∼N(ut,σ2I)v_t \sim \mathcal{N}(u_t, \sigma^2I)：

\text{softmax}(W_{LLM}(z_t + v_t) + b_{LLM}), & v_t \sim \mathcal{N}(u_t, \sigma^2I), \text{ for } 1 \leq t \leq \tau \\ \text{softmax}(W_{LLM}z_t + b_{LLM}), & \text{for } \tau < t \end{cases}$$ 其中： - $z_t = \phi_{LLM}(y_{<t})$ 是pre-logit(倒数第二层输出) - $u_t$ 是待优化的扰动均值 - $\sigma^2I$ 是固定的协方差矩阵 - $\tau$ 是控制时间窗口 #### 2. 输入轨迹分布 输入轨迹 $V = [v_1, ..., v_\tau]$ 服从联合高斯分布： $$q(V|U, \sigma^2) = \frac{1}{(2\pi\sigma^2)^{d\tau/2}} \exp\left(-\frac{1}{2\sigma^2}\sum_{t=1}^\tau (v_t - u_t)^\top(v_t - u_t)\right)$$ 基础分布为零均值高斯：$p(V|0, \sigma^2)$ #### 3. 最优分布推导 通过自由能： $$F(r, p, x, \lambda) = \log\left(\mathbb{E}_{V\sim P}\left[\exp\left(\frac{1}{\lambda}r(x,y(V))\right)\right]\right)$$ **定理3.1** 证明了最优密度函数为： $$q^*(V) = \frac{1}{\eta}\exp\left(\frac{1}{\lambda}r(x,y(V))\right)p(V)$$ 其中 $\eta$ 是归一化常数。 #### 4. 自适应重要性采样 由于最优分布难以直接计算，使用重要性采样近似。 **定理3.2** 证明了最优均值为： $$u_t^* = \mathbb{E}_{V\sim Q^*}[v_t] = \mathbb{E}_{V\sim Q_{\hat{U},\sigma^2}}[w(V)v_t]$$ 权重函数为： $$\tilde{w}^i = \frac{\exp\left(\frac{1}{\lambda}r(x,y(V^i)) - \frac{1-\alpha}{\sigma^2}\sum_{t=1}^\tau \hat{u}_t^\top v_t^i\right)}{\sum_j \exp\left(\frac{1}{\lambda}r(x,y(V^j)) - \frac{1-\alpha}{\sigma^2}\sum_{t=1}^\tau \hat{u}_t^\top v_t^j\right)}$$ 其中引入了松弛参数 $\alpha \in (0,1)$ 以增强数值稳定性。 #### 5. 迭代更新 通过 $\kappa$ 次迭代，每次生成 $n$ 个样本： $$\hat{u}_t^{k+1} = \sum_{i=1}^n \tilde{w}^i v_t^{i,k}, \quad v_t^{i,k} \sim \mathcal{N}(\hat{u}_t^k, \sigma^2I)$$ 最终选择所有样本中奖励最高的响应。 ### 技术创新点 #### 1. Pre-logit空间vs Token空间 - **优势**：pre-logit分布可用闭式高斯分布表示，而token序列分布难以建模 - **可计算性**：权重函数易于计算，无需复杂的归一化流等技术 #### 2. 高斯假设的合理性 论文从理论上分析了高斯假设与softmax层的联系： 如果 $p(z_t|y_t=y_i) = \mathcal{N}(\mu_{y_i}, \Sigma)$，则由贝叶斯定理： $$P(y_t=y_i|z_t) = \frac{\exp(\mu_{y_i}^\top\Sigma^{-1}z - \frac{1}{2}\mu_{y_i}\Sigma^{-1}\mu_{y_i} + \ln P(y_t=y_i))}{\sum_j \exp(\mu_{y_j}^\top\Sigma^{-1}z - \frac{1}{2}\mu_{y_j}\Sigma^{-1}\mu_{y_j} + \ln P(y_t=y_j))}$$ 这正对应于softmax函数形式，说明高斯假设与神经语言模型的隐含假设一致。 #### 3. 与BoN的理论联系 **定理3.3** 证明：当 $\lambda \to 0^+$ 且 $\kappa=1$ 时，AISP退化为BoN。 这表明AISP是BoN的连续近似和泛化，提供了更灵活的优化框架。 #### 4. 固定控制窗口 与MPPI的滑动窗口不同，AISP使用固定窗口 $t \in [1, \tau]$，避免固定前缀token导致的多样性损失。 ## 实验设置 ### 数据集 1. **Anthropic HH-RLHF**：用于对齐LLM的有用性和无害性 2. **Stanford Human Preferences (SHP)**：人类偏好数据集 3. **规模**：从测试集中随机选择1000个样本(受计算资源限制) ### 基础模型 - **LLMs**：Llama-3-8B, Vicuna-7B-v1.5, Gemma3-4B - **奖励模型**：UltraRM-13b, Eurus-RM-7b ### 评价指标 1. **奖励值**：使用UltraRM评估 $r(x,y)$ 2. **多样性(Diversity)**：$\sum_{n=2}^4 \frac{\text{unique n-gram}(y)}{\text{total n-gram}(y)}$，评估响应中的重复程度 3. **连贯性(Coherence)**：使用SimCSE计算提示和响应嵌入的余弦相似度 4. **胜率(Win Rate)**：使用GPT-4评估AISP相对于BoN的胜率 ### 对比方法 1. **BoN (top-p)**：使用nucleus采样的Best-of-N，N=1024 (= κn) 2. **RE-Control**：基于训练值函数的控制方法 3. **ARGS-greedy**：在logit上添加加权奖励的方法 ### 实现细节 - **AISP参数**：$n=32$, $\kappa=32$，总样本数1024 - **超参数调优**：在10个训练样本上进行网格搜索 - $\lambda \in [0.1, 0.3, 0.5, 0.7]$ (UltraRM), $[60, 120, 240, 480]$ (Eurus) - $\sigma^2 \in [0.1, 0.3, 0.5, 0.7]$ - $\alpha \in [0.99, 0.999, 0.9999, 0.99999]$ - **生成设置**：最大新token长度128，半精度(bfloat16) - **硬件**：NVIDIA A100 (40GB) 和 H100 (80GB) ## 实验结果 ### 主要结果 #### 平均奖励对比(表1) 在6种模型-奖励模型组合、2个数据集上的结果显示： **SHP数据集**： - **Llama3 & UltraRM**：AISP (-1.39) vs BoN (-2.38)，提升**41.6%** - **Vicuna & UltraRM**：AISP (-1.46) vs BoN (-1.78)，提升18.0% - **Gemma3 & UltraRM**：AISP (-2.39) vs BoN (-3.43)，提升30.3% **HH-RLHF数据集**： - **Llama3 & UltraRM**：AISP (-5.02) vs BoN (-5.074)，提升1.1% - **Vicuna & UltraRM**：AISP (-4.73) vs BoN (-4.85)，提升2.5% **关键发现**： - AISP在所有设置下均达到或超过BoN的平均奖励 - 相比需要训练的RE-Control，AISP在多数情况下表现更优(如Llama3 & UltraRM: -1.39 vs -9.28) - ARGS在本实验中表现不佳，可能因为轨迹级奖励模型不适合token级评估 #### 胜率分析(表2) 使用GPT-4评估100对样本： **SHP数据集**： - Llama & UltraRM：AISP 51.3% vs BoN 42.0% - Gemma3 & UltraRM：AISP 53.0% vs BoN 41.3% - 平均胜率显著高于BoN **HH-RLHF数据集**： - 结果更均衡，但AISP在多数设置下仍保持优势 - 部分设置(如Vicuna)显示较高的平局率(27.7%-36.0%) ### 样本效率分析(图3) **收敛曲线**显示AISP的关键优势： - **初期**：BoN在前几次迭代中表现更好(因为直接采样多样性高) - **中期**：AISP迅速赶上，约在k=10-15次迭代后超越BoN - **后期**：AISP持续改进，最终显著优于BoN **三条曲线分析**： 1. **AISP (Mean at k)**：$\frac{1}{n}\sum_i r(x,y(V^{i,k}))$，随迭代稳步上升 2. **AISP (Best at k)**：$\max_i r(x,y(V^{i,k}))$，单次迭代最佳 3. **AISP (Best so far)**：$\max_{i,1\leq j\leq k} r(x,y(V^{i,j}))$，全局最佳 **重要洞察**：AISP不仅优化单个响应，还优化响应分布，Mean曲线的上升证明了分布优化的有效性。 ### Batched AISP实验(图4) 在相同迭代次数下比较(BoN N=128 vs AISP κ=b, n=N/b)： **设置对比**： - AISP1: (b=8, n=16) - AISP2: (b=16, n=8) - AISP3: (b=32, n=4) - AISP4: (b=64, n=2) **结果**： - 所有AISP设置均优于BoN(-4.2至-4.4 vs BoN约-4.7) - 只要每次迭代至少4个样本，AISP就能超越BoN - 证明了AISP在时间约束下的实用性 ### KL散度分析(表3) **不同超参数下的KL散度**： - AISP (λ=0.1, α=0.9999)：KL=140.9, Reward=-2.15 - AISP (λ=10.0, α=0.99)：KL=2.98, Reward=-3.37 - RE-Control：KL=0.172, Reward=-9.30 - ARGS：KL=78.8, Reward=-5.11 **关键发现**： - 通过调节λ和α，AISP可灵活控制与基础LLM的偏离程度 - 即使KL散度小于ARGS(18.9 vs 78.8)，AISP仍获得更高奖励(-2.75 vs -5.11) - 证明了AISP在奖励提升和保持基础LLM特性之间的良好权衡 ### 消融实验 #### 超参数敏感性(附录D.1，图6-7) **λ的影响**： - 小λ(0.1)：均值不增长，优化失效 - 大λ(0.7)：均值增长率提高，但需保持数值稳定性 - 最终奖励在λ∈[0.1, 0.7]范围内均优于BoN **σ的影响**： - 小σ(0.1)：探索空间受限，奖励早期饱和 - 大σ(0.7)：探索充分但略有不稳定 - 最佳值约为σ=0.5 **α的影响**： - 小α(0.5-0.8)：过度惩罚偏离，奖励改进受限 - 大α(0.999-0.9999)：允许足够探索，奖励稳步提升 **总体评价**：超参数行为符合直觉，调优相对容易 ### 实验发现 1. **样本效率**：AISP在相同样本数下获得更高奖励，在迭代过程中展现更快的改进速度 2. **无需训练优势**：无需预先收集数据集或训练值函数，即可超越RE-Control 3. **分布优化**：不仅优化单个响应，还优化整体响应分布 4. **灵活性**：通过超参数可控制奖励提升与基础LLM保真度的权衡 5. **并行化潜力**：Batched AISP在时间约束下仍能保持性能优势 6. **跨模型泛化**：在多种LLM(Llama3, Vicuna, Gemma3)和奖励模型上均有效 ## 相关工作 ### 测试时对齐方法分类 #### 1. 基于训练的方法 - **RE-Control** (Kong et al., 2024)：训练值函数优化pre-logit - **Critic-Guide Decoding** (Kim et al., 2023)：训练批评网络预测状态值 - **Controlled Decoding** (Mudgal et al., 2024)：训练值函数进行块级生成 - **局限性**：需要大规模数据集(如RE-Control使用349,000样本)和训练成本 #### 2. 采样型方法 - **Best-of-N (BoN)**：简单有效，但样本效率低 - Yang et al. (2024)证明BoN渐近优化KL约束RL目标 - Beirami et al. (2024)证明BoN胜率上界为N/(N+1) - **Soft Reasoning** (Zhu et al., 2025)：基于贝叶斯优化，但仅扰动初始token嵌入 - **Importance Sampling方法** (Loula et al., 2025)：在token空间使用重要性采样，需任务特定势函数 #### 3. Logit操作方法 - **ARGS** (Khanov et al., 2024)：在logit上添加加权奖励 - **局限性**：需要token级奖励模型 ### 本文优势 1. **vs BoN**：主动探索最优响应，样本效率更高 2. **vs RE-Control**：无需训练，避免数据收集和训练成本 3. **vs Soft Reasoning**：优化完整pre-logit序列，而非仅初始嵌入 4. **vs Loula et al.**：在pre-logit空间使用易处理的高斯分布 ### 理论基础 **控制理论视角**： - 传统最优控制(如Pontryagin最大值原理)不适用于非线性大规模LLM - **MPPI** (Williams et al., 2017, 2018)：采样型模型预测控制，利用GPU并行计算 - AISP将MPPI应用于LLM对齐，引入自适应重要性采样 ## 结论与讨论 ### 主要结论 1. **方法有效性**：AISP作为无需训练的测试时对齐方法，在奖励优化上显著优于BoN和RE-Control 2. **理论贡献**：建立了pre-logit空间随机控制框架，证明了最优分布可通过自适应重要性采样近似 3. **样本效率**：AISP在样本使用效率上优于BoN，在相同样本数下获得更高奖励 4. **实用性**：Batched AISP在时间约束下仍能保持性能，适合实际应用 5. **可控性**：通过超参数可灵活调节奖励提升与基础LLM保真度的权衡 ### 局限性 #### 1. 计算复杂度 - **顺序迭代**：需要κ次顺序迭代，时间复杂度为O(κ) - **额外计算**：权重函数需要计算$\sum_{t=1}^\tau \hat{u}_t^\top v_t^i$，虽然开销为O(τd)相对可忽略 #### 2. 高斯假设 - **假设限制**：pre-logit分布的高斯假设可能不完全准确 - **简化代价**：为获得可处理的闭式解而做的简化 #### 3. 超参数调优 - **三个超参数**：λ, σ², α需要调优 - **数据集依赖**：不同奖励模型(UltraRM vs Eurus)需要不同的λ范围 #### 4. 实验规模 - **样本限制**：由于计算资源，仅使用1000个测试样本 - **模型规模**：主要在7B-13B规模模型上测试，更大模型的表现未知 #### 5. 多样性和连贯性 - 在某些设置下，AISP的多样性和连贯性不如BoN - 可能因为奖励模型未优先考虑这些维度 ### 未来方向 1. **结合微调**：探索AISP与参数高效微调(如LoRA)的结合 2. **不同采样技术**：研究其他重要性采样变体(如sequential Monte Carlo) 3. **更复杂的分布**：使用归一化流等技术建模更复杂的pre-logit分布 4. **多目标优化**：同时优化奖励、多样性和连贯性 5. **更大规模模型**：在更大规模LLM(如70B+)上验证方法 6. **理论分析**：提供收敛速率和样本复杂度的理论保证 ## 深度评价 ### 优点 #### 1. 创新性 - **跨学科融合**：首次将MPPI控制理论应用于LLM对齐，开辟新研究方向 - **Pre-logit空间**：在pre-logit而非token空间操作，利用高斯分布的可处理性 - **理论完备**：提供完整的理论推导(定理3.1-3.3)和闭式解 #### 2. 实用性 - **无需训练**：相比RE-Control节省大量数据收集和训练成本 - **即插即用**：可直接应用于预训练LLM，无需修改模型结构 - **Batched版本**：提供并行化方案，适应实际部署需求 #### 3. 实验充分性 - **多维度评估**：奖励、多样性、连贯性、胜率、KL散度 - **多种设置**：3种LLM × 2种奖励模型 × 2个数据集 = 12种组合 - **消融实验**：详细的超参数敏感性分析(附录) - **收敛分析**：展示样本效率优势的动态过程 #### 4. 理论洞察 - **高斯假设合理性**：从softmax层推导pre-logit高斯分布的合理性 - **与BoN的联系**：证明AISP是BoN的泛化，提供统一框架 - **自由能边界**：利用变分推断思想，建立优雅的理论框架 #### 5. 写作质量 - 结构清晰，从问题定义到理论推导再到实验验证层次分明 - 提供详细的算法伪代码(Algorithm 1)和实现细节 - 附录包含完整证明和额外实验 ### 不足 #### 1. 方法局限 - **计算开销**：虽然无需训练，但推理时需要κn次前向传播，对于κ=32, n=32，共1024次前向传播 - **顺序依赖**：κ次迭代必须顺序执行，限制了并行化潜力 - **内存需求**：需要存储n个样本的pre-logit轨迹，空间复杂度O(nτd) #### 2. 实验设计 - **样本规模**：仅1000个测试样本，统计显著性可能不足 - **token长度限制**：由于内存限制，对prompt和生成长度有严格限制(128 tokens) - **缺少大模型实验**：未在更大规模模型(如Llama-70B)上验证 #### 3. 对比公平性 - **BoN设置**：BoN使用top-p采样，而AISP内部使用贪婪解码，可能不完全公平 - **RE-Control训练**：RE-Control在测试集上训练值函数，可能存在过拟合 #### 4. 理论分析不足 - **收敛保证**：缺少自适应重要性采样的收敛速率分析 - **有效样本数**：未分析重要性采样的有效样本数(ESS) - **高斯假设验证**：缺少对实际pre-logit分布的经验验证 #### 5. 多样性问题 - 在某些设置下，AISP的多样性和连贯性不如BoN - 缺少对这一现象的深入分析和解决方案 ### 影响力 #### 1. 学术贡献 - **新范式**：为测试时对齐提供控制理论视角，可能启发后续研究 - **理论桥梁**：连接控制理论、变分推断和LLM对齐 - **方法论**：自适应重要性采样在pre-logit空间的成功应用可推广到其他生成任务 #### 2. 实用价值 - **成本效益**：无需训练的特性在资源受限场景下有重要价值 - **灵活性**：可与不同LLM和奖励模型组合，适应性强 - **可扩展性**：Batched AISP提供实际部署路径 #### 3. 可复现性 - **代码可用性**：论文未明确提及代码开源，但提供了详细的算法和超参数 - **实现复杂度**：算法相对简单，基于标准重要性采样，易于复现 - **计算需求**：需要GPU资源(H100 80GB或A100 40GB)，对个人研究者有门槛 #### 4. 局限性 - **适用场景**：主要适用于有明确奖励模型的场景 - **扩展性**：在更大模型或更长序列上的表现未知 - **工业应用**：1024次前向传播的推理成本在生产环境可能不可接受 ### 适用场景 #### 最适合的场景 1. **有明确奖励模型**：如安全性检测、事实准确性评估 2. **中等规模模型**：7B-13B参数量的LLM 3. **离线批处理**：可容忍κ次顺序迭代的延迟 4. **资源受限**：无法承担微调成本但有推理资源 #### 不太适合的场景 1. **实时交互**：需要低延迟响应的对话系统 2. **超大规模模型**：内存和计算成本可能过高 3. **无奖励模型**：依赖明确的奖励信号 4. **极端长序列**：控制窗口τ过大会显著增加计算 #### 潜在扩展 1. **多模态生成**：将方法扩展到图像-文本生成 2. **强化学习**：作为exploration策略 3. **主动学习**：用于不确定性采样 4. **对抗鲁棒性**：探索最坏情况下的响应 ## 参考文献 ### 核心引用 1. **Williams et al. (2017, 2018)**: Model Predictive Path Integral Control - AISP的理论基础 2. **Kong et al. (2024)**: RE-Control - 主要对比方法 3. **Yang et al. (2024)**: BoN的理论分析 4. **Lee et al. (2018)**: 高斯假设在神经网络中的应用 ### 相关工作 5. **Ouyang et al. (2022)**: RLHF原始论文 6. **Snell et al. (2024)**: 测试时计算的最优分配 7. **Beirami et al. (2024)**: BoN的理论保证 8. **Khanov et al. (2024)**: ARGS方法 --- ## 总结 本文提出的AISP方法通过将控制理论引入LLM对齐，提供了一个理论优雅且实用有效的测试时对齐方案。其核心创新在于在pre-logit空间施加高斯扰动，通过自适应重要性采样优化扰动分布，实现了无需训练即可超越现有方法的性能。 **主要优势**是样本效率高、无需训练、理论完备；**主要局限**是推理成本较高、需要顺序迭代、对超大模型的扩展性未知。该方法为测试时对齐提供了新的研究方向，特别是在资源受限但有明确奖励模型的场景下具有重要应用价值。 未来研究可在减少推理成本、扩展到更大模型、结合微调方法等方向进一步改进。整体而言，这是一篇高质量的研究工作，兼具理论深度和实用价值。