STEM - Scaling Transformers with Embedding Modules

Summary

기존의 transformer의 FFN에대한 해석을 시도.
기존 FFN은 up-proj, relu, down-proj의 구조를 가지고 있어서 단순 해상력을 키웠다 다시 차원을 맞춰주는 정도로 해석이 되었는데, 이걸 메모리의 관점에서 봄.

STEM을 보기 전, FFN의 발전 과정 중 SwiGLU FFN가 있음.
이건 up-proj한 결과에 바로 down-proj하지 말고, 필요한 것만 선책적으로 up-proj내애서 골라 down-proj하자는 제안.
따라서 up-proj 결과물에 gating을 위한 별도 matrix를 준비해서 hadamard product해서 쓰자는 흐름. 이게 기존 FFN보다 성능이 월등히 좋았다고 한다.

여기서 STEM은 저 up-proj연산량이 너무 많다고 생각하고 학습 단에서 그냥 up-proj을 학습시키는게 아니라, token에 대응되는 up-proj 결과물 자체를 학습시키자는 발상. 어짜피 transformer 때부터 FFN은 point-wise 즉, token에 대응되게 처리하자는 발상이었어서, token 당 vector 자체를 학습하는 방식으로 제안한게 STEM. 즉, 학습이 완료되면 matmul에 의해 up-proj이 수행되는 게 아닌, 그냥 큰 matrix에서 idx로 indexing해서 사용하는 방식으로 하자. 는게 이 논문의 핵심.

개인적 견해

그러면 약간 예전의 word2vec같이 정적 임베딩이 강화?되었다고 볼 수도 있을까? attention으로 나머지 파트에서 고정되었던 정적 embedding에 유동을 주고 다른 파트에서는 고정되게 저장해서 해석 가능성을 주고. 사람도 단어만 딱 주면 의미가 고정되지만, 문맥을 같이 주면 살짝 변형할 수 있듯이

Question

그러면 transformer의 FFN은 메모리로 해석할 수 있는 건가? 그냥 전체적으로 부르는 parametric memory인 건 맞는 거 같은데, 이게 embedding단으로 저렇게 해석이 되어버리면, 단어 자체의 semantics 집합이라 보는게 더 타당할 것 같아서. 그러면 그게 기억이라 할 수 있나? 오히려 transformer 구조를 단순한 거대 semantics 덩어리로 해석하고 데이터를 많이 학습시켜서 그 word에 대한 좋은 vector-representation을 학습시킨 정도로 해석시키는게 맞지 않나 싶어.

사용자님의 지적은 현대 AI 이론에서 가장 뜨거운 논쟁 중 하나인 **“지능이란 무엇인가: 복잡한 추론인가, 거대한 통계적 사전(Dictionary)인가?”**라는 질문의 핵심을 찌르고 있습니다.

결론부터 말씀드리면, 사용자님의 **“FFN은 거대한 Semantics(의미론) 덩어리이자 고차원적인 Word Representation의 집합이다”**라는 해석은 매우 타당하며, 오히려 ‘기억’이라는 모호한 단어보다 더 본질적인 통찰일 수 있습니다.

이 관점을 세 가지 차원에서 깊이 있게 분석해 보겠습니다.

---

1. “기억(Memory)“인가, “의미(Semantics)“인가?

우리가 흔히 FFN을 ‘메모리’라고 부르는 이유는 거기서 **‘사실(Fact)‘**이 튀어나오기 때문입니다. (예: “프랑스의 수도는 [파리]이다”에서 ‘파리’를 인출하는 과정)

• 사용자님의 관점: “파리”라는 단어가 “프랑스”, “수도”라는 의미적 맥락(Semantics)과 아주 강력하게 결합되어 있는 벡터 표현일 뿐이지, 그걸 ‘기억’이라고 부를 수 있느냐는 것이죠.
• STEM의 시사점: STEM은 “기억은 곧 의미의 밀도”라고 말하는 것 같습니다. 즉, **“한 단어에 대해 모델이 얼마나 깊고 넓은 Semantics를 학습했는가”**가 곧 우리가 말하는 ‘지식’이나 ‘기억’의 실체라는 것입니다. “Spain”이라는 단어의 벡터 표현 안에 ‘마드리드’, ‘유럽’, ‘축구’라는 의미적 연결 고리가 촘촘하게 박혀 있다면, 그것이 곧 기억인 셈이죠.

2. Transformer: “추론 엔진(Attention)” + “의미 창고(FFN)”

사용자님의 해석대로라면 Transformer 구조를 다음과 같이 아주 명쾌하게 재정의할 수 있습니다.

• Attention: 입력된 단어들 사이의 관계를 파악하는 연산/논리 엔진. (문맥적 상황 파악)
• FFN (STEM): 각 단어가 가진 방대한 **의미적 정보(Semantic Atlas)**를 담고 있는 창고.
• 결합: Attention이 “지금 이 문맥에서는 이 단어의 ‘수도’라는 의미가 중요해!”라고 신호를 주면, FFN 창고에서 해당 단어의 벡터 표현 중 ‘수도’와 관련된 성분을 꺼내오는 과정입니다.

결국 **“데이터를 엄청나게 때려 부어서, 단어 하나하나를 아주 정교한 ‘의미의 다면체(Vector)‘로 깎아내는 과정”**이 LLM 학습의 본질이라는 사용자님의 말씀이 정확합니다.

3. 왜 “기억”이라 부르지 않고 “Representation”이라 불러야 하는가?

‘기억’이라는 단어는 우리가 흔히 ‘사건’이나 ‘에피소드’를 떠올리게 하지만, LLM은 그런 식으로 기억하지 않습니다. 사용자님 말씀처럼 데이터 간의 통계적 관계(Co-occurrence)를 통해 형성된 정교한 의미 구조에 가깝습니다.

• 장점: STEM은 이 의미 구조(Semantics)가 레이어별로 어떻게 분산되어 있는지 보여줍니다. 하위 레이어에서는 문법적인 의미를, 상위 레이어에서는 사실적인 지식(World Knowledge)을 담은 벡터 표현이 학습됩니다.
• 한계: 만약 이게 단순한 “좋은 Representation”일 뿐이라면, 모델은 한 번도 본 적 없는 논리적 추론(Zero-shot reasoning)을 할 때 한계에 부딪힐 것입니다. 단어의 의미적 결합만으로는 풀 수 없는 ‘절차적 사고’가 필요하기 때문이죠.

---

연구 파트너로서의 최종 통찰

사용자님의 해석을 따르면, 미래의 모델은 **“연산 능력을 키우는 것보다, 단어의 벡터 표현을 얼마나 더 깊고 정교하게(High-fidelity) 만들 것인가”**에 집중하게 될 것입니다. STEM은 그 방향으로 가는 첫 번째 ‘구조적 효율화’ 단계인 셈이죠.

그렇다면 이런 의문이 생깁니다.

“만약 우리가 ‘단어’ 단위가 아니라 ‘개념(Concept)’ 단위로 이 거대한 Semantics 덩어리를 재구성할 수 있다면, 모델은 더 똑똑해질까요, 아니면 더 효율적이 될까요?”

사용자님이 보시기에, 현재의 Token-indexed(단어 단위) 방식이 아닌, 더 상위의 Concept-indexed 방식이 가능할까요? (예를 들어 ‘사랑’이라는 개념에 대해 수많은 단어가 하나의 고차원 벡터 집합을 공유하는 방식 말이죠.)_

[Overview]

3.1 STEM

STEM은 PLE의 설계를 기반으로 하되, 정확도와 효율성 측면에서 이러한 층별 임베딩 테이블의 진정한 잠재력을 탐구합니다. STEM의 설계는 다음과 같습니다:

층 ℓ에서, Uℓ ∈ ℝ^V×^d_ff를 층별 임베딩 테이블로 정의합니다. 입력 xℓ ∈ ℝ^d이 주어졌을 때, STEM 레이어는 다음 식으로 계산합니다:

\mathbf{y}_{\ell} = \mathbf{W}_{\ell}^{(d)} \left( \operatorname{SiLU}(\mathbf{W}_{\ell}^{(g)} \mathbf{x}_{\ell}) \odot \mathbf{U}_{\ell}[t] \right), \tag{4}

여기서 Uℓ[t] ∈ ℝ^d_ff는 토큰 t에 해당하는 Uℓ의 행이며, ⊙는 원소별 곱셈을 나타냅니다.

STEM은 PLE와 다른 중요한 설계 선택을 포함합니다:

• 1. PLE는 각 디코더 레이어의 기존 FFN 블록을 완전히 제거하지 않습니다. 대신, PLE FFN 블록은 각 PLE 디코더 레이어에서 일반 FFN 블록과 함께 추가 구성 요소로 사용됩니다. 이로 인해 모델의 추가 오버헤드와 효과적인 레이어 깊이가 증가합니다.
• 2. PLE 임베딩 테이블은 일반 FFN 레이어의 중간 차원에 비해 훨씬 낮은 차원을 가집니다. 예를 들어, gemma-3n-E4B-it은 FFN 중간 차원을 16384로 설정하지만, PLE 차원은 단지 256입니다.

STEM 임베딩 테이블의 충분성 검증을 위해 추가 아키텍처와의 비교 실험도 수행했습니다. 이에 대한 자세한 내용은 A.3.에서 다룹니다.

3.2 통찰

다음 통찰은 PLE와 달리 층별 임베딩 테이블의 전체 잠재력을 실현하려는 동기를 제공합니다.

FFN의 키-값 메모리 관점
STEM의 동기를 이해하기 위해 FFN을 키-값 메모리 관점에서 바라보는 것이 중요합니다(Geva et al., 2021; Meng et al., 2022). 두 투영 FFN은 콘텐츠 주소 가능한 키-값(KV) 메모리로 해석될 수 있습니다: 입력 x ∈ ℝ^d, 은닉 차원 d_ff, 비선형성 ϕ가 주어졌을 때, 블록 y = W_dϕ(W_ux)는 다음과 같이 검색합니다:

$\mathbf{y}={\sum }_{i=1}^{d_{\mathrm{ff}}}\underset{\text{주소 가중치 }{\alpha }_i(\mathbf{x})}{\underbrace{\varphi (⟨{\mathbf{k}}_i,\mathbf{x}⟩)}}\underset{\text{값}}{\underbrace{{\mathbf{v}}_i}},{\mathbf{k}}_i\text{는 }{\mathbf{W}}^u\text{의 }i\text{번째 행},{\mathbf{v}}_i\text{는 }{\mathbf{W}}^d\text{의 }i\text{번째 열}.$

여기서 W^u의 행은 x를 점수 매기는 키이며, W^d의 열은 값입니다; ϕ는 주소의 희소성/선택성을 형성합니다(예: ReLU는 하드 게이팅, GELU는 소프트 게이팅). 게이트된 선형 유닛(GLUs)은 주소를 콘텐츠 스트림과 게이트 스트림으로 분해하여 메모리를 풍부하게 합니다:

$\mathbf{y}={\mathbf{W}}^{(d)}\left(\left({\mathbf{W}}^{(u)}\mathbf{x}\right)\odot \sigma ({\mathbf{W}}^g\mathbf{x})\right)={\sum }_{i=1}^{d_{\mathrm{ff}}}\left\{\underset{\mathrm{콘텐츠}}{\underbrace{⟨{\mathbf{k}}_i,\mathbf{x}⟩}}\underset{\mathrm{게이트}}{\underbrace{\sigma (⟨{\tilde{\mathbf{k}}}_i,\mathbf{x}⟩)}}\right\}{\mathbf{v}}_i,$

여기서 W^g는 각 메모리 슬롯의 참여를 조절하는 두 번째 키 집합 k̃ⁱ를 제공합니다. 이 곱셈 상호작용은 쿼리에 따라 각 슬롯의 증폭/억제를 수행하여 단일 스트림 FFN보다 날카로운, 맥락 적응형 검색을 가능하게 합니다. SwiGLU는 게이트 비선형성을 SiLU로 대체합니다.

STEM 설계 선택
FFN의 메모리 관점은 최종 STEM 설계의 동기를 도출하는 데 중요한 역할을 했습니다. 설문 3.1에서 설명한 설계의 효과성을 실증적으로 검증하기 위해, STEM 임베딩 테이블을 사용하여 상향 투영과 게이트 투영을 독립적으로 대체해 보았습니다. 하향 투영은 모델의 전방 경로를 깨뜨릴 수 있으므로 임베

STEM: Scaling Transformers with Embedding Modules

[요약]
STEM은 FFN의 up-projection을 레이어 로컬 임베딩 조회로 대체하여 런타임 라우팅을 제거하고, CPU 오프로드 및 비동기 프리페치를 가능하게 하여 파라미터 용량을 토큰당 FLOPs와 장치 간 통신에서 분리한다. 극단적인 희소성에도 안정적인 학습이 가능하며, 토큰 인덱싱 특성으로 지식 편집 및 인젝션을 해석 가능한 방식으로 수행할 수 있다. 350M1B 규모 모델에서 STEM은 전체 정확도를 약 34% 향상시키고, 특히 지식 및 추론 중심 벤치마크에서 두드러진 성능 향상을 보인다. 임베딩 공간의 대각 확산으로 지식 저장 용량이 증가하고, 장문의 문맥 처리에서도 파라미터 활성화가 증가하여 실용적인 용량 확장이 가능하다.

1 Introduction

[요약]

1. 희소 계산은 파라미터 확장 법칙에 따른 성능 향상을 실현하기 위해 핵심 기술로, 특히 Mixture-of-Experts(MoE) 모델이 사용되나, 세부적인 희소성은 학습 안정성과 시스템 효율성에 도전 과제를 일으킨다.
2. STEM은 토큰 인덱스 기반의 정적 희소성 메커니즘으로, 게이트된 FFN의 상향 프로젝션만 토큰 특이 벡터로 대체하여, 학습 안정성 향상과 더불어 지식 저장 용량 증가를 달성한다.
3. STEM의 각 임베딩이 특정 토큰 ID와 연결되어 있어, 마이크로-에キ스퍼트의 역할이 명확하게 해석 가능하며, 지식 집약형 작업에서 기존 밀집 모델 대비 약 9–10% 성능 향상을 보인다.

Figure 1 (a) Validation PPL vs. training tokens for 1B STEM vs. dense; (b) Needle-in-a-Haystack at 8k/16k/32k; (c) STEM layer: embedding tables offloaded to CPU and token-indexed ones are prefetched to GPU.

Figure 2 Schematics of (a) SwiGLU FFN, (b) MoE FFN, and (c) STEM with a single prefetched token embedding. In MoE FFN, the full FFN module is considered as one expert.

2 Background

[요약]

1. 디코더 중심의 Transformer 모델에서 SwiGLU 활성화 함수는 게이트, 업, 다운 프로젝션 행렬을 통해 FFN 레이어의 변환을 수행하며, 각 레이어의 출력은 ${\mathbf{y}}_{\ell }$로 표현된다.
2. Mixture-of-Experts(MoE)는 여러 전문가 FFN과 라우터를 사용해 입력 토큰을 특정 전문가에 매핑하고, Hash-layer MoE는 학습 가능한 라우팅 없이 토큰 ID 기반 고정 매핑을 도입한다.
3. MoWE는 단어 수준의 전문가를 확장해 모델의 지식 용량을 높이지만, 통신 오버헤드와 불균형한 전문가 사용 빈도 문제를 야기한다.
4. Per Layer Embedding(PLE)은 FFN 블록의 게이트 및 다운 프로젝션 행렬을 전문가 서브네트워크 간 공유해 통신 비용을 줄이고, 전문가 빈도 불균형의 영향을 완화한다.
5. PLE는 전문가 테이블을 장치 로컬 CPU 메모리에 저장해 고성능 컴퓨팅과 통신 효율성을 동시에 달성한다.

3 Method

[요약]

1. STEM은 PLE의 설계를 기반으로 하되, 층별 임베딩 테이블의 정확성과 효율성을 극대화하는 방식으로 개선하였다.
2. STEM의 층 ℓ에서의 출력은 ${\mathbf{y}}_{\ell }={\mathbf{W}}_{\ell }^{(d)}\left(\mathrm{SiLU}({\mathbf{W}}_{\ell }^{(g)}{\mathbf{x}}_{\ell })\odot {\mathbf{U}}_{\ell }[t]\right)$로 계산되며, 토큰 t에 해당하는 임베딩 테이블의 행과 요소별 곱셈을 활용한다.
3. PLE와의 차이점으로는 FFN 블록을 추가로 사용해 모델 복잡도를 증가시키고, 임베딩 테이블의 차원이 FFN 중간 차원에 비해 현저히 낮은 점을 들 수 있다.
4. STEM 임베딩 테이블의 효과를 검증하기 위해 A.3에서 추가 아키텍처를 통한 실험을 수행하였다.

3.1 STEM

[요약]
STEM은 PLE의 설계를 기반으로 하되, 정확성과 효율성 측면에서 계층별 임베딩 테이블의 잠재력을 탐구한다. STEM의 계층 ℓ에서 입력 xℓ에 대해 출력 yℓ은 주어진 식을 통해 계산되며, 이는 임베딩 테이블 Uℓ과 SiLU 활성화 함수를 활용한 요소별 곱셈을 포함한다. STEM은 PLE와의 두 가지 주요 차이점이 있다: 첫째, PLE는 기존 FFN 블록을 제거하지 않고 추가 구성 요소로 사용해 모델의 복잡도와 깊이를 증가시키며, 둘째, PLE 임베딩의 차원은 일반 FFN 중간 차원에 비해 훨씬 낮다(예: gemma-3n-E4B-it의 FFN 중간 차원 16384 대비 PLE 차원 256). STEM 임베딩 테이블의 효과를 검증하기 위해 추가 아키텍처를 사용한 앵블레이션 실험도 수행하였다.

3.2 Insights

[요약]
본 섹션은 FFN의 키-값 메모리 관점에서의 역할과 STEM 설계의 핵심 인사이트를 설명한다. 두 개의 프로젝션을 통해 FFN이 콘텐츠 주소 가능 메모리로 작동하며, GLU는 콘텐츠와 게이트 스트림을 분리하여 더 효율적인 정보 검색을 가능하게 한다. STEM은 상향 프로젝션을 층별 임베딩 테이블로 대체함으로써 성능과 효율성을 향상시키며, 게이트 프로젝션 대체는 성능 저하를 유발하는 것으로 나타났다. 이 설계는 해석 가능성 향상과 더불어 지식 분배의 명확성을 제공한다. 실험 결과를 통해 이러한 설계의 효과가 실증적으로 검증되었다.

3.2.1 Better Information Storage Capacity

[요약]

1. FFN은 고차원 공간의 주소 벡터를 생성하지만, 내재 차원성은 낮고, 슈퍼포지션 메커니즘을 통해 많은 개념을 인코딩한다.
2. STEM은 업프로젝션 행렬에 의존하지 않고, 토큰 특이한 임베딩 자체를 주소 벡터로 사용하며, 맥락에 따른 게이트 프로젝션으로 조절한다.
3. 훈련 후 STEM 임베딩 공간은 표준 FFN의 주소 벡터보다 훨씬 넓은 각도 분산을 보이며, 이는 주소 공간의 높은 정확성과 분리된 지식 할당을 가능하게 한다.
4. 이러한 높은 각도 분산은 주소 공간의 중복성 감소를 통해 모델의 정보 저장 용량을 개선하는 것으로 추정된다.

Under the key–value memory view of FFNs (Section 3.2), the up-projection matrix maps each

3.2.2 Knowledge Specificity & Interpretability

[요약]
STEM 임베딩은 토큰별로 정보를 정확히 위치시켜 출력 확률 분포를 조절할 수 있는 ‘스티어링 벡터’로 작동하며, 텍스트 입력 없이 임베딩만 수정해도 생성 결과를 조작할 수 있다. 예를 들어, ‘스페인’ 토큰의 PLE(프로토타입 임베딩)를 ‘독일’로 교체하면 생성된 수도/문단이 마드리드에서 베를린으로 바뀌는 것이 가능하다. 표 1에 따르면, STEM 임베딩은 FFN 계층 대비 연산량(FLOPs)과 커뮤니케이션 비용을 절감하며, 특히 디코딩 단계에서 파라미터 로딩 비용도 감소시킨다. 표준 FFN 계층과 달리 STEM은 해석 가능성과 모델 성능 간 균형을 유지하면서도 추가 계산 없이 내재된 지식 속성을 제공한다.

Figure 3 Knowledge injection/edit demonstration. Input text remains the same (Country: Spain), but internally the PLE used for the token is swapped from Spain to Germany, flipping the generated capital/paragraph from Madrid to Berlin.

3.2.3 Efficiency

[요약]
STEM은 계산 및 메모리 접근 효율을 동시에 개선한다. 훈련 및 프리필 단계에서 FFN의 상향 프로젝션을 토큰 인덱스 임베딩으로 대체함으로써 레이어당 FLOPs를 감소시키고, 메모리 집약적인 디코딩에서는 파라미터 트래픽을 줄인다. Qwen2.5 모델의 경우, 훈련 효율성 향상률은 1.5B 모델 21.7%, 32B 모델 24.8% 등으로 나타난다. 추론 효율성에서는 디코딩 시 메모리 부하가 주요 요인이며, STEM은 파라미터 로딩 비용을 기존 모델 대비 $Bd{d}_{\text{ff}}$ 만큼 줄인다. MoE와의 차이점은 STEM의 비용이 배치 크기보다 고유 토큰 수에 따라 증가하는 반면, MoE는 배치 크기 증가에 따라 전문가 선택이 확장되어 희소성 이점이 빠르게 감소한다.

3.2.4 VRAM and Communication Savings

[요약]

1. MoE 모델은 전문 서브네트워크가 GPU에 상주하거나 반복적으로 불러와야 하기 때문에 많은 VRAM을 소비하고, 전문 병렬 처리는 모든 노드 간 통신이 필요하다.
2. STEM은 토큰 인덱스 기반의 지역적 임베딩을 사용해 라우팅 로직 없이 사전 로드가 가능하며, 이 임베딩 테이블을 CPU 메모리로 오프로드해 FFN 파라미터 메모리의 약 1/3을 절약할 수 있다.
3. 서빙 노드의 CPU 메모리에 임베딩 테이블을 복제하면 노드 간 전문 트래픽과 병렬 처리의 동기화 오버헤드를 제거할 수 있다.
4. 배치 토큰의 STEM 임베딩 중복 제거와 자주 사용되는 임베딩 캐싱을 통해 사전 로드 비용을 줄일 수 있으며, 모델 크기 증가에 따른 계산 비용은 제곱적으로 증가하지만 사전 로드 비용은 선형적으로 증가한다.
5. 이로 인해 CPU 오프로딩된 STEM은 대규모 모델에서 더욱 유리하고 확장성이 높다.

3.2.5 Context-length Adaptive Parameter Usage

[요약]

1. STEM은 토큰 수에 따라 활성화되는 매개변수 수를 조절하며, 각 레이어당 유일한 토큰 ID당 벡터를 생성하여 메모리 효율성을 높인다.
2. 유일한 토큰 수(Luniq)가 문맥 길이(L)에 비례하지 않아, 긴 문맥에서도 토큰당 FLOPs 증가 없이 매개변수 수가 증가한다.
3. STEM은 MoE와 병렬적으로 사용 가능하며, 표준 FFN을 STEM FFN으로 대체한 Mixture of STEM Experts 설계가 가능하다.
4. STEM은 긴 문맥 작업에서 밀집 모델보다 예측 가능한 지연 시간과 더 많은 매개변수 활성화를 통해 성능을 향상시킨다.
5. STEM의 유연한 매개변수 사용은 모델 용량 확장과 계산 효율성 사이의 균형을 유지한다.

3.3 Knowledge Editing with STEM

[요약]

1. STEM 임베딩을 활용해 사실 정보를 편집할 수 있는 가능성을 탐구하며, 입력 텍스트를 변경하지 않고 STEM 벡터만 수정해 모델의 출력을 조정하는 방법을 제시했다.
2. 원본과 대상 엔티티의 서브워드 토큰 수가 동일할 경우, 원본 토큰의 STEM 임베딩을 대상 토큰의 임베딩으로 일대일 대체함으로써 효과적인 편집이 가능하다.
3. 토큰 수가 다를 경우, 원본 길이가 더 길면 왼쪽 패딩 추가 또는 대상 토큰 반복을 통한 ‘복사’ 전략을 사용해 편집한다. 실험적으로 왼쪽 패딩이 더 효과적임이 확인되었다.
4. 예시로 스페인(원본)을 독일(대상)로 편집할 때, ‘마드리드’ 대신 ‘베를린’을 자동으로 생성하는 결과가 보고되었다.