Needle: 將 Gemini 工具調用蒸餾至 26M 參數模型

大型語言模型 (LLM) 開發的主流趨勢一直是追求更大的規模。然而，對於特定的代理任務——特別是工具調用 (tool calling)——龐大的模型往往是大材小用。Cactus Compute 透過發布 Needle 解決了這個問題，這是一個開源的 2600 萬參數模型，專門從 Gemini 蒸餾而成，旨在於手機、手錶和眼鏡等預算型消費級裝置上處理函數調用 (function calling)。

Needle 代表了一種哲學轉變：不再將工具調用視為複雜的推理任務，而是將其視為檢索與組裝的問題。透過專注於將查詢與工具名稱進行匹配，並將參數值提取到 JSON 中，開發者創建了一個極其快速且輕量化的模型，且不會犧牲代理行為的核心效用。

架構：超越 MLP

Needle 背後最重要的技術洞察之一是意識到在工具調用任務中，Feed-Forward Networks (FFNs) 在這種規模下很大程度上是被浪費的。在傳統的 Transformer 架構中，FFNs 用於存儲事實性知識。然而，在工具調用場景中，「事實」 (工具定義和可用函數) 是在輸入上下文中提供的。

為了優化這一點，Needle 使用了 Simple Attention Networks。整個模型由注意力機制和門控機制組成，完全移除了 MLPs。這種架構選擇是基於以下觀察：交叉注意力 (cross-attention) 是將用戶查詢與結構化工具定義進行匹配的理想原語 (primitive)。

這一發現對 AI 開發具有更廣泛的意義：對於任何模型可以訪問外部結構化知識的任務——例如檢索增強生成 (RAG) 或工具使用——模型不需要在其權重中記憶事實，這允許了顯著更小、更高效的架構。

性能與訓練

Needle 專為消費級硬體上的極致效率而設計，擁有 6,000 tokens/s 的 prefill 速度和 1,200 tokens/s 的 decode 速度。其訓練過程非常精簡：

  • Pretraining: 透過 16 個 TPU v6e 在 27 小時內完成 200B tokens 的預訓練。
  • Post-training: 透過 45 分鐘內完成 2B tokens 的合成函數調用數據的後訓練。
  • Data Synthesis: 數據集是透過 Gemini 生成的，涵蓋了 15 種工具類別，包括導航、智慧家庭控制、訊息傳遞和定時器。

在單次 (single-shot) 函數調用中，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 可以作為複雜代理系統中的「第一道程序」。它可以處理初始的工具選擇和參數提取，然後將結果移交給更大、更強大的模型進行最終的綜合與摘要。

批判性觀點與考量

儘管令人興奮，該項目也引發了關於法律和效用的重要討論：

  • 服務條款: 一些用戶指出，蒸餾 Gemini 可能會違反 Google 的服務條款，這些條款通常禁止使用其服務來開發競爭性模型。
  • 失敗模式: 關於模型如何處理歧義（例如，兩個相似的工具）或無法履行的請求，仍存在疑問。開發者鼓勵用戶透過其 playground 測試這些失敗模式。
  • 數據結構化: 有些人對訓練數據的特定結構化方式提出了疑問——例如，使用人類可讀的時間字符串而非電腦標準化格式——這表明可能需要第二個支架 (harness) 來完成輸出。

Needle 是更廣泛的 Cactus 專案的一部分，這是一個為移動和穿戴式硬體構建的推理引擎。透過證明 26M 參數模型可以有效地處理工具調用，Cactus Compute 正在挑戰業界，重新思考代理式 AI 所需的規模。

Sources