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

社区文章发布于2024年2月23日

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

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

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

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

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

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

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

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

图1。左上：正则表达式 $\texttt{a*b*}$ 的字符级DFA，右上： $\texttt{Outlines}$ 为同一正则表达式返回的分词级DFA，下：解码时的分词掩码。

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

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

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

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

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

图2. $\texttt{Outlines}$ 为正则表达式 $\texttt{"boolean: ((true)|(false))"}$ 构建的DFA

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

图3. $\texttt{DirectMerge}$ 和 $\texttt{CartesianMerge}$ 为正则表达式 $\texttt{"boolean: ((true)|(false))"}$ 构建的DFA。

挑战 #1：概率分布偏差

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

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

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

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

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

挑战 #2：自我中毒

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

挑战 #3：错失加速机会

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

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

确保符合性与正确分词

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


初始字符串	$\texttt{aaaabaacac}$
转换为单字母分词序列	$\texttt{a a a a b a a c a c}$
$\texttt{merge}(\texttt{a}, \texttt{b})$ 后的分词序列	$\texttt{a a a }\underline{\texttt{ab}}\texttt{ a a c a c}$
$\texttt{merge}(\texttt{a}, \texttt{a})$ 后的分词序列	$\underline{\texttt{aa}}\texttt{ a ab }\underline{\texttt{aa}}\texttt{ c a c}$
$\texttt{merge}(\texttt{a}, \texttt{c})$ 后的最终分词序列	$\texttt{aa a ab aa c }\underline{\texttt{ac}}$

表1. 字符串 $\texttt{aaaabaacac}$ 使用合并表 $[(\texttt{a}, \texttt{b}), (\texttt{a}, \texttt{a}), (\texttt{a}, \texttt{c})]$ 的分词。合并操作的顺序和字符串中字母的顺序都会影响最终的分词序列。

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

$S$ 是否有一个 $x\:(x \neq a)$ 的传入转换，或者 $S$ 是否是起始状态？
$S$ 是否有 $y\:(y \neq b)$ 的出边转换，或者 $S$ 是否是接受状态？

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

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

图 4. 在 $(a, b)$ 合并操作且 $a \neq b$ 的情况下应用的转换示例。左： $\texttt{DirectMerge}$ 之前的情况，右： $\texttt{DirectMerge}$ 之后的情况。

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

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

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

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

图 5. 目标语言是两种正则语言的交集，我们已经知道如何构建它们的 DFA。

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

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

速度够快吗？

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

步骤	活动	平均持续时间（秒）	评论
1	使用 $\texttt{DirectMerge}$ 计算识别正确标记序列的 DFA（预处理）	20.9	对于给定的分词器仅执行一次
2	使用 $\texttt{Outlines}$ 计算识别解码为匹配正则表达式的字符串的标记序列的 DFA（预处理）	0.687	对于给定的正则表达式仅执行一次
3	识别相关状态对（预处理）	27.3	对于给定的正则表达式仅执行一次
4	构建标记掩码（解码）	0.00153	在每个解码步骤

表 2. $\texttt{CartesianMerge}$ 在 $\texttt{glaive-function-calling-v2}$ 数据集中 1000 个 JSON 字符串样本上的各个步骤的执行时间。

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

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

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

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

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

图 6. 与 Vegas 配置对应的 DFA，即由不能合并的两个分隔符包围的重复部分，只需计算并存储一次。上图： $\texttt{CartesianMerge}$ 的朴素实现；下图：利用 Vegas 配置的 $\texttt{CartesianMerge}$ 实现。

结论

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

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

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

感谢上述论文的作者，特别是 Luca Beurer-Kellner 的有益讨论，以及所使用的各种资源的贡献者（其中包括 $\texttt{outlines}$ 、 $\texttt{interegular}$ 、 $\texttt{transformers}$ 和 $\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}
}

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

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录以评论