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} }