本文对常见的几种神经网络组件进行了简单的实现和注释，便于理解。包括：
- 层归一化（Layer Normalization，LN）
- RMSNorm（Root Mean Square Layer Normalization）
- 批次归一化（Batch Normalization，BN）
- Dropout
LayerNorm
Layer Normalization (LN) 是一种归一化技术，旨在改善神经网络的训练稳定性和性能。LN 的基本思想是对每个样本在特征维度上进行归一化，而不是在批次维度上。对于每个输入 $ $ ，LN 的公式如下：
\[ \mu = \frac{1}{H} \sum_{i=1}^{H} x_i \]
\[ \sigma^2 = \frac{1}{H} \sum_{i=1}^{H} (x_i - \mu)^2 \]
\[ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \]
其中，$ H $ 是特征维度的大小， $ $ 和 $ ^2 $ 分别是特征维度上的均值和方差。归一化后，会应用可学习的缩放参数 $ $ 和偏移参数 $ $ ：
\[ y_i = \gamma \hat{x}_i + \beta \]
代码实现如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
| import torch
from torch import nn
class LayerNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.hidden_size = hidden_size
self.eps = eps
self.gamma = nn.Parameter(torch.ones(hidden_size))
self.beta = nn.Parameter(torch.zeros(hidden_size))
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
variance = x.var(dim=-1, keepdim=True, unbiased=False)
x_normalized = (x - mean) / torch.sqrt(variance + self.eps)
output = self.gamma * x_normalized + self.beta
return output
def test_layer_norm():
batch_size = 2
seq_len = 4
hidden_size = 8
x = torch.randn(batch_size, seq_len, hidden_size)
layer_norm = LayerNorm(hidden_size)
output = layer_norm(x)
print("Input shape:", x.shape)
print("Output shape:", output.shape)
if __name__ == "__main__":
test_layer_norm()
|
RMSNorm
LayerNorm 需同时计算输入数据的均值和方差。研究发现 LayerNorm 的优势在于缩放不变性（即方差的计算，让模型对于输入和权重上的偏移噪声不敏感），而不是重新居中（即均值的计算，让模型在输入和权重都被随机缩放时保持输出表示不变）。
所以 RMSNorm 省略了归一化过程中的均值计算，仅计算特征的均方根（Root Mean Square），使得算法更加简洁，而效果不减，且运算效率显著提升。RMSNorm 的公式如下： \[ RMS=\sqrt{\frac{1}{H} \sum_{i=1}^{H} x_i^2 + \epsilon} \]
\[ \hat{x}_i = \frac{x_i}{RMS} \]
此外，通常仅保留缩放参数 $ $，省略偏移参数 $ $。实验表明，中心化（减均值）对性能影响有限，移除后仍能保持模型表达能力。
\[ y_i = \gamma \hat{x}_i \] 代码实现如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
| import torch
from torch import nn
class RMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.hidden_size = hidden_size
self.eps = eps
self.gamma = nn.Parameter(torch.ones(hidden_size))
def forward(self, x):
rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps)
x_normalized = x / rms
output = self.gamma * x_normalized
return output
def test_rms_norm():
batch_size = 2
seq_len = 4
hidden_size = 8
x = torch.randn(batch_size, seq_len, hidden_size)
rms_norm = RMSNorm(hidden_size)
output = rms_norm(x)
print("Input shape:", x.shape)
print("Output shape:", output.shape)
print("Params:", list(rms_norm.parameters()))
if __name__ == "__main__":
test_rms_norm()
|
BatchNorm
Batch Normalization (BN) 是另一种归一化技术，主要用于加速神经网络的训练。BN 对每个 mini-batch 的数据在批次维度上进行归一化。对于所有输入 $ $ ，BN 的公式如下：
\[ \mu_B = \frac{1}{m} \sum_{i=1}^{m} x_i \]
\[ \sigma_B^2 = \frac{1}{m} \sum_{i=1}^{m} (x_i - \mu_B)^2 \]
\[ \hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}} \]
其中， $ m $ 是 mini-batch 的大小， $ _B $ 和 $ _B^2 $ 分别是批次维度上的均值和方差。归一化后，也会应用可学习的缩放参数 $ $ 和偏移参数 $ $ ：
\[ y_i = \gamma \hat{x}_i + \beta \]
LN 和 BN 的区别：
- 归一化的维度：LN 在特征维度上进行归一化，而 BN 在批次维度上进行归一化。
- 应用场景：LN 更适用于 Recurrent Neural Networks (RNNs) 和 Transformer 等序列模型，而 BN 通常用于 Convolutional Neural Networks (CNNs)。
- 批量大小依赖性：LN 不依赖于批量大小，因此在小批量甚至单样本的情况下也能很好地工作。BN 依赖于较大的批量大小以稳定均值和方差的估计。
为什么 Transformer 使用 LN 而不是 BN：
- Transformer 模型通常处理变长序列数据，其批次大小可能会变化或者在推理阶段可能只有一个样本。
- BN 依赖于批次维度的均值和方差估计，因此在这种情况下表现可能不稳定。LN 则对每个样本独立进行归一化，不依赖于批次大小，因此更适合于 Transformer 这种模型。
- 此外，LN 对序列模型的时间步长无关的归一化方式有助于保持输入数据的顺序特性，从而提高模型的性能和稳定性。
在 BatchNorm 的实现中，通常会区分训练（training）和推理（inference）阶段，这是因为在这两个阶段中，BN 的行为有所不同：
训练阶段（training）：
- 在训练阶段，BN 会计算当前 mini-batch 的均值和方差，并使用这些统计量对数据进行归一化。具体公式同前文所述。
- 训练过程中，BN 还会使用移动平均的方法更新运行时均值和方差，这样在推理阶段就可以使用这些全局统计量来代替 mini-batch 的统计量。更新规则如下： \[ \text{running\_mean} = (1 - \text{momentum}) \times \text{running\_mean} + \text{momentum} \times \mu_B \] \[ \text{running\_var} = (1 - \text{momentum}) \times \text{running\_var} + \text{momentum} \times \sigma_B^2 \]
推理阶段（inference）：
- 在推理阶段，BN 不再使用 mini-batch 的均值和方差，而是使用训练阶段累积的运行时均值和方差。这是因为在推理阶段，通常批量大小很小甚至为 1，使用 mini-batch 的统计量会导致不稳定的输出。推理阶段的归一化公式如下： \[ \hat{x}_i = \frac{x_i - \text{running\_mean}}{\sqrt{\text{running\_var} + \epsilon}} \]
代码实现如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
| import torch
from torch import nn
class BatchNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-5, momentum=0.1):
super().__init__()
self.hidden_size = hidden_size
self.eps = eps
self.momentum = momentum
self.gamma = nn.Parameter(torch.ones(hidden_size))
self.beta = nn.Parameter(torch.zeros(hidden_size))
self.running_mean = torch.zeros(hidden_size)
self.running_var = torch.ones(hidden_size)
def forward(self, x):
if self.training:
batch_mean = x.mean(dim=(0, 1), keepdim=False)
batch_var = x.var(dim=(0, 1), keepdim=False, unbiased=False)
self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean
self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
mean = batch_mean
variance = batch_var
else:
mean = self.running_mean
variance = self.running_var
x_normalized = (x - mean) / torch.sqrt(variance + self.eps)
output = self.gamma * x_normalized + self.beta
return output
def test_batch_norm():
batch_size = 2
seq_len = 4
hidden_size = 8
x = torch.randn(batch_size, seq_len, hidden_size)
batch_norm = BatchNorm(hidden_size)
output = batch_norm(x)
print("Input shape:", x.shape)
print("Output shape:", output.shape)
if __name__ == "__main__":
test_batch_norm()
|
注意：这里的实现其实不太合理，因为 BN 通常不用于变长序列模型的输入，因此这里的 seq_len 维度只是摆设。
Dropout
Dropout 是一种正则化技术，用于防止神经网络的过拟合。通过在训练过程中随机「丢弃」一部分神经元，Dropout 使得模型不会过度依赖某些特定的神经元，从而增强模型的泛化能力。
在推理阶段，所有神经元都被激活，并根据 Dropout 概率进行缩放，主要目的是为了在训练和推理阶段保持一致的输出期望值。
具体而言，假设在训练阶段，输入神经元的激活值为 $ x $，Dropout 的概率为 $ p $。每个神经元以 $ 1 - p $ 的概率被保留（即不被丢弃）。因此，每个神经元的激活值期望为：
\[ E[\text{激活值}] = x \cdot (1 - p) \]
为了使得训练和推理阶段的输出期望一致，我们需要在训练阶段对保留的神经元进行缩放，即乘以 \(\frac{1}{1 - p}\)，这样可以抵消丢弃神经元带来的期望值减少。
代码实现如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
| import torch
from torch import nn
class Dropout(nn.Module):
def __init__(self, dropout_prob=0.1):
super().__init__()
self.dropout_prob = dropout_prob
def forward(self, x):
if self.training:
mask = (torch.rand(x.shape) > self.dropout_prob).float()
output = mask * x / (1.0 - self.dropout_prob)
else:
output = x
return output
def test_dropout():
batch_size = 2
seq_len = 4
hidden_size = 8
x = torch.randn(batch_size, seq_len, hidden_size)
dropout = Dropout(dropout_prob=0.1)
dropout.train()
output_train = dropout(x)
dropout.eval()
output_eval = dropout(x)
print("Input shape:", x.shape)
print("Output shape during training:", output_train.shape)
print("Output shape during evaluation:", output_eval.shape)
if __name__ == "__main__":
test_dropout()
|