Needle - A 26M-Parameter Model Distilled from Gemini Tool Calling

Needle is an experimental model that distills the tool calling of Gemini 3.1 into a 26 million parameter Simple Attention Network. It allows local fine-tuning on Mac/PC, and aims to be a small AI for consumer devices like phones, watches, and glasses.

The production implementation runs on top of Cactus and delivers prefill 6000 toks/sec and decode 1200 speeds. The weights are fully open on Hugging Face at Cactus-Compute/needle, and the dataset generation code has also been released.

Training was carried out in two stages.

• 27 hours of pretraining on 16 TPU v6e with 200B tokens
• 45 minutes of follow-up training on a single-run function calling dataset of 2B tokens

On single tool call tasks, it is presented as outperforming models like FunctionGemma-270m, Qwen-0.6B, and LFM2.5-350m. However, it is also noted that these models have strengths in broader scope and conversational settings.

The needle playground web UI opens at http://127.0.0.1:7860, where it automatically downloads the weights, lets you attach local tools for testing, and allows custom fine-tuning with a button click. In Python, you can use SimpleAttentionNetwork, load_checkpoint, generate, and get_tokenizer to feed in queries and tool schemas to generate JSON tool calls like get_weather, and the CLI provides playground, finetune, run, train, pretrain, eval, tokenize, generate-data, and tpu.

The model configuration is d=512, 8H/4KV, BPE=8192, 12-layer encoder, 8-layer decoder, and it uses GQA+RoPE, cross attention, gated residual, tied linear, and shared embedding.