STEP 04 / 18

토큰마다 혼자 생각하기: Feed-Forward

원본: 위치별 피드포워드 네트워크

어텐션이 토큰끼리 정보를 주고받는 단계라면, 피드포워드(FFN)는 받은 정보를 토큰마다 혼자 소화하는 단계다.

$$\mathrm{FFN}(x)=\max(0, xW_1 + b_1) W_2 + b_2$$

512차원을 2048차원으로 넓혔다가(ReLU) 다시 512로 줄인다. 모든 위치에 같은 가중치를 쓰지만, 위치끼리는 섞지 않는다. 그래서 이름이 "position-wise"다.

코드는 nn.Linear 두 개와 ReLU, dropout이 전부다.

실행 결과에서 볼 것

문장 전체를 넣었을 때 4번째 토큰의 출력과 그 토큰 하나만 넣었을 때의 출력이 같다. 토큰 사이에 정보가 오가지 않는다는 뜻이다.
파라미터가 약 210만 개로, 어텐션 블록(105만 개)의 두 배다. Transformer 파라미터의 상당 부분이 FFN에 있다.

코드 변화

+13 −0 · 파일 전체 86줄

이전 단계의 transformer.py와 비교했다. 초록 줄이 새로 생긴 줄, 빨간 줄이 사라진 줄이다.

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import torch
2
3	from transformer import PositionwiseFeedForward
4
5	torch.manual_seed(0)
6
7	ffn = PositionwiseFeedForward(d_model=512, d_ff=2048, dropout=0.0)
8	x = torch.randn(2, 10, 512)
9	print("입력:", tuple(x.shape), "→ w_1:", tuple(ffn.w_1(x).shape), "→ 출력:", tuple(ffn(x).shape))
10	n = sum(p.numel() for p in ffn.parameters())
11	print(f"파라미터 수: {n:,} (512×2048+2048) + (2048×512+512)")
12	print()
13
14	print("'위치별(position-wise)' = 토큰마다 따로, 같은 가중치로 계산한다")
15	alone = ffn(x[:, 3:4]) # 4번째 토큰만 넣어도
16	print(" 전체 문장에서 4번째 토큰 출력과 혼자 넣었을 때 출력의 차이:",
17	f"{(ffn(x)[:, 3:4] - alone).abs().max().item():.1e}")
18	print(" → 토큰끼리 정보를 섞는 건 어텐션뿐이고, FFN은 각 토큰을 혼자 변환한다")
19	print()
20
21	h = ffn.w_1(x).relu()
22	print(f"ReLU 뒤 0이 된 비율: {(h == 0).float().mean().item():.0%}")

입력: (2, 10, 512) → w_1: (2, 10, 2048) → 출력: (2, 10, 512)
파라미터 수: 2,099,712  (512×2048+2048) + (2048×512+512)

'위치별(position-wise)' = 토큰마다 따로, 같은 가중치로 계산한다
  전체 문장에서 4번째 토큰 출력과 혼자 넣었을 때 출력의 차이: 5.4e-07
  → 토큰끼리 정보를 섞는 건 어텐션뿐이고, FFN은 각 토큰을 혼자 변환한다

ReLU 뒤 0이 된 비율: 50%


1	1		import copy
2	2		import math
3	3
4	4		import torch
5	5		import torch.nn as nn
6	6
7	7
8	8		def attention(query, key, value, mask=None, dropout=None):
9	9		"Compute 'Scaled Dot Product Attention'"
10	10		d_k = query.size(-1)
11	11		scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
12	12		if mask is not None:
13	13		scores = scores.masked_fill(mask == 0, -1e9)
14	14		p_attn = scores.softmax(dim=-1)
15	15		if dropout is not None:
16	16		p_attn = dropout(p_attn)
17	17		return torch.matmul(p_attn, value), p_attn
18	18
19	19
20	20		def subsequent_mask(size):
21	21		"Mask out subsequent positions."
22	22		attn_shape = (1, size, size)
23	23		subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
24	24		torch.uint8
25	25		)
26	26		return subsequent_mask == 0
27	27
28	28
29	29		def clones(module, N):
30	30		"Produce N identical layers."
31	31		return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
32	32
33	33
34	34		class MultiHeadedAttention(nn.Module):
35	35		def __init__(self, h, d_model, dropout=0.1):
36	36		"Take in model size and number of heads."
37	37		super(MultiHeadedAttention, self).__init__()
38	38		assert d_model % h == 0
39	39		# We assume d_v always equals d_k
40	40		self.d_k = d_model // h
41	41		self.h = h
42	42		self.linears = clones(nn.Linear(d_model, d_model), 4)
43	43		self.attn = None
44	44		self.dropout = nn.Dropout(p=dropout)
45	45
46	46		def forward(self, query, key, value, mask=None):
47	47		"Implements Figure 2"
48	48		if mask is not None:
49	49		# Same mask applied to all h heads.
50	50		mask = mask.unsqueeze(1)
51	51		nbatches = query.size(0)
52	52
53	53		# 1) Do all the linear projections in batch from d_model => h x d_k
54	54		query, key, value = [
55	55		lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
56	56		for lin, x in zip(self.linears, (query, key, value))
57	57		]
58	58
59	59		# 2) Apply attention on all the projected vectors in batch.
60	60		x, self.attn = attention(
61	61		query, key, value, mask=mask, dropout=self.dropout
62	62		)
63	63
64	64		# 3) "Concat" using a view and apply a final linear.
65	65		x = (
66	66		x.transpose(1, 2)
67	67		.contiguous()
68	68		.view(nbatches, -1, self.h * self.d_k)
69	69		)
70	70		del query
71	71		del key
72	72		del value
73	73		return self.linears[-1](x)
	74	+
	75	+
	76	+	class PositionwiseFeedForward(nn.Module):
	77	+	"Implements FFN equation."
	78	+
	79	+	def __init__(self, d_model, d_ff, dropout=0.1):
	80	+	super(PositionwiseFeedForward, self).__init__()
	81	+	self.w_1 = nn.Linear(d_model, d_ff)
	82	+	self.w_2 = nn.Linear(d_ff, d_model)
	83	+	self.dropout = nn.Dropout(dropout)
	84	+
	85	+	def forward(self, x):
	86	+	return self.w_2(self.dropout(self.w_1(x).relu()))

1	import copy
2	import math
3
4	import torch
5	import torch.nn as nn
6
7
8	def attention(query, key, value, mask=None, dropout=None):
9	"Compute 'Scaled Dot Product Attention'"
10	d_k = query.size(-1)
11	scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
12	if mask is not None:
13	scores = scores.masked_fill(mask == 0, -1e9)
14	p_attn = scores.softmax(dim=-1)
15	if dropout is not None:
16	p_attn = dropout(p_attn)
17	return torch.matmul(p_attn, value), p_attn
18
19
20	def subsequent_mask(size):
21	"Mask out subsequent positions."
22	attn_shape = (1, size, size)
23	subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
24	torch.uint8
25	)
26	return subsequent_mask == 0
27
28
29	def clones(module, N):
30	"Produce N identical layers."
31	return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
32
33
34	class MultiHeadedAttention(nn.Module):
35	def __init__(self, h, d_model, dropout=0.1):
36	"Take in model size and number of heads."
37	super(MultiHeadedAttention, self).__init__()
38	assert d_model % h == 0
39	# We assume d_v always equals d_k
40	self.d_k = d_model // h
41	self.h = h
42	self.linears = clones(nn.Linear(d_model, d_model), 4)
43	self.attn = None
44	self.dropout = nn.Dropout(p=dropout)
45
46	def forward(self, query, key, value, mask=None):
47	"Implements Figure 2"
48	if mask is not None:
49	# Same mask applied to all h heads.
50	mask = mask.unsqueeze(1)
51	nbatches = query.size(0)
52
53	# 1) Do all the linear projections in batch from d_model => h x d_k
54	query, key, value = [
55	lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
56	for lin, x in zip(self.linears, (query, key, value))
57	]
58
59	# 2) Apply attention on all the projected vectors in batch.
60	x, self.attn = attention(
61	query, key, value, mask=mask, dropout=self.dropout
62	)
63
64	# 3) "Concat" using a view and apply a final linear.
65	x = (
66	x.transpose(1, 2)
67	.contiguous()
68	.view(nbatches, -1, self.h * self.d_k)
69	)
70	del query
71	del key
72	del value
73	return self.linears[-1](x)
74
75
76	class PositionwiseFeedForward(nn.Module):
77	"Implements FFN equation."
78
79	def __init__(self, d_model, d_ff, dropout=0.1):
80	super(PositionwiseFeedForward, self).__init__()
81	self.w_1 = nn.Linear(d_model, d_ff)
82	self.w_2 = nn.Linear(d_ff, d_model)
83	self.dropout = nn.Dropout(dropout)
84
85	def forward(self, x):
86	return self.w_2(self.dropout(self.w_1(x).relu()))