让大模型提前预判输出属性,实现精准可控生成。
TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
- 用隐马尔可夫模型+轻量分类器估算未来属性概率,实现全局控制。
- 毒性和个性化生成效果领先,解码开销仅增加20%。
- 支持快速适配新属性,适合需要灵活控制的生成场景。
随着大语言模型的发展,控制其输出以符合人类价值观(如去毒化)或特定属性(如个性化、主题)的需求日益增长。然而,自回归模型仅关注下一词预测,难以处理需前瞻性的全局属性。现有方法要么为每个新属性单独微调模型(成本高、不灵活),要么通过采样或训练近似未来序列的期望属性概率(EAP),但对稀有属性效率低且不可靠。本文提出TRACE(可计算的概率推理用于可适应的可控生成),一种高效计算EAP并实现轻量级控制的新框架。TRACE从语言模型中蒸馏出隐马尔可夫模型(HMM),并搭配一个小型分类器,精确计算基于HMM预测未来的属性概率。该概率用于重加权语言模型的下一步词概率,从而生成全局合规的文本。实验表明,TRACE在去毒化任务上达到当前最佳表现,解码开销仅增加20%,可在数秒内生成76个低资源个性化模型,并无缝扩展至复合属性。代码已开源:https://github.com/yidouweng/trace。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。