Needle: 소비자 디바이스를 위한 초소형 AI 모델과 Simple Attention Network 구조 정리

이 글은 초소형 AI 모델 Needle이 무엇인지, 어떤 구조와 배경에서 만들어졌는지, 그리고 실제로 어떻게 사용하고 파인튜닝할 수 있는지를 정리한 기술 블로그입니다. Needle은 대규모 언어 모델을 증류해 만든 2,600만 파라미터 규모의 모델로, 로컬 PC나 맥에서도 직접 실행과 학습이 가능하다는 점이 핵심 특징입니다. 본문에서는 Needle의 아키텍처, 성능 특성, 활용 방식, 그리고 왜 이 모델이 ‘소비자 디바이스용 AI’라는 방향성을 갖는지 차근차근 살펴봅니다.

반응형

---

Needle 프로젝트 개요와 등장 배경

Needle은 대형 모델인 Gemini 3.1을 기반으로 증류(distillation)해 만든 Simple Attention Network(SAN) 실험 모델입니다. 목표는 명확합니다.
스마트폰, 워치, 글래스 같은 소비자 디바이스에서도 실행 가능한 작고 빠른 개인 AI를 정의하는 것입니다.

기존 소형 모델들은 파라미터 수는 작지만, 여전히 리소스 부담이 크거나 특정 작업(예: 함수 호출)에서는 성능이 불안정한 경우가 많았습니다. Needle은 이 지점을 정면으로 다루며, 구조 단순화와 학습 전략을 통해 실사용에 가까운 성능을 목표로 설계되었습니다.

---

Simple Attention Network 구조적 특징

Needle의 가장 큰 특징은 전통적인 트랜스포머 구조에서 불필요한 요소를 과감히 제거하고, Attention 중심 구조로 재구성했다는 점입니다.

핵심 아키텍처 요소

• 파라미터 수: 26M
• 차원(d): 512
• 헤드 구성: 8H / 4KV
• 토크나이저: BPE 8192
• Encoder 12층, Decoder 8층
• Encoder에는 FFN이 존재하지 않음

구조적 특징 요약

• Encoder는 Self-Attention과 Gated Residual만 사용
• Decoder는 Masked Self-Attention과 Cross-Attention을 결합
• ZCRMSNorm을 사용해 안정적인 정규화 수행
• RoPE(Position Embedding) 적용
• Embedding과 Linear(T) 가중치 공유로 파라미터 효율 극대화

이 구조는 계산량을 줄이면서도 함수 호출과 같은 단일 샷 태스크에서 높은 정확도를 유지하는 데 초점이 맞춰져 있습니다.

---

학습 방식과 성능 특성

Needle은 짧은 시간 안에 대규모 데이터를 활용해 학습됐다는 점도 인상적입니다.

학습 이력

• 사전 학습: TPU v6e 16개에서 2,000억 토큰, 약 27시간
• 후처리 학습: 단일 샷 함수 호출 데이터 20억 토큰, 약 45분

실행 성능

• Cactus 환경 기준
  • Prefill: 초당 6,000 토큰
  • Decode: 초당 1,200 토큰

이 수치는 초소형 모델 기준으로 매우 공격적인 성능이며, 실제 디바이스 내 추론이나 온디바이스 AI 시나리오에 적합한 수준입니다.

Needle은 FunctionGemma-270m, Qwen-0.6B, Granite-350m, LFM2.5-350m 등과 비교해 단일 샷 함수 호출(single-shot function call) 영역에서 더 나은 성능을 보입니다. 다만, 대화형 시나리오 전반에서는 여전히 더 큰 모델들이 유리하다는 점도 명확히 언급되고 있습니다.

---

오픈 가중치와 실험 중심 접근

Needle은 실험용 프로젝트라는 성격을 분명히 하고 있습니다.
가중치와 데이터 생성 방식 모두 공개되어 있으며, 사용자가 직접 테스트하고 파인튜닝하는 것을 전제로 설계되었습니다.

• 가중치 및 코드 공개
• 데이터 생성 파이프라인 공개
• 로컬 환경에서 파인튜닝 가능

이는 단순히 “작은 모델”을 제공하는 것을 넘어, 개인 AI를 직접 설계하고 최적화할 수 있는 실험 도구에 가깝습니다.

프로젝트와 코드는 다음 저장소에서 확인할 수 있습니다.
Needle GitHub 저장소

---

Needle 빠른 시작: 로컬 실행과 테스트

설치 및 실행

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

위 명령을 실행하면 로컬에서 웹 UI가 열리며, 브라우저(http://127.0.0.1:7860)를 통해 모델 테스트와 파인튜닝을 바로 진행할 수 있습니다. 가중치는 자동으로 다운로드됩니다.

---

Python에서의 기본 사용 예제

Needle은 함수 호출 기반 개인 AI 시나리오에 최적화돼 있습니다.

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","description":"Get current weather for a city.","parameters":{"location":{"type":"string","description":"City name.","required":true}}}]',
    stream=False,
)

print(result)

출력 결과는 자연어 응답이 아닌, 다음과 같은 구조화된 함수 호출입니다.

[{ "name": "get_weather", "arguments": { "location": "San Francisco" } }]

이 방식은 개인 비서, 스마트홈, 디바이스 제어 같은 시나리오에 매우 적합합니다.

---

파인튜닝과 데이터 구성 방식

Needle은 사용자의 도구(tool)에 맞게 모델을 쉽게 적응시킬 수 있도록 파인튜닝 워크플로우를 제공합니다.

데이터 포맷

• JSONL 형식
• 각 라인은 query, tools, answers 필드 포함
• 도구별 최소 120개 예제 권장
  (학습 100 / 검증 10 / 테스트 10)

파인튜닝 실행

needle finetune data.jsonl

또는 Playground UI를 사용하면 데이터 생성, 학습, 평가, 결과 번들링까지 한 번에 수행할 수 있습니다.

---

Needle이 의미하는 방향성

Needle은 “가장 똑똑한 모델”을 목표로 하지 않습니다. 대신 다음 질문에 답하려는 프로젝트입니다.

• AI를 디바이스 안으로 가져올 수 있는 최소 단위는 어디인가
• 개인화된 도구 호출 AI를 얼마나 단순한 구조로 구현할 수 있는가
• 사용자가 직접 학습시키고 조정할 수 있는 AI는 어떤 형태여야 하는가

이 실험은 Cactus Compute를 중심으로 진행되고 있으며, 초소형 AI와 온디바이스 AI 설계에 관심 있는 개발자에게 매우 의미 있는 참고 사례가 됩니다.

---

728x90

Needle은 단순히 “작은 모델”이 아니라, 작은 구조로 무엇까지 가능한지 보여주는 실험적 기준점입니다.
초고속 추론, 완전 공개 가중치, 쉬운 파인튜닝, 그리고 함수 호출 중심의 설계는 개인 AI와 소비자 디바이스 AI의 현실적인 방향을 제시합니다.

앞으로 스마트폰, 웨어러블, 임베디드 환경에서 AI를 직접 다루고 싶은 개발자라면, Needle은 충분히 실험해볼 가치가 있는 프로젝트입니다. 대규모 모델 시대 이후를 준비하는 하나의 실마리로서, 이 접근 방식이 어떤 진화를 이어갈지 주목해볼 만합니다.

300x250

https://github.com/cactus-compute/needle

GitHub - cactus-compute/needle: 26m function call model that runs on incredibly small devices