TL;DR

Needle은 Gemini 3.1을 증류해 만든 26M 파라미터의 경량 인코더-디코더 모델로서 순수 어텐션 구조와 GQA·RoPE·gated residual을 결합해 함수 호출과 도구 사용 출력에 초점을 맞췄다. 학습은 JAX/Flax로 200B 토큰을 16x TPU v6e에서 수행했으며 bfloat16과 INT4 QAT를 적용해 양자화 친화적인 배포 준비를 마쳤고 후처리로 2B 토큰의 function-calling 데이터를 추가 학습했다. 가중치와 데이터셋 생성 파이프라인이 MIT 라이선스로 공개되어 로컬 미세조정이 가능하며 프로덕션에서는 Cactus 런타임에서 prefill 6000 toks/sec, decode 1200 toks/sec의 처리량을 보고했다. 이러한 설계는 온디바이스 환경에서 구조화된 도구 호출 워크플로우를 구현할 때 유의미한 처리량과 배포 유연성을 제공하지만 문서상에는 표준 벤치마크 성능 지표(예: MMLU 등)가 제시되지 않아 일반적인 언어 이해·생성 벤치마크와의 직접 비교는 불가능하다.

핵심 역량

• 구조화된 도구 호출을 생성하여 외부 API나 툴 호출 이름과 인자 형태의 출력을 생산할 수 있다. 이 출력은 JSON 유사 포맷으로 직렬화되어 후속 시스템에서 즉시 실행 가능하다. 온디바이스 환경에서 짧은 대화나 쿼리 기반 툴 연동에 적합하도록 설계되었다.
• 경량 모델로서 로컬 환경(예: Mac/PC)에서의 미세조정과 빠른 배포가 가능하다. 체크포인트와 데이터셋 생성 파이프라인이 공개되어 사용자가 직접 파인튜닝 데이터를 만들고 재학습 흐름을 운영할 수 있다. INT4 QAT를 활용해 낮은 비트폭 환경에서도 추론을 유지하도록 설계되었다.
• 인코더-디코더 구조에서 cross-attention을 통해 외부 문맥이나 툴 임베딩을 통합하여 응답을 생성할 수 있다. Encoder의 GQA와 RoPE는 장문 문맥에서 위치·질의 관계를 유지하는 역할을 수행한다. Decoder는 masked self-attn과 cross-attn을 결합해 조건부 생성 품질을 확보한다.
• 프로덕션 환경에서는 Cactus 런타임을 통해 높은 처리량을 달성하며 prefill과 decode 단계에서 각각 다른 속도 특성을 보인다. 공개된 수치에 따르면 prefill은 초당 6000토큰, decode는 초당 1200토큰 수준으로 운영되었다. 이러한 처리량은 온디바이스 중심의 실시간 도구 호출 시나리오에서 유의미한 성능 이점을 제공한다.

강점

• 매우 작은 파라미터 수(26M)로 Gemini 3.1 기반의 증류 효과를 유지하여 온디바이스 미세조정과 배포가 가능하다. 경량화 덕분에 Mac/PC 수준의 로컬 환경에서도 실험과 튜닝이 실용적으로 이루어질 수 있다. 가중치와 데이터 생성 파이프라인이 공개되어 재현성과 확장성이 확보되었다.
• 추론 처리량 측면에서 Cactus 런타임을 통해 높은 실시간 처리 성능을 달성하여 prefill 6000 toks/sec와 decode 1200 toks/sec을 보고했다. 이러한 처리량은 도구 호출 중심의 실시간 응답 파이프라인에서 지연을 줄이고 처리량을 높이는 데 유리하게 작용한다. 런타임 최적화와 모델 경량화가 함께 작용한 결과이다.
• 아키텍처적 선택으로 FFN을 제거한 순수 어텐션 구조와 GQA·RoPE·gated residual의 결합으로 작은 모델 규모에서도 기능 특화된 표현력을 확보했다. INT4 QAT를 통해 양자화된 환경에서도 성능 유지 가능성을 높였고 bfloat16을 통해 학습 안정성을 확보했다. 이러한 기술적 조합이 온디바이스 환경에서의 실용성을 뒷받침한다.

훈련 방식

Gemini 3.1을 증류해 Simple Attention Network 구조로 축소한 뒤 JAX/Flax로 사전학습을 수행했으며, 200B 토큰을 16x TPU v6e에서 약 27시간 동안 학습했다. 학습 중 bfloat16 정밀도를 사용했고 INT4 QAT를 적용하여 양자화에 견디는 가중치를 확보했다. 사전학습 이후 function-calling 역량 보강을 위해 별도 2B 토큰의 도구 호출 데이터로 후처리 학습을 45분간 진행했다.

알아두면 좋은 것

• 모델 가중치와 데이터셋 생성 파이프라인이 완전히 공개되어 있으며 라이선스는 MIT로 배포되었다. 공개된 체크포인트와 코드로 사용자가 로컬에서 재현 및 미세조정이 가능하도록 설계되었다. 이로 인해 연구와 커스텀 도구 연동 실험이 수월해졌다.
• 아키텍처는 순수 어텐션 기반의 Simple Attention Network로서 인코더에서 FFN을 제거했고 gated residual과 RoPE를 결합하여 안정적인 표현 학습을 유지했다. 이 설계는 파라미터 절감과 연산 간결화를 통해 26M 규모에서 실용적 성능을 확보하기 위해 선택되었다. Encoder는 12층, Decoder는 8층 구성을 채택했다.
• 학습 스택은 JAX/Flax 기반이며 사전학습에 200B 토큰을 사용했고 16x TPU v6e에서 약 27시간이 소요되었다. 사전학습 이후 function-calling 성능 보강을 위해 별도로 2B 토큰의 도구 호출 데이터로 후처리 학습을 45분간 수행했다. 학습 중에는 bfloat16을 기본으로 사용하고 INT4 QAT를 적용하여 최종 배포 효율을 높였다.
• 프로덕션 배포는 Cactus 런타임에서 운영되며 공개된 수치로 prefill 6000 toks/sec, decode 1200 toks/sec의 처리량이 보고되었다. 웹 UI와 CLI를 통해 로컬 미세조정 및 도구 통합 실험이 가능하며, README가 제시한 링크들에서 훈련 및 서빙 코드를 확인할 수 있다. 이 구조는 온디바이스 도구 연동 워크플로우를 빠르게 구성하려는 목적에 부합한다.

기술적 특징

• 순수 어텐션 기반의 Simple Attention Network 구조가 채택되어 Feed-Forward Network를 제거하고 어텐션·정규화·잔차만으로 표현력을 확보했다. 이 설계는 파라미터와 연산량을 줄여 26M 규모에서도 조건부 생성 품질을 유지하도록 설계되었다. 결과적으로 메모리 접근 패턴이 단순해지고 온디바이스 추론 비용을 낮추는 효과가 있다.
• 인코더는 12층 구조에 GQA(8H/4KV), RoPE, gated residual을 조합하여 쿼리·키·값 처리와 위치 정보를 효율적으로 다루도록 구성되었다. GQA는 쿼리 중심의 연산 분할로 계산 비용을 줄이고 RoPE는 상대적 위치 정보를 보존하여 긴 문맥에서의 위치 추론을 개선한다. gated residual은 신호 흐름을 제어해 작은 모델에서의 학습 안정성과 표현 분리를 돕는다.
• 디코더는 8층으로 구성되며 masked self-attn과 cross-attn을 결합해 조건부 생성을 수행하고, 임베딩과 소프트맥스 계층을 weight tying으로 공유함으로써 파라미터 효율을 높였다. Decoder는 인코더의 정보를 cross-attn으로 통합하여 도구 호출 시퀀스와 문맥을 매핑하는 역할을 수행한다. 이 구조는 구조화된 function call 출력 품질을 유지하면서 모델 크기를 억제한다.
• 정규화로는 ZCRMSNorm을 사용하여 초기화를 0으로 두고 RMS 기반 스케일링을 적용해 초기 학습 단계의 수치적 불안정을 완화했다. 이 기법은 작은 모델이 초기 단계에서 발산하거나 편향된 출력을 보이는 위험을 낮추기 위해 선택되었으며 학습 초기에 안정적인 스케일을 확보한다. ZCRMSNorm은 pure-attention 설계와 결합되어 모델의 수렴 특성을 개선했다.
• 학습과정에서는 bfloat16 정밀도를 기본으로 사용했으며 INT4 QAT를 훈련 단계에 통합해 양자화 손실을 보정하도록 설계되었다. QAT는 학습 중 양자화를 모사하여 낮은 비트폭으로 배포했을 때 정확도 저하를 완화하며, Needle은 이를 통해 온디바이스 저정밀 추론에서 실용적 성능을 유지한다. 이러한 접근은 추론 메모리와 대역폭 제약이 있는 엣지 환경에서 유리하다.

차별점

• Gemini 3.1을 증류해 26M로 극단적으로 축소한 점이 핵심 차별점이며, 이로 인해 대형 기반 모델의 능력을 경량 환경에 이식할 수 있는 구조적 기반이 마련되었다. 증류 과정은 사전학습과 후처리 function-calling 데이터로 보강되어 실제 도구 호출 품질을 높이는 데 초점이 맞춰졌다. 결과적으로 대형 모델 없이도 도구 연동 워크플로우를 실현할 수 있다.
• 인코더에서 FFN을 제거한 순수 어텐션 설계를 통해 파라미터 효율과 연산 단순화를 달성한 점이 구체적 차별점이다. 이 설계는 메모리·연산 비용을 낮춰 온디바이스 배포를 용이하게 했으며, 대신 GQA·gated residual·RoPE 같은 구조적 보완을 통해 표현력을 유지했다. 따라서 동일 규모 대비 도구 호출 특화 작업에서 효율적인 처리 성능을 보인다.
• 학습 파이프라인에 INT4 QAT를 통합하고 bfloat16을 사용한 점이 배포 효율성 측면에서 차별화되었다. 이 조합은 낮은 비트폭 환경에서도 성능 저하를 최소화하도록 설계되어 엣지 디바이스에서의 실행 가능성을 높였다. 또한 가중치와 데이터셋이 공개되어 로컬 미세조정과 재학습이 가능한 점이 실용적 유연성을 제공한다.

벤치마크