在 🤗 Transformers 中使用约束集束搜索引导文本生成

引言

本博文假设读者熟悉使用不同集束搜索变体的文本生成方法，相关解释请参阅博文：“如何生成文本：在 Transformers 中使用不同解码方法进行语言生成”

与普通的集束搜索不同，约束集束搜索允许我们对文本生成的输出进行控制。这很有用，因为我们有时确切地知道我们希望输出中包含什么内容。例如，在神经机器翻译任务中，我们可能通过查字典知道最终译文必须包含哪些词。有时，对于语言模型来说几乎同样可能的生成输出，由于特定上下文的原因，对最终用户来说可能并非同样合意。这两种情况都可以通过允许用户告诉模型最终输出必须包含哪些词来解决。

为何如此困难

然而，这实际上是一个非常不简单的问题。这是因为任务要求我们在最终输出的某个位置、在生成的某个时间点强制生成某些子序列。

假设我们想生成一个句子 S，它必须包含短语 $p_1=\{ t_1, t_2 \}$，其中按顺序包含词元 $t_1, t_2$。我们将期望的句子 $S$ 定义为

$$S_{expected} = \{ s_1, s_2, ..., s_k, t_1, t_2, s_{k+1}, ..., s_n \}$$

问题在于集束搜索是逐词元生成序列的。虽然不完全准确，但可以将集束搜索看作函数 $B(\mathbf{s}_{0:i}) = s_{i+1}$，它查看当前从 $0$ 到 $i$ 已生成的词元序列，然后预测在 $i+1$ 位置的下一个词元。但是这个函数如何能在任意步骤 $i < k$ 时，知道必须在未来的某个步骤 $k$ 生成这些词元呢？或者当它在步骤 $i=k$ 时，它如何能确定这是强制生成这些词元的最佳位置，而不是在未来的某个步骤 $i>k$ 呢？

如果你有多个具有不同要求的约束怎么办？如果你想强制生成短语 $p_1=\{t_1, t_2\}$ 和短语 $p_2=\{ t_3, t_4, t_5, t_6\}$ 呢？如果你想让模型在这两个短语之间选择一个呢？如果我们想强制生成短语 $p_1$，并从短语列表 $\{p_{21}, p_{22}, p_{23}\}$ 中强制生成一个短语呢？

以上这些例子实际上是非常合理的用例，下面将会展示，而新的约束集束搜索功能允许所有这些操作！

本文将快速介绍新的约束集束搜索功能能为你做什么，然后深入探讨其内部工作原理。

示例 1：强制生成一个词

假设我们想把 "How old are you?" 翻译成德语。

"Wie alt bist du?" 是非正式场合的说法，而 "Wie alt sind Sie?" 是正式场合的说法。

根据上下文，我们可能希望使用一种礼貌形式而不是另一种，但我们如何告诉模型这一点呢？

传统的集束搜索

以下是我们在传统集束搜索设置中进行文本翻译的方式。

!pip install -q git+https://github.com/huggingface/transformers.git

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt bist du?

使用约束集束搜索

但如果我们知道我们想要正式的输出而不是非正式的呢？如果我们根据先验知识知道生成结果必须包含什么，并且可以将其注入到生成过程中呢？

以下是现在可以使用 model.generate() 的 force_words_ids 关键字参数实现的功能

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

force_words = ["Sie"]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids

outputs = model.generate(
    input_ids,
    force_words_ids=force_words_ids,
    num_beams=5,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

如你所见，我们能够利用关于期望输出的先验知识来引导生成。以前我们必须生成一堆可能的输出，然后筛选出符合我们要求的。现在我们可以在生成阶段就做到这一点。

示例 2：析取约束

我们上面提到了一个用例，我们知道最终输出中要包含哪些词。一个例子可能是在神经机器翻译中使用字典查找。

但如果我们不知道该使用哪种词形，比如我们希望 ["raining", "rained", "rains", ...] 这样的输出同样可能呢？更普遍地说，总有我们不想要逐字逐句完全相同的词，而可能对其他相关的可能性持开放态度的情况。

允许这种行为的约束是析取约束，它允许用户输入一个词列表，其目的是引导生成，使得最终输出必须包含该词列表中的至少一个。

以下是一个混合使用上述两种约束类型的例子

from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

force_word = "scared"
force_flexible = ["scream", "screams", "screaming", "screamed"]

force_words_ids = [
    tokenizer([force_word], add_prefix_space=True, add_special_tokens=False).input_ids,
    tokenizer(force_flexible, add_prefix_space=True, add_special_tokens=False).input_ids,
]

starting_text = ["The soldiers", "The child"]

input_ids = tokenizer(starting_text, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    force_words_ids=force_words_ids,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(tokenizer.decode(outputs[1], skip_special_tokens=True))

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.


Output:
----------------------------------------------------------------------------------------------------
The soldiers, who were all scared and screaming at each other as they tried to get out of the
The child was taken to a local hospital where she screamed and scared for her life, police said.

如你所见，第一个输出使用了 "screaming"，第二个输出使用了 "screamed"，并且两者都逐字使用了 "scared"。可供选择的列表 ["screaming", "screamed", ...] 不必是词形变化；这可以满足任何我们需要从词列表中选择一个的用例。

传统的集束搜索

以下是传统集束搜索的一个例子，摘自之前的一篇博文

与贪婪搜索不同，集束搜索通过保留更长的假设列表来工作。在上图中，我们展示了在生成的每个可能步骤中的三个可能的下一个词元。

在 num_beams=3 的情况下，以下是看待上述示例集束搜索第一步的另一种方式

集束搜索不会像贪婪搜索那样只选择 "The dog"，而是允许对 "The nice" 和 "The car" 进行进一步考虑。

在下一步中，我们为上一步创建的三个分支中的每一个考虑下一个可能的词元。

尽管我们最终考虑的输出远多于 num_beams，但在每一步结束时，我们会将它们减少到 num_beams。我们不能一直分支下去，否则我们需要跟踪的 beams 数量将是 $\text{beams}^{n}$ 对于 $n$ 步来说，这个数字会很快变得非常大（$10$ 个 beam 经过 $10$ 步后是 $10,000,000,000$ 个 beam！）。

在剩下的生成过程中，我们重复上述步骤，直到满足结束条件，例如生成 <eos> 词元或达到 max_length。分支、排序、缩减，然后重复。

约束集束搜索

约束集束搜索试图通过在生成的每一步注入所需的词元来满足约束。

假设我们想在生成的输出中强制包含短语 "is fast"。

在传统的集束搜索设置中，我们在每个分支找到概率最高的 k 个下一个词元，并将它们加入考虑范围。在约束设置中，我们做同样的事情，但同时也会添加那些能让我们更接近满足约束的词元。下面是一个演示

除了像 "dog" 和 "nice" 这样常见的高概率下一个词元之外，我们强制加入词元 "is"，以便更接近满足我们 "is fast" 的约束。

在下一步中，下面分支出的候选词元与传统集束搜索的大部分相同。但和上面的例子一样，约束集束搜索通过在每个新分支上强制约束来扩展现有候选词元

组 (Banks)

在讨论下一步之前，我们需要考虑上一步中可能出现的不良行为。

简单地在输出中强制加入所需短语 "is fast" 的问题在于，大多数情况下，你会得到像上面 "The is fast" 这样无意义的输出。这实际上是使这个问题变得不简单的原因。关于解决这个问题的复杂性的更深入讨论，可以在 huggingface/transformers 中提出的原始功能请求 issue 中找到。

组 (Banks) 通过在满足约束和创建合理输出之间建立平衡来解决这个问题。

组 $n$ 指的是在满足约束方面取得了 $n$ 步进展的 beam 列表。将所有可能的 beam 分配到各自的组后，我们进行轮询选择。在上面的例子中，我们会从组 2 中选择概率最高的输出，然后从组 1 中选择概率最高的，再从组 0 中选择一个，然后是组 2 中概率第二高的，组 1 中概率第二高的，依此类推。由于我们使用 num_beams=3，我们只需将上述过程重复三次，最终得到 ["The is fast", "The dog is", "The dog and"]。

这样，即使我们强制模型考虑我们手动添加了所需词元的分支，我们仍然会跟踪其他可能更有意义的高概率序列。尽管 "The is fast" 完全满足了我们的约束，但它并不是一个非常合理的短语。幸运的是，我们在未来的步骤中有 "The dog is" 和 "The dog and" 可用，希望它们能引导出更合理的输出。

这种行为在上述示例的第三步中得到了展示

请注意 "The is fast" 不需要任何手动添加约束词元，因为它已经满足了约束（即，已包含短语 "is fast"）。另外，请注意像 "The dog is slow" 或 "The dog is mad" 这样的 beam 实际上在组 0 中，因为虽然它包含了词元 "is"，但它必须从头开始才能生成 "is fast"。在 "is" 后面加上 "slow" 这样的词，实际上是重置了它的进展。

最后请注意，我们最终得到了一个包含我们约束短语的合理输出："The dog is fast"！

我们最初担心盲目添加所需词元会导致像 "The is fast" 这样无意义的短语。然而，通过使用从组中轮询选择的方法，我们隐式地排除了无意义的输出，而偏向于更合理的输出。

更多关于 Constraint 类和自定义约束

这个解释的主要 takeaways 可以总结如下。在每一步，我们不断地“纠缠”模型，让它考虑满足我们约束的词元，同时跟踪那些不满足约束的 beam，直到我们最终得到包含我们期望短语的、概率相当高的序列。

因此，设计这个实现的一个原则性方法是将每个约束表示为一个 Constraint 对象，其目的是跟踪其进展并告诉集束搜索接下来要生成哪些词元。虽然我们为 model.generate() 提供了关键字参数 force_words_ids，但实际上在后端发生的是以下情况

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, PhrasalConstraint

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

constraints = [
    PhrasalConstraint(
        tokenizer("Sie", add_special_tokens=False).input_ids
    )
]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids


outputs = model.generate(
    input_ids,
    constraints=constraints,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

你可以自己定义一个并将其输入到 constraints 关键字参数中，以设计你独特的约束。你只需创建一个 Constraint 抽象接口类的子类，并遵循其要求。你可以在这里找到 Constraint 的定义以获取更多信息。

一些独特的想法（尚未实现；也许你可以试试！）包括像 OrderedConstraints、TemplateConstraints 这样的约束，这些可能会在未来加入。目前，生成是通过在输出的任何位置包含序列来完成的。例如，前面的一个例子中，一个序列是 scared -> screaming，另一个是 screamed -> scared。OrderedConstraints 可以允许用户指定这些约束被满足的顺序。

TemplateConstraints 可以允许该功能更小众的用法，其目标可以是这样的

starting_text = "The woman"
template = ["the", "", "School of", "", "in"]

possible_outputs == [
   "The woman attended the Ross School of Business in Michigan.",
   "The woman was the administrator for the Harvard School of Business in MA."
]

或者

starting_text = "The woman"
template = ["the", "", "", "University", "", "in"]

possible_outputs == [
   "The woman attended the Carnegie Mellon University in Pittsburgh.",
]
impossible_outputs == [
  "The woman attended the Harvard University in MA."
]

或者如果用户不关心两个词之间可以有多少个词元，那么可以直接使用 OrderedConstraint。

结论

约束集束搜索为我们提供了一种灵活的方式，将外部知识和要求注入到文本生成中。以前，没有简单的方法告诉模型 1. 包含一个序列列表，其中 2. 一些是可选的，一些不是，以便 3. 它们在序列的某个位置以各自合理的位置生成。现在，我们可以通过混合使用不同子类的 Constraint 对象来完全控制我们的生成！

这个新功能主要基于以下论文

• 《使用约束集束搜索的引导式开放词汇图像描述》(Guided Open Vocabulary Image Captioning with Constrained Beam Search)
• 《用于神经机器翻译的带动态束分配的快速词汇约束解码》(Fast Lexically Constrained Decoding with Dynamic Beam Allocation for Neural Machine Translation)
• 《用于翻译和单语重写的改进词汇约束解码》(Improved Lexically Constrained Decoding for Translation and Monolingual Rewriting)
• 《因果关系的引导式生成》(Guided Generation of Cause and Effect)

像上面这些一样，许多新的研究论文正在探索使用外部知识（例如知识图谱、知识库）来引导大型深度学习模型输出的方法。希望这个约束集束搜索功能能成为实现这一目标的又一个有效方法。

感谢所有为这个功能贡献提供指导的人：Patrick von Platen，从最初的 issue 到最终的 PR 都参与其中，以及 Narsil Patry，为代码提供了详细的反馈。

本帖缩略图使用了归属于以下来源的图标：由 Freepik - Flaticon 创建的速记图标