TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation

📄 arXiv: 2504.18535v2 📥 PDF

作者: Gwen Yidou Weng, Benjie Wang, Guy Van den Broeck

分类: cs.CL, cs.LG

发布日期: 2025-04-25 (更新: 2025-09-29)

备注: Proceedings of the 42nd International Conference on Machine Learning (ICML). 2025

🔗 代码/项目: GITHUB

💡 一句话要点

TRACE：提出一种基于概率推理的可控语言生成框架，提升生成质量与效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 可控语言生成 概率推理 隐马尔可夫模型 期望属性概率 文本去毒化 个性化语言模型 语言模型控制

📋 核心要点

  1. 现有自回归语言模型难以控制生成文本的全局属性，后训练成本高昂，采样近似EAP效率低且不稳定。
  2. TRACE框架通过从语言模型中提取HMM，并结合分类器估计属性概率，实现高效的期望属性概率（EAP）计算。
  3. 实验表明，TRACE在去毒化、个性化等任务上取得了SOTA结果，且解码开销低，能快速生成多个个性化模型。

📝 摘要（中文）

随着大型语言模型（LM）的发展，越来越需要控制它们的输出，使其与人类价值观（例如，去毒化）或期望的属性（例如，个性化、主题）对齐。然而，自回归模型专注于下一个token的预测，难以处理需要前瞻的全局属性。现有的解决方案要么为每个新属性对LM进行后训练（成本高且不灵活），要么通过抽样或训练来近似未来序列的期望属性概率（EAP），这对于稀有属性来说速度慢且不可靠。我们引入TRACE（用于自适应可控生成的易处理概率推理），这是一个新颖的框架，它通过易处理的概率推理和轻量级控制来有效地计算EAP并适应新属性。TRACE从LM中提取一个隐马尔可夫模型（HMM），并将其与一个小分类器配对以估计属性概率，从而能够精确计算HMM预测的未来序列的EAP。然后，该EAP用于重新加权LM的下一个token概率，以实现全局兼容的延续。实验表明，TRACE仅以20%的解码开销实现了最先进的去毒化结果，在几秒钟内产生了76个低资源个性化LM，并无缝扩展到复合属性。我们的代码可在https://github.com/yidouweng/trace获得。

🔬 方法详解

问题定义：论文旨在解决大型语言模型生成文本时难以控制全局属性的问题，例如去毒化、个性化等。现有方法主要存在两个痛点：一是针对每个属性进行后训练，成本高且缺乏灵活性；二是采用采样或训练的方式近似期望属性概率（EAP），对于稀有属性而言，效率低且结果不稳定。

核心思路：论文的核心思路是利用可处理的概率推理来高效计算EAP，从而指导语言模型的生成过程。具体而言，通过从预训练语言模型中提取一个隐马尔可夫模型（HMM），并结合一个小型分类器来估计属性概率，从而实现对未来序列EAP的精确计算。

技术框架：TRACE框架主要包含两个阶段：1) HMM提取与属性概率估计：首先，从预训练语言模型中提取一个HMM，该HMM能够捕捉语言模型的序列生成特性。然后，训练一个小型的分类器，用于估计给定token序列的属性概率。2) EAP计算与重加权：利用提取的HMM和训练好的分类器，计算未来序列的EAP。然后，使用计算得到的EAP来重新加权语言模型的下一个token概率，从而引导生成过程，使其更符合期望的属性。

关键创新：TRACE的关键创新在于利用HMM进行可处理的概率推理，从而高效且精确地计算EAP。与现有方法相比，TRACE避免了昂贵的后训练过程和不稳定的采样近似，能够快速适应新的属性，并实现对生成文本全局属性的有效控制。

关键设计：HMM的提取方式（例如，状态数量、转移概率的计算方法）以及分类器的选择（例如，网络结构、损失函数）是关键的设计因素。此外，如何有效地利用EAP来重新加权下一个token的概率，也是影响生成质量的重要因素。论文中可能涉及对这些参数和方法的具体选择和优化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TRACE在去毒化任务上取得了SOTA结果，仅以20%的解码开销就超越了现有方法。此外，TRACE还能够在几秒钟内生成76个低资源个性化语言模型，展示了其高效性和可扩展性。实验结果表明，TRACE能够有效地控制生成文本的全局属性，并快速适应新的属性。

🎯 应用场景

TRACE框架具有广泛的应用前景，可用于各种需要控制生成文本属性的场景，例如：1) 内容审核：用于生成无毒、无害的内容；2) 个性化推荐：根据用户偏好生成个性化的文本；3) 创意写作：辅助生成特定风格或主题的文章。该研究有助于提升语言模型的可控性和实用性，促进其在更多领域的应用。

📄 摘要（原文）

As large language models (LMs) advance, there is an increasing need to control their outputs to align with human values (e.g., detoxification) or desired attributes (e.g., personalization, topic). However, autoregressive models focus on next-token predictions and struggle with global properties that require looking ahead. Existing solutions either post-train LMs for each new attribute--expensive and inflexible--or approximate the Expected Attribute Probability (EAP) of future sequences by sampling or training, which is slow and unreliable for rare attributes. We introduce TRACE (Tractable Probabilistic Reasoning for Adaptable Controllable gEneration), a novel framework that efficiently computes EAP and adapts to new attributes through tractable probabilistic reasoning and lightweight control. TRACE distills a Hidden Markov Model (HMM) from an LM and pairs it with a small classifier to estimate attribute probabilities, enabling exact EAP computation over the HMM's predicted futures. This EAP is then used to reweigh the LM's next-token probabilities for globally compliant continuations. Empirically, TRACE achieves state-of-the-art detoxification results with only 20% decoding overhead, yields 76 low-resource personalized LMs within seconds, and seamlessly extends to composite attributes. Our code is available at: https://github.com/yidouweng/trace.