CactusはNeedleをオープンソース化しました。これは2600万パラメータの関数呼び出しモデルで、格安スマートフォン、スマートウォッチ、メガネ上での動作を想定しています。独自の推論エンジンCactusを使用することで、コンシューマデバイス上で6000 tok/sのプリフィルと1200 tok/sのデコードを達成します。

アーキテクチャ: Simple Attention Networks

NeedleはSimple Attention Networkを採用しており、MLPは一切使用していません。モデル全体はアテンション層とゲート層のみで構成されています。主要な設計は以下の通りです: d=512, 8H/4KV, BPE=8192。エンコーダ・デコーダ構造（12エンコーダ層、8デコーダ層）で、クロスアテンション、RoPEを使用したマスク付き自己アテンション、重み共有の埋め込みを採用しています。

トレーニング詳細

• 16個のTPU v6eを用いて200Bトークンで事前トレーニング（27時間）
• 合成された2Bトークンの関数呼び出しデータで追加トレーニング（45分）
• データはGeminiを用いて15のツールカテゴリ（タイマー、メッセージング、ナビゲーション、スマートホームなど）で合成

ベンチマーク結果

Needleはシングルショット関数呼び出しにおいて、FunctionGemma-270M、Qwen-0.6B、Granite-350M、LFM2.5-350Mを上回ります。ただし、これらのモデルはより広いスコープと容量を持ち、会話型設定で優れています。

クイックスタート

git clone https://github.com/cactus-compute/needle.git
cd needle && source ./setup
needle playground

http://127.0.0.1:7860でWeb UIが開き、テストや独自ツールでのファインチューニングが可能です。

使用例（Python）

from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer

params, config = load_checkpoint("checkpoints/needle.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()
result = generate(
    model, params, tokenizer,
    query="What's the weather in San Francisco?",
    tools='[{"name":"get_weather","parameters":{"location":"string"}}]',
    stream=False
)
print(result)
[{"name":"get_weather","arguments":{"location":"San Francisco"}}]

ローカルでのファインチューニング

# playground経由（Geminiで自動データ生成）
needle playground
または独自データを提供
needle finetune data.jsonl

提供情報

重みはHugging Faceで公開されています: Cactus-Compute/needle。すべてMITライセンスです。