STEM: Scaling Transformers with Embedding Modules

STEM (Scaling Transformers with Embedding Modules)은 2026년 1월 Carnegie Mellon University와 Meta AI 연구진이 발표한 논문입니다.

Transformer 내부의 거대한 행렬 곱셈(Matrix Multiplication) 하나를, 단순한 사전 찾기(Lookup)로 바꿔버렸습니다.

1. 무엇을 교체했나? "Up-Projection" 행렬

Transformer의 한 블록은 크게 Attention과 FFN으로 나뉩니다. 보통 FFN은 입력 벡터를 뻥튀기(Up)했다가 다시 줄이는(Down) 구조를 가집니다. LLaMA 등 최신 모델에서 쓰는 SwiGLU 구조를 수식으로 보면 이렇습니다.

$\text{FFN}(x) = (xW_{gate} \odot \text{Swish}(xW_{up})) W_{down}$

여기서 STEM이 주목한 건 바로 $xW_{up}$ (Up-Projection)입니다.

• 기존 방식:
  • 현재의 입력 벡터(문맥이 섞인 복잡한 값)와 거대한 행렬 $W_{up}$을 곱합니다.
  • 문제점: $W_{up}$은 모델 파라미터의 상당 부분을 차지하며, FLOPs 연산량이 큽니다.
  • 의미: 문맥에 따라 매번 새로운 정보를 계산해냅니다.

• STEM 방식:
  • 행렬을 제거하고, 대신 거대한 O(1) 임베딩 테이블(Lookup Table)을 둡니다.
  • 계산 대신, 이 레이어에 들어온 토큰의 원래 ID(예: "Apple")를 보고 테이블에서 해당 벡터를 꺼내옵니다(Lookup).

$xW_{up} \quad \Rightarrow \quad \text{EmbeddingTable}[\text{TokenID}]$

2. 작동 원리: "정적인 지식" vs "동적인 문맥"

이 논문의 핵심 원리는 "지식(Knowledge)과 문맥(Context)을 분리하자"는 것입니다.

(1) 지식은 "정적"이어도 된다 (Lookup)

"사과(Apple)"라는 단어가 가진 본질적인 속성(빨갛다, 과일이다, IT 기업이다 등)은 문장이 바뀐다고 해서 매번 새로 계산할 필요가 없다는 것입니다. 그래서 $W_{up}$을 행렬 곱셈으로 매번 계산하는 대신, "사과"라는 단어에 배정된 전용 벡터를 그냥 꺼내 쓰는 방식(Lookup)으로 바꿨습니다. 이것이 Embedding Module입니다.

(2) 문맥은 "Gate"가 담당한다 (Dynamic)

하지만 "사과"가 과일인지, IT 기업인지는 문맥에 따라 다릅니다. 이걸 어떻게 구분할까요? STEM은 $W_{up}$은 Lookup으로 바꿨지만, $W_{gate}$ (Gate Projection)는 그대로 둡니다.

• Gate의 역할: Lookup 해온 "사과" 벡터(과일 정보 + 기업 정보가 다 섞여 있음) 중에서, 현재 문맥에 필요한 정보만 통과시키고 나머지는 0으로 만듭니다.
  • 문장이 "맛있는 사과"라면 기업 정보 차단, 과일 정보 통과
  • 문장이 "Apple의 주가"라면 과일 정보 차단, 기업 정보 통과

3. 왜 이렇게 바꿨을까?

이 단순한 교체가 가져오는 이득은 생각보다 큽니다.

특징	기존 Transformer (Dense FFN)	STEM	이점
연산 비용	행렬 곱셈	메모리 접근	FLOPs 대폭 감소 (행렬 곱셈 하나가 사라짐)
파라미터	모든 토큰이 전체 공유	각 토큰이 독립된 벡터를 가짐	희소성(Sparsity) 확보 (필요한 것만 로드 가능)
해석 가능성	내부를 알기 어려움	"사과" 벡터가 명확히 존재	지식 수정(Knowledge Editing)이 쉬워짐

핵심: 지식 수정(Knowledge Editing)이 가능하다

기존 모델에서 "미국의 수도는?"이라고 물었을 때 틀린 답을 내놓으면, 수억 개의 파라미터를 재학습해야 했습니다. 하지만 STEM에서는 해당 레이어의 "미국" 토큰에 해당하는 임베딩 벡터만 찾아서 수정하면 됩니다. 즉, 모델 전체를 건드리지 않고 특정 지식만 수술하듯 고칠 수 있게 된 것입니다.

4. 요약: 셰프의 비유

• 기존 방식 (Dense FFN):
  셰프(모델)가 손님이 주문한 요리(문맥)를 보고, 매번 밀가루와 향신료를 배합해(행렬 곱셈) 소스를 처음부터 만듭니다. (힘들고 느림)
• STEM 방식:
  셰프가 주문서의 메뉴 이름(Token ID)을 보고, 선반에서 미리 만들어진 소스 병을 꺼냅니다(Lookup). 그리고 현재 요리 상태를 맛본 뒤(Gate), 소스를 얼마나 넣을지 결정합니다. (빠르고 효율적이며, 소스 맛이 이상하면 그 병만 바꾸면 됨)

결론

STEM은 Transformer의 가장 무거운 연산 중 하나를 토큰 기반의 사전 조회(Dictionary Lookup)로 대체하여, 계산 효율성과 해석 가능성 두 마리 토끼를 잡은 구조입니다.

https://arxiv.org/abs/2601.10639