和田山先生の「モデルベース深層学習と深層展開」 の第2章に記載のあった勾配法のFluxによる実装を試す．実行環境は以下の通り
Julia Version 1.11.1
Commit 8f5b7ca12ad (2024-10-16 10:53 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (arm64-apple-darwin22.4.0)
CPU: 8 × Apple M1
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
ここでは２次式の最小化問題
$$
\underset{(x_1, x_2) \in \mathbb{R}^2}{\text{minimize}} \quad 4x_1^2 + x_2^2
$$
を勾配法で解くことを考える．$x_1 = x_2 = 0$ が最小値を与えることは明らかであるが，この関数の勾配は
$$
\frac{\partial f}{\partial x_1} = 8x_1, \quad
\frac{\partial f}{\partial x_2} = 2x_2.
$$
であることから，勾配降下ステップは
$$
\begin{pmatrix}
x_1^{(t+1)} \\
x_2^{(t+1)}
\end{pmatrix}
=
\begin{pmatrix}
x_1^{(t)} \\
x_2^{(t)}
\end{pmatrix}
– 0.1
\begin{pmatrix}
8x_1^{(t)} \\
2x_2^{(t)}
\end{pmatrix}.
$$
となる．ただし $0.1$ は勾配法におけるステップサイズである．この勾配法を直接Juliaで実装したのが以下のコードである．
# パッケージの読み込み
using LinearAlgebra
using PyPlot
using Random
Random.seed!(1)
using Flux
# 2変数の2次式 f(x)
f(x) = 4*x[1]^2 + x[2]^2
f (generic function with 1 method)
# 適当な初期値
x0 = randn(2)
2-element Vector{Float64}:
-0.07058313895389791
0.5314767537831963
# 勾配法による x の更新
x = copy(x0)
xseq1 = [x]
println(x)
for i in 1:15
x = x - 0.1*[8*x[1], 2*x[2]]
println(x)
push!(xseq1, x)
end
[-0.07058313895389791, 0.5314767537831963]
[-0.01411662779077958, 0.425181403026557]
[-0.0028233255581559154, 0.3401451224212456]
[-0.0005646651116311828, 0.2721160979369965]
[-0.00011293302232623652, 0.2176928783495972]
[-2.2586604465247296e-5, 0.17415430267967774]
[-4.5173208930494585e-6, 0.1393234421437422]
[-9.034641786098914e-7, 0.11145875371499377]
[-1.8069283572197825e-7, 0.08916700297199501]
[-3.6138567144395633e-8, 0.071333602377596]
[-7.227713428879127e-9, 0.0570668819020768]
[-1.4455426857758253e-9, 0.04565350552166144]
[-2.8910853715516494e-10, 0.03652280441732915]
[-5.7821707431032983e-11, 0.02921824353386332]
[-1.1564341486206594e-11, 0.023374594827090655]
[-2.312868297241318e-12, 0.018699675861672524]
等高線プロットで最小値に収束する様子を確認してみる．
function plot_contour(f, xseq, xlim=1)
# データの生成
X = range(-xlim, xlim, length=100)
xaxis = [[x1, x2] for x1 in X, x2 in X]
x_opt = [0, 0]
# 1つめのsubplotには数値をつける
fig, ax = subplots(1, 1, figsize=(4, 4))
cs = ax.contour(X, X, f.(xaxis))
ax.clabel(cs, inline=true)
ax.scatter(x_opt[1], x_opt[2],color="red", marker="o")
ax.scatter([x[1] for x in xseq], [x[2] for x in xseq], color="blue", marker="x")
display(fig)
end
plot_contour(f, xseq1)
すると，確かに最小値 $(0, 0)$ に収束していることがわかる．なお，等高線プロットの赤い点が最小値 $(0, 0)$ で，青い点が勾配法による更新の様子を表している．
次に，Fluxの自動微分機能を使用してこの問題を解く．自動微分により，勾配ベクトルを直接記述することなく勾配法を実装することが可能となる．
x = copy(x0)
xseq2 = [x]
println(x)
for i in 1:15
ps = Flux.params(x)
gs = Flux.gradient(ps) do
f(x)
end
x = x - 0.1*gs[x]
println(x)
push!(xseq2, x)
end
plot_contour(f, xseq2)
[-0.07058313895389791, 0.5314767537831963]
[-0.01411662779077958, 0.425181403026557]
[-0.0028233255581559154, 0.3401451224212456]
[-0.0005646651116311828, 0.2721160979369965]
[-0.00011293302232623652, 0.2176928783495972]
[-2.2586604465247296e-5, 0.17415430267967774]
[-4.5173208930494585e-6, 0.1393234421437422]
[-9.034641786098914e-7, 0.11145875371499377]
[-1.8069283572197825e-7, 0.08916700297199501]
[-3.6138567144395633e-8, 0.071333602377596]
[-7.227713428879127e-9, 0.0570668819020768]
[-1.4455426857758253e-9, 0.04565350552166144]
[-2.8910853715516494e-10, 0.03652280441732915]
[-5.7821707431032983e-11, 0.02921824353386332]
[-1.1564341486206594e-11, 0.023374594827090655]
[-2.312868297241318e-12, 0.018699675861672524]
先の勾配を手計算で求め，勾配降下ステップを直接実装した結果と同様になることが確認できる．
続いて，Flux.jl で提供される最適化関数 Descent() とパラメータ更新関数 update! を使用して最適化を行う．
なお，update! でパラメータを更新した場合，x 自身の値がイタレーション毎に更新されてしまい，先と同様に push!(xseq3, x) で値を保存しようとすると， xseq3 の値がすべて同じになってしまう．そのため，copy(x) で x のコピーを保存することで，x の値が更新されても xseq3 には更新前の値が保存されるようにしている．
x = copy(x0)
xseq3 = [copy(x)]
println(x)
opt = Flux.Descent(0.1)
for i in 1:15
ps = Flux.params(x)
gs = Flux.gradient(ps) do
f(x)
end
Flux.Optimise.update!(opt, ps, gs)
println(x)
push!(xseq3, copy(x))
end
plot_contour(f, xseq3)
[-0.07058313895389791, 0.5314767537831963]
[-0.01411662779077958, 0.425181403026557]
[-0.0028233255581559154, 0.3401451224212456]
[-0.0005646651116311828, 0.2721160979369965]
[-0.00011293302232623652, 0.2176928783495972]
[-2.2586604465247296e-5, 0.17415430267967774]
[-4.5173208930494585e-6, 0.1393234421437422]
[-9.034641786098914e-7, 0.11145875371499377]
[-1.8069283572197825e-7, 0.08916700297199501]
[-3.6138567144395633e-8, 0.071333602377596]
[-7.227713428879127e-9, 0.0570668819020768]
[-1.4455426857758253e-9, 0.04565350552166144]
[-2.8910853715516494e-10, 0.03652280441732915]
[-5.7821707431032983e-11, 0.02921824353386332]
[-1.1564341486206594e-11, 0.023374594827090655]
[-2.312868297241318e-12, 0.018699675861672524]
これまでと同様，確かに最小値 $(0, 0)$ に収束していることがわかる．以上，Fluxを使用した勾配法のデモを行った．
Leave a Reply