Needle: Geminiのツール呼び出しを26Mパラメータモデルに蒸留する

大規模言語モデル(LLM)開発の主流な傾向は、より大規模なモデルへの競争となってきました。しかし、特定の自律的なタスク、特にツール呼び出しにおいては、巨大なモデルはしばしば過剰です。Cactus Computeは、スマートフォン、時計、メガネなどの低価格な消費者向けデバイスで関数呼び出しを処理するために、Geminiから特別に蒸留されたオープンソースの2600万(26M)パラメータモデルであるNeedleをリリースすることで、この問題に対処しました。

Needleは、哲学的な転換を象徴しています。ツール呼び出しを複雑な推論タスクとしてではなく、検索と組み立ての問題として扱うという転換です。クエリをツール名に一致させ、引数の値をJSONに抽出することに焦点を当てることで、開発者は、エージェント的な動作の核となる有用性を損なうことなく、非常に高速で軽量なモデルを作成しました。

アーキテクチャ: MLPを超えて

Needleの背後にある最も重要な技術的な洞察の一つは、フィードフォワードネットワーク(FFN)が、ツール呼び出しタスクにおけるこの規模では大部分が浪費されているという気づきです。従来のTransformerアーキテクチャでは、FFNは事実的な知識を保存するために使用されます。しかし、ツール呼び出しのシナリオでは、「事実」は入力コンテキスト内で提供されます。

これを最適化するために、NeedleはSimple Attention Networksを利用しています。モデル全体がアテンションとゲーティングメカニズムで構成されており、MLPを完全に排除しています。このアーキテクチャの選択は、ユーザーのクエリを構造化されたツール定義に一致させるための理想的なプリミティブ(基本要素)はクロスアテンションであるという観察に基づいています。

この発見は、AI開発におけるより広範な示唆を与えています。Retrieval-Augmented Generation(RAG)やツール使用のように、モデルが外部の構造化された知識にアクセスできるタスクにおいては、モデルは自身の重みに事実を記憶させる必要はなく、これにより大幅に小型化し、より効率的なアーキテクチャが可能になります。

パフォーマンスとトレーニング

Needleは消費者向けハードウェアでの極限の効率性を実現するために設計されており、プリフィル速度は6,000 tokens/s、デコード速度は1,200 tokens/sを誇ります。そのトレーニングプロセスは驚くほど軽量でした。

  • Pretraining: 16台のTPU v6eを使用して27時間で200B tokens。
  • Post-training: 45分間で2B tokensの合成された関数呼び出しデータ。
  • Data Synthesis: データセットはGeminiを介して生成され、ナビゲーション、スマートホーム制御、メッセージング、タイマーなど15のツールカテゴリをカバーしています。

シングルショットの関数呼び出しにおいて、NeedleはFunctionGemma-270M、Qwen-0.6B、Granite-350Mを含むいくつかのより大きなモデルを凌駕しています。これらの大きなモデルは、会話の流暢さと一般的な能力において優位性を維持していますが、Needleはツールを起動する特定の動作に最適化されています。

実世界でのアプリケーションとユースケース

コミュニティは、この規模のモデルに対して、特にレイテンシとプライバシーが重要となる環境において、いくつかの影響力の大きいアプリケーションを特定しています。

1. デバイス上での仮想アシスタント

Needleは低価格なハードウェアで動作できるため、「Siriのような」コアとして理想的な候補です。書き起こされた音声テキストとシステムアクション(例:「10分間のタイマーをセットして」)の間の架け橋として機能することができます。

2. 自然言語コマンドラインインターフェース (CLIs)

開発者は、Needleを使用して、引数を自然言語で指定できるCLIツールを作成することを提案しています。複雑なフラグを覚える代わりに、ユーザーは toolcli add tom to teamfutz group と入力し、Needleはこれを toolcli --gadd teamfutz tom に解析します。

3. スマートホーム統合

Home Assistantのようなプラットフォームでは、Needleは、クラウドサーバーに音声データを送信したり、ホームサーバーで巨大な9Bパラメータモデルを動かしたりすることなく、照明を切り替えたり家電を制御したりするための、低レイテンシでローカルな方法を提供できます。

4. エージェント的なオーケストレーション

Needleは、複雑なエージェントシステムにおける「最初のパス」として機能できます。最初のツール選択と引数の抽出をハンドルし、その結果を結果を最終的な合成と要約のために、より大規模でより有能なモデルに渡すことができます。

批判的な視点と検討事項

期待が高まる一方で、プロジェクトは、合法性と有用性に関する重要な議論を呼び起こしています。

  • Terms of Service: 一部のユーザーは、Geminiを蒸留することがGoogleの利用規約(Terms of Service)に違反する可能性があると指摘しています。これは、通常、彼らのサービスを使用して競合モデルを開発することを禁止しています。
  • Failure Modes: モデルが曖昧さ(例:2つの類似したツール)や、実行できないリクエストをどのように処理するかについては、疑問が残っています。開発者は、ユーザーがプレイグラウンドを通じてこれらの失敗モードをテストすることを推奨しています。
  • Data Structuring: 一部のユーザーは、トレーニングデータの特定の構造化方法(例:コンピュータ標準形式ではなく、人間が読みやすい時間文字列を使用すること)に疑問を疑問を呈しており、出力を最終化するために二次的なハーネスが必要になる可能性があることを示唆しています。

Needleは、モバイルおよびウェアラブルハードウェア向けにゼロから構築された推論エンジンである、より広範なCactusプロジェクトの一部です。26Mパラメータモデルがツール呼び出しを効果的に処理できることを証明することで、Cactus Computeは、業界に対して、エージェント的なAIに必要とされる規模の検討を再考させるよう促しています。

Sources