Cactus가 Needle을 오픈소스로 공개했습니다. Needle은 저가형 휴대폰, 시계, 안경에서 구동되도록 설계된 2600만 파라미터 함수 호출 모델입니다. 자체 추론 엔진인 Cactus를 사용하여 소비자 기기에서 6000 tok/s 프리필과 1200 tok/s 디코드를 달성합니다.

아키텍처: Simple Attention Networks

Needle은 Simple Attention Network를 사용합니다. MLP는 전혀 없으며, 모델 전체가 어텐션과 게이팅 계층으로 구성됩니다. 주요 설계: d=512, 8H/4KV, BPE=8192, 인코더-디코더 구조(12개 인코더 계층, 8개 디코더 계층)로 크로스 어텐션, RoPE를 적용한 마스크드 셀프 어텐션, 그리고 결합된 임베딩을 사용합니다.

학습 세부사항

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

벤치마크 결과

Needle은 단일 호출 함수 실행에서 FunctionGemma-270M, Qwen-0.6B, Granite-350M, LFM2.5-350M을 능가합니다. 하지만 이 모델들은 더 넓은 범위와 용량을 가지며 대화형 환경에서 뛰어납니다.

빠른 시작

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

웹 UI가 http://127.0.0.1:7860에서 열리며, 자체 도구로 테스트 및 파인튜닝할 수 있습니다.

사용법 (Python)

from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer

params, config = load_checkpoint("checkpoints/needle.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()
result = generate(
    model, params, tokenizer,
    query="What's the weather in San Francisco?",
    tools='[{"name":"get_weather","parameters":{"location":"string"}}]',
    stream=False
)
print(result)
[{"name":"get_weather","arguments":{"location":"San Francisco"}}]

로컬 파인튜닝

# playground 사용 (Gemini로 데이터 자동 생성)
needle playground
또는 자체 데이터 제공
needle finetune data.jsonl

다운로드

가중치는 Hugging Face에서 확인 가능: Cactus-Compute/needle. 모든 것은 MIT 라이선스입니다.