Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models

📄 arXiv: 2603.12248v1 📥 PDF

作者: Samy Jelassi, Mujin Kwun, Rosie Zhao, Yuanzhi Li, Nicolo Fusi, Yilun Du, Sham M. Kakade, Carles Domingo-Enrich

分类: cs.LG

发布日期: 2026-03-12


💡 一句话要点

提出能量基微调(EBFT),通过特征匹配优化语言模型序列级行为。

🎯 匹配领域: 支柱六:视频提取与匹配 (Video Extraction)

关键词: 语言模型微调 特征匹配 能量基模型 序列生成 策略梯度 强化学习 自然语言处理

📋 核心要点

  1. 现有交叉熵训练侧重于token预测,忽略了模型rollout的序列级行为,导致与实际应用场景存在差距。
  2. 论文提出能量基微调(EBFT),通过特征匹配优化序列级统计信息,提供密集语义反馈,无需任务特定验证器。
  3. 实验表明,EBFT在问答编码、非结构化编码和翻译任务上,性能与RLVR相当,优于SFT,且验证交叉熵更低。

📝 摘要(中文)

交叉熵(CE)训练为语言模型提供了密集且可扩展的监督,但它优化的是教师强制下的下一个token预测,而不是模型rollout下的序列级行为。本文提出了一种用于语言模型微调的特征匹配目标,该目标针对completion分布的序列级统计信息,提供密集的语义反馈,而无需特定于任务的验证器或偏好模型。为了有效地优化这个目标,我们提出了能量基微调(EBFT),它使用步长块并行采样从嵌套前缀并发生成多个rollout,批量提取这些rollout的特征,并使用生成的嵌入来执行on-policy策略梯度更新。我们提出了一个理论视角,将EBFT与kl正则化特征匹配和能量基建模联系起来。在问答编码、非结构化编码和翻译等任务中,EBFT与RLVR相当,并且在下游精度上优于SFT,同时实现了比这两种方法更低的验证交叉熵。

🔬 方法详解

问题定义:现有语言模型微调方法,如交叉熵训练,主要关注下一个token的预测,而忽略了模型在实际生成序列时的整体行为。这种训练方式导致模型在序列层面的表现与预期不符,尤其是在需要长文本生成或复杂推理的任务中,效果不佳。现有方法缺乏对序列级别语义信息的有效利用,需要人工标注或复杂的奖励函数设计。

核心思路:论文的核心思路是通过特征匹配来优化语言模型。具体来说,就是让模型生成的序列的统计特征与目标分布的统计特征尽可能接近。这种方法避免了直接预测token,而是关注整个序列的语义信息,从而更好地提升模型在序列生成任务中的表现。能量基模型(Energy-Based Model)提供了一种有效的框架来建模这种特征匹配,通过定义能量函数来衡量序列的质量。

技术框架:EBFT的整体框架包括以下几个主要步骤: 1. Rollout生成:使用步长块并行采样,从嵌套前缀并发生成多个rollout。这种方法可以高效地探索模型的生成空间。 2. 特征提取:对生成的rollout进行特征提取,得到序列的嵌入表示。这些特征可以捕捉序列的语义信息。 3. 能量计算:使用能量函数来衡量序列的质量。能量函数可以基于序列的特征进行设计。 4. 策略梯度更新:使用策略梯度方法来更新模型参数,使得模型生成的序列的能量更低,从而更接近目标分布。

关键创新:EBFT的关键创新在于使用能量基模型进行特征匹配,从而优化语言模型的序列级行为。与传统的交叉熵训练相比,EBFT关注整个序列的语义信息,而不是单个token的预测。此外,EBFT使用步长块并行采样,可以高效地生成多个rollout,从而加速训练过程。EBFT避免了对任务特定验证器或偏好模型的依赖,降低了微调的成本。

关键设计:EBFT的关键设计包括: 1. 能量函数:能量函数的设计至关重要,需要能够有效地衡量序列的质量。论文中可能使用了预训练模型的embedding作为特征,并设计了基于这些特征的能量函数。 2. 步长块并行采样:这种采样方法可以高效地生成多个rollout,从而加速训练过程。具体实现可能涉及到对序列进行分块,并并行地生成每个块的多个候选序列。 3. 策略梯度算法:选择合适的策略梯度算法对于训练的稳定性和收敛速度至关重要。论文中可能使用了TRPO或PPO等算法。

📊 实验亮点

EBFT在Q&A编码、非结构化编码和翻译任务上表现出色,与RLVR性能相当,并在下游精度上超越了SFT。更重要的是,EBFT在所有任务中都实现了比SFT和RLVR更低的验证交叉熵,表明其具有更好的泛化能力和更稳定的训练过程。这些实验结果充分证明了EBFT在优化语言模型序列级行为方面的有效性。

🎯 应用场景

该研究成果可广泛应用于各种需要高质量文本生成的场景,例如机器翻译、文本摘要、对话系统、代码生成等。通过优化序列级行为,可以提升生成文本的流畅性、连贯性和语义准确性。该方法无需人工标注或复杂的奖励函数设计,降低了微调成本,具有很高的实际应用价值和潜力。

📄 摘要(原文)

Cross-entropy (CE) training provides dense and scalable supervision for language models, but it optimizes next-token prediction under teacher forcing rather than sequence-level behavior under model rollouts. We introduce a feature-matching objective for language-model fine-tuning that targets sequence-level statistics of the completion distribution, providing dense semantic feedback without requiring a task-specific verifier or preference model. To optimize this objective efficiently, we propose energy-based fine-tuning (EBFT), which uses strided block-parallel sampling to generate multiple rollouts from nested prefixes concurrently, batches feature extraction over these rollouts, and uses the resulting embeddings to perform an on-policy policy-gradient update. We present a theoretical perspective connecting EBFT to KL-regularized feature-matching and energy-based modeling. Empirically, across Q&A coding, unstructured coding, and translation, EBFT matches RLVR and outperforms SFT on downstream accuracy while achieving a lower validation cross-entropy than both methods.