Needle: Gemini의 Tool Calling 능력을 26M 파라미터 모델로 증류

TL;DR Highlight

Gemini의 함수 호출(tool calling) 능력만 뽑아서 26M(2600만) 파라미터짜리 초경량 모델로 만든 프로젝트로, 폰/워치/스마트글라스 같은 엣지 디바이스에서 직접 실행 가능하다.

Who Should Read

스마트폰이나 IoT 기기처럼 리소스가 제한된 환경에서 AI 기반 함수 호출 기능을 구현하고 싶은 모바일/임베디드 개발자, 또는 에이전트 파이프라인에서 tool routing 비용을 줄이고 싶은 AI 서비스 개발자.

Core Mechanics

Needle은 Gemini 3.1의 tool calling 능력만을 knowledge distillation(큰 모델의 지식을 작은 모델에 전달하는 기법)으로 뽑아낸 26M(2600만) 파라미터짜리 모델이다. 모델 파일 크기가 약 14MB 수준이라 일반 앱에 번들로 포함할 수 있다.
아키텍처는 'Simple Attention Network'라고 부르는 인코더-디코더 구조로, 인코더 12개 + 디코더 8개 레이어로 구성된다. FFN(Feed-Forward Network) 레이어 없이 Attention만으로 설계해서 파라미터 수를 극단적으로 줄였다.
d=512 차원, 8헤드 어텐션/4 KV헤드(GQA), BPE 어휘 크기 8192로 설계되어 있고, ZCRMSNorm, RoPE 위치 인코딩, Gated Residual 같은 최신 기법들을 조합했다.
사전학습은 TPU v6e 16개로 200B(2000억) 토큰을 27시간 동안 학습했고, 이후 single-shot 함수 호출 데이터셋 2B(20억) 토큰으로 45분간 파인튜닝했다.
실제 배포 환경(Cactus 플랫폼)에서 prefill 속도 6000 tokens/sec, decode 속도 1200 tokens/sec로 동작한다. 이 정도면 엣지 디바이스에서도 실시간 응답이 충분히 가능한 속도다.
single-shot 함수 호출 벤치마크에서 FunctionGemma-270m, Qwen-0.6B, Granite-350m, LFM2.5-350m 같은 더 큰 모델들을 앞섰다. 단, 이 모델들은 대화 능력까지 포함한 범용 모델이라 직접 비교는 공정하지 않다.
Needle은 오직 'single-shot tool call' 즉 사용자 입력 → 함수 호출 한 번 매핑만 잘한다. 자유로운 대화나 멀티턴 추론은 지원하지 않으며, 그 용도로는 쓰면 안 된다.
가중치와 데이터셋 생성 코드가 완전 오픈소스로 공개되어 있고, Mac/PC에서 로컬로 파인튜닝할 수 있는 UI도 제공한다. 자신만의 툴셋에 맞게 추가 학습이 가능하다.

Evidence

CLI 도구에 자연어 인터페이스를 붙이는 데 쓸 수 있겠다는 아이디어가 나왔다. 예를 들어 `toolcli add tom to teamfutz group`이라고 입력하면 내부적으로 `toolcli --gadd teamfutz tom`을 실행하는 식인데, 모델 크기가 14MB밖에 안 되니 CLI 앱에 번들로 포함하는 게 현실적으로 가능하다는 의견이었다.
실제로 알람 설정과 쇼핑 목록 추가를 테스트해봤더니 Siri보다 잘 동작했다는 사용 후기가 있었다. Home Assistant 연동에도 쓰면 좋겠다는 의견도 여럿 달렸다.
이 모델이 정확히 어디에 쓰이는 건지 모르겠다는 질문이 많았다. Claude Haiku나 Gemini Flash 같은 경량 LLM의 대체재인지, 아니면 에이전트 파이프라인에서 tool routing만 담당하는 특수 모델인지 포지셔닝이 불분명하다는 지적이었다. 이에 대해 '더 큰 모델에 넘기기 전 첫 번째 tool call 패스로 쓰면 좋겠다'는 댓글이 달리기도 했다.
Gemini 대신 왜 굳이 tool calling이 약한 모델(Gemini)을 distillation 대상으로 선택했냐는 날카로운 질문도 있었다. GPT-4o나 Claude처럼 tool calling이 더 강한 모델에서 증류했다면 결과가 더 좋지 않았겠냐는 의문이었다.
HuggingFace 레포에서 `needle-tokenizer` 데이터셋 접근이 막혀 있어서 README의 설치 가이드를 그대로 따르면 에러가 난다는 버그 리포트가 있었다. 아직 README와 실제 저장소 접근 권한이 맞지 않는 상태인 것으로 보인다.

How to Apply

에이전트 파이프라인에서 tool routing(어떤 함수를 호출할지 판단)만 Needle에 맡기고, 실제 응답 생성은 GPT-4o나 Claude 같은 큰 모델에 넘기는 구조로 쓰면, API 비용을 크게 줄이면서도 함수 선택 정확도를 유지할 수 있다.
Home Assistant나 음성 비서처럼 '전사된 텍스트 + 사용 가능한 툴 목록 → 함수 호출' 패턴이 반복되는 IoT/스마트홈 프로젝트에 Needle을 임베딩하면, 클라우드 API 없이 기기 자체에서 intent 파악이 가능해진다.
자신만의 CLI 도구에 자연어 인터페이스를 붙이고 싶다면, Needle을 앱에 번들로 포함(~14MB)하고 자신의 커맨드 스펙으로 파인튜닝하면 된다. 저장소에 Mac/PC에서 버튼 하나로 파인튜닝하는 UI가 포함되어 있어 ML 지식 없이도 시도해볼 수 있다.
Qwen이나 다른 오픈소스 LLM에 tool calling 능력을 추가하거나 개선하고 싶은 경우, Needle의 데이터셋 생성 파이프라인과 학습 코드를 참고해서 비슷한 방식으로 자체 모델을 증류하는 실험을 해볼 수 있다.

Code Example

snippet

# 설치 및 quickstart
git clone https://github.com/cactus-compute/needle
cd needle
pip install -r requirements.txt

# 모델 아키텍처 핵심 설정
# d=512, 8 heads / 4 KV heads (GQA), BPE vocab=8192
# Encoder x 12 (Self Attn + GQA + RoPE, no FFN)
# Decoder x 8 (Masked Self Attn + Cross Attn + RoPE)
# ZCRMSNorm + Gated Residual 사용
# Tied embedding (입력 임베딩과 출력 Linear 가중치 공유)

# 파인튜닝 실행
bash launch_train.sh

Terminology

Knowledge Distillation선생님(큰 모델)이 학생(작은 모델)을 가르치는 방식. 큰 모델의 출력 분포를 정답 레이블로 삼아 작은 모델을 학습시켜서, 작은 모델이 큰 모델의 행동을 흉내 낼 수 있게 된다.

GQA (Grouped Query Attention)어텐션 헤드 수를 줄이는 메모리 절약 기법. 여러 Query 헤드가 하나의 Key/Value 헤드를 공유해서 추론 시 메모리 사용량과 속도를 개선한다.

RoPE (Rotary Position Embedding)토큰의 위치 정보를 인코딩하는 방식 중 하나. 절대 위치 대신 회전 행렬로 상대적 위치 관계를 표현해서 긴 문맥에서도 성능이 잘 유지된다.

Single-shot function call사용자 입력 하나를 받아 함수 호출 하나를 출력하는 단순한 패턴. 멀티턴 대화나 복잡한 추론 없이 '이 텍스트 → 이 함수 호출'만 하는 것.

PrefillLLM이 사용자 입력 전체를 한 번에 처리하는 단계. 이 단계가 빠를수록 첫 번째 토큰이 빨리 나온다.

Tied Embedding입력 임베딩 행렬과 출력 선형 레이어의 가중치를 같은 것으로 공유하는 기법. 파라미터 수를 줄이면서도 성능을 유지하는 데 효과적이다.