Quail 如何用 Trie 树执行 AI.CLASSIFY 分类

返回博客 目录

1. 引言。
2. Quail 中的 AI.CLASSIFY。
3. How Quail Executes AI.CLASSIFY.
   1. trie_decode：贪心地生成标签。
   2. trie_tree：一次性为所有标签打分。
   3. letters：每个标签只打一个 token 的分。
   4. 为后续算子保留文档 KV。
4. Pricing and Picking a Method.
   1. 工作负载与符号约定。
   2. trie_tree 和 letters 的成本。
   3. trie_decode 的成本。
   4. 选择一种方法。
   5. 在 QUAIL-B 上的实验。
5. 附注：决策模型。
6. 结论。

1. 引言

我们最近发布了 Quail，一个用于 AI-SQL 的执行引擎。最常见的 AI-SQL 操作之一是分类：给定一个文档和一个标签列表，选出最合适的标签。例如，支持团队可能希望将每条收到的消息归类为“账单”、“退款请求”、“配送延迟”等等。

大多数系统对文档进行分类的方式与回答其他任何问题相同。它们把文档和标签放进一个提示词中，让 LLM 生成一个响应，然后解析该响应得到一个标签。BigQuery 的 AI.CLASSIFY 和 Snowflake 的 AI_CLASSIFY 的文档描述了类似的接口：每个函数都会向输入添加一个提示词，将其发送给 LLM，并返回最匹配的类别。但让 LLM 自由生成既浪费又脆弱——LLM 可能在响应上花费许多 token，而且响应甚至可能不包含有效的标签！

由于 Quail 在自己的推理引擎内运行 LLM，它不必让 LLM 自由生成。相反，Quail 可以将 LLM 限制在标签范围内，而事实证明有几种方法可以做到这一点，每种方法的成本各不相同。在本文中，我们讨论：对于给定的 LLM，Quail 应如何执行 AI.CLASSIFY？

我们将依次介绍：

• AI.CLASSIFY 在 Quail 中如何工作（第 2 节）。
• Quail 执行 AI.CLASSIFY 的三种方式：贪心地生成标签、一次性为所有标签打分，以及每个标签只打一个字母的分（第 3 节）。
• 如何估算每种方式的成本（基于我们关于 AI 驱动过滤器的博客文章），以及 Quail 如何选择最便宜的一种（第 4 节）。
• 关于决策模型的附注，决策模型是经过微调以在答案中进行选择的 LLM（第 5 节）。

2. Quail 中的 AI.CLASSIFY

AI.CLASSIFY 接收一个文档和一个标签列表。它返回其中一个标签。

以下是我们在整篇文章中使用的查询，使用 Qwen3-4B-fp8 作为 LLM。它将支持消息分类为 10 个主题：

SELECT m.id,
       AI.CLASSIFY(
         m.body,
         ARRAY['billing', 'refund request', 'refund status',
               'shipping delay', 'shipping damage',
               'shipping address change', 'product defect',
               'two-factor authentication issue', 'praise', 'other']
       ) AS topic
FROM messages m;

该函数有三个参数：

• input：要分类的文档。这里，输入是一条支持消息的文本。
• categories：2 到 255 个标签。每个标签可以有一个最多 25 个词的简短描述。
• options (not in the example query): an object with optional settings.
  • task_description：给 LLM 的额外指令，最多 50 个词。
  • probabilities：当 true 时，Quail 还会返回列 topic_probabilities，其中每个标签对应一个概率。这些概率之和为 1。
  • selectivity：对标签上的条件所保留行比例的估计，例如 WHERE topic = 'billing'。优化器使用该估计来规划查询。该估计不会改变标签。

Quail 始终返回其中一个标签。

AI.CLASSIFY 也可以出现在 WHERE 子句中。例如，以下查询只保留关于退款的消息：

SELECT m.id
FROM messages m
WHERE AI.CLASSIFY(m.body, ARRAY[...]) IN ('refund request', 'refund status');

对标签的条件使用 = 或 IN。当同一个分类同时出现在 SELECT 和 WHERE 中时，Quail 只运行一次该分类。

3. Quail 如何执行 AI.CLASSIFY

在数据库术语中，AI.CLASSIFY 是一个逻辑运算符。Quail 支持该运算符的三种物理实现，优化器会为每个查询选择一种实现。

这三种实现都会为每个文档向 LLM 发送一个提示。提示以 DOCUMENT: 和文档开头，然后是一个带有标签列表的问题。提示之后是 token ANSWER:，LLM 在此为答案的第一个 token 打分。三种实现的提示在问题和标签列表上略有不同，如 图 1 所示。

这些实现在选择标签的方式上也各不相同。trie_decode 逐个 token 生成标签，trie_tree 在一个请求中为每个标签打分，而 letters 为每个标签的一个字母打分。

3.1 trie_decode：贪心地生成标签

trie_decode 是最接近先前系统做法的实现。它像 LLM 生成文本一样，一次一个 token 地生成标签。¹ 在每个解码步骤，LLM 选择最有可能延续某个标签的 token。

从高层来看，我们可以把标签集合组装成一棵 trie（图 2）。trie 也称为前缀树，它将每个标签存储为从根开始的一条 token 路径。以相同 token 开头的标签共享路径的开头，例如两个退款标签都经过 refund。trie 的根是 ANSWER: token。在本文余下部分，我们用引号写标签，用代码格式写 token，例如标签“refund request”的 token 是 refund 和 request。

生成标签的简单方式是每个解码步骤生成一个 token。在每个解码步骤，LLM 读取到目前为止的 token，并选择 trie 中当前节点最有可能的子节点。例如，要生成“shipping address change”，LLM 可以经过 3 个解码步骤。它会在 ANSWER: 处选择 shipping，然后选择 address，再选择 change。

现在，你可能会意识到——如果答案确实是“shipping address change”，那么当我们处于第二个解码步骤（address）时，实际上根本不需要解码 change token，因为没有其他有效的 token！

更一般地说，LLM 只需要在 trie 中标签分叉的地方做出选择，即在有多个子节点的节点处。在我们的运行示例中，只有 3 个这样的节点：ANSWER:（第一个词是什么？）、refund（request 还是 status？）和 shipping（delay、damage 还是 address？）。在其他任何地方，下一个 token 都已经是确定的。因此 Quail 只在 3 个分叉节点处运行 LLM，并且一旦只剩下一个标签就停止。例如，如果 LLM 在 ANSWER: 处选择 product，答案必定是“product defect”，所以一步就完成了！

当我们一次对许多文档进行分类时，这会是什么样子？图 3 展示了一批 10 个文档的解码步骤。Quail 在一次前向传播中为批次中的所有文档一起计算每个步骤。步骤 0 与提示词在同一个请求中发生：LLM 读取文档并在 ANSWER: 处选择第一个 token。然后，之后的每个解码步骤再发送一个请求，但只针对那些标签尚未确定的文档。在我们的运行示例中，只有选择了 refund 或 shipping 的 5 个文档需要步骤 1，另外 5 个文档已经完成。

所以 trie_decode 相当便宜：如果每个标签等可能，我们运行示例中的一个文档平均只需要 1.5 个请求。但有一个问题——由于它是贪婪的，trie_decode 只查看自己路径上的 token。它从不给其他标签打分，因此无法返回标签概率。所以，如果用户要求概率，Quail 必须使用其他方法之一。²

3.2 trie_tree：一次性为所有标签打分

如果用户想要概率呢？一种简单的方法是在 LLM 下为每个标签计算其完整概率。一个标签的概率是其各个 token 概率的乘积。如果标签 \(m\) 的 token 为 \(x_1, \dots, x_t\)，那么

\[P(m) = \prod_{k=1}^{t} P(x_k \mid \text{prompt}, x_1, \dots, x_{k-1}).\]

例如，对于“refund request”，

\[P(\text{refund request}) = P(\texttt{refund} \mid \text{prompt}) \times P(\texttt{request} \mid \text{prompt}, \texttt{refund}).\]

然而，这种简单方法会不公平地惩罚较长的标签。每个额外的 token 都会将分数乘以一个小于 1 的数，即使标签已经确定。例如，“shipping address change”的分数会包含 change 的概率，尽管 change 是唯一能跟在 address 后面的 token。

所以在 trie_tree 中，我们只对每个标签打分，直到我们知道它是哪个标签为止。更准确地说，我们为 token 打分，直到该标签与所有其他标签分叉的那个词的末尾。设 \(d_m\) 为标签 \(m\) 中该词最后一个 token 的位置。那么

\[\text{score}(m) = \prod_{k=1}^{d_m} P(x_k \mid \text{prompt}, x_1, \dots, x_{k-1}).\]

为了说明，以下是运行示例中三个标签的分数。为了保持表格简短，我们在每个概率中省略了“prompt”：

“shipping address change”在“address”处与其他标签分叉，因为“delay”和“damage”也跟在 shipping 后面。“product defect”在“product”处分叉，因为没有其他标签以 product 开头。“two-factor authentication issue”在“two-factor”处分叉，这是一个包含两个 token 的词，所以我们对两个 token 都打分。

接下来，我们需要把分数转换为概率。这 10 个分数加起来不会等于 1，因为在每个 token 处，LLM 都会把它的概率分布到词表中的全部 151,936 个 token 上，包括那些不属于任何标签的 token，比如 The。所以 Quail 会把每个分数除以这 10 个分数之和：

\[\hat{P}(m) = \frac{\text{score}(m)}{\sum_{m'} \text{score}(m')}.\]

这里，\(\hat{P}(m)\) 是 Quail 为标签 \(m\) 返回的概率。

最后，我们必须处理一种情况：一个标签是另一个标签的开头。当前示例中没有这样的标签。假设我们有“refund”和“refund request”两个标签。在 refund 之后，LLM 可以停止，这意味着答案是“refund”，也可以继续输出 request。所以“refund”的分数等于 refund 的概率乘以不继续输出 request 的概率：

\[\text{score}(\text{refund}) = P(\texttt{refund} \mid \text{prompt}) \times \bigl(1 - P(\texttt{request} \mid \text{prompt}, \texttt{refund})\bigr).\]

否则，refund 的概率会被同时计入两个标签。

现在，回到当前示例。就像在 trie_decode 中那样，我们可以在每个有子节点的 trie 节点上运行 LLM，也就是 8 个节点。然而，与 trie_decode 类似，在每个这样的节点上运行 LLM 并没有达到它本可以达到的效率。你可能以为 trie_tree 只需要 3 个分支节点，就像 trie_decode 那样。但它还需要多一个 token，two！否则，“two-factor authentication issue”的分数也会包含其他以 two 开头的文本的概率，比如“two days ago”。所以我们要乘以 two 之后出现 -factor 的概率。因此，trie_tree 会在 4 个节点（ANSWER:、refund、shipping 和 two）上运行 LLM，所有这些都在与提示相同的请求中完成。

自定义注意力内核。最后，考虑如何在 GPU 上执行上述算法。一种简单的方法是向 LLM 发送 4 个请求，每个请求都包含相同的前缀（提示和标签描述），后面跟着一个不同的 trie 节点，用于计算其概率。当然，前缀的 KV 可以在 4 个请求之间共享，但每个请求仍然需要从 GPU 内存中读取整个前缀的 KV。相反，Quail 使用自定义的树注意力内核，将全部 4 个节点放在一个请求中发送给 LLM。³ 每个节点只关注提示以及自身路径上的节点，如图 4 所示。

总体而言，trie_tree 处理的 token 数总是至少与 trie_decode 一样多，因为它会在每个标签的分支节点上运行 LLM。然而，trie_tree 每个文档只需要一个请求。

3.3 letters：为每个标签评分一个 token

还有一种更简单的方法来获得概率：把每个标签都变成一个单独的 token。⁴ letters 在提示中给每个标签前面加一个字母（“A: billing”、“B: refund request”等等），并要求 LLM 返回字母而不是标签，如图 1 右侧所示。现在 Quail 只需要一步。在 ANSWER: 处，它读取这 10 个字母的概率，对它们重新缩放使其总和为 1，然后选择字母概率最高的标签。因此，不需要 trie，也不需要额外的请求。

缺点是提示更长，因为现在每个标签行都有一个字母。准确率也可能发生变化，因为 LLM 必须用字母而不是标签本身来回答。⁵

3.4 为后续算子保留文档 KV

到目前为止，我们只单独看了一次 AI.CLASSIFY 调用。但真实查询往往会对同一批文档运行多个 AI 算子。例如，一个查询可能先对每条消息的主题进行分类，只保留退款消息，然后再对它们的紧急程度进行分类。每个算子都要读取同一条消息，因此每次都从头处理这条消息会很浪费。

相反，Quail 将每个文档的 KV 存储在其 KV 缓存中，这样下一个算子只需处理自己的问题，即使这些算子具有不同的语义（例如过滤、分类和连接）。当然，GPU 内存是有限的，所以 Quail 无法永久保留每个文档的 KV。当内存填满时，Quail 会驱逐那些它预计每页节省工作量最少的 KV。⁶

4. 定价与选择方法

那么 Quail 应该使用哪种方法呢？简而言之，当 trie_decode 可以运行时，它通常是最快的，前提是 Quail 让每次前向传播都填满 token。但 trie_decode 并不总是最快的，而且用户可能还要求概率，这是 trie_decode 无法给出的。因此，这里我们将逐一分析每种方法的成本，看看差异来自哪里。

与 Quail 中的所有成本模型一样，我们的分类成本模型基于光速（SoL）估计，即 GPU 上该工作可能的最低耗时。我们关于 AI 驱动过滤器的博客文章详细解释了 SoL 估计，我们沿用其符号：\(N\) 是文档数量，\(\ell_i\) 是文档 \(i\) 中的 token 数，\(K\) 是前向传播的次数。

4.1 工作负载与符号

表 1. 分类成本模型中的符号。

让我们先建立一些符号。回顾第 3 节，Quail 向 LLM 为每个文档发送一个请求。请求包含前导部分（文本 DOCUMENT:）、文档、指令（问题和标签列表），然后是答案位置。答案位置是指令之后的一个 token，Quail 在此读取 LLM 对下一个 token（如 ANSWER: token）给出的分数。

设 \(q_{\text{pre}}\) 和 \(q_{\text{tail}}\) 分别为前导部分和指令中的 token 数。设 \(a\) 为每个文档的预期答案位置数，设 \(s\) 为 Quail 读取分数的预期答案位置数。那么请求长度、总 token 数和前向传播次数分别为

\[\begin{aligned} r_i &= q_{\text{pre}} + \ell_i + q_{\text{tail}} + a, \\ n_{\text{tok}} &= L_1 + N(q_{\text{pre}} + q_{\text{tail}} + a), \\ K &= \left\lceil \frac{n_{\text{tok}}}{C} \right\rceil. \end{aligned}\]

这里，\(L_1 = \sum_i \ell_i\) 是文档 token 的总数，而 \(C\) 是一次前向传播中 token 的最大数量。

这三种方法仅在 \(q_{\text{tail}}\) 和 \(a\) 上有所不同。表 2 给出了我们运行示例中的值。

表 2. 运行示例中每个文档的指令长度和预期答案位置。

对于投影、注意力和 MLP，我们可以复用我们关于 AI 驱动过滤器的博客文章中的公式，代入新的 \(n_{\text{tok}}\) 和 \(r_i\)。然而，分类与过滤器有两点不同：(1) 每种方法的 KV 读取次数不同，即加载的缓存 KV token 数量不同；(2) 分类需要整个输出头，即整个词汇表的分数，而不仅仅是 TRUE 和 FALSE。

KV 读取。 每个答案位置都会关注提示，因此 GPU 必须从 HBM 加载提示的 KV。letters 和 trie_tree 每个文档加载一次，因为它们的所有答案位置都在一个请求中。trie_decode 每个额外的解码步骤加载一次。步骤 0 共享提示的请求，但之后的每个解码步骤都是一个新请求。因此，如果 \(b\) 是 trie_decode 每个文档的预期请求数（包括步骤 0），则加载的 KV token 数量为

\[\begin{aligned} R_{\text{letters}} = R_{\text{tree}} &= \sum_i (q_{\text{pre}} + \ell_i + q_{\text{tail}}), \\ R_{\text{decode}} &\approx (b - 1) \sum_i (q_{\text{pre}} + \ell_i + q_{\text{tail}}). \end{aligned}\]

将 \(R\) 代入我们关于 AI 驱动过滤器的博客文章中的公式，注意力会移动 \(B_{\text{attn}} = B_{\text{kv}}(W + R)\) 字节。这里，\(B_{\text{kv}}\) 是所有层中一个 token 的 KV 字节数，而 \(W\) 是新 KV 注意力写入的 token 数量。

输出头。 在最后一层之后，每个答案位置都有一个隐藏状态，它是一个包含 \(d_{\text{model}}\) 个值的向量。LLM 通过再一次矩阵乘法将隐藏状态转换为词汇表中每个 token 的分数。我们将该矩阵称为输出头。⁷ 输出头为每个词汇 token 有一行，因此一个 token 的分数就是隐藏状态与该 token 所在行的点积。

设 \(V_{\text{out}}\) 为 Quail 在每个答案位置计算的行数。如果每个标签都是单个 token（如 letters 中的字母），Quail 只计算标签 token 的行。否则，Quail 计算所有 \(\text{vocab}\) 行，这是将分数转换为概率所必需的。由于输出头采用 BF16，

\[\begin{aligned} F_{\text{head}} &= 2\,d_{\text{model}}\,V_{\text{out}}\,N a, \\ B_{\text{head}} &= 2\,d_{\text{model}}\,V_{\text{out}}\,K, \\ T_{\text{head}} &= \max\!\left(\frac{F_{\text{head}}}{\Pi_{bf16}}, \frac{B_{\text{head}}}{\beta}\right). \end{aligned}\]

这里，\(F_{\text{head}}\) 计算每个答案位置上每个计算行的每个值两次 FLOPs，而 \(B_{\text{head}}\) 计算每次前向传播中每个计算行的每个值两个字节。\(\Pi_{bf16}\) 和 \(\beta\) 分别是 GPU 的峰值 BF16 吞吐量和 HBM 带宽。对于 Qwen3-4B，\(d_{\text{model}} = 2{,}560\) 且 \(\text{vocab} = 151{,}936\)。

一次分类的 SoL 估计为

\[T_{\text{classify}} = T_{\text{proj}} + T_{\text{attn}} + T_{\text{mlp}} + T_{\text{head}}.\]

这里，\(T_{\text{proj}}\)、\(T_{\text{attn}}\) 和 \(T_{\text{mlp}}\) 分别是投影、注意力和 MLP 的估计延迟。

4.2 trie_tree 和 letters 的成本

trie_tree 和 letters 都将它们所有的答案位置放在与提示相同的请求中，因此它们只在每个文档添加的 token 数量上有所不同。

对于 trie_tree，每个答案位置是额外的一个 token，会经过整个 LLM，这在投影和 MLP 中花费约 72 亿 FLOPs，就像文档 token 一样。由于某些标签有多个 token，Quail 还会在每个答案位置计算完整的输出头，这增加了约 7.8 亿 FLOPs。在我们的运行示例中，trie_tree 有 4 个答案位置，因此每个文档增加约 320 亿 FLOPs（\(a = 4\) 且 \(V_{\text{out}} = \text{vocab}\)）。

对于 letters，额外的 token 在提示中。letters 向指令添加 \(2M + 3\) 个 token，其中 \(M\) 是标签数量。每个标签行因其字母获得 2 个 token，问题因“the letter of”这些词获得 3 个 token。在我们的运行示例中，letters 添加 23 个 token，每个文档花费约 1660 亿 FLOPs。输出头很便宜，因为 Quail 只计算 letters 的 10 行（\(a = 1\) 且 \(V_{\text{out}} = M\)）。因此，在我们的运行示例中，letters 的额外工作量约为 trie_tree 额外工作量的 5 倍。然而，两者与文档本身相比都很小，对于 1000 token 的文档，文档本身花费约 7 万亿 FLOPs。

4.3 trie_decode 的成本

与其他两种方法不同，trie_decode 的成本取决于 LLM 选择哪个标签，因为某些标签需要比其他标签更多的解码步骤。Quail 在运行查询之前不知道标签，因此成本模型假设每个标签的可能性相同。⁸

标签路径上的每个分支节点花费一个请求。因此，如果 \(c_m\) 是标签 \(m\) 路径上的分支节点数，包括决定该标签的那个节点，则每个文档的预期请求数为

\[b = \mathbb{E}[\text{requests per document}] = \frac{1}{M} \sum_{m=1}^{M} c_m.\]

在我们的运行示例中，5 个标签在 ANSWER: 处直接决定，另外 5 个需要一个额外的分支节点（refund 或 shipping），如 图 3 所示。因此

\[b = \frac{5 \cdot 1 + 5 \cdot 2}{10} = 1.5.\]

在我们的运行示例中，每个解码步骤在提示后添加一个 token：步骤 0 中的 ANSWER:，以及步骤 1 中的 refund 或 shipping。因此 trie_decode 在提示之外，平均每个文档运行 LLM 处理 1.5 个 token。与 trie_tree 类似，每个 token 花费约 80 亿 FLOPs（投影和 MLP 中 72 亿，加上完整输出头的 7.8 亿）。因此 trie_decode 每个文档仅增加约 120 亿 FLOPs，而 trie_tree 约为 320 亿。

注意，该估计假设了最优打包，即 Quail 将许多文档打包在一起，使每次前向传播都能充分利用 GPU。⁹ 如果只有少量文档，后续的解码步骤可能在几乎空的前向传播中运行，此时 trie_decode 的开销会高于估计值。

4.4 选择方法

综合来看，优化器会为每个可运行的方法计算 \(T_{\text{classify}}\)，并选择估计值最低的方法。

在我们的运行示例中，trie_decode 胜出，因为它对每篇文档运行的 LLM 额外 token 最少：1.5 个，而 trie_tree 为 4 个，letters 为 24 个。letters 仅在少数标签共享较长起始部分时胜出，例如 2 个标签仅最后一个词不同，因为两种 trie 方法都必须对每个共享 token 运行 LLM。而 trie_tree 主要在查询要求概率时胜出，因为此时 trie_decode 无法运行。

注意，差异很小，对于 1,000 篇 100 token 的文档，差异在 2% 到 11% 之间，因为所有三种方法都必须预填充前导、文档和指令，而预填充占 \(T_{\text{classify}}\) 的大部分。

4.5 QUAIL-B 上的实验

为了端到端地观察这一切如何运作，我们向 QUAIL-B（我们用于 AI-SQL 的基准）添加了 12 个分类查询。这些查询跨越五个数据集（电影评论、不良药物事件报告、事实核查声明、法律引用和智能体轨迹），每个 AI.CLASSIFY 运算符有 4 到 27 个标签。有些查询只有一个运算符（AI.CLASSIFY）；其他查询还有过滤器、连接或多个不同的 AI.CLASSIFY 运算符。

我们在单张 H100 上以缩放因子 0.5 使用 Qwen3-4B-fp8 运行每个查询，分别在 Quail 和 vLLM 基线（v0.26.0，带前缀缓存）上运行。在 vLLM 基线中，LLM 生成标签，然后我们解析它。图 6 显示了结果。

Quail 在所有 12 个查询上都更快：总速度快 2.1 倍，每个查询快 1.8 倍（几何平均）。 最大的优势来自多次读取相同文档的查询。例如，AGENT-5 对相同的智能体轨迹有三个不同的 AI.CLASSIFY 运算符。vLLM 的前缀缓存会驱逐最近最少使用的 KV，因此当下一个运算符读取轨迹时，其 KV 通常已被驱逐，vLLM 必须重新计算。在 AGENT-5 上，vLLM 计算的 token 中有 64% 是重新计算。而 Quail 会为后续运算符保留轨迹的 KV（第 3.4 节）。Quail 在每个查询上的运行速度也在其 SoL 估计的 1.8 倍到 2.8 倍以内。

5. 附注：决策模型

上个月，TypeSafe AI 发布了 Jev，一个决策模型。Jev 不生成文本，而是回答一个带类型的问题（例如是或否，或从几个标签中选择一个），并为每个答案返回一个概率。Jev 席卷了社区，开源后续项目在几周内就出现了：Jared Palmer 的 Kev、Cloudflare 的 Clef，以及 vLLM Semantic Router 团队的 Decision models。Databricks 甚至在 SQL 中运行开源决策模型。Quail 支持其中一个开源模型 Decision 2.0 Kai 0.6B，作为上述三种方法的替代方案。

Decision 2.0 Kai 是 Qwen3-0.6B 的微调版本。它从不生成文本；它只是在一次前向传播中为每个答案返回一个概率。

让我们简要介绍一下它的工作原理，如 图 5 所示。Quail 没有在指令中列出标签，而是将每个标签写入文档之后各自的选项块中。LLM 像往常一样在整个提示上运行，每个 token 从最后一层输出为隐藏状态。然后，一个额外的决策头读取每个选项块 \(m\) 最后一个 token 的隐藏状态 \(h_m\)，以及提示最后一个 token 的隐藏状态 \(q\)。

决策头首先重新缩放 \(h_m\) 和 \(q\)，使它们的值处于相近的尺度，得到 \(\bar{h}_m\) 和 \(\bar{q}\)。然后，它为每个标签 \(m\) 计算一个分数 \(s_m\)，并通过 softmax 将分数转换为概率：

\[\begin{aligned} s_m &= \frac{(W_h \bar{h}_m) \cdot (W_q \bar{q})}{\sqrt{d}} + w \cdot \operatorname{GELU}\!\left(U_h \bar{h}_m + U_q \bar{q} + b\right), \\ P(m) &= \frac{e^{s_m}}{\sum_{j=1}^{M} e^{s_j}}. \end{aligned}\]

第一项是点积，类似于注意力分数，第二项是一个小型神经网络。\(W_h\)、\(W_q\)、\(U_h\)、\(U_q\)、\(w\) 和 \(b\) 是微调期间学习到的权重。\(W\) 和 \(U\) 是 \(d \times d_{\text{model}}\) 矩阵，将隐藏状态缩小到 \(d\) 个值，而 \(w\) 和 \(b\) 是 \(d\) 个值的向量，将 GELU 输出转换为标量。¹⁰ 对于 Decision 2.0 Kai，整个决策头约有 100 万个权重，而 LLM 中约有 6 亿个，因此其成本几乎可以忽略不计。

与 trie_tree 类似，Decision 2.0 Kai 在一次请求中给出概率。但 trie_tree 必须在每个答案位置对词汇表中的每个 token 进行评分，而 Decision 2.0 Kai 只对 \(M\) 个标签进行评分。

有人可能会注意到，选项块增加了大量 token：在我们的运行示例中，每个标签约 16 个 token，每篇文档约 190 个 token，而 letters 为 24 个。但额外的 token（更具体地说，是训练好的决策头）带来了大量准确率提升。我们在一个 H100 上对 QUAIL-B 中的 12 个分类查询运行了 Decision 2.0 Kai 和基础 Qwen3-0.6B。Decision 2.0 Kai 平均在 63.1% 的文档上与参考标签一致，而 Qwen3-0.6B 仅为 39.0%！我们对这些决策模型感到非常兴奋，并相信我们可以为 AI-SQL 任务微调我们自己的决策模型（甚至可能是更小的模型）。

6. 结论

在本文中，我们介绍了执行 AI.CLASSIFY 的多种方式。由于 Quail 在自己的推理引擎内运行 LLM，它可以将 LLM 限制在标签上：它可以贪婪地生成标签（trie_decode）、一次性对所有标签评分（trie_tree），或为每个标签评分一个字母（letters），并选择对给定 LLM 和标签来说最便宜的方式。大多数情况下，答案是 trie：默认使用 trie_decode，当查询要求概率时使用 trie_tree。

接下来，我们正在努力支持更大的决策模型（基于混合架构构建）以及像 AI.EXTRACT 这样的新算子。敬请期待更多精彩内容！同时，试试 Quail，如果你喜欢我们正在构建的东西，请在 GitHub 上给它一个 star！

致谢

我们感谢 Modal 赞助本研究所使用的算力。

注释

1. 此前来自学术界的开源 LLM 驱动的数据处理系统，例如 DocETL、LOTUS 和 Palimpzest，同样让 LLM 逐 token 生成答案。

2. 此外，如果某个标签是另一个标签的精确前缀，我们还必须使用其他方法之一。例如，假设标签是“refund”和“refund request”。在 refund 之后，该标签可能已经完整，也可能继续为 request。trie_decode 选择下一个 token，但我们没有定义用于停止的 token。因此 trie_decode 永远无法选择标签“refund”。

3. 我们介绍 Quail 的博客文章描述了用于 AI 连接的树注意力内核。在一次连接中，许多伙伴文档共享一个锚文档。在 trie_tree 中，许多 trie 节点共享一个提示。

4. vLLM 中 DiffusionGemma 的结构化模式使用了相同的思路。它只接受单 token 选项，并将较长的选项映射为字母，例如将“moderation_spam”映射为“A”（vLLM PR #57250）。

5. 我们不知道哪种方法在总体上最准确。在我们的两个 QUAIL-B 查询上，trie_tree 和 trie_decode 与参考标签的一致频率大致相同（在 BIO-5 上为 72.1% 对 72.0%，在 AGENT-5 上相差不超过 0.2 个百分点）。就本文而言，我们假设这三种方法的准确率大致相同，并且我们预计强大的开源 LLM 会使这三种方法对大多数查询都足够好。

6. Quail 将某一页的节省估计为后续算子读取该文档的概率，乘以再次处理该文档的时间，再除以该文档的页数。参见 Quail 关于 KV 保留的文档。

7. 我们关于 AI 驱动过滤器的博客文章在其成本模型中忽略输出头的成本是可以的，因为过滤器只需要 TRUE 和 FALSE 的行，当我们把 True 这类拼写也算上时最多 8 行，而不是全部 151,936 行。这 8 行每文档约花费 41,000 FLOPs，而每个文档 token 在投影和 MLP 中约花费 72 亿 FLOPs。

8. 真实标签通常并非等概率，因此对 \(b\) 的估计可能过高或过低。例如，假设大多数文档得到的标签在 ANSWER: token 处就已确定。那么 trie_decode 发送的请求会少于估计值，因此 trie_decode 的 SoL 估计值过高。

9. 这与 vLLM 等在线服务引擎相当不同。vLLM 关心每个请求的延迟，因此它会尽快运行解码步骤，即使是在大部分为空的前向传播中。Quail 不关心每个单独文档的延迟，只关心整个查询的延迟，因此它总能填满前向传播。

10. 我们找到了 decision_model.py 中描述的决策头，即 Hugging Face 上随模型一起发布的代码。

引用本文

@misc{shankar2026aiclassify,
  title = {To Classify, Try a Trie Tree},
  author = {Shankar, Shreya and Frye, Charles},
  year = {2026},
  month = oct,
  url = {https://fsdatalab.github.io/blog/ai-classification/}
}