活性化関数と重みの初期値の関係

『ゼロから作るDeep Learning』を久々に読み返すと，活性化関数と重みの初期値の関係についての記載がありました．
恥ずかしながら初期値についてはあまり意識せず適当に決めていたので，意識付けのためにも，これらの関係について検証を行ってみました（n番煎じ）．

本書曰く，

活性化関数にReLUを使う場合は「Heの初期値」，sigmoidやtanhなどのS字カーブのときは「Xavierの初期値」を使う

ということらしいです．
前層のノード数を $n$ とするとき，Heの初期値では $\sqrt{\frac{2}{n}}$ を，Xavierの初期値では $\sqrt{\frac{1}{n}}$ を標準偏差とするガウス分布によって初期化を行います．

モデルの定義

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class TestModel(nn.Module):
    def __init__(self, hidden_num, layer_num, init_func=None, act_func=None):
        super(TestModel, self).__init__()

        self.act_dist = []
        def forward_func(module, input, outputs):
            self.act_dist.append(outputs.view(-1).detach().clone().numpy())

        if init_func == 'he_normal':
            init = nn.init.kaiming_normal_
        elif init_func == 'he_uniform':
            init = nn.init.kaiming_uniform_
        elif init_func == 'xavier':
            init = nn.init.xavier_normal_
        else:
            init = nn.init.normal_

        layers = []
        for _ in range(layer_num):
            l = nn.Linear(hidden_num, hidden_num)
            init(l.weight)
            layers.append(l)
            
            if act_func == 'relu':
                act = nn.ReLU()
            elif act_func == 'sigmoid':
                act = nn.Sigmoid()
            elif act_func == 'tanh':
                act = nn.Tanh()
            else:
                act = nn.Identity()

            act.register_forward_hook(forward_func)
            layers.append(act)

        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        x = self.layers(x)
        return x

順伝播の実行と分布の出力

hidden_num = 100
layer_num = 5
init_func = 'he_normal'
act_func = 'relu'
input = torch.randn(1000, 1, hidden_num)
model = TestModel(hidden_num, layer_num, init_func=init_func, act_func=act_func)

outputs = model(input)

fig = plt.figure(figsize=(24, 4))
fig.suptitle(f'Initialization: {init_func}, Activation: {act_func}', fontsize=20)
for i, ad in enumerate(model.act_dist):
    plt.subplot(1,len(model.act_dist), i+1)
    plt.hist(ad, bins=50)
    if act_func == 'relu':
        plt.xlim([0.0, 1.0])
        plt.ylim([0, 10000])
    elif act_func == 'sigmoid':
        plt.xlim([0.0, 1.0])
        plt.ylim([0, 6000])
    elif act_func == 'tanh':
        plt.xlim([-1.0, 1.0])
        plt.ylim([0, 5000])

検証

隠れ層の次元が100の線形結合を5層つなげて活性化関数の出力の分布をヒストグラムにしました．

ReLU

Xavierの初期値
f:id:spookyboogie:20220202055353p:plain

Heの初期値
f:id:spookyboogie:20220202055510p:plain
確かに「Xavierの初期値」では分布に偏りが見られます．

Sigmoid

Xavierの初期値
f:id:spookyboogie:20220202060304p:plain

Heの初期値
f:id:spookyboogie:20220202060514p:plain
「Heの初期値」では，「Xavierの初期値」に比べて僅かですが3層目以降の分布が歪になる傾向が見られました．
また，「Heの初期値」では1層目の分布が，山を均したような形状になっています．『活性化関数の出力の分布が広い方が表現力が高い』と本書に記載がありましたが，Sigmoidでは0.5から離れるにつれて傾きが小さくなるため，誤差消失が起きる可能性が高まります．一概に広い→良いとは言えないのかもしれません．

Tanh

Xavierの初期値
f:id:spookyboogie:20220202061708p:plain

Heの初期値
f:id:spookyboogie:20220202061732p:plain
「Xavierの初期値」では，程よく0.0付近に集中している一方で，「Heの初期値」では，出力が広く分布しており，1層目では-1.0と1.0の度数が高いことがわかります．これも先程のSigmoid同様に，Tanhの外形を考えると誤差消失の観点から好ましくないのだと思われます．

これまでは重みの初期値はPyTorchのデフォルトのままでしたが，今後は気を配ろうと思います．

チラ裏備忘録

情報整理