ニューラルネットワークにおける初期値とバッチサイズの影響の図示

データセット

以下のようなデータセットを考える。

x_train = Matrix{Float32}([
    0.09291784, 0.46809093, 0.93089486, 0.67612654, 0.73441752, 0.86847339,
    0.49873225, 0.51083168, 0.18343972, 0.99380898, 0.27840809, 0.38028817,
    0.12055708, 0.56715537, 0.92005746, 0.77072270, 0.85278176, 0.05315950,
    0.87168699, 0.58858043
] |> permutedims)
y_train = Matrix{Float32}([
    -0.25934537, 0.18195445, 0.651270150, 0.13921448, 0.09366691, 0.30567674,
    0.372291170, 0.20716968, -0.08131792, 0.51187806, 0.16943738, 0.3994327,
    0.019062570, 0.55820410, 0.452564960, -0.1183121, 0.02957665, -1.24354444,
    0.248038840, 0.26824970
] |> permutedims)

データはudlbookのノートブックからお借りした。

モデル

１層、隠れユニット３のニューラルネットワークをモデルとする。

using Flux
model = Chain(Dense(1 => 3, relu), Dense(3 => 1))

描画用関数を定義する。

using CairoMakie
function plot_univariate_regression(x_model, y_model, x_data=nothing, y_data=nothing, sigma_model=nothing, title="")
    x_model = dropdims(x_model; dims=1)
    y_model = dropdims(y_model; dims=1)
    fig = Figure()
    ax = Axis(fig[1, 1], title=title)
    if !isnothing(sigma_model)
        band!(ax, x_model, y_model .+ sigma_model, y_model .- sigma_model, color=:lightgray)
    end
    lines!(ax, x_model, y_model)
    if !isnothing(x_data) && !isnothing(y_data)
        x_data = dropdims(x_data; dims=1)
        y_data = dropdims(y_data; dims=1)
        scatter!(ax, x_data, y_data)
    end
    fig
end

sigma = 0.2
x_model = Matrix{Float32}(0:0.01:1 |> permutedims)
y_model = model(x_model)
plot_univariate_regression(x_model, y_model, x_train, y_train, sigma)

初期値の影響

ランダムな初期値

最急降下法(フルバッチ勾配降下法)でモデルを訓練してみる。

optim = Flux.setup(Flux.Descent(0.01), model)
loader = Flux.DataLoader((x_train, y_train), batchsize=length(x_train))
losses = Float32[]
for epoch in 1:10
    for (x, y) in loader
        loss, grads = Flux.withgradient(model) do m
            y_hat = m(x)
            Flux.mse(y_hat, y)
        end
        Flux.update!(optim, model, grads[1])
        push!(losses, loss)
    end
end
y_model = model(x_model)
plot_univariate_regression(x_model, y_model, x_train, y_train, sigma)

10周で左側に変化が見られる。
200周すると

全体的に平均値が下がってきているが形は変わらない。この後いくら訓練数を増やしても大きくは変わらなくなる。

初期値を指定する

モデルの初期値を指定してみる。

function create_model()
    model = Chain(Dense(1 => 3, relu), Dense(3 => 1))
    model[1].bias .= [0.3, -1.0, -0.5]
    model[1].weight .= [-1.0, 1.8, 0.65]
    model[2].bias .= 0.1
    model[2].weight .= [-2.0 -1.0 7.0]
    model
end
model = create_model()

最急降下法(フルバッチ勾配降下法)でモデルを訓練すると、10周で

位置が調整され、200周で

こんな感じ。この後1000周ほどしたが、位置の調整はあるが、形自体が変わることはなかった。

２つを見比べると初期値の影響が非常に大きいことがわかる。1つ目の例は左端の外れ値に過度に適合しようとしないが、２つ目の例は訓練を重ねると外れ値に引きずられるように下がってくる。この２つの間には損失のギャップがあり、訓練を繰り返せば同一の損失関数の最小値に行けるというものではないということがわかる。

バッチサイズの影響

確率的勾配降下法

２つ目のモデルを確率的勾配降下法で訓練する。

optim = Flux.setup(Flux.Descent(0.01), model)
loader = Flux.DataLoader((x_train, y_train), batchsize=1)
losses = Float32[]
for epoch in 1:10
    for (x, y) in loader
        loss, grads = Flux.withgradient(model) do m
            y_hat = m(x)
            Flux.mse(y_hat, y)
        end
        Flux.update!(optim, model, grads[1])
        push!(losses, loss)
    end
end

batchsize=1にした。10周すると

200周すると

1000周すると

となる。
確率的勾配降下法による汎化がわかりやすく見てとれる。

今回の結果は浅いニューラルネットワークで、入力次元も1次元の非常に単純なモデルであるため、注意は必要である。多層の効果は未だに完全に理解されておらず、多次元は人の想像力を容易に上回ってくる。が、深層学習というツールを使うのに全くブラックボックスというわけにもいかないので、イメージとして持っておいてもいいと思う。