Needle - Gemini tool calling을 증류한 2,600만 파라미터 모델

Needle은 Gemini 3.1의 tool calling을 2,600만 파라미터 Simple Attention Network로 증류한 실험적 모델이다. Mac/PC에서 로컬 파인튜닝이 가능하며, 휴대폰·시계·안경 같은 소비자 기기용 소형 AI를 목표로 한다.

프로덕션 구현은 Cactus 위에서 돌아가며 prefill 6000 toks/sec, decode 1200 속도를 제시한다. 가중치는 Hugging Face Cactus-Compute/needle에 완전 공개됐고, 데이터셋 생성 코드도 함께 열었다.

학습은 두 단계로 진행됐다.

• 16 TPU v6e에서 200B 토큰으로 27시간 pretraining
• 단일 실행 함수 호출 데이터셋 2B 토큰으로 45분 후속 학습

단일 tool call 과제에서는 FunctionGemma-270m, Qwen-0.6B, LFM2.5-350m 등보다 낫다고 제시한다. 다만 이들 모델은 더 넓은 범위와 대화형 설정에서 강점이 있다는 점도 함께 언급한다.

needle playground 웹 UI는 http://127.0.0.1:7860에서 열리며, 가중치를 자동으로 내려받아 로컬 도구를 붙여 테스트하고 버튼 클릭으로 맞춤 파인튜닝할 수 있다. Python에서는 SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer로 쿼리와 tool schema를 넣어 get_weather 같은 JSON tool call을 생성할 수 있고, CLI는 playground, finetune, run, train, pretrain, eval, tokenize, generate-data, tpu를 제공한다.

모델 구성은 d=512, 8H/4KV, BPE=8192, 12-layer encoder, 8-layer decoder이며, GQA+RoPE, cross attention, gated residual, tied linear, shared embedding을 사용한다.