Transformer FFN

Transformer Positionwise Feed Forward Network(FFN) 완벽 이해

Transformer 구조를 공부하다 보면 Self-Attention만큼 자주 등장하는 것이 바로 Positionwise Feed Forward Network(FFN)입니다.

Transformer의 Encoder와 Decoder에서는 Attention 연산이 끝난 후 FFN을 적용합니다.
FFN은 각 토큰의 정보를 독립적으로 변환하면서 특징을 더욱 풍부하게 만들어 주는 역할을 합니다.

이번 글에서는 PyTorch로 구현한 PositionwiseFeedForward 클래스를 통해 FFN의 구조와 d_model, d_ff, Linear의 동작을 쉽게 살펴보겠습니다.

1. Positionwise Feed Forward Network란?

Transformer에서 FFN은 일반적으로 다음과 같은 구조를 사용합니다.

입력
 ↓
Linear(d_model → d_ff)
 ↓
ReLU
 ↓
Dropout
 ↓
Linear(d_ff → d_model)
 ↓
출력

수식으로 표현하면 다음과 같습니다.

FFN(x)=W2(ReLU(W1x+b1))+b2FFN(x)=W_2(ReLU(W_1x+b_1))+b_2

여기서 중요한 점은 각 토큰에 대해 동일한 FFN을 독립적으로 적용한다는 것입니다.

예를 들어 문장이 다음과 같다고 해보겠습니다.

나는 / 오늘 / 학교에 / 간다

각 단어에 대한 벡터가 존재한다면,

나는      → FFN
오늘      → FFN
학교에    → FFN
간다      → FFN

와 같이 각각 처리됩니다.

토큰끼리 정보를 교환하는 역할은 주로 Self-Attention이 담당하고, FFN은 Attention을 통해 얻어진 각 토큰의 표현을 다시 변환하고 강화하는 역할을 합니다.

2. PyTorch로 FFN 구현하기

다음은 간단한 Positionwise Feed Forward Network 구현입니다.

import torch
import torch.nn as nn
import torch.nn.functional as F


# Positionwise Feed Forward Network 클래스
class PositionwiseFeedForward(nn.Module):

    def __init__(self, d_model, d_ff, dropout=0.1):

        # d_model:
        # 입력 벡터의 차원이며,
        # 첫 번째 Linear의 입력 차원과
        # 두 번째 Linear의 출력 차원
        # 을 결정합니다.

        # d_ff:
        # 첫 번째 Linear의 출력 차원이며,
        # 두 번째 Linear의 입력 차원입니다.

        # dropout:
        # Dropout에서 뉴런을 무작위로 제거할 확률입니다.

        super().__init__()

        # 첫 번째 Linear
        self.w1 = nn.Linear(d_model, d_ff)

        # 두 번째 Linear
        self.w2 = nn.Linear(d_ff, d_model)

        # Dropout
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):

        # 1. 첫 번째 Linear
        # 2. ReLU 활성화 함수
        # 3. Dropout
        # 4. 두 번째 Linear
        return self.w2(
            self.dropout(
                F.relu(
                    self.w1(x)
                )
            )
        )

3. d_model은 무엇인가?

먼저 d_model을 이해해야 합니다.

d_model = 512

d_model은 Transformer에서 사용하는 토큰 벡터의 기본 차원입니다.

예를 들어 하나의 토큰이 다음과 같이 표현된다고 생각해 보겠습니다.

[0.12, 0.35, 0.87, ..., 0.21]

이 벡터의 크기가 512라면,

토큰 벡터
   ↓
[0.12, 0.35, ..., 0.21]
   ↓
512차원

이 됩니다.

따라서 FFN의 입력과 최종 출력도 다시 512차원이 됩니다.

입력
512차원
   ↓
FFN
   ↓
출력
512차원

즉,

nn.Linear(d_model, d_ff)

와

nn.Linear(d_ff, d_model)

를 연결하여 입력 차원과 출력 차원을 동일하게 유지합니다.

4. d_ff는 무엇인가?

코드에서는 다음과 같이 설정했습니다.

d_ff = 64

그러면 FFN은 다음과 같은 구조가 됩니다.

512
 ↓
64
 ↓
512

즉,

self.w1 = nn.Linear(512, 64)
self.w2 = nn.Linear(64, 512)

입니다.

첫 번째 Linear에서 벡터의 차원을

512→64512 \rightarrow 64

로 변경하고,

두 번째 Linear에서 다시

64→51264 \rightarrow 512

로 변경합니다.

5. 그런데 왜 중간 차원을 변경할까?

FFN의 핵심 아이디어 중 하나입니다.

단순히

512 → 512

로 처리하는 것보다 중간에서 다른 차원으로 변환하고 활성화 함수를 적용하면서 비선형적인 특징 변환을 수행할 수 있습니다.

구조적으로 보면,

Input
  │
  │ 512
  ▼
Linear
  │
  │ 64
  ▼
ReLU
  │
  ▼
Dropout
  │
  ▼
Linear
  │
  │ 512
  ▼
Output

과 같습니다.

여기서 첫 번째 Linear가 특징을 변환하고, ReLU가 비선형성을 추가하며, 두 번째 Linear가 다시 Transformer의 기본 차원인 d_model로 복원합니다.

6. nn.Linear는 실제로 무엇을 하는가?

PyTorch의

nn.Linear(in_features, out_features)

는 기본적으로 다음과 같은 연산을 수행합니다.

y=xWT+by=xW^T+b

따라서

nn.Linear(512, 64)

라면 내부적으로 다음과 같은 Parameter를 갖습니다.

weight
shape = (64, 512)

bias
shape = (64)

따라서 입력이

x
shape = (batch, sequence, 512)

라면,

self.w1(x)

의 결과는

(batch, sequence, 64)

가 됩니다.

그리고 두 번째 Linear는

nn.Linear(64, 512)

이므로

(batch, sequence, 64)
        ↓
      Linear
        ↓
(batch, sequence, 512)

가 됩니다.

7. 실제 Shape으로 확인해 보기

예를 들어 Transformer에 다음과 같은 입력이 들어왔다고 가정하겠습니다.

x.shape

결과가

(2, 10, 512)

라고 해보겠습니다.

각 숫자의 의미는 다음과 같습니다.

2   → Batch Size
10  → Sequence Length
512 → d_model

즉,

2개의 문장
각 문장에는 10개의 토큰
각 토큰은 512차원

입니다.

첫 번째 Linear를 통과하면,

(2, 10, 512)
      ↓
Linear(512 → 64)
      ↓
(2, 10, 64)

가 됩니다.

ReLU와 Dropout을 거쳐도 Shape은 그대로입니다.

(2, 10, 64)

마지막 Linear를 통과하면,

(2, 10, 64)
      ↓
Linear(64 → 512)
      ↓
(2, 10, 512)

가 됩니다.

따라서 FFN을 통과한 후에도 전체 Tensor의 Shape은 처음과 동일합니다.

8. 왜 Sequence 차원은 그대로 유지되는가?

이 부분은 Transformer를 공부할 때 매우 중요합니다.

FFN은 일반적으로 마지막 차원에 대해서 Linear 연산을 수행합니다.

예를 들어,

(batch, sequence, d_model)

형태의 입력이 있다면,

(2, 10, 512)

에서 Linear는 512 부분을 변환합니다.

(2, 10, 512)
        ↓
(2, 10, 64)

즉,

Batch       → 그대로
Sequence    → 그대로
Feature     → 변경

입니다.

따라서 Positionwise라는 이름처럼 각 위치의 토큰에 동일한 네트워크를 적용할 수 있습니다.

9. ReLU는 왜 사용하는가?

코드에서는 다음과 같이 작성했습니다.

F.relu(self.w1(x))

ReLU는 다음과 같은 함수입니다.

ReLU(x)=max(0,x)ReLU(x)=max(0,x)

즉,

음수 → 0
양수 → 그대로

입니다.

예를 들어,

[-2.0, 1.5, -0.3, 4.2]

가 ReLU를 통과하면,

[0, 1.5, 0, 4.2]

가 됩니다.

ReLU의 중요한 역할은 비선형성(non-linearity)을 추가하는 것입니다.

Linear Layer만 여러 개 연결하면 결과적으로 하나의 Linear Transformation으로 표현할 수 있습니다.

하지만 그 사이에 ReLU와 같은 활성화 함수를 넣으면 보다 복잡한 특징을 학습할 수 있습니다.

10. Dropout은 어떤 역할을 하는가?

코드에서는 다음과 같이 설정했습니다.

dropout = 0.2

그리고

self.dropout = nn.Dropout(p=dropout)

로 사용합니다.

즉,

Dropout 확률 = 20%

입니다.

학습 과정에서 일부 값을 무작위로 0으로 만들어 모델이 특정 특징에 지나치게 의존하는 것을 방지합니다.

구조는 다음과 같습니다.

Linear
  ↓
ReLU
  ↓
Dropout
  ↓
Linear

다만 중요한 점은 평가 모드에서는 Dropout이 비활성화된다는 것입니다.

model.train()

에서는 Dropout이 적용되고,

model.eval()

에서는 Dropout이 적용되지 않습니다.

11. 전체 코드 실행하기

다음과 같이 FFN을 생성할 수 있습니다.

d_model = 512
d_ff = 64
dropout = 0.2

ffn = PositionwiseFeedForward(
    d_model=d_model,
    d_ff=d_ff,
    dropout=dropout
)

그리고 임의의 입력을 만들어 보겠습니다.

x = torch.randn(2, 10, 512)

output = ffn(x)

print("Input shape :", x.shape)
print("Output shape:", output.shape)

결과는 다음과 같습니다.

Input shape : torch.Size([2, 10, 512])
Output shape: torch.Size([2, 10, 512])

중간 과정을 직접 확인하면 더 이해하기 쉽습니다.

x1 = ffn.w1(x)

print(x1.shape)

결과:

torch.Size([2, 10, 64])

ReLU를 적용하면,

x2 = F.relu(x1)

print(x2.shape)

결과:

torch.Size([2, 10, 64])

Dropout을 적용해도,

x3 = ffn.dropout(x2)

print(x3.shape)

결과:

torch.Size([2, 10, 64])

마지막 Linear를 적용하면,

x4 = ffn.w2(x3)

print(x4.shape)

결과:

torch.Size([2, 10, 512])

따라서 전체적인 Shape 변화는 다음과 같습니다.

Input
(2, 10, 512)
      ↓
Linear
512 → 64
      ↓
(2, 10, 64)
      ↓
ReLU
      ↓
(2, 10, 64)
      ↓
Dropout
      ↓
(2, 10, 64)
      ↓
Linear
64 → 512
      ↓
Output
(2, 10, 512)

12. Transformer에서 FFN의 역할

Transformer에서는 대략 다음과 같은 흐름으로 데이터가 처리됩니다.

Token Embedding
      ↓
Positional Encoding
      ↓
Self-Attention
      ↓
Add & Norm
      ↓
Feed Forward Network
      ↓
Add & Norm

Self-Attention은 토큰 간의 관계를 학습합니다.

예를 들어,

나는 오늘 학교에 간다

에서 간다라는 토큰이 학교에와 어떤 관계를 갖는지 학습하는 것이 Attention의 중요한 역할입니다.

반면 FFN은 Attention 결과를 각각의 위치에서 독립적으로 처리하여 각 토큰의 특징 표현을 더욱 변환하고 강화합니다.

따라서 간단하게 정리하면,

Self-Attention
→ 토큰과 토큰 사이의 관계를 학습

FFN
→ 각 토큰의 특징을 변환하고 강화

라고 이해할 수 있습니다.

13. d_model과 d_ff 설정에서 주의할 점

실제 Transformer에서는 일반적으로 d_ff가 d_model보다 훨씬 크게 설정되는 경우가 많습니다.

예를 들어 Transformer의 전형적인 설정에서는,

d_model = 512
d_ff   = 2048

처럼 사용할 수 있습니다.

즉,

512 → 2048 → 512

와 같은 구조입니다.

반면 이번 예제에서는

d_model = 512
d_ff = 64

이므로,

512 → 64 → 512

가 됩니다.

따라서 이 코드는 FFN의 구조를 이해하기 위한 간단한 예제로 보는 것이 좋습니다.

d_ff를 크게 설정하면 더 많은 특징을 표현할 수 있지만, 그만큼 Parameter 수와 계산량도 증가합니다.

14. Parameter 개수도 직접 확인해 보기

PyTorch에서는 다음과 같이 Parameter를 확인할 수 있습니다.

for name, param in ffn.named_parameters():
    print(name, param.shape)

결과는 다음과 비슷합니다.

w1.weight torch.Size([64, 512])
w1.bias   torch.Size([64])

w2.weight torch.Size([512, 64])
w2.bias   torch.Size([512])

Parameter 개수를 계산하면,

첫 번째 Linear:

512×64+64=32,832512\times64+64=32,832

두 번째 Linear:

64×512+512=33,28064\times512+512=33,280

따라서 전체 Parameter는

32,832+33,280=66,11232,832+33,280=66,112

개입니다.

즉, 이 FFN에는 총 66,112개의 학습 가능한 Parameter가 존재합니다.

15. 핵심 정리

Positionwise Feed Forward Network는 Transformer의 중요한 구성 요소입니다.

전체 구조는 다음과 같이 기억하면 됩니다.

          d_model
             │
             ▼
      Linear(d_model → d_ff)
             │
             ▼
           ReLU
             │
             ▼
          Dropout
             │
             ▼
       Linear(d_ff → d_model)
             │
             ▼
          d_model

핵심은 다음 네 가지입니다.

① d_model

Transformer에서 사용하는 임베딩 토큰 표현의 기본 차원입니다.

② d_ff

FFN 내부에서 사용하는 중간 차원입니다.

③ 두 개의 Linear

d_model → d_ff → d_model

형태로 특징을 변환합니다.

④ Positionwise

각 토큰 위치에 동일한 FFN을 독립적으로 적용합니다.

결국 Transformer에서 Attention이 “토큰 간 관계를 이해하는 역할”을 한다면, FFN은 Attention을 통해 얻은 각 토큰의 표현을 더욱 풍부하게 변환하는 역할을 한다고 이해하면 좋습니다.

태그

댓글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다