Needle: 将 Gemini 工具调用能力蒸馏至 26M 参数模型

大型语言模型 (LLM) 开发的主流趋势一直是追求更大的规模。然而,对于特定的智能体任务——特别是工具调用——庞大的模型往往是大材小用。Cactus Compute 通过发布 Needle 解决了这一问题,这是一个开源的、拥有 2600 万参数的模型,专门从 Gemini 蒸馏而来,旨在处理手机、手表和眼镜等廉价消费级设备上的函数调用。

Needle 代表了一种理念的转变:不再将工具调用视为复杂的推理任务,而是一个检索与组装的问题。通过专注于将查询与工具名称进行匹配并将参数值提取为 JSON,开发者创建了一个极其快速且轻量级的模型,且没有牺牲智能体行为的核心效用。

架构:超越 MLP

Needle 背后的一个最重要的技术洞察是:意识到在工具调用任务中,这种规模下的前馈网络 (FFNs) 在很大程度上是浪费的。在传统的 Transformer 架构中,FFNs 用于存储事实性知识。然而,在工具调用场景中,“事实” (工具定义和可用函数) 是在输入上下文中提供的。

为了优化这一点,Needle 利用了 Simple Attention Networks。整个模型由注意力机制和门控机制组成,完全移除了 MLP。这种架构选择基于这样一个观察:交叉注意力 (cross-attention) 是将用户查询与结构化工具定义进行匹配的理想原语。

这一发现对 AI 开发具有更广泛的启示:对于任何模型可以访问外部结构化知识的任务——例如检索增强生成 (RAG) 或工具使用——模型不需要在其权重中记忆事实,从而允许使用显著更小、更高效的架构。

性能与训练

Needle 专为消费级硬件上的极致效率而设计,其 prefill 速度达到 6,000 tokens/s,decode 速度达到 1,200 tokens/s。其训练过程非常精简:

  • 预训练 (Pretraining): 在 16 个 TPU v6e 上花费 27 小时处理了 200B tokens。
  • 训练后处理 (Post-training): 在 45 分钟内使用 2B tokens 的合成函数调用数据进行训练。
  • 数据合成 (Data Synthesis): 数据集是通过 Gemini 生成的,涵盖了包括导航、智能家居控制、消息传递和定时器在内的 15 个工具类别。

在单次函数调用 (single-shot function calling) 中,Needle 的表现优于几个更大的模型,包括 FunctionGemma-270M、Qwen-0.6B 和 Granite-350M。虽然这些较大的模型在对话流畅度和通用能力方面保持优势,但 Needle 针对特定的工具触发行为进行了优化。

现实世界应用与用例

社区已经为这种规模的模型确定了几个具有高影响力的应用,特别是在对延迟和隐私要求极高的环境中:

1. 设备端虚拟助手

由于 Needle 可以运行在廉价硬件上,它是“类 Siri”核心的理想候选者。它可以作为转录语音文本与系统操作(例如,“设置一个 10 分钟的定时器”)之间的桥梁。

2. 自然语言命令行界面 (CLIs)

开发者建议使用 Needle 来创建 CLI 工具,其中参数可以通过自然语言指定。用户无需记忆复杂的 flag,只需输入 toolcli add tom to teamfutz group,Needle 就会将其解析为 toolcli --gadd teamfutz tom

3. 智能家居集成

对于 Home Assistant 等平台,Needle 可以提供一种低延迟、本地化的方式来切换灯光或控制家电,而无需将语音数据发送到云端服务器,也无需在家庭服务器上运行庞大的 9B 参数模型。

4. 智能体编排 (Agentic Orchestration)

Needle 可以作为复杂智能体系统中的“第一道工序”。它可以处理初始的工具选择和参数提取,然后将结果移交给一个更大、更强大的模型进行最终的合成与总结。

批判性视角与考量

尽管令人兴奋,但该项目引发了关于法律和效用的重要讨论:

  • 服务条款 (Terms of Service): 一些用户指出,蒸馏 Gemini 可能违反了 Google 的服务条款,这些条款通常禁止使用其服务来开发竞争性模型。
  • 失败模式 (Failure Modes): 关于模型如何处理歧义(例如,两个相似的工具)或无法完成的请求,仍存在疑问。开发者鼓励用户通过其 playground 进行测试。
  • 数据结构化 (Data Structuring): 有些人对训练数据的特定结构化方式提出了质疑——例如,使用人类可读的时间字符串而非计算机标准格式——这表明可能需要一个辅助工具来完成最终输出。

Needle 是更广泛的 Cactus 项目的一部分,这是一个专为移动和可穿戴硬件构建的推理引擎。通过证明一个 26M 参数的模型可以有效地处理工具调用,Cactus Compute 正在挑战行业,要求重新思考智能体 AI 所需的规模。

Sources