在 Hugging Face Transformers 中使用约束束搜索引导文本生成

Hugging Face 已在 transformers 库中实现了约束束搜索（constrained beam search），使用户能够通过强制包含特定词汇或短语来对模型输出进行精确控制。这种能力对于神经机器翻译等任务至关重要，例如在这些任务中，词典查询可以规定强制性术语，或者需要根据上下文强制执行特定的正式程度（例如，德语中的正式与非正式）。

约束生成的挑战

在文本生成任务中强制加入特定的子序列并非易事，因为束搜索是逐 token 进行的。标准的解码函数根据当前序列预测下一个 token，但没有内在机制来了解某个要求的 token 是否应该出现在当前步骤或未来的某个步骤。

此外，管理多个约束——例如同时要求两个不同的短语，或允许模型从多个可能的短语列表中选择一个——为解码过程增加了显著的复杂性。

约束束搜索的能力

约束束搜索允许用户在生成阶段注入先验知识，而不是在输出生成后进行过滤。该实现支持两种主要的约束类型：

强制词约束

通过在 model.generate() 中使用 force_words_ids 关键字参数，用户可以指定必须出现在输出中的确切 token。例如，在“How old are you?”的英德翻译中，强制使用单词“Sie”可以确保模型生成正式的翻译（“Wie alt sind Sie?”）而不是非正式的翻译（“Wie alt bist du?”）。

析取约束

析取约束允许模型通过包含所提供列表中的至少一个单词来满足要求。这在多种词形（例如，“raining”, ""raining", "rained", "rains"）均可接受时非常有用。在混合约束场景中，模型可以被强制逐字包含一个特定的单词，同时从相关的术语灵活列表中选择一个选项。

技术实现：Banks 与轮询选择

约束束搜索通过在生成的每一步注入所需的 token，使序列更接近于满足约束。然而，盲目地强制加入 token 往往会导致无意义的结果（例如，“The is fast”）。为了解决这个问题，Hugging Face 利用“Banks”系统来平衡约束的满足与语言概率。

Bank 系统

束（Beams）根据其在满足约束方面的进度被分入不同的 banks：

• Bank $n$：包含在满足约束方面已取得 $n$ 步进度的 beams。

轮询选择

为了防止模型仅选择受约束最严（但可能无意义）的 beams，算法在 banks 之间执行轮询选择。如果 num_beams=3，系统会先从 Bank 2 中选择概率最高的输出，然后是 Bank 1，最后是 Bank 0。这确保了即使在模型被推向约束时，高概率且合理的序列仍能被保留。

如果一个 beam 偏离了约束（例如，在约束为“is fast”时生成了“slow”），其进度将被重置，并移回 Bank 0。

通过 Constraint 类扩展约束

虽然 force_words_ids 是主要的面向用户接口，但后端利用 Constraint 对象来跟踪进度并建议下一个 token。开发者可以通过继承 Constraint 抽象接口类来创建自定义约束。

提到的潜在未来扩展包括：

• OrderedConstraints：确保约束按特定顺序得到满足。
• TemplateConstraints：强制模型遵循具有待填充缺口的特定结构模板。

研究基础

该功能基于多篇专注于词汇约束解码和引导生成的研究论文，包括关于开放词汇图像描述和神经机器翻译的工作。