Needle: Gemini의 Tool Calling 기능을 26M 모델로 증류(Distilled)했습니다
요약
Needle은 26M 파라미터 규모의 경량 함수 호출(tool calling) 모델을 오픈 소스로 공개했습니다. 이 모델은 소비자용 기기에서 높은 속도로 실행 가능하며, 도구 사용 과정을 근본적으로 검색 및 조립 과정으로 정의하여 Cross-attention만을 사용하여 FFN(Feed-Forward Network) 파라미터를 제거한 것이 특징입니다. Needle은 온디바이스 AI 환경을 목표로 하며, RAG나 도구 사용 등 외부 구조화된 지식 접근 작업 전반에 적용 가능한 원시 요소임을 제시합니다.
핵심 포인트
- Needle은 26M 파라미터의 경량 함수 호출 모델로, 소비자용 기기(스마트폰, 워치 등)에서의 온디바이스 AI 구현을 목표로 합니다.
- 도구 호출 과정을 추론이 아닌 검색 및 조립 과정으로 정의하고, Cross-attention만을 사용하여 FFN 레이어를 제거함으로써 효율성을 극대화했습니다.
- Needle은 단일 샷 함수 호출 성능에서 여러 경쟁 모델(FunctionGemma-270M 등)을 능가하지만, 대화형 환경에서의 활용이 권장됩니다.
- 모델 가중치와 코드는 MIT 라이선스 하에 공개되었으며, 모바일 및 웨어러블용 추론 엔진인 Cactus도 함께 제공됩니다.
저희는 26M 파라미터의 함수 호출 (function-calling, 도구 사용) 모델인 Needle을 오픈 소스로 공개합니다. 이 모델은 소비자용 기기에서 prefill 시 6000 tok/s, decode 시 1200 tok/s의 속도로 실행됩니다.
저희는 저가형 스마트폰에서 실행 가능한 에이전트 모델 (agentic models)을 구축하려는 노력이 부족하다는 점에 항상 답답함을 느껴왔으며, 이에 대한 조사를 진행한 결과 한 가지 관찰을 얻었습니다: 에이전트 경험은 도구 호출 (tool calling)을 기반으로 구축되지만, 이를 위해 거대 모델을 사용하는 것은 과잉 (overkill)이라는 점입니다. 도구 호출은 근본적으로 검색 및 조립 (retrieval-and-assembly; 쿼리를 도구 이름과 매칭하고, 인자 값을 추출하여, JSON을 출력하는 과정)이지, 추론 (reasoning)이 아닙니다. 이를 위한 적절한 원시 요소 (primitive)는 교차 주의 집중 (Cross-attention)이며, 이 규모에서는 FFN (Feed-Forward Network) 파라미터가 낭비됩니다.
단순 주의 집중 네트워크 (Simple Attention Networks): 모델 전체가 오직 attention과 gating으로만 구성되어 있으며, 어디에도 MLP (Multi-Layer Perceptron)가 없습니다. Needle은 소비자용 기기 (스마트폰, 워치, 안경 등)를 위한 단발성 함수 호출 (single-shot function calling)을 위한 실험적 시도입니다.
학습 (Training):
-
16개의 TPU v6e에서 200B 토큰으로 사전 학습 (Pretrained) (27시간)
-
합성된 함수 호출 데이터 2B 토큰으로 사후 학습 (Post-trained) (45분)
-
15가지 도구 카테고리 (타이머, 메시징, 내비게이션, 스마트 홈 등)를 포함하여 Gemini를 통해 합성된 데이터셋 사용
지금 바로 테스트하고 Mac/PC에서 미세 조정 (finetune)할 수 있습니다: https://github.com/cactus-compute/needle
아키텍처에 대한 전체 기술 문서는 여기에서 확인할 수 있습니다: https://github.com/cactus-compute/needle/blob/main/docs/simple_attention_networks.md
저희는 "FFN 없음"이라는 발견이 함수 호출을 넘어 모델이 외부의 구조화된 지식에 접근할 수 있는 모든 작업 (RAG, 도구 사용, 검색 증강 생성 (retrieval-augmented generation))으로 일반화된다는 것을 발견했습니다. 사실 관계가 입력값으로 제공된다면, 모델은 FFN 가중치에 해당 사실을 암기할 필요가 없습니다. 실험 결과는 추후 발표될 예정입니다.
Needle는 단일 샷 함수 호출 (single-shot function calling)에서 FunctionGemma-270M, Qwen-0.6B, Granite-350M, LFM2.5-350M을 능가하지만, 해당 모델들은 더 넓은 범위/용량 (scope/capacity)을 가지고 있으며 대화형 환경 (conversational settings)에서 뛰어난 성능을 보입니다. 저희는 여러분이 플레이그라운드 (playground)를 통해 직접 도구들을 테스트하고 그에 맞춰 미세 조정 (finetune)해 보시기를 권장합니다.
Needle는 온디바이스 AI (on-device AI)를 실용적으로 만들기 위한 광범위한 노력의 일환입니다. 저희는 또한 모바일 및 웨어러블을 위한 오픈 소스 추론 엔진 (inference engine)인 Cactus (https://github.com/cactus-compute/cactus)를 구축하고 있습니다.
모든 것은 MIT 라이선스 하에 제공됩니다. 가중치 (Weights): https://huggingface.co/Cactus-Compute/needle
AI 자동 생성 콘텐츠
본 콘텐츠는 Reddit AI Engineering의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기