[LLM 해석가능성] 피드포워드 네트워크(FFN)를 키-값 메모리로 이해하기
#ML#LLM Interpretability

[LLM 해석가능성] 피드포워드 네트워크(FFN)를 키-값 메모리로 이해하기

0. Introduction

Week 1에서 남은 질문

이전 글에서 Transformer의 구조를 따라갔습니다. Self-Attention이 토큰 간 정보를 교환하는 메커니즘이라는 것, Residual Stream이 정보의 공유 버스라는 것을 확인했습니다.

그런데 Transformer 블록에는 Attention 말고도 또 하나의 핵심 구성요소가 있습니다 — FFN(Feed-Forward Network). FFN은 Residual Stream에서 어떤 역할을 하고 있을까요?

Geva et al.의 논문 "Transformer Feed-Forward Layers Are Key-Value Memories" (EMNLP 2021)은 이 질문에 대한 놀라운 답을 제시합니다. FFN이 단순한 비선형 변환이 아니라, 사실적 지식을 저장하는 키-값 메모리라는 것입니다.

이 글의 여정
1. FFN의 수식 구조
2. Key-Value Memory 대응
3. Attention K-V와의 비교
4. FFN이 실제로 저장하는 것
5. ROME으로의 연결

1. FFN의 수식 구조

차원(shape)까지 정확히 이해하기

Transformer 블록 안에서 Attention 이후에 실행되는 FFN은 다음과 같은 수식을 따릅니다:

FFN(x)=W2⋅ReLU(W1x+b1)+b2\text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x + b_1) + b_2

각 구성요소의 shape을 살펴봅시다:

차원 정리
입력 x∈Rdmodelx \in \mathbb{R}^{d_{model}} (예: 768)W1∈Rdff×dmodelW_1 \in \mathbb{R}^{d_{ff} \times d_{model}} (예: 3072 × 768) — 차원 확장ReLU 적용 → 음수를 0으로W2∈Rdmodel×dffW_2 \in \mathbb{R}^{d_{model} \times d_{ff}} (예: 768 × 3072) — 차원 축소출력: 다시 Rdmodel\mathbb{R}^{d_{model}} → Residual Stream에 더해짐
Q1
W1xW_1 x를 계산할 때, W1W_1의 각 행(row) 벡터는 입력 xx와 어떤 연산을 수행하게 될까요?
내적(dot product)입니다. W1W_1의 i번째 행과 xx의 내적이 출력 벡터의 i번째 원소가 됩니다.
(W1x)i=wi(1)⋅x(W_1 x)_i = w_i^{(1)} \cdot x
즉, W1W_1의 각 행은 입력 벡터와의 유사도를 측정하는 것입니다.

2. Key-Value Memory 대응

FFN의 수식을 재해석하다

Section 1에서 우리는 W1W_1의 각 행이 입력과 내적을 수행하여 매칭 여부를 판단한다는 것을 확인했습니다. 이것은 무엇과 비슷할까요?

데이터베이스에 query를 던져서 매칭되는 key를 찾고, 대응하는 value를 가져오는 구조 — 바로 Key-Value Memory입니다.

Key-Value 대응
W1W_1의 i번째 행 = key kik_iki⋅xk_i \cdot x = key와 입력의 매칭 점수ReLU → 매칭 점수가 양수인 것만 활성화 → mim_iW2W_2의 i번째 열 = value viv_iW2m=∑imi⋅viW_2 m = \sum_i m_i \cdot v_i = 활성화된 value의 가중합
FFN(x)=∑i=1dffReLU(ki⋅x+bi)⏟mi (활성화 강도)⋅vi⏟value 벡터\text{FFN}(x) = \sum_{i=1}^{d_{ff}} \underbrace{\text{ReLU}(k_i \cdot x + b_i)}_{m_i \text{ (활성화 강도)}} \cdot \underbrace{v_i}_{\text{value 벡터}}
핵심
이 대응은 비유가 아니라 수식적으로 정확히 성립합니다. FFN의 각 뉴런(i)은 하나의 key-value 쌍이며, 입력과 매칭된 key에 대응하는 value를 가중합하여 출력합니다.
Q1
W1W_1에 3072개의 행이 있다는 것은, 이 FFN 층에 메모리 슬롯이 몇 개 있다는 뜻일까요?
3072개입니다. 각 행이 하나의 key이고, 대응하는 W2W_2의 열이 그 key의 value이므로, 3072개의 key-value 쌍이 이 층 하나에 저장되어 있습니다.

3. Attention K-V vs FFN K-V

같은 이름, 다른 메커니즘

Attention에도 Key와 Value가 있고, FFN에도 Key와 Value가 있습니다. 이 둘은 같은 것일까요?

핵심 차이는 동적(dynamic) vs 고정(static)입니다.

Attention — 동적 Key-Value
K, V가 현재 입력 시퀀스의 토큰으로부터 매번 새로 계산됩니다. "지금 이 문장에서 어떤 단어의 정보를 가져올까" → 문맥 의존적
FFN — 고정된 Key-Value
K (W1W_1 행), V (W2W_2 열)가 훈련 후 고정됩니다. 문장이 바뀌어도 W1W_1, W2W_2는 동일 → 고정된 지식 저장소

비유하자면: Attention은 대화 중에 상대방이 방금 한 말을 참조하는 것이고, FFN은 학교에서 배워서 머릿속에 저장해둔 지식을 꺼내 쓰는 것입니다.

핵심 정리
Attention은 문맥을 읽고, FFN은 지식을 꺼낸다.
Q1
Attention과 FFN이 Transformer 블록 안에서 순서대로 실행된다는 것을 떠올려보세요 (Attention → FFN). 이 순서가 의미적으로 자연스러운 이유는 무엇일까요?
먼저 Attention이 문맥 정보를 모아서 '지금 무엇에 대해 이야기하고 있는가'를 파악합니다. 그 다음 FFN이 파악된 subject에 대한 저장된 지식을 꺼냅니다.

예를 들어 "The Eiffel Tower is in..."에서 Attention이 "Eiffel Tower"라는 subject를 모은 후, FFN이 "Eiffel Tower → Paris"라는 사실을 recall하는 것입니다.

4. FFN이 실제로 저장하는 것

Geva et al.의 실험적 발견

수학적으로 key-value memory 구조라는 건 알겠는데, 실제로 훈련된 모델의 FFN에서 key와 value는 구체적으로 무엇을 저장하고 있을까요?

Geva et al.은 실험을 통해 두 가지를 밝혔습니다:

Key 측 분석
특정 key(뉴런)는 특정 입력 텍스트 패턴에 의해 활성화됩니다. 예를 들어 어떤 key는 "수도 이름이 나오는 문맥"에서만 활성화되고, 다른 key는 "연도가 나오는 문맥"에서 활성화됩니다. 상위 layer로 갈수록 더 추상적/의미적 패턴에 반응합니다.
Value 측 분석
활성화된 key에 대응하는 value를 unembedding으로 변환하면, 특정 출력 토큰 분포에 대응합니다. 예를 들어 "수도" key에 대응하는 value는 도시 이름 토큰들 ("Paris", "London", "Tokyo" 등)의 확률을 높입니다.
핵심 통찰
FFN의 각 뉴런은 "이런 패턴이 나타나면 → 이런 토큰을 예측하라"는 if-then 규칙으로 해석할 수 있습니다.

또한 층(layer)에 따라 패턴의 성격이 다릅니다:

하위 층 (Layer 1-4)
구문적/표면적 패턴
예: "대문자 뒤에 대문자가 올 가능성"
상위 층 (Layer 8-12)
의미적/사실적 패턴
예: "에펠탑에 대한 문맥 → Paris"
Q1
FFN의 value를 unembedding matrix로 변환하면 토큰 분포가 나온다고 했습니다. 이것은 이전 포스트에서 배운 어떤 기법과 비슷할까요?
Logit Lens입니다! Logit Lens는 residual stream의 중간 상태에 unembedding을 적용하여 '현재 시점의 예측'을 관찰하는 기법이었습니다. FFN의 value에 같은 것을 적용하면, 각 메모리 슬롯이 어떤 토큰을 '기억하고 있는지' 엿볼 수 있는 것입니다.

5. 요약 + ROME으로의 다리

핵심 정리와 다음 여정

FFN 수식 구조
x → W₁(확장) → ReLU → W₂(축소) → 출력
Key-Value 대응
W₁ 행 = key (패턴 매칭), W₂ 열 = value (정보 출력)
Attention과의 차이
Attention은 동적 문맥, FFN은 고정된 지식
계층적 패턴
하위 층은 구문적, 상위 층은 의미적 패턴
다음 포스트 예고

FFN이 고정된 사실 저장소라면, 자연스럽게 두 가지 질문이 떠오릅니다:

어디에?특정 사실이 수십 개의 FFN 층 중 정확히 어디에 저장되어 있는가? → Causal Tracing
편집?찾은 위치의 가중치를 수정하여 사실을 바꿀 수 있는가? → ROME

다음 포스트에서 Meng et al.의 ROME 논문을 따라가며, 이 두 질문에 답합니다.