[24’ ACL Findings] Finding and Editing Multi-Modal Neurons in Pre-Trained Transformers

Date: 2024.08.23 Updated: 2024.08.23

카테고리: Multimodal

🔍 Abstract

저자들은 LLaVA와 같은 Large Multimodal Model에서 Multimodal Neuron을 찾고, 이를 편집하여 모델의 행동을 수정하여 Multimodal Neuron의 존재를 증명한다. 지금까지 LLM의 FFN(Feed Forward Network)에서는 여러 역할을 하는 Neuron이 발견되면서 이를 흔히 key neuron이라고 불렀다. 이를 Multimodal World로 확장한 것이다.

1. Background

기존에 LLM에서는 FFN을 Key-Value Memory Retrieval로 해석하였다. 즉, Input은 Key로 변환되고 이를 통해 LLM Parameter에 저장된 Value 중 어떠한 것을 가져올지 결정한다. 이를 통하여 특정 word token에 FFN이 기여하게 되고, 이때 이러한 word token에 가장 큰 영향을 미치는 Neuron을 찾는 연구들이 진행되었다.

2. Previous Work: Mmns

2.1. Introduction

Multimodal에서도 이러한 연구가 Analogous하게 진행되었다. 즉, LLM 내부에서 Image Concept를 Text로 변환하는 특정한 Neuron을 찾는 것이다. 이를 Multimodal Neuron이라고 부른다.

ICCV CLVL (Closing the Loop Between Vision and Language) Workshop 2023에서는 이를 LMM에서 처음으로 제안한 논문인 Multimodal Neurons in Pretrained Text-Only Transformers가 발표되었다. (기존에 CLIP과 같은 Vision-Language Pretraining에서는 이미 제안되었다.)

해당 논문에서는 굉장히 간단한 Setting을 사용하였다. 즉, Visual Encoder + LLM에 추가적인 Instruction Text를 넣는 일반적인 LMM Setting이 아니라 Visual Encoder로 얻은 Embedding만 LLM에 넣어 Captioning을 진행하는 일종의 Toy Setting을 만든 것이다.

저자들은 해당 논문을 통해 Image Embedding은 그대로 Text Semantic을 가지고 있는 것이 아니며, FFN의 Multimodal Neuron을 통해 이해할 수 있는 Text Semantic으로 Translation이 일어난다고 주장하였다.

2.2. Gradient-based Approach

저자들은 이러한 Neuron을 찾기 위해 가장 Output Token Probability에 영향을 많이 미치는 Neuron을 계산하고자 하였다. 저자들은 Gradient-based Approach를 사용하여 다음과 같이 Attribution Score를 계산하였다.

\[g_ {k, c} = Z^ k _ p \frac{\partial y^ c}{\partial Z ^ k _ p}\]

이때 Attribution Score $g_ {k, c}$는 FFN Neuron $u_ k$가 Output Token $y^ c$에 미치는 영향을 나타낸다. 이때 Output Token은 모델이 처음으로 뱉는 Noun Token을 선택하였다.

여기서 $p$는 Image Patch Index이고, $Z^ k _ p$는 Pre-activation Output이다. 즉 FFN에서 $W_ {in} ^ \ell h_ p ^ {\ell - 1} \in \mathbb{R} ^ {\vert U ^ \ell \vert}$ 일 때 이 중 하나의 Neuron $u_ k$에서의 Pre-activation 값을 $Z ^ k _ p$라고 한다. 따라서 정확하게 이야기하자면 Attribution Score는 $g_ {k, c, p}$로 작성하는 것이 바람직하다.

직관적으로 이 값이 크기 위해서는 Neuron의 Output, 즉 $Z ^ k _ p$도 크고 $y^ c$에 대한 Neuron의 Effect $\frac{\partial y^ c}{\partial Z ^ k _ p}$도 커야 한다. 수학적으로 이는 Neuron의 First-order effect에만 해당하지만, 근사적으로 Output Token에 미치는 영향을 계산하기에는 충분하다.

여기서 $g_ {k, c}$를 계산하면, Top-k Neuron을 찾아내고, 이를 Multimodal Neuron이라고 부른다. 이때 이들의 의미를 해석하기 위하여, FFN의 Key-Value Memory Retrieval 해석을 적용한다. 찾아낸 Neuron이 $u_ k$라고 하자. 여기서 Key는 Non-linear activation function $g$를 적용한 Activation $A_ i ^ k = g( Z ^ k _ p)$가 되고, Value는 $W_ {out} ^ k$가 된다. 여기서 Layer $\ell$은 표기를 간단하게 하기 위하여 생략하였다. 따라서 Residual Stream에 더해지는 값은 $W_ {out} ^ k A_ i ^ k$가 되고, $A_ i ^ k$는 상수값, 즉 일종의 가중치로 해석할 수 있다.

이를 정리하면, Neuron $u_ k$가 Output에 미치는 영향은 정확히 $\text{softmax}(W_ u W_ {out} ^ k A_ i ^ k)$가 되며, $A_ i ^ k$는 단순 가중치이므로 softmax에서 무시된다. 여기서 $W_ u$는 Unembedding Matrix이다. 이를 통해 Multimodal Neuron이 Output Token에 미치는 영향을 계산할 수 있으며, 따라서 $\text{softmax}(W_ u W_ {out} ^ k)$에서 Top-k Vocabulary를 찾아 해당 Neuron의 의미를 해석할 수 있다.

2.3. Interpretation

저자들은 이를 사용하여 Multimodal Neuron을 찾아내었다. 위는 이를 찾아낸 결과이다. 참고로 Patch는 총 이미지 196개, A image of 3개로 199개이므로, Patch 199는 Last Token Prediction을 의미한다. 저자들은 이 중에서도 해석 가능한 몇 개의 Neuron만을 선택하여 분석하였다.

Image Activation의 경우, Image Prompt $x_ i \in [x_ 1, \cdots, x_ P]$에 대하여 Activation $A_ i ^ k$를 계산하고, 이 값의 상위 5%만 Visualization한 것이다. 이를 통하여 Multimodal Neuron이 어떠한 Image Patch에 대하여 Activation이 높은지를 확인할 수 있으며, Relevant Image Patch에서 높은 Activation을 보이는 것으로 보아 Semantic Relationship을 실제로 반영하고 있는 것을 알 수 있다. 또한, 이러한 Neuron은 Sensitivity, Specificity를 어느 정도 가지고 있었다고 한다. 즉, 동일한 Semantic에 대해서는 항상 잘 반응하고, 다른 Semantic에 대해서는 반응하지 않는다는 것이다.

2.4. Layer-wise Analysis

저자들은 추가로 이러한 Neuron의 분포를 분석하였으며, Early-to-mid Layer에서 높은 비율로 분포한다는 것을 알게 되었다. 저자들은 LLM에서 Factual Recall을 할 때, Early-to-mid Layer의 FFN이 중요한 역할을 한다는 Locating and Editing Factual Associations in GPT 논문과도 일맥상통한다고 주장한다. 다만 이 주장에는 몇 가지 허점이 있다.

Gradient-based Approach는 그 정의에 의해 Late Layer보다 Early Layer의 Contribution을 과대평가하는 경향이 있다.
LLM의 Factual Recall은 Parametric Knowledge를 통해 새로운 정보를 얻는 과정인 반면, Image의 정보를 해석하는 것은 Non-parametric Knowledge인 Image의 정보를 해석하는 것에 가깝다. 즉 Factual Recall이 아니라 Grounding Task에 가깝다.
완전한 LMM Setting이 아니고 거의 Image만 Input으로 들어가기 때문에 Image Processing 의존성이 높을 수 있다.

실제 $W_ u$를 사용하여 해석하였을 때도 실제 Output Token과 $W^ k _ {out}$ 사이에는 약한 Semantic Relationship이 존재할 뿐, 여기서 온전한 Knowledge Emergence가 일어나는 것으로 보이지는 않는다.

본 논문의 저자들은 이러한 Gradient-based Approach의 문제를 지적하였다. 추가로, Gradient-based Method는 backward pass가 필요하기 때문에 상당히 비용이 많이 든다. 따라서 저자들은 Forward pass에서 분석이 가능한 Contribution 방법론을 제안하였고, 이는 비용을 줄일 뿐만 아니라 해석 가능성을 높일 수 있다.

3.1. Forward Pass Contribution

저자들은 다음과 같이 Layer $l$의 Neuron $u_ i$가 Token $t$에 미치는 Contribution을 계산하였다.

\[s_ {i, t} ^ l = \mathbf{Q}^ l (i, t)\]

이때 Contribution $\mathbf{Q}$의 식을 유도하여 보자. 모델의 Hidden Dimension을 $d$, FFN의 Dimension을 $d_ m$, Vocabulary Size를 $v$라고 하자. 그리고 Unembedding Matrix $W_ u \in \mathbb{R} ^ {v \times d}$, FFN Output MLP Layer Matrix $W_ {out} \in \mathbb{R} ^ {d \times d_ m}$이라 하자. 이때 Layer $l$의 FFN이 Token에 미치는 영향은 $W_ u W_ {out} ^ l \mathbf{O}^ l _ {-1}$로 나타낼 수 있고, 이때 Activation Function Output $\mathbf{O}^ l _ {-1}$은 다음과 같이 정의된다.

\[\mathbf{O}^ l _ {-1} = \sigma \left( W_ {in} ^ l (\mathbf{a} ^ l _ {-1} + \mathbf{h} ^ {l - 1} _ {-1}) \right) \in \mathbb{R} ^ {d_ m}\]

이때 $\mathbf{a} ^ l _ {-1}$은 Attention Output, $\mathbf{h} ^ {l - 1} _ {-1}$은 Hidden State Output이다. 그러나 이는 Layer $l$의 FFN 전체가 Token에 미치는 영향을 계산한 것이므로, Neuron $u_ i$가 Token $t$에 미치는 영향을 계산하기 위해서는 이를 Neuron 단위로 쪼개어야 한다.

Transpose $\mathcal{T}(\cdot)$를 사용하면, $W_ u W_ {out} ^ l \mathbf{O}^ l _ {-1}$을 다음과 같이 재구성할 수 있다.

\[W_ u W_ {out} ^ l \mathbf{O}^ l _ {-1} = \sum_ {d_m} \mathcal{T} \left( W_ u W_ {out} ^ l \odot \mathcal{T}(\mathbf{O}^ l _ {-1}) \right)\]

즉, $W_ u W_ {out} ^ l \in \mathbb{R} ^ {v \times d_ m}$와 $\mathcal{T}(\mathbf{O}^ l _ {-1}) \in \mathbb{R} ^ {1 \times d_ m}$이 Element-wise Multiplication에 의해 $\mathbb{R} ^ {v \times d_ m}$로 확장된다. 이때 각 Column $\mathbb{R} ^ {v}$은 Neuron $u_ i$가 Vocabulary Space에 미치는 영향을 개별적으로 나타낸다. 따라서 Contribution $\mathbf{Q} ^ l$는 다음과 같이 정의된다.

\[\mathbf{Q}^ l = \mathcal{T} \left( W_ u W_ {out} ^ l \odot \mathcal{T}(\mathbf{O}^ l _ {-1}) \right) \in \mathbb{R} ^ {d_ m \times v}\]

여기서 Indexing을 통해 $s_ {i, t} ^ l = \mathbf{Q}^ l (i, t)$로 정의하여 Layer $l$의 Neuron $u_ i$가 Token $t$에 미치는 Contribution을 계산할 수 있다. 저자들은 Noun Token $t$에 대해서만 이러한 Contribution Score를 계산하였고, 여기서 Top-k Neuron을 찾아 Multimodal Neuron이라고 부르기로 하였다.

저자들은 추가로 학습을 통한 Model Editing도 진행하였는데, 여기서는 그 과정은 생략하기로 한다.

\[\text{loss} = o_ i ^ l (\mathbf{w}^ \prime \mathbf{v}_ 0 - \mathbf{w}^ \prime \mathbf{v}_ 1) + \beta \Vert \Delta \mathbf{w} \Vert _ 2\]

3.2. Analysis

저자들은 Multimodal Neuron이 (1) Sensitivity, (2) Specificity를 가지고 있으며, (3) Causal-Effect를 가지고 있다는 것을 보였다. 이를 통해 실제로 유의미한 Neuron을 찾아냈다는 것을 보였다.

3.2.1. Sensitivity

Sensitivity는 동일한 Concept에 대하여 항상 높은 Activation을 보이는 것을 의미한다. 저자들은 이를 (1) Visual, (2) Textual, (3) Similar & Different Image에 대하여 확인하였다.

먼저 Visual에 대하여 확인하기 위하여 저자들은 찾은 Neuron $u_ i$에 대하여 Image Patch에 대한 Activation을 계산하고, 이를 Top-5%만 Visualization한 결과를 보였다. 이를 통해 Semantic Relationship이 실제로 반영되어 있음을 확인하였다.

다음으로 Textual에 대하여 확인하기 위하여 저자들은 찾은 Neuron $u_ i$이 활성화하는 $W^ l _ {out} [:, i]$에 Unembedding Matrix $W_ u$를 곱하여 Top-k Vocabulary를 찾아내었다. 위에서 설명한 Gradient-based Approach(Mmns)는 이러한 Neuron 간의 일관성이 없었다. 그러나 이러한 Forward Pass Contribution을 통해 찾은 Neuron은 Semantic Information이 일관되게 나타나는 것을 확인할 수 있었다.

마지막으로 저자들은 Region Invariance, Cross-Image Invariance를 측정하였다. Region Invariance는 Image Patch를 Shuffling하여도 Activation이 유지되는지를 측정한 것이고, Cross-Image Invariance는 Neuron이 동일한 Concept을 가진 다른 Image에 대해서도 Activation이 유지되는지를 측정한 것이다. 이를 통해 Sensitivity가 유지되는 것을 확인하였다.

3.2.2. Specificity

Specificity는 Neuron이 하나의 Concept에 대해서만 Activation을 보이고, 다른 Concept에 대해서는 Activation을 보이지 않는 것을 의미한다. 위 그림을 통하여, 각 Neuron이 다른 Concept에 대해서는 Activation을 보이지 않는 것을 확인할 수 있다.

정량적으로도 연관된 Concept에 대하여 Contribution 값이 더 높은 것을 알 수 있었다.

3.2.3. Causal-Effect

그렇다면 실제로 이러한 Neuron에 Perturbation을 가하면 모델이 제대로 된 대답을 하지 못할까? 저자들은 이러한 Causal-Effect를 측정하였다. 그 결과, Perturbation을 가한 경우 해당 Token이 생성되지 않았으며, 새로운 답변을 생성하거나 무의미한 답변을 생성하였다. 이를 통해 Causal-Effect가 존재함을 확인하였다.

저자들은 이러한 관찰을 바탕으로 Multimodal Neuron Editing도 진행하였으며, 여기서는 간단한 결과만 관찰해 보자.

3.3. Layer-wise Analysis

저자들은 Layer-wise Analysis를 통해 Deep Layer에서 Multimodal Neuron이 더 많이 발견된다는 것을 확인하였다. 이는 Factual Recall과는 다른 결과이다. 이를 통해 Factual Recall과 Grounding Task의 차이를 확인할 수 있다. 한편, 이러한 분석에도 한계점이 있다.

Forward Contribution을 이용한 방법은 해당 FFN의 Parametric Knowledge가 반드시 Vocabulary Space와 Align하는 경우에만 유효하다. 즉, 실제 Output Token이 생성되기 위해서는 Syntactic Concept이 추상화되며 Semantic Concept이 도출되어야 하는데, 필연적으로 Early Layer에서 발생하는 Syntactic Concept은 Vocabulary Space와 Align되어 있지 않을 가능성이 높아 Multimodal Neuron으로 검출되지 않을 가능성이 높다. 따라서 어떻게 보면 Output Token과 가까운 위치에 있는 Late Layer에서 Multimodal Neuron이 더 많이 발견되는 것은 당연한 결과이다.

💡 Summary

여기에서는 두 개의 논문을 리뷰하였다. 해당 논문들의 내용을 정리하면 다음과 같다.

[24’ ICCV-WS] Multimodal Neurons in Pretrained Text-Only Transformers
- Gradient-based Attribution 계산을 통해 처음으로 Multimodal Neuron을 발견하였음; Early-to-Mid Layer에서 이러한 FFN이 발견되나, 이는 Gradient-based Approach의 Bias 때문일 수 있음
- Image Embedding은 그대로 Text Semantic을 가지고 있는 것이 아니며, FFN의 Multimodal Neuron을 통해 이해할 수 있는 Text Semantic으로 Translation이 일어난다고 주장
[24’ ACL Findings] Finding and Editing Multi-Modal Neurons in Pre-Trained Transformers
- LLaVA와 같은 LMM Setting에서 Forward Contribution을 통해 Multimodal Neuron을 찾아내고, Late Layer에서 더 많이 발견되는 것을 확인
- 이러한 Neuron의 Sensitivity, Specificity, Causal-Effect를 확인

I'm rubatoyeong