니들: FFN 전혀 없이 구축된 2600만 파라미터 도구 호출 모델

Needle은 단일 호출 함수 호출을 위해 설계된 26M 파라미터 모델입니다. FFN이 전혀 없는 교차 어텐션과 게이팅 레이어를 사용하며, 이는 도구 호출이 추론이 아니라 검색 및 조립(쿼리를 도구 이름에 매칭하고, 인자 값을 추출하며, JSON을 생성)이라는 통찰에 기반합니다. 이 모델은 소비자 기기에서 6000 tok/s 프리필과 1200 tok/s 디코드 속도로 실행됩니다.

훈련 세부 사항

• 16개의 TPU v6e에서 200B 토큰으로 사전 훈련(27시간)
• 합성 함수 호출 데이터 2B 토큰으로 후속 훈련(45분)
• 데이터는 15개의 도구 카테고리(타이머, 메시징, 내비게이션, 스마트 홈 등)로 Gemini를 통해 합성됨

아키텍처: 단순 어텐션 네트워크

전체 모델은 어텐션과 게이팅만으로 구성되며, MLP는 어디에도 없습니다. 저자들은 FFN 파라미터가 이 규모의 도구 호출에는 낭비라고 주장하며, 'FFN 없는' 발견이 모델이 외부 구조화된 지식(RAG, 도구 사용, 검색 증강 생성)에 접근할 수 있는 모든 작업에 일반화된다고 설명합니다. 모델은 입력에 사실이 제공되면 FFN 가중치에 사실을 기억할 필요가 없습니다.

벤치마크

Needle은 단일 호출 함수 호출에서 FunctionGemma-270M, Qwen-0.6B, Granite-350M, LFM2.5-350M을 능가하지만, 이들 모델은 대화 환경에서 더 많은 용량을 갖습니다.

사용 방법

# 플레이그라운드를 통해 모델을 테스트하거나 Mac/PC에서 미세 조정
git clone https://github.com/cactus-compute/needle

• GitHub: github.com/cactus-compute/needle
• 가중치: huggingface.co/Cactus-Compute/needle
• 아키텍처 문서: Simple Attention Networks 문서
• 모바일/웨어러블용 추론 엔진(Cactus): github.com/cactus-compute/cactus

모든 것은 MIT 라이선스로 제공됩니다.