通过模型蒸馏优化 AI 函数费用

本文档介绍了如何在 BigQuery 中使用代管式 AI 函数的优化模式。您可以使用优化模式处理包含数千甚至数十亿行的大规模数据集,与标准逐行 LLM 推理相比,大语言模型 (LLM) 词元消耗和查询延迟时间显著降低。此优化仅适用于 AI.IF 和 AI.CLASSIFY 函数。

如果您想设置总 token 消耗量上限以防止预算超支,而不是优化查询的 token 效率,请参阅使用 token 配额控制费用。

为了更好地了解您的 token 消耗情况,您可以在Google Cloud 控制台中查看查询使用的 token 数量。如需在运行查询之前估算此用量,请使用 AI.COUNT_TOKENS 函数。

以下示例演示了如何将 AI.IF 函数与优化模式搭配使用,以识别有关自然灾害的新闻报道,并使用 text-embedding-005 作为嵌入模型:

SELECT
  title,
  body,
  AI.IF(
    ('The following news story is about a natural disaster: ', body),
    embeddings => AI.EMBED(body, endpoint => 'text-embedding-005', task_type => 'CLASSIFICATION').result,
    -- Optional, 'MINIMIZE_COST' is the default when embeddings are provided.
    optimization_mode => 'MINIMIZE_COST'
   ) AS is_natural_disaster
FROM
  `bigquery-public-data.bbc_news.fulltext`;

optimization_mode => 'MINIMIZE_COST' 参数会启用优化模式。这是提供嵌入时的默认设置,因此您可以省略此参数。

在此示例中,嵌入是实时生成的。在实践中,我们建议您具体化嵌入,以便重复使用。

优化模式的运作方式

托管式 AI 函数 AI.IF 和 AI.CLASSIFY 通常会针对数据集中的每一行调用远程 LLM。使用优化模式时,BigQuery 会在查询执行期间自动训练轻量级蒸馏模型。

该过程的工作原理如下:

启用优化模式时的 AI 函数工作流程

  • 抽样和添加标签:BigQuery 会选择一小部分代表性样本数据,并调用 Gemini 来提供标签。
  • 蒸馏模型训练:使用 LLM 标签和数据嵌入作为特征,实时训练本地蒸馏模型。
  • 质量检查:BigQuery 会根据 LLM 的结果评估提炼模型的准确率。默认情况下,如果提炼模型未能达到所需的质量阈值,查询会失败,并显示错误,说明模型被舍弃的原因。如果模型质量可接受,BigQuery 仍可能会针对特定行回退到远程 LLM,以保持一致的质量,或针对缺少有效嵌入的行回退到远程 LLM。
  • 推理:蒸馏模型会处理大部分行,从而显著减少 Gemini 调用次数。

限制

优化模式具有以下限制:

  • 行数下限:AI 函数的输入必须包含大约 3,000 行,以确保有足够的数据用于模型训练。
  • 数据类型:对于引用多个列的提示,仅支持对字符串列进行优化。
  • 多标签分类:优化模式不支持使用 output_mode => 'multi' 的 AI.CLASSIFY。
  • 函数支持:只有 AI.IF 和 AI.CLASSIFY 函数支持优化模式;但是,如果将优化模式与 AI.CLASSIFY 搭配使用,当提炼模型质量不足时,查询可能会失败。
  • 错误率:优化模式不支持 max_error_ratio 参数。

准备工作

如需获取在 BigQuery 中运行托管式 AI 函数所需的权限,请参阅为调用 Gemini Enterprise Agent Platform LLM 的生成式 AI 函数设置权限。

选择嵌入模型

如需使用优化模式,您必须为数据计算嵌入,并将其提供给 AI 函数。为了使输入列具有关联的嵌入,所有行都必须具有一致的嵌入维度,并由同一嵌入模型生成。

为了获得最佳的性价比和可伸缩性,我们建议您使用嵌入模型(例如 text-embedding-005 或 Gemini 嵌入)为英语或多语言任务计算数据的嵌入。对于多模态数据(文本和图片),我们建议使用多模态嵌入模型,例如 Gemini Embedding 2 (gemini-embedding-2)。

生成嵌入

您可以使用由 BigQuery 管理的自主生成功能或通过手动创建嵌入列来计算数据的嵌入。以下部分介绍了如何将这两种方法与 AI.CLASSIFY 和 AI.IF 函数结合使用。

自主嵌入生成

如果您使用自主嵌入生成,BigQuery 会在调用 AI.IF 或 AI.CLASSIFY 时自动使用嵌入。这是推荐的方法,但仅限于每个表一个嵌入列。

以下示例使用 text-embedding-005 作为嵌入模型,创建一个具有自主生成的嵌入列的表,然后使用 AI.CLASSIFY 函数对数据进行分类:

-- Create a table with an autonomously generated embedding column
CREATE TABLE my_dataset.bbc_news (
  title STRING,
  body STRING,
  body_embedding STRUCT<result ARRAY<FLOAT64>, status STRING>
    GENERATED ALWAYS AS (
      AI.EMBED(
        body,
        connection_id => '<my_connection_id>',
        task_type => 'CLASSIFICATION',
        endpoint => 'text-embedding-005')
    ) STORED
    OPTIONS(asynchronous = TRUE)
);

-- Insert data into the table
INSERT INTO my_dataset.bbc_news (title, body)
SELECT title, body FROM `bigquery-public-data.bbc_news.fulltext`;

-- Run the optimized query.
-- Wait for the background job to finish generating embeddings before running.
SELECT
  title,
  body,
  AI.CLASSIFY(
    body,
    categories => ['tech', 'sport', 'business', 'other']
  ) AS category
FROM
  my_dataset.bbc_news;

手动列规范

如果您已有嵌入列,请在 AI.IF 或 AI.CLASSIFY 的 embeddings 实参中指定该列。您可以使用 AI.EMBED 函数生成此值。

以下示例演示了如何使用 text-embedding-005 作为嵌入模型创建具有嵌入列的表,然后在 AI.CLASSIFY 查询中使用该列:

-- Create a table with an embedding column
CREATE TABLE my_dataset.bbc_news AS
SELECT
  title,
  body,
  AI.EMBED(
    body,
    endpoint => 'text-embedding-005',
    task_type => 'CLASSIFICATION'
  ).result AS body_embedding
FROM
  `bigquery-public-data.bbc_news.fulltext`;

-- Run the optimized query
SELECT
  title,
  body,
  AI.CLASSIFY(
    body,
    categories => ['tech', 'sport', 'business', 'other'],
    embeddings => body_embedding,
  ) AS category
FROM
  my_dataset.bbc_news;

如果提示引用了多个列,请在 embeddings 参数中提供列名称及其对应嵌入的列表。例如:embeddings => [('body', body_embedding), ('title', title_embedding)]。

监控查询优化

如需验证在查询执行期间优化了多少行,您可以在 Google Cloud 控制台中或通过 API 查看执行统计信息:

控制台

如需查看已优化的行数以及有关优化状态的系统消息,请执行以下操作:

  1. 在 Google Cloud 控制台中,前往 BigQuery 页面。

    转到 BigQuery

  2. 在导航菜单中,点击作业探索器。

  3. 点击任务 ID 以查看作业详情窗格。

  4. 点击作业信息标签页,然后在生成式 AI 函数优化字段中查看指标和状态。

    “作业信息”标签页中的“生成式 AI 函数优化”字段

API

检查作业元数据的 GenAIFunctionStats 对象中的 FunctionGenAiCostOptimizationStats。此对象包含通过优化后的工作流推断出的行数,以及提供优化状态数据洞见的系统生成的消息。

问题排查

以下部分介绍了如何诊断和解决使用优化模式时的常见问题。

数据规模过小

问题:模型训练数据不足。您可能会看到以下出错提示: Fail to apply cost optimizations because the number of rows to process is too small.

解决方案:将输入的大小至少增加到出错提示中显示的行数,并验证是否已为所有行正确生成有效的嵌入。

某些类别中的样本很少或没有样本

问题:在采样阶段,某些类别或用例的样本数量不足,导致无法进行模型训练。

  • 对于 AI.CLASSIFY 函数,您可能会看到以下出错提示:Failed to apply cost optimizations because the random training sample did not contain enough examples of all categories.
  • 对于 AI.IF 函数,您可能会看到以下出错提示:Failed to apply cost optimizations because the random training sample did not include enough examples of both true/false cases.

解决方案:

  • 对于 AI.CLASSIFY:
    • 从 AI.CLASSIFY 函数调用中移除罕见或空类别。
    • 将罕见类别归入更广泛的类别,以增加样本量。 您可以使用 OTHER 类别对未涵盖在更具体类别中的商品进行分组。不过,如果您的类别列表已经完整,请勿添加 OTHER,因为此字词含义模糊,可能会引起混淆。
  • 对于 AI.IF:更新提示的过滤条件,以便在输入行中充分体现 true 和 false 两种情况。如果条件几乎匹配所有行或几乎不匹配任何行,则可能会导致此错误。

嵌入的维度不一致

问题:各行之间的嵌入维度不一致。您可能会看到一条出错提示:Failed to apply cost optimizations because the embeddings have inconsistent dimensions.

解决方案:验证所有嵌入是否由同一模型生成,并且具有相同的嵌入向量长度。您可以使用类似于以下内容的 SQL 查询来检查列中的嵌入是否具有相同的长度:

SELECT ARRAY_LENGTH(body_embedding.result), COUNT(*)
FROM `PROJECT_ID.DATASET.TABLE_NAME`
GROUP BY 1;

提示过于复杂

问题:蒸馏模型无法达到高精确度阈值。您可能会看到以下某条出错提示:

  • Fail to apply cost optimizations because the distilled model does not approximate LLM behavior well enough for some categories.
  • Fail to apply cost optimization because the prompt complexity is too high.

解决方案:

  • 如果出错提示列出了 F1 得分较低的类别,请先重点关注这些类别。请优化其名称和说明,以便与其他类别明确区分。

  • 使用一组构成一个分区的类别。确保类别之间的重叠最小,并涵盖所有可能的输入。

    • 避免类别重叠,即一个输入可能同时属于多个类别。例如,避免使用 ['terrible', 'bad', 'okay', 'good', 'excellent'] 等类别。
    • 避免出现不适用任何类别的情况。例如,类别列表 ['bad', 'average'] 不涵盖表达赞美的评价。
    • 提供类别说明,以引导 LLM 消除类别之间的歧义。例如:

      AI.CLASSIFY(
        review,
        categories => [
          ('terrible', 'Review where customer was not happy and the message indicates they will never try this product again'),
          ('bad', 'Review where customer was not happy but suggested improvements to the product'),
          ('okay', 'Review where customer was neutral about the product. Short reviews qualify for this category'),
          ('good', 'Review where customers were happy using this product but had minor critiques'),
          ('excellent', 'Review where customers were very happy using this product and will recommend others to try it too')],
        embeddings => review_embeddings)
      
  • 尝试使用更高级的嵌入模型,例如 text-embedding-005 或 gemini-embedding-2。

  • 如需其他调试帮助,请联系 bqml-feedback@google.com。

LLM 处理的行数不符合预期

问题:查询执行统计信息显示,远程 LLM 处理的行数(而非提炼模型处理的行数)出乎意料地高。这可能是由于以下原因:

  • 提炼模型已成功训练,但某些行缺少嵌入。这些行由远程 LLM 处理。
  • 提炼模型无法应用于每一行,必须回退到远程 LLM 才能保持一致的质量。

解决方案:验证是否已正确生成嵌入,并且嵌入对数据中的所有行都有效。如果问题仍然存在,请联系 bqml-feedback@google.com 进行调试。

未检测到自主嵌入列

问题:BigQuery 无法检测到自主嵌入列。如果您的脚本使用临时表,并且对原始表的引用丢失,则可能会发生这种情况。

解决方案:使用 embeddings 参数显式传递自主嵌入列(例如 embeddings => content_embedding.result),这会触发费用优化。

后续步骤