26M参数小模型让AI助手直接调用工具，还能本地微调

我们将Gemini 3.1蒸馏为一个2600万参数的“简单注意力网络”，你甚至可以在Mac/PC上本地微调它。在生产环境中，Needle在Cactus上以每秒6000个token的预填充速度和1200个token的解码速度运行。权重完全开源在Cactus-Compute/needle，数据集生成代码也已公开。d=512，8H/4KV，BPE=8192。

┌──────────────┐
│ 工具调用 │
└──────┬───────┘
┌┴──────────┐
│ Softmax │
└─────┬─────┘
┌─────┴─────┐
│ 线性层 (T)│ ← 权重共享
└─────┬─────┘
┌─────┴─────┐
│ ZCRMSNorm │
└─────┬─────┘
┌────────┴────────┐
│ 解码器 x 8 │
│┌───────────────┐│
││ ZCRMSNorm ││
││ 掩码自注意力 ││
││ + RoPE ││
││ 门控残差连接 ││
│├───────────────┤│
┌──────────────┐ ││ ZCRMSNorm ││
│ 编码器 x 12 │──────────────────────▶交叉注意力 ││
│ │ ││ 门控残差连接 ││
│ ┌──────────┐ │ │└───────────────┘│
│ │ZCRMSNorm │ │ └────────┬────────┘
│ │自注意力 │ │ ┌─────┴─────┐
│ │GQA+RoPE │ │ │ 嵌入层 │ ← 共享
│ │门控残差 │ │ └─────┬─────┘
│ │ │ │ ┌───────┴───────-┐
│ │ (无FFN) │ │ │[EOS]<tool_call>│
│ └──────────┘ │ │ + 答案 │
│ │ └───────────────-┘
└──────┬───────┘
│
┌────┴──────┐
│ 嵌入层 │
└────┬──────┘
│
┌────┴──────┐
│ 文本 │
│ 查询 │
└───────────┘

在16个TPU v6e上预训练了2000亿个token（耗时27小时）。在20亿个token的单次函数调用数据集上进行后训练（耗时45分钟）。

Needle是简单注意力网络的一次实验性尝试，旨在重新定义面向消费设备（手机、手表、眼镜等）的微型AI。虽然它在单次函数调用任务上超越了FunctionGemma-270m、Qwen-0.6B、Graninte-350m和LFM2.5-350m，适用于个人AI场景，但这些模型具有更广的范围和容量，在对话场景中表现更优。此外，小模型可能较为敏感。请使用下一节中的UI测试你自己的工具，并一键进行相应微调。

## 快速开始
git clone https://github.com/cactus-compute/needle.git
cd needle && source ./setup
needle playground
打开一个Web UI，地址为http://127.0.0.1:7860，你可以在其中测试和微调自己的工具。权重会自动下载。

## 使用（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="旧金山的天气怎么样？",
tools='[{"name":"get_weather","parameters":{"location":"string"}}]',
stream=False,
)
print(result)
# [{"name":"get_weather","arguments":{"location":"San Francisco"}}]

## 微调
# 沙盒（通过Gemini生成数据，训练，评估，打包结果）
needle playground

# 命令行界面（如果本地没有权重，会自动下载）
needle finetune data.jsonl

## 命令行界面
needle playground 通过Web UI测试和微调
needle finetune <data.jsonl> 在自定义数据上微调
needle run --query "..." --tools 单次推理
needle train 完整训练运行
needle pretrain 在PleIAs/SYNTH上预训练
needle eval --checkpoint <path> 评估检查点
needle tokenize 对数据集进行分词
needle generate-data 通过Gemini合成训练数据
needle tpu <action> TPU管理（参见docs/tpu.md）

@misc{ndubuaku2026needle,
title={Needle},
author={Henry Ndubuaku, Jakub Mroz, Karen Mosoyan, Roman Shemet, Parkirat Sandhu, Satyajit Kumar, Noah Cylich, Justin H. Lee},
year={2026},
url={https://github.com/cactus-compute/needle}
}