요약

Needle는 Gemini의 도구 호출(Tool Calling) 기능을 26M 파라미터 규모의 초소형 모델로 증류한 Simple Attention Network 모델입니다. 소비자용 기기에 최적화된 초소형 AI를 목표로 하며, 특정 함수 호출 작업에서 더 큰 규모의 모델들을 능가하는 성능을 보여줍니다.

핵심 포인트

Gemini의 도구 호출 기능을 26M 규모의 초소형 모델로 성공적으로 증류
Cactus 위에서 초당 6000 토큰의 프리필 및 1200 토큰의 디코드 속도 구현
FunctionGemma-270m, Qwen-0.6B 등 유사 규모 모델 대비 싱글샷 함수 호출 성능 우위
가중치와 데이터셋 생성 방법이 Hugging Face에 완전히 공개됨
휴대폰, 시계, 안경 등 소비자용 기기를 위한 Tiny AI 설계

Needle

우리는 Gemini 3.1을 Mac/PC에서 로컬로 파인튜닝 (Finetune)할 수 있는 26M 파라미터 규모의 "Simple Attention Network"로 증류 (Distilled)했습니다.
실제 운영 환경에서 Needle은 Cactus 위에서 초당 6000 토큰 (toks/sec)의 프리필 (Prefill) 속도와 1200 토큰의 디코드 (Decode) 속도로 실행됩니다.
가중치 (Weights)와 데이터셋 생성 방법은 Cactus-Compute/needle에 완전히 공개되어 있습니다.

d=512, 8H/4KV, BPE=8192
                                  ┌──────────────┐
                                  │  Tool Call   │
...

16개의 TPU v6e에서 200B 토큰으로 사전 학습 (Pretrained) (27시간).
2B 토큰의 싱글샷 (Single-shot) 함수 호출 (Function call) 데이터셋으로 사후 학습 (Post-trained) (45분).

Needle는 소비자용 기기(휴대폰, 시계, 안경 등)를 위한 초소형 AI (Tiny AI)를 재정의하기 위해 설계된 Simple Attention Networks의 실험적 실행 결과입니다.
따라서 개인용 AI를 위한 싱글샷 함수 호출 성능에서는 FunctionGemma-270m, Qwen-0.6B, Graninte-350m, LFM2.5-350m을 능가하지만,
해당 모델들은 더 넓은 범위/용량 (Scope/Capacity)을 가지고 있으며 대화형 설정 (Conversational settings)에서 뛰어난 성능을 발휘합니다. 또한, 소형 모델은 까다로울 수 있습니다.
다음 섹션의 UI를 사용하여 직접 도구를 테스트하고, 버튼 클릭 한 번으로 그에 맞춰 파인튜닝 (Finetune)해 보시기 바랍니다.

빠른 시작 (Quickstart)

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

http://127.0.0.1:7860 에서 웹 UI가 열리며, 여기서 자신만의 도구를 테스트하고 파인튜닝 (Finetune)할 수 있습니다. 가중치 (Weights)는 자동으로 다운로드됩니다.

사용법 (Usage - Python)

from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer

params, config = load_checkpoint("checkpoints/needle.pkl")
...

파인튜닝 (Finetuning)

# Playground (Gemini를 통해 데이터를 생성하고, 학습하며, 평가하고, 결과를 묶음)
needle playground

...

데이터 형식 (Data format)

JSONL 파일의 각 라인은 query, tools, answers 세 가지 필드를 가집니다.

도구 스키마 (Tool schema):

{
  "name": "get_weather",
  "description": "도시의 현재 날씨를 가져옵니다.",
...

응답 스키마 (Answer schema):

{ "name": "get_weather", "arguments": { "location": "Paris" } }

전체 JSONL 예시 (각 라인은 하나의 학습 예시이며, tools와 answers는 JSON 인코딩된 문자열입니다):

{"query": "파리 날씨가 어때?", "tools": "[{\"name\":\"get_weather\",\"description\":\"도시의 현재 날씨를 가져옵니다.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"도시 이름.\",\"required\":true}}}]", "answers": "[{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\"}}]"}
{"query": "불을 꺼줘", "tools": "[{\"name\":\"get_weather\",\"description\":\"도시의 현재 날씨를 가져옵니다.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"도시 이름.\",\"required\":true}}},{\"name\":\"toggle_lights\",\"description\":\"스마트 조명을 켜거나 끄는 기능을 전환합니다.\",\"parameters\":{\"state\":{\"type\":\"string\",\"description\":\"on 또는 off.\",\"required\":true}}}]", "answers": "[{\"name\":\"toggle_lights\",\"arguments\":{\"state\":\"off\"}}]"}

각 도구당 최소 120개의 예시를 제공해야 합니다 (학습 100개 / 검증 10개 / 테스트 10개). 예시가 적으면 과적합(overfit)될 수 있습니다. 완벽한 학습 지표가 나타나더라도 모델이 일반화되지 못할 수 있으니, 질의어 표현을 다양하게 하고 여러 도구가 사용 가능한 예시를 포함하세요.

파인튜닝된 모델 사용하기

파인튜닝은 최적의 체크포인트(checkpoint)를 checkpoints/needle_finetuned_<id>_best.pkl로 저장합니다:

needle run --checkpoint checkpoints/needle_finetuned_*_best.pkl \
  --query "날씨가 어때?" --tools '[{"name":"get_weather","description":"도시의 현재 날씨를 가져옵니다.","parameters":{"location":{"type":"string","description":"도시 이름.","required":true}}}]'

params, config = load_checkpoint("checkpoints/needle_finetuned_<id>_best.pkl")
model = SimpleAttentionNetwork(config)
result = generate(model, params, get_tokenizer(), query="...", tools='[...]', stream=False)

CLI

needle playground                  웹 UI를 통해 테스트 및 파인튜닝
needle finetune <data.jsonl>       자체 데이터를 이용한 파인튜닝
needle run --query "..." --tools   단일 추론
...

@misc{ndubuaku2026needle,
  title={Needle},
  author={Henry Ndubuaku, Jakub Mroz,  Karen Mosoyan, Roman Shemet, Parkirat Sandhu, Satyajit Kumar, Noah Cylich, Justin H. Lee},
...

AI 자동 생성 콘텐츠

원문 바로가기

Insights

Show HN: Needle: Gemini의 도구 호출 (Tool Calling) 기능을 26M 모델로 증류 (Distilled)

요약