STEM(Scaling Transformers with Embedding Modules)은 카네기멜론대(CMU)와 Meta AI가 공동 개발한 트랜스포머 스케일링 기법이다. Feed-Forward Network(FFN)의 up-projection을 레이어 로컬 임베딩 룩업(layer-local embedding lookup)으로 교체해, Mixture-of-Experts(MoE)의 복잡한 런타임 라우팅 없이도 파라미터 용량을 효율적으로 늘릴 수 있다.
핵심 아이디어
기존 파인그레인드 희소성(fine-grained sparsity, 대표적으로 MoE) 접근법은 학습 불안정성, 부하 균형(load balancing), 통신 오버헤드 문제가 있다. STEM은 이를 정적(static), 토큰 인덱스 기반(token-indexed) 방식으로 해결한다:
- FFN의 up-projection을 레이어 로컬 임베딩 룩업으로 대체
- gate와 down-projection은 dense 유지
- 런타임 라우팅 없음 → 학습 안정성 확보
- CPU 오프로드 + 비동기 프리페치(async prefetch) 지원
MoE와의 비교
| 항목 | MoE | STEM |
|---|---|---|
| 라우팅 | 런타임 동적 라우팅 | 정적 토큰 인덱스 |
| 학습 안정성 | 부하 균형 튜닝 필요 | 안정적 |
| 통신 오버헤드 | 디바이스 간 통신 | CPU 오프로드 가능 |
| 해석 가능성 | 낮음 | 지식 편집·주입 가능 |
주요 장점
지식 저장 용량 증가
임베딩 공간이 큰 각도 스프레드(angular spread)를 가져 지식 저장 용량이 향상된다.
해석 가능성
토큰 인덱스 방식이어서 특정 임베딩 엔트리가 어떤 지식을 담는지 파악하고, 입력 텍스트나 추가 연산 없이 지식 편집(knowledge editing) 및 지식 주입(knowledge injection)이 가능하다.
장문 컨텍스트 성능
시퀀스 길이가 늘어날수록 더 많은 고유 파라미터가 활성화되어 테스트 시 실질적인 용량 확장(test-time capacity scaling)이 이뤄진다.
효율성
FFN 파라미터 약 1/3을 제거하면서 토큰당 FLOPs와 파라미터 접근 수를 줄인다.
성능 결과
350M 및 1B 모델 스케일에서 dense 베이스라인 대비 약 3–4% 정확도 향상을 달성했다. 지식·추론 집약적 벤치마크에서 특히 개선됐다:
- ARC-Challenge
- OpenBookQA
- GSM8K
- MMLU
논문 및 코드
- arxiv: 2601.10639 — 2026년 1월 게재
- GitHub: Infini-AI-Lab/STEM
참고 자료
- STEM: Scaling Transformers with Embedding Modules — arXiv (2026-01-15)
- Infini-AI-Lab/STEM — GitHub 공식 저장소
관련 문서
- llm-fine-tuning — LLM 파인튜닝 전체 가이드 (SFT, PEFT, LoRA)