torch.nn.Linear が kaiming_uniform_ を使っているのに Kaiming 初期化ではない話

PyTorch の torch.nn.Linear はデフォルトで torch.nn.init.kaiming_uniform_ によって初期化されていて、Kaiming 初期化 (Hu 初期化) をしているようにみえますが、そういうわけではないです。

参考文献

1. https://github.com/pytorch/pytorch/blob/v2.12.0/torch/nn/modules/linear.py#L121-L124
2. https://docs.pytorch.org/docs/2.12/nn.init.html#torch.nn.init.kaiming_uniform_
3. https://arxiv.org/abs/1502.01852
4. https://github.com/pytorch/pytorch/issues/57109#issuecomment-828847575
5. https://soumith.ch/files/20141213_gplus_nninit_discussion.htm
   • There must be some other reason why making the target output variance smaller than 1 somehow works a bit better, and the ideal scale factor probably isn't sqrt(3). I really want to know what's going on here, "it just works a bit better" is rather unsatisfying :)

torch.nn.Linear のデフォルト初期化

torch.nn.Linear は、デフォルトで torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) によって初期化されます。結果的に、$U([-\sqrt{1/h_{\mathrm{in}}}, \sqrt{1/h_{\mathrm{in}}}])$ にしたがう乱数で初期化されます。ソースコードのコメントにも結果そうだとあります。[1]

ではこの torch.nn.init.kaiming_uniform_ [2] は何なのかというと、

• 出力を活性化しないなら $U([-\sqrt{3/h_{\mathrm{in}}}, \sqrt{3/h_{\mathrm{in}}}])$
• 出力を ReLU 活性化するなら $U([-\sqrt{6/h_{\mathrm{in}}}, \sqrt{6/h_{\mathrm{in}}}])$
• 出力を Leaky ReLU 活性化するなら $U([-\sqrt{6/(h_{\mathrm{in}}\cdot(1 + a^2))}, \sqrt{6/(h_{\mathrm{in}}\cdot(1 + a^2))}])$
  • ただし $a$ は Leaky ReLU の負側の傾き

にしたがう乱数で重みパラメータを初期化します。そうすることで、層を経ても分散を保とうとします。[3] これは提唱者にちなんで Kaiming 初期化 (Hu 初期化) とよばれます。

となると、torch.nn.Linear のデフォルト初期化 $U([-\sqrt{1/h_{\mathrm{in}}}, \sqrt{1/h_{\mathrm{in}}}])$ は $a=\sqrt{5}$ の Leaky ReLU で分散が保たれるような初期化になりますが、Leaky ReLU で $a=\sqrt{5}$ にすることは考えにくいです ($a=\sqrt{5}$ だと負側がちょっと Leak どころか正側より流れてくる)。

ではなぜ $U([-\sqrt{1/h_{\mathrm{in}}}, \sqrt{1/h_{\mathrm{in}}}])$ で初期化しているのかは、ソースコード [1] にある Issue リンクをたどると、Torch 開発者の Collobert らが経験的にその方がよいと見出したからのようです。[4][5]

結局何なのかというと、torch.nn.Linear のデフォルト初期化は一見 Kaiming 初期化を呼び出しているようで、Kaiming 初期化をしているわけではないので注意してください。

(おまけ) 検証スニペット

以下のスニペットは torch.nn.Linear の重みパラメータがデフォルトで $U([-\sqrt{1/h_{\mathrm{in}}}, \sqrt{1/h_{\mathrm{in}}}])$ にしたがう乱数で初期化されることを確かめるだけです。

import torch


def main(seed=0, in_features=100):
    torch.manual_seed(seed)
    layer = torch.nn.Linear(in_features, out_features=1, bias=False)
    w = layer.weight
    _fmt = lambda x: f'{x:+2.6f}'
    print(
        '\nshape         :', str(w.shape),
        '\nmin (expected):', _fmt(-1 / in_features ** 0.5),
        '  min (actual)  :', _fmt(w.min().item()),
        '\nmax (expected):', _fmt(1 / in_features ** 0.5),
        '  max (actual)  :', _fmt(w.max().item()),
        '\nvar (expected):', _fmt(1 / (3 * in_features)),
        '  var (actual)  :', _fmt(w.var(unbiased=False).item()),
    )


if __name__ == '__main__':
    main(seed=0, in_features=100)
    main(seed=1, in_features=100)
    main(seed=2, in_features=100)

shape         : torch.Size([1, 100])
min (expected): -0.100000   min (actual)  : -0.099753
max (expected): +0.100000   max (actual)  : +0.099413
var (expected): +0.003333   var (actual)  : +0.003435

shape         : torch.Size([1, 100])
min (expected): -0.100000   min (actual)  : -0.097755
max (expected): +0.100000   max (actual)  : +0.099605
var (expected): +0.003333   var (actual)  : +0.002867

shape         : torch.Size([1, 100])
min (expected): -0.100000   min (actual)  : -0.099572
max (expected): +0.100000   max (actual)  : +0.096172
var (expected): +0.003333   var (actual)  : +0.002789