メインコンテンツまでスキップ

手計算したLayerNormの値が合わない？

· 約6分
Zephyr
Zephyr
Engineer

今日はふと LayerNorm の値を計算してみたくなりました。

LayerNorm の公式は次の通りです：

LayerNorm(x)=x−μVar[x]+ϵ×γ+β\text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\text{Var[}x\text{]} + \epsilon}} \times \gamma + \beta

ここで、μ\muはxxの平均値、Var\text{Var}はxxの分散です。

この情報をもとに、γ\gammaとβ\betaを無視して実際に計算してみます：

import torch

x = torch.rand(16, 768)
mu = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True)
eps = 1e-5
y = (x - mu) / (var + eps).sqrt()

以下のような結果が得られました：

# tensor([[ 0.1219, -0.0222, -1.4742,  ...,  0.1738, -0.6124, -0.3001],
#         [-1.6009, -1.5814,  1.5357,  ...,  0.1917,  1.3787, -0.2772],
#         [ 0.3738,  1.0520,  0.4403,  ...,  1.1353, -0.7488, -0.9137],
#         ...,
#         [ 0.8823, -1.5427,  0.4725,  ..., -1.2544, -1.5354, -0.4305],
#         [ 1.4548,  0.3059, -0.6732,  ..., -0.7109,  0.4908, -1.2447],
#         [-0.4067,  0.5974, -0.9113,  ..., -0.2511, -0.2279, -0.9675]])

次に、この結果を PyTorch のtorch.nn.LayerNormと比較してみます：

layer_norm = torch.nn.LayerNorm(768, elementwise_affine=False, bias=False)

y_ln = layer_norm(x)

結果は以下の通りです：

# tensor([[ 0.1220, -0.0222, -1.4752,  ...,  0.1739, -0.6128, -0.3003],
#         [-1.6020, -1.5824,  1.5367,  ...,  0.1918,  1.3796, -0.2774],
#         [ 0.3741,  1.0527,  0.4406,  ...,  1.1360, -0.7493, -0.9143],
#         ...,
#         [ 0.8829, -1.5437,  0.4728,  ..., -1.2552, -1.5364, -0.4308],
#         [ 1.4557,  0.3061, -0.6736,  ..., -0.7113,  0.4911, -1.2455],
#         [-0.4069,  0.5978, -0.9119,  ..., -0.2513, -0.2281, -0.9681]])

見比べてみると、値が異なります。

不偏推定​

少し調べてみたところ、torch.varにはcorrectionというパラメータがあり、デフォルト値は1、つまり不偏推定を使用していました。

これは分母をNではなくN-1で割ることを意味します。一方で、torch.nn.LayerNormはNを使用しています。

そこで、torch.varのcorrectionパラメータを0に設定してみます：

var = x.var(dim=-1, correction=0, keepdim=True)
ヒント

correctionはunbiasedの別名であり、PyTorch 2.0.0 で導入されました。

古いバージョンでは、次のように設定します：

var = x.var(dim=-1, unbiased=False, keepdim=True)

再度結果を比較すると：

# tensor([[ 0.1220, -0.0222, -1.4752,  ...,  0.1739, -0.6128, -0.3003],
#         [-1.6020, -1.5824,  1.5367,  ...,  0.1918,  1.3796, -0.2774],
#         [ 0.3741,  1.0527,  0.4406,  ...,  1.1360, -0.7493, -0.9143],
#         ...,
#         [ 0.8829, -1.5437,  0.4728,  ..., -1.2552, -1.5364, -0.4308],
#         [ 1.4557,  0.3061, -0.6736,  ..., -0.7113,  0.4911, -1.2455],
#         [-0.4069,  0.5978, -0.9119,  ..., -0.2513, -0.2281, -0.9681]])

今度は値が一致しました！

なぜ LayerNorm は不偏推定を使用しないのか？​

簡単にまとめると、安定性と計算の簡略化が理由です。

この点について詳しく知りたい場合は以下をご覧ください：

  • 小規模データでの計算安定性

    LayerNorm は、通常、バッチ全体ではなく個々のサンプルの特徴次元（例えば各ニューロンや特徴量）に適用されます。各サンプルの特徴数はサンプル数よりも大きいことが多いため、サンプル分散よりも母分散のほうが安定で正確な推定が可能です。特にサンプル数が少ない場合に効果的です。

  • 不偏推定の重要性が低い

    サンプル分散の不偏性（分母を n-1 とする理由）は、統計学では母集団のパラメータを推定する際に重要です。しかし、LayerNorm のような深層学習における正規化操作では、不偏性の影響は小さいです。母分散を用いることで計算が簡略化され、結果への影響はほとんどありません。

  • 勾配計算の安定性

    逆伝播では、安定した勾配が重要です。母分散を使用することで、サンプル数が少ない場合に発生する追加のノイズを回避し、勾配計算を滑らかで安定したものにします。これにより、ネットワークの収束性と学習効果が向上します。

  • 計算の簡略化

    計算の観点では、母分散の計算はサンプル分散よりも若干簡単です（分母が n-1 ではなく n になるため、1 つの減算操作が省略されます）。これは決定的な理由ではありませんが、設計時に考慮された可能性があります。

まとめ​

この疑問をふと思いつき、この記事を書いてみました。

もしこの問題の解答が役立つことを願っています。