使用正则表达式约束实现快速、高保真LLM解码

约束LLM生成的文本对于开发健壮的LLM应用至关重要。例如，开发者可能要求LLM的响应遵循特定的JSON或YAML模式，以确保其内容全面且可可靠解析。

去年，已经发布了强大的工具和方法（Beurer-Kellner 等，2022；Lundberg & Ribeiro，2023；Willard & Louf，2023；Zheng 等，2023）来满足这一需求。值得注意的是，开源库 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 实现了一种高效的方法（Willard & Louf，2023），可确保符合任意正则表达式。

在这篇博客文章中，我提出了两种新颖、快速、高保真的替代方法。具体来说，我将：

• 介绍 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 的原理，并讨论允许不当分词序列（即与字符串分词不对应的分词序列）如何扭曲LLM的概率分布并阻碍解码速度；
• 介绍 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$，一种在仅生成正确分词序列的同时保持保证合规性的方法。$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 适用于所有基于合并表的分词器，尤其是非常流行的字节对编码分词器（Sennrich 等，2016）；
• 提出 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$，它是 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 的一个变体，可以更好地随正则表达式的复杂性而扩展；
• 讨论 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 的执行时间，并展示在JSON或YAML模式约束的情况下如何显著缩短执行时间。

这篇博客文章附带一份技术附录，其中对算法进行了形式化和证明。此外，还提供了一个包含 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 和 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 实现的Python notebook，以及额外的实证结果。

确定性有限自动机与约束解码

作为预处理步骤，$\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 将测试字符串是否匹配正则表达式的字符级确定性有限自动机 (DFA)（参见图 1a）转换为测试分词序列是否被解码成匹配正则表达式的字符串的分词级DFA（参见图 1b）

• 状态，包括起始状态和接受状态，保持不变；
• 分词级DFA的字母表是分词器的词汇表；
• 分词级DFA的转换函数是字符级DFA的扩展转换函数，限制在分词器的词汇表内。例如，在图 1b 的示例中，从 $\mathtt{\text{0}}\backslash texttt\{0\}$ 到 $\mathtt{\text{1}}\backslash texttt\{1\}$ 存在一个 $\mathtt{\text{abb}}\backslash texttt\{abb\}$ 转换，因为如果我们在字符级DFA上从 $\mathtt{\text{0}}\backslash texttt\{0\}$ 开始并连续看到 $\mathtt{\text{a}}\backslash texttt\{a\}$、$\mathtt{\text{b}}\backslash texttt\{b\}$ 和 $\mathtt{\text{b}}\backslash texttt\{b\}$，我们将最终到达状态 $\mathtt{\text{1}}\backslash texttt\{1\}$。

在解码时，DFA用于确定每个新分词允许哪些潜在分词。为此，从DFA的初始状态开始，我们通过当前状态的出边转换来确定允许的分词，将相应的掩码应用于下一个分词概率并重新归一化这些概率。然后我们可以采样一个新的分词并更新DFA的状态。

例如，假设令牌 $\mathtt{\text{aaaa}}\backslash texttt\{aaaa\}$、$\mathtt{\text{aa}}\backslash texttt\{aa\}$ 和 $\mathtt{\text{ab}}\backslash texttt\{ab\}$ 已经生成。从状态 $\mathtt{\text{0}}\backslash texttt\{0\}$ 开始，我们随着 $\mathtt{\text{aaaa}}\backslash texttt\{aaaa\}$ 和 $\mathtt{\text{aa}}\backslash texttt\{aa\}$ 停留在状态 $\mathtt{\text{0}}\backslash texttt\{0\}$，但随着 $\mathtt{\text{ab}}\backslash texttt\{ab\}$ 转换到状态 $\mathtt{\text{1}}\backslash texttt\{1\}$。现在允许的令牌是那些从状态 $\mathtt{\text{1}}\backslash texttt\{1\}$ 出发的转换，即 $\mathtt{\text{b}}\backslash texttt\{b\}$ 和 $\mathtt{\text{bb}}\backslash texttt\{bb\}$。然后我们考虑这两个令牌的下一个令牌概率并将其重新归一化，使其总和为1。

由于不当分词序列导致的三大挑战

使用 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 潜在生成的分词序列，正是那些被解码成符合正则表达式的字符串的序列。这包括并非由字符串分词产生的序列。我们称之为不当分词序列。

考虑正则表达式：$\mathtt{\text{"boolean:}}\mathtt{\text{((true)}}\mid \mathtt{\text{(false))"}}\backslash texttt\{"boolean:\; ((true)\}\backslash mid\backslash texttt\{(false))"\}$。虽然只有两个字符串与此正则表达式匹配，但 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 返回的相应DFA有17个状态和48个转换（图2）。例如，第一个分词是从 $\mathtt{\text{boolean}}\backslash texttt\{boolean\}$、$\mathtt{\text{bool}}\backslash texttt\{bool\}$、$\mathtt{\text{bo}}\backslash texttt\{bo\}$ 和 $\mathtt{\text{b}}\backslash texttt\{b\}$ 中选择的。

相比之下，正如我们将在以下部分中看到的，由 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 和 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 构建的DFA仅接受由符合字符串分词产生的序列，它只有5个状态和4个转换（图3）。

挑战 #1：概率分布偏差

$\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 对不当分词序列的宽容性可能会显著扭曲LLM的概率分布，超出强制正则表达式约束所严格必需的范围。

为了说明这一点，我们假设我们生成一个对提示 Question: What is the first name of the US president who was a member of the Sigma Alpha Epsilon fraternity?\nAnswer: 的续写，并加上正则表达式约束 $\mathtt{\text{"(}}\mathtt{\text{William)}}\mid \mathtt{\text{(}}\mathtt{\text{Theodore)"}}\backslash texttt\{"(\; William)\}\backslash mid\; \backslash texttt\{(\; Theodore)"\}$ （正确答案是 William，就像 William McKinley，西奥多·罗斯福的前任）。

当文本生成不受约束时，我们可以计算所有导致两个可能答案之一的分词序列的概率。如果使用 $\mathtt{\text{Mistral-7B}}\backslash texttt\{Mistral-7B\}$ 和温度等于1，我们得到 ，并且注意到这两个答案几乎所有的概率质量都来自正确的令牌序列（$\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$ 为 99.8586%，$\mathtt{\text{"}}\mathtt{\text{William"}}\backslash texttt\{"\; William"\}$ 为 99.9997%）。

如果我们使用 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 约束文本生成并采用多项式采样作为解码策略，$\mathtt{\text{"}}\mathtt{\text{William"}}\backslash texttt\{"\; William"\}$ 和 $\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$ 的任何前缀都将被接受为第一个标记并决定答案（如果此第一个标记不是空格）。例如，标记 $\mathtt{\text{"}}\mathtt{\text{The"}}\backslash texttt\{"\; The"\}$ 可以被选作第一个标记，在这种情况下，答案必然是 $\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$。结果，生成 $\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$ 的概率将是生成 $\mathtt{\text{"}}\mathtt{\text{T"}}\backslash texttt\{"\; T"\}$、$\mathtt{\text{"}}\mathtt{\text{Th"}}\backslash texttt\{"\; Th"\}$、$\mathtt{\text{"}}\mathtt{\text{The"}}\backslash texttt\{"\; The"\}$ 等作为第一个标记的概率之和，尽管除了 $\mathtt{\text{"}}\mathtt{\text{Theod"}}\backslash texttt\{"\; Theod"\}$ 之外，这些第一个标记几乎从不导致 $\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$，并且 $\mathtt{\text{"}}\mathtt{\text{Theodore"}}\backslash texttt\{"\; Theodore"\}$ 将仅仅因为有一个非常常见的标记 $\mathtt{\text{"}}\mathtt{\text{The"}}\backslash texttt\{"\; The"\}$ 作为前缀而显著提升。事实上，在这种情况下 。

使用 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 和 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$，只会考虑与可能答案分词相对应的令牌（例如 $\mathtt{\text{"}}\mathtt{\text{The"}}\backslash texttt\{"\; The"\}$ 被忽略）。条件概率 ，与非约束情况非常接近，这很合理，因为LLM在大多数情况下自发生成适当的令牌序列。

挑战 #2：自我中毒

另一个问题是，生成的令牌会反馈给LLM以选择后续令牌。在其预训练期间，根据是否使用了BPE-Dropout（Provilkov 等，2020）等子词正则化技术，LLM很少或从未接触过不当分词序列，因此LLM可能无法正确解释新添加的令牌，这可能会影响其输出的质量。

挑战 #3：错失加速机会

正如最近两篇博客文章所建议的，$\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 提供了加速解码的机会。

让我们回到 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 和 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 在图3中提供的DFA。从状态 $\mathtt{\text{0}}\backslash texttt\{0\}$ 和 $\mathtt{\text{1}}\backslash texttt\{1\}$ 只有一个出边转换。这意味着当我们到达这些状态时，甚至不需要计算下一个令牌的概率，因为只有一个可能的令牌。然后我们可以在一次LLM调用中生成整个文本。同时，用 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 为图3中相同的正则表达式构建的DFA有16个状态，但只有3个状态具有一个出边转换（状态6、7和13）。因此，跳过计算成本高昂的LLM调用的机会要少得多。此外，有许多潜在路径比三个转换长得多。这也会减慢解码速度。

确保符合性与正确分词

现在我们可以专注于 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$。尽管其目标略有不同，$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 与 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 共享相同的总体方法：它从识别正则表达式的字符级DFA派生出令牌级DFA，并且这个令牌级DFA用于在解码时识别允许的令牌。$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 适用于基于合并表的分词器，例如字节对编码分词器（Sennrich 等，2016）。使用此类分词器对字符串进行编码包括将字符串转换为单字母令牌序列，然后执行一系列有序的合并操作。每个合并操作都由两个令牌定义（例如 $\mathtt{\text{a}}\backslash texttt\{a\}$ 和 $\mathtt{\text{b}}\backslash texttt\{b\}$），并从当前令牌序列的左侧开始合并所有相应的连续令牌对。

使用 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$，我们从字符级DFA开始，并为每个合并操作应用一系列简单的转换。让我们考虑 $(a,b)(a,\; b)$ 合并操作且 $a\mathrm{\ne }ba\; \backslash neq\; b$ 的情况。$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 影响每个具有一个或多个传入 $aa$ 转换和一个传出 $bb$ 转换的状态 $SS$，其影响取决于两个标准：

1. $SS$ 是否有一个 $x(x\mathrm{\ne }a)x\backslash :(x\; \backslash neq\; a)$ 的传入转换，或者 $SS$ 是否是起始状态？
2. $SS$ 是否有 $y(y\mathrm{\ne }b)y\backslash :(y\; \backslash neq\; b)$ 的出边转换，或者 $SS$ 是否是接受状态？

更精确地说，如果从 $SS$ 出发的 $bb$ 转换导致 $S_{b}S\_b$，则对状态 $SS$ 应用的转换如下：

• 将每个传入的 $aa$ 转换替换为到 $S_{b}S\_b$ 的 $abab$ 转换。
• 如果两个条件都为假（参见图 4a）
  • 移除 $SS$
• 如果标准1为假且标准2为真
  • 移除出站 $bb$ 转换
• 如果条件1为真且条件2为假
  • 删除所有传入的 $aa$ 转换
• 如果两个条件都为真（参见图 4c）
  • 添加一个状态 $S^{\mathrm{\prime }}S\text{'}$
  • 将所有传入的 $aa$ 转换重定向到 $S^{\mathrm{\prime }}S\text{'}$
  • 复制 $SS$ 的所有 $y(y\mathrm{\ne }b)y\backslash :(y\; \backslash neq\; b)$ 转换并将其添加到 $S^{\mathrm{\prime }}S\text{'}$
  • 如果 $SS$ 是接受状态，则使 $S^{\mathrm{\prime }}S\text{'}$ 成为接受状态

包含两个相同标记的合并操作更具挑战性，因为并非所有匹配标记对都应该合并。例如，在表1的第四行中，第二个和第三个 $aa$ 未合并，因为合并操作是左到右应用的，第二个 $aa$ 已经与第一个 $aa$ 合并。您可以在技术附录中找到 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 在 $a=ba\; =\; b$ 情况下的描述以及两种情况的正确性证明。

随着正则表达式复杂度的扩展

$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 保证了对正则表达式约束的遵守和正确的标记化。然而，$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 会移除或添加状态和转换，在某些不利情况下，状态和转换的数量可能会变得异常庞大。

幸运的是，通过 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$，我们可以在不显式构建它的情况下，高效地模拟由 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 创建的 DFA。我们确实可以注意到，目标语言是所有解码为匹配正则表达式的字符串的标记序列集合与所有正确标记序列集合的交集。这两个集合都是正则语言，我们已经知道如何构建相应的 DFA，对于第一个集合使用 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$，对于第二个集合则对 $\mathtt{\text{".*"}}\backslash texttt\{".*"\}$ 正则表达式应用 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$。此外，与 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 返回的 DFA 不同，它们不会随着正则表达式变得更复杂而变得无法控制地大。$\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 不会向字符级 DFA 添加任何状态，而 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 结合 $\mathtt{\text{".*"}}\backslash texttt\{".*"\}$ 返回的 DFA 具有有趣的特性，使其易于计算和压缩（详见技术附录）。

一旦我们构建了这两个 DFA，我们就可以在解码时跟踪它们各自的状态，并且允许的标记是它们各自允许标记集合的交集。

然而，仅仅知道哪些标记在两个 DFA 中都是有效转换是不够的，因为我们还需要检查生成的标记最终是否能导致解码成匹配正则表达式的字符串的标记序列。这以前不是问题，因为 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 和 $\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 生成的标记级 DFA 的状态都是相关的（即，在起始状态和接受状态之间的一条路径上），只要初始字符级 DFA 只有相关状态。有效状态对可以通过从起始状态对开始，然后从接受状态对以及反向转换进行广度优先探索来计算。

速度够快吗？

表 2 报告了 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 各个步骤的执行时间，使用了 $\mathtt{\text{glaive-function-calling-v2}}\backslash texttt\{glaive-function-calling-v2\}$ 数据集中的 JSON 字符串样本。第一步的延迟可以忽略不计，因为此活动只需为给定的分词器执行一次，其结果将对所有未来的正则表达式都有用。第二步和第四步也只增加了很少的开销，而第三步是性能瓶颈。

如果正则表达式约束被大量使用，则第三步的延迟仍然有限，但当预计仅在解码时提供新的正则表达式时，它可能会成为一个问题。这通常发生在 LLM 提供商提供类似于 OpenAI 函数调用的功能时。然而，当真正的目标是遵守 JSON 或 YAML 模式时，我们可以显著缩短预处理时间。

在实际情况中，正则表达式的某些部分通常会被两个分隔符包围，并且无法与这些分隔符合并。这对于我们来说尤其有趣，因为我们知道表示这些部分的 DFA 部分不受正则表达式所有其他部分的影响。我们称这些情况为 Vegas 配置（因为分隔符之间发生的事情就保留在分隔符之间）。

例如，对于 Mistral-7B 分词器，任何以空格开头并由冒号和换行符分隔的字符串都是 Vegas 配置，这对于 YAML 模式当然非常有用。

Vegas 配置可以通过两种方式利用，如图 6 所示：

• 如果我们发现相同的 Vegas 配置在正则表达式中多次出现，我们只需计算并存储相应的 DFA 部分一次；
• 如果我们经常使用的正则表达式的可变部分对应于有限数量的 Vegas 配置，我们可以构建这些 Vegas 配置的库，存储相应的 DFA 部分，并在需要时直接使用它们。对于 JSON 或 YAML 模式，该库通常会包含常见类型（字符串、数字、日期等）的 Vegas 配置。

结论

在这篇博客文章中，我介绍了 $\mathtt{\text{Outlines}}\backslash texttt\{Outlines\}$ 的两种替代方案。$\mathtt{\text{DirectMerge}}\backslash texttt\{DirectMerge\}$ 和 $\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 也保证遵守正则表达式约束，同时加速解码并降低语言模型分布扭曲的风险。

这种扭曲发生的频率尚不清楚。然而，很容易找到它们的例子，如果它们在实践中发生，很可能不会被注意到。考虑到这种不确定性以及解码速度的提升，$\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 对于大多数受限解码需求来说，应该是一个无悔的选择。

除了受限解码，评估模型有时生成不正确标记序列是否会影响其性能也很有趣。如果确实如此，$\mathtt{\text{CartesianMerge}}\backslash texttt\{CartesianMerge\}$ 可用于强制执行正确的标记化。

感谢上述论文的作者，特别是 Luca Beurer-Kellner 的有益讨论，以及所使用的各种资源的贡献者（其中包括 $\mathtt{\text{outlines}}\backslash texttt\{outlines\}$、$\mathtt{\text{interegular}}\backslash texttt\{interegular\}$、$\mathtt{\text{transformers}}\backslash texttt\{transformers\}$ 和 $\mathtt{\text{glaive-function-calling-v2}}\backslash texttt\{glaive-function-calling-v2\}$）。

如果您想引用这篇博客文章，欢迎使用以下 BibTeX 条目

@misc{Tran-Thien_2024,
title={Fast, High-Fidelity LLM Decoding with Regex Constraints},
url={https://vivien000.github.io/blog/journal/llm-decoding-with-regex-constraints.html},
journal={Unsupervised Thoughts (blog)},
author={Tran-Thien, Vivien},
year={2024}
}

这篇博客文章最初发布在我的个人博客上.