Transformer 제대로 이해하자

Transformer 하면 “Attention is all you need”를 떠올리며 대충 Self-Attention과 MLP가 있고, 이를 반복해서 Encoder와 Decoder가 존재하며 예측한다고 겉핡기로 아는 경우가 많은데 작년에 나온, 3B1B Transformer 영상을 보면 self-attention과 MLP에 대해 더 깊은 의미를 생각해볼 수 있고 후속 연구들도 이런 해석을 지지하고 있어, 이런 내용을 포함해 Transformer의 작동 원리를 말로서 정리해보고자 한다.

 

1. Encoder-Only 구조로 큰 틀에서 이해하기

먼저 핵심 원리를 파악하기 위해 Encoder-Only 구조에 대해 먼저 살펴보겠다. Transformer는 입력 문장 내 단어 임베딩들 간의 상호작용을 위해 self-attention 메커니즘을 적용한다. 각 토큰은 이 과정을 통해 문맥 정보를 반영한 합성 embedding으로 재구성된다. 특히multi-head attention을 사용함으로써 서로 다른 관점에서 독립적으로 임베딩을 만들어 다양한 문맥적 특성이 동시에 반영된다.

 

그 후 MLP 블록에서는 첫 번째 선형 변환을 통해 semantic probing을 수행한다. 구체적으로, self-attention에 의해 형성된 embedding E에 대해 가중치 행렬 W₁의 각 row는 기존 embedding과 다른 서로 다른 semantic direction을 내포하며, W₁와 E의 곱셈을 통해 각 semantic 과의 유사도를 얻어서 입력 embedding이 latent space에서 어느 정도 활성화되는지를 평가한다. 이후 첫 번째 선형 변환의 결과에 GELU나 ReLU와 같은 비선형 활성화 함수를 적용하여 게이팅 효과를 부여한다. 이 과정으로 기존 embedding에 대해 중요한 semantic 정보만을 남기도록 조정된다. 마지막으로, 이 유사도 정보를 바탕으로 두 번째 선형 변환을 수행하고 가중치 행렬 W₂와 선형 결합하여 새로운 embedding을 생성한다. W₂는 각 의미 성분을 재구성하는 템플릿 역할을 하며, 첫 번째 단계에서 얻은 정보를 기반으로 의미적 합성을 수행한다고 해석할 수 있다.

 

정리하면, Transformer의 self-attention은 단어 임베딩 간의 내부 소통을 통해 문맥 정보를 반영하고, multi-head attention은 이를 다양한 관점에서 독립적으로 처리한다. 이후 MLP의 두 계층은 입력 embedding과 잠재 의미 방향 간의 내적을 통해 각 semantic component의 활성화 정도를 평가한 후, 비선형 활성화로 선택적으로 게이팅하고, semantic 유사도에 기반한 템플릿을 적용하여 최종적으로 새로운 embedding을 얻어낸다. 이와 같이 MLP는 단순한 Feed-Forward Network 이상의 역할을 수행하며, Self-Attention으로 획득한 문맥 정보를 더욱 정제 및 확장하여 의미적 재구성(semantic synthesis)을 이루는 모듈로 해석할 수 있다.

 

또한, Transformer는 문장 길이에 상관없이 충분히 scale up이 가능하다. n개의 토큰을 입력받아도 Embedding matrix는 각 토큰의 vector 크기만 고려하여 동일하게 적용되고, Self-Attention에서 사용하는 W_q, W_k, W_v는 모든 토큰에 동일하게 곱해지므로 입력 길이 n과 관계없이 처리된다. 이후 MLP도 각 token별로 독립적으로 적용되므로 전체 구조는 문장 길이에 관계없이 안정적으로 작동한다.

 

2. 위 해석이 과대해석인가?

Geva et al. (2021)은 Transformer의 FFN 계층이 단순한 비선형 변환(projection) 이상의 역할을 수행한다고 주장한다. 이 논문에서는 FFN 계층이 마치 key-value memory처럼 작동하여, 각 뉴런이 입력 토큰에 내재된 discrete한 정보를 저장하고, 필요할 때 이를 검색하는 역할을 한다고 설명한다. 또한 Rogers et al. (2020)의 “A Primer in BERTology: What We Know About How BERT Works”에서는 BERT의 FFN 계층이 self-attention을 통해 집계된 문맥 정보를 단순히 전달하는 것이 아니라, 이를 bottleneck 역할로 필터링하고 재가중치화하여 더 정제된 표현을 생성하는 데 기여한다고 분석한다. 그리고 Clark et al. (2019)은 “What Does BERT Look At? An Analysis of BERT’s Attention”을 통해, FFN 계층이 단순한 선형 변환을 넘어서 중요한 semantic signal을 선택적으로 증폭하여 토큰의 표현을 재구성하는 데 기여한다고 분석한다.

 

위와 같은 Transformer 관련 후속 논문들은 Transformer의 MLP (FFN) 블록이 Self-Attention으로 집계한 문맥 정보를 단순히 전달하는 역할을 넘어서, semantic probing(각 의미 방향에 대한 활성화 정도 측정), semantic gating(중요 정보를 선택적으로 강화) 및 semantic synthesis(선택된 정보를 재구성)와 같은 복합적인 역할을 수행한다고 해석할 수 있음을 보여준다.

3. Decoder 를 사용하는 Transformer구조

다음으로 Decoder-Only 구조를 살펴보자. Decoder-Only 구조는 (x₁, x₂, …, xₙ)을 입력받아 그 다음 토큰 xₙ₊₁을 예측하는 문제에 사용된다. 모델에 (x₁, x₂, …, xₙ)을 입력하면 Embedding과 Positional Encoding을 거쳐 각 토큰에 대한 초기 embedding을 생성하고, 이후 Masked Self-Attention과 MLP를 통해 (E₁, E₂, …, Eₙ)의 embedding을 얻는다. 여기서 각 Eₖ는 x₁부터 xₖ까지의 정보만 담으며, prediction head(linear projection 후 softmax)를 통해 word dictionary와의 유사도를 계산하여 확률 분포를 만든다. 학습 시에는 전체 시퀀스의 shifted 버전, 즉 (x₂, …, xₙ₊₁)을 예측하도록 loss를 계산하며, 실제 사용할 때는 autoregressive하게 마지막 임베딩 토큰 Eₙ만을 사용해 다음 토큰 xₙ₊₁을 예측하고 가장 확률이 높은 단어를 선택한다.

 

마지막으로 Encoder-Decoder 구조에 대해 설명한겠다.. Encoder-Decoder 구조는 (x₁, …, xₙ)이라는 문장을 입력받아 (y₁, …, yₘ)의 번역 결과나 문맥 전환(context switching) 등 다양한 Seq2Seq 작업에 적용할 수 있으며, 학습 시에는 Teacher Forcing 기법을 사용한다. 구체적으로, 타깃 문장 (y₁, …, yₘ)이 있을 때 Decoder의 입력으로 (BOS, y₁, y₂, …, yₘ₋₁)을 넣는다. 입력 문장 (x₁, …, xₙ)은 Encoder를 통해 self-attention과 MLP로 처리되어 내부 정보를 섞고 semantic 재구성을 거친 embedding E = (E₁, …, Eₙ)을 생성한다. Decoder는 미래 토큰 정보가 섞이지 않도록 masked self-attention을 적용하여 (E’₀, E’₁, E’₂, …, E’ₘ₋₁)와 같은 embedding을 생성한다. 이후 Decoder는 Encoder의 출력 E와 Decoder의 embedding E’ 사이에서 attention 연산을 수행하여 두 정보를 융합하고, 이 융합된 정보를 MLP를 통해 재구성한다. 이 과정을 반복하여 생성된 k번째 token embedding Eₖ는 k번째까지의 실제 정보와 기존 입력 (x₁, …, xₙ)의 정보를 모두 담고 있으며, Decoder only 때처럼 마찬가지로 prediction head를 통해 word dictionary와의 유사도 기반 확률 분포를 바탕으로 실제 정답 (y₁, y₂, …, yₘ)에 가까워지도록 학습을 진행한다. 실제 사용할 때는 Encoder에 전체 문장 (x₁, …, xₙ)을 입력하고, Decoder는 처음에 BOS만 넣어 다음 토큰을 예측한 후 (BOS, y₁), (BOS, y₁, y₂) 등의 방식으로 이전 예측 결과를 추가하면서 autoregressive하게 전체 문장을 생성하면 된다.