STEP 06 / 18

깊게 쌓아도 무너지지 않게: LayerNorm과 잔차 연결

원본: 인코더

이제 블록을 6층씩 쌓을 준비를 한다. 층이 깊어지면 신호가 점점 작아지거나 커져서 학습이 망가진다. 이를 막는 장치가 두 개 있다.

LayerNorm: 토큰 벡터 하나(512차원)를 평균 0, 표준편차 1로 맞춘 뒤 학습 가능한 스케일 a_2와 이동 b_2를 적용한다.
SublayerConnection: 서브층(어텐션이나 FFN)을 감싸서 x + dropout(sublayer(norm(x)))를 돌려준다. 서브층의 결과를 입력에 더하기 때문에 서브층이 아무것도 못 해도 입력은 그대로 통과한다.

논문 수식은 LayerNorm(x + Sublayer(x))(Post-LN)인데 코드는 정규화를 먼저 한다(Pre-LN). 원본 코드 주석에 "for code simplicity the norm is first"라고 적혀 있다. Pre-LN은 학습이 더 안정적이어서 요즘 LLM도 대부분 이 순서를 쓴다.

sublayer를 함수로 받는다는 점도 눈여겨보자. 다음 단계에서 어텐션은 lambda x: self.self_attn(x, x, x, mask) 형태로, FFN은 모듈 그대로 넘긴다.

실행 결과에서 볼 것

정규화 뒤 평균 0, 표준편차 1.
서브층이 0만 내놓아도 출력이 입력과 같다.
FFN 20개를 그냥 쌓으면 신호가 0.05 수준으로 쪼그라들지만, 잔차 연결과 정규화로 감싸면 크기가 유지된다.

코드 변화

+31 −0 · 파일 전체 150줄

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

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import torch
2
3	from transformer import LayerNorm, SublayerConnection, PositionwiseFeedForward
4
5	torch.manual_seed(0)
6
7	x = torch.randn(2, 3, 8) * 5 + 3
8	ln = LayerNorm(8)
9	y = ln(x)
10	print("LayerNorm은 토큰 하나(마지막 차원)마다 평균 0, 표준편차 1로 맞춘다")
11	print(" 전 평균 :", [f"{v:.2f}" for v in x.mean(-1)[0].tolist()])
12	print(" 후 평균 :", [f"{abs(v):.2f}" for v in y.mean(-1)[0].tolist()])
13	print(" 후 표준편차:", [f"{v:.2f}" for v in y.std(-1)[0].tolist()])
14	print(" 학습 파라미터:", [n for n, _ in ln.named_parameters()], "(스케일 γ, 이동 β)")
15	print()
16
17	sub = SublayerConnection(8, dropout=0.0)
18	print("SublayerConnection: x + Sublayer(LayerNorm(x))")
19	out = sub(x, lambda t: torch.zeros_like(t))
20	print(" 서브층이 0만 내놓으면 출력 = 입력?", torch.allclose(out, x))
21	print(" → 서브층이 아무것도 못 배워도 신호는 그대로 흐른다 (잔차 연결)")
22	print()
23
24	print("층을 깊게 쌓았을 때 신호 크기 (FFN 20개)")
25	torch.manual_seed(0)
26	ffns = [PositionwiseFeedForward(64, 256, dropout=0.0) for _ in range(20)]
27	subs = [SublayerConnection(64, dropout=0.0) for _ in range(20)]
28	h_plain = h_res = torch.randn(1, 5, 64)
29	for f, s in zip(ffns, subs):
30	h_plain = f(h_plain)
31	h_res = s(h_res, f)
32	print(f" 그냥 쌓기 : 표준편차 {h_plain.std():.2e}")
33	print(f" 잔차 + 정규화: 표준편차 {h_res.std():.2e}")

LayerNorm은 토큰 하나(마지막 차원)마다 평균 0, 표준편차 1로 맞춘다
  전 평균  : ['0.59', '4.22', '2.29']
  후 평균  : ['0.00', '0.00', '0.00']
  후 표준편차: ['1.00', '1.00', '1.00']
  학습 파라미터: ['a_2', 'b_2'] (스케일 γ, 이동 β)

SublayerConnection: x + Sublayer(LayerNorm(x))
  서브층이 0만 내놓으면 출력 = 입력? True
  → 서브층이 아무것도 못 배워도 신호는 그대로 흐른다 (잔차 연결)

층을 깊게 쌓았을 때 신호 크기 (FFN 20개)
  그냥 쌓기    : 표준편차 4.59e-02
  잔차 + 정규화: 표준편차 1.50e+00


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	74
75	75
76	76		class PositionwiseFeedForward(nn.Module):
77	77		"Implements FFN equation."
78	78
79	79		def __init__(self, d_model, d_ff, dropout=0.1):
80	80		super(PositionwiseFeedForward, self).__init__()
81	81		self.w_1 = nn.Linear(d_model, d_ff)
82	82		self.w_2 = nn.Linear(d_ff, d_model)
83	83		self.dropout = nn.Dropout(dropout)
84	84
85	85		def forward(self, x):
86	86		return self.w_2(self.dropout(self.w_1(x).relu()))
87	87
88	88
89	89		class Embeddings(nn.Module):
90	90		def __init__(self, d_model, vocab):
91	91		super(Embeddings, self).__init__()
92	92		self.lut = nn.Embedding(vocab, d_model)
93	93		self.d_model = d_model
94	94
95	95		def forward(self, x):
96	96		return self.lut(x) * math.sqrt(self.d_model)
97	97
98	98
99	99		class PositionalEncoding(nn.Module):
100	100		"Implement the PE function."
101	101
102	102		def __init__(self, d_model, dropout, max_len=5000):
103	103		super(PositionalEncoding, self).__init__()
104	104		self.dropout = nn.Dropout(p=dropout)
105	105
106	106		# Compute the positional encodings once in log space.
107	107		pe = torch.zeros(max_len, d_model)
108	108		position = torch.arange(0, max_len).unsqueeze(1)
109	109		div_term = torch.exp(
110	110		torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
111	111		)
112	112		pe[:, 0::2] = torch.sin(position * div_term)
113	113		pe[:, 1::2] = torch.cos(position * div_term)
114	114		pe = pe.unsqueeze(0)
115	115		self.register_buffer("pe", pe)
116	116
117	117		def forward(self, x):
118	118		x = x + self.pe[:, : x.size(1)].requires_grad_(False)
119	119		return self.dropout(x)
	120	+
	121	+
	122	+	class LayerNorm(nn.Module):
	123	+	"Construct a layernorm module (See citation for details)."
	124	+
	125	+	def __init__(self, features, eps=1e-6):
	126	+	super(LayerNorm, self).__init__()
	127	+	self.a_2 = nn.Parameter(torch.ones(features))
	128	+	self.b_2 = nn.Parameter(torch.zeros(features))
	129	+	self.eps = eps
	130	+
	131	+	def forward(self, x):
	132	+	mean = x.mean(-1, keepdim=True)
	133	+	std = x.std(-1, keepdim=True)
	134	+	return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
	135	+
	136	+
	137	+	class SublayerConnection(nn.Module):
	138	+	"""
	139	+	A residual connection followed by a layer norm.
	140	+	Note for code simplicity the norm is first as opposed to last.
	141	+	"""
	142	+
	143	+	def __init__(self, size, dropout):
	144	+	super(SublayerConnection, self).__init__()
	145	+	self.norm = LayerNorm(size)
	146	+	self.dropout = nn.Dropout(dropout)
	147	+
	148	+	def forward(self, x, sublayer):
	149	+	"Apply residual connection to any sublayer with the same size."
	150	+	return x + self.dropout(sublayer(self.norm(x)))

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()))
87
88
89	class Embeddings(nn.Module):
90	def __init__(self, d_model, vocab):
91	super(Embeddings, self).__init__()
92	self.lut = nn.Embedding(vocab, d_model)
93	self.d_model = d_model
94
95	def forward(self, x):
96	return self.lut(x) * math.sqrt(self.d_model)
97
98
99	class PositionalEncoding(nn.Module):
100	"Implement the PE function."
101
102	def __init__(self, d_model, dropout, max_len=5000):
103	super(PositionalEncoding, self).__init__()
104	self.dropout = nn.Dropout(p=dropout)
105
106	# Compute the positional encodings once in log space.
107	pe = torch.zeros(max_len, d_model)
108	position = torch.arange(0, max_len).unsqueeze(1)
109	div_term = torch.exp(
110	torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
111	)
112	pe[:, 0::2] = torch.sin(position * div_term)
113	pe[:, 1::2] = torch.cos(position * div_term)
114	pe = pe.unsqueeze(0)
115	self.register_buffer("pe", pe)
116
117	def forward(self, x):
118	x = x + self.pe[:, : x.size(1)].requires_grad_(False)
119	return self.dropout(x)
120
121
122	class LayerNorm(nn.Module):
123	"Construct a layernorm module (See citation for details)."
124
125	def __init__(self, features, eps=1e-6):
126	super(LayerNorm, self).__init__()
127	self.a_2 = nn.Parameter(torch.ones(features))
128	self.b_2 = nn.Parameter(torch.zeros(features))
129	self.eps = eps
130
131	def forward(self, x):
132	mean = x.mean(-1, keepdim=True)
133	std = x.std(-1, keepdim=True)
134	return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
135
136
137	class SublayerConnection(nn.Module):
138	"""
139	A residual connection followed by a layer norm.
140	Note for code simplicity the norm is first as opposed to last.
141	"""
142
143	def __init__(self, size, dropout):
144	super(SublayerConnection, self).__init__()
145	self.norm = LayerNorm(size)
146	self.dropout = nn.Dropout(dropout)
147
148	def forward(self, x, sublayer):
149	"Apply residual connection to any sublayer with the same size."
150	return x + self.dropout(sublayer(self.norm(x)))