Needle: Gemini의 도구 호출 기능을 26M 모델로 증류하다
Needle: We Distilled Gemini Tool Calling Into a 26M Model
핵심 요약
소비자용 기기에서 초고속으로 작동하는 26M 파라미터 규모의 경량 도구 호출 모델 'Needle'이 공개되었습니다.
- 경량 모델 아키텍처 — FFN 없이 어텐션과 게이팅만 사용하여 소비자용 기기에서 초고속 추론을 구현함.
- 도구 호출 최적화 — 복잡한 추론 대신 외부 데이터 검색 및 조립에 집중하여 효율성을 극대화함.
- 합성 데이터 학습 — Gemini로 생성한 200B 토큰의 사전 학습 데이터와 2B 토큰의 도구 호출 데이터를 활용함.
- 오픈소스 생태계 — MIT 라이선스로 공개되어 누구나 로컬 환경에서 미세 조정 및 테스트가 가능함.
우리는 26M 파라미터의 함수 호출(도구 사용) 모델인 Needle을 오픈소스로 공개했습니다. 이 모델은 소비자용 기기에서 6000 tok/s의 프리필(prefill) 속도와 1200 tok/s의 디코드(decode) 속도로 작동합니다.
우리는 저가형 휴대폰에서 실행되는 에이전트 모델을 만드는 데 노력이 부족하다는 점에 항상 좌절감을 느꼈고, 다음과 같은 관찰을 통해 연구를 진행했습니다. 에이전트 경험은 도구 호출을 기반으로 구축되며, 이를 위해 거대 모델을 사용하는 것은 과도하다는 것입니다. 도구 호출은 본질적으로 검색 및 조립(쿼리를 도구 이름과 매칭하고, 인자 값을 추출하며, JSON을 출력하는 것)이지 추론이 아닙니다. 이 규모에서는 크로스 어텐션(Cross-attention)이 올바른 기본 요소이며, FFN 파라미터는 낭비됩니다.
단순 어텐션 네트워크(Simple Attention Networks): 전체 모델은 어텐션과 게이팅으로만 구성되며, MLP는 어디에도 없습니다. Needle은 소비자용 기기(휴대폰, 시계, 안경 등)를 위한 단일 샷 함수 호출 실험 모델입니다.
학습:
- 16개의 TPU v6e에서 200B 토큰으로 사전 학습 (27시간)
- 2B 토큰의 합성 함수 호출 데이터로 사후 학습 (45분)
- 데이터셋은 15개의 도구 카테고리(타이머, 메시징, 내비게이션, 스마트 홈 등)를 사용하여 Gemini로 합성됨
지금 바로 테스트하고 Mac/PC에서 미세 조정할 수 있습니다: https://github.com/cactus-compute/needle
아키텍처에 대한 전체 설명은 여기 있습니다: https://github.com/cactus-compute/needle/blob/main/docs/simple_attention_networks.md
우리는 "FFN 없음"이라는 발견이 함수 호출을 넘어 모델이 외부 구조화된 지식(RAG, 도구 사용, 검색 증강 생성)에 접근할 수 있는 모든 작업으로 일반화된다는 것을 발견했습니다. 입력으로 사실이 제공된다면 모델은 FFN 가중치에 사실을 암기할 필요가 없습니다. 실험 결과는 곧 발표될 예정입니다.
이 모델은 단일 샷 함수 호출에서 FunctionGemma-270M, Qwen-0.6B, Granite-350M, LFM2.5-350M보다 뛰어난 성능을 보이지만, 해당 모델들은 더 넓은 범위와 용량을 가지고 있으며 대화형 환경에서 뛰어납니다. 여러분의 도구를 사용하여 플레이그라운드에서 테스트하고 그에 따라 미세 조정해 보시길 권장합니다.
Needle은 온디바이스 AI를 실용적으로 만들기 위한 광범위한 노력의 일환입니다. 우리는 또한 모바일 및 웨어러블을 위한 오픈소스 추론 엔진인 Cactus(https://github.com/cactus-compute/cactus)를 구축하고 있습니다.
모든 것은 MIT 라이선스입니다. 가중치: https://huggingface.co/Cactus-Compute/needle
