07e366e56b20dbfd50de.md

title	tags	author	slide
Hamiltonian Descent Methods の Chainer 実装を改善した	Python 機械学習 Chainer Julia Flux.jl	omi_UT	false

この記事は Hamiltonian Descent Methods を chainer で実装した の続きになります。

数学的な話は特にしないので、引き続き

$$p_{i+1}=\delta p_i-\epsilon\delta\nabla f(x_i)\\\ x_{i+1}=x_i+\epsilon\nabla k(p_{i+1})$$

がわかっていれば大丈夫です。

改善点

k を使わない

$k = \sqrt{\|p\|^2+1}-1$ ですが、

relativistic kinetic と呼ばれ，微分すると， $$\nabla k(p)=\frac{p}{\sqrt{\|p\|^2+1}}$$ となり，様々な適応的勾配降下法で使われている形になる．

これ ¹ をすっかり忘れていました。更新に必要なのは $ \nabla k$ であり、 $k$ ではありません。 よって更新処理の中に Variable を出現させる必要はありませんでした。

cuda.elementwise の使用

使い方を調べました。$\nabla k(p_{i+1})$ には $\|p_{i+1}\|^2$ が必要なため、$p$ の更新と $x$ の更新とで二段階に分ける必要があります。

実装

    def update_core_gpu(self, param):
        grad = param.grad
        if grad is None:
            return
        hp = self.hyperparam
        p = self.state['p']
        
        if HamiltonianRule._kernel is None:
            HamiltonianRule._kernel = cuda.elementwise(
                'T delta, T epsilon, T grad', 'T p',
                'p *= delta; p -= epsilon * delta * grad;',
                'Hamiltonian_p')
        HamiltonianRule._kernel(hp.delta, hp.epsilon, grad, p)

        sqsum = cuda.cupy.vdot(p,p)
        if HamiltonianRule._kernel_2 is None:
            HamiltonianRule._kernel_2 = cuda.elementwise(
                'T epsilon, T p, T sqsum', 'T param', 
                'param += epsilon * p / (1.0 + sqsum)', 'Hamiltonian_q')
        
        HamiltonianRule._kernel_2(hp.epsilon, p, sqsum, param.data)

結果

実験の条件等は前回 から一切変更していません。

最初のほう滑らかに行ってくれないのが気に食わないのですが、途中からちゃんと収束が始まったのはよさげです。 また Adam と同じくらいのスピードで動いてくれるようです。

まとめ

実装がまともになった

おまけ

前回断念した Julialang による実装を行います

使うもの

• julialang
• Flux.jl
• その他種々のパッケージ
  • CuArrays (たぶん Flux で入る…？)
  • PyPlot
  • JSON
  • DataFrames
  • Lazy

実装

本体部分のみ。全部は GitHub レポジトリの方 に上げてます。

function Flux.Optimise.apply!(o::Hamiltonian, x, Δ)
    ϵ, δ = o.epsilon, o.delta
    p = get!(o.momentum, x, zero(x))
    @. p *= δ
    @. p -= ϵ * δ * Δ

    o.momentum[x] = p
    sqsum = p ⋅ p
    @. Δ = - ϵ / (1.0 + sqsum) * p
    return Δ
end

$\epsilon$ とか $\delta$ とか使えると楽ですね apply! はデータから 引く 値を返すべき関数らしいので、 $\Delta$ の正負を逆にしています。

モデルは python のときと一緒にしました。 softmax だったのかは謎だけど

m = Chain(
  Dense(28^2, 500, relu),
  Dense(500, 500, relu),
  Dense(500, 10),
  softmax) |> gpu

あとは流すだけなのですが CUDA 周りに深刻なバグが山ほど埋まってるようです。特に broadcast 関係が厳しいらしく、それは julia 本体が悪いのでは…？ という感じ

# バッチに分ける（手動）
dataset = [ (X[:,100i-99:100i],Y[:,100i-99:100i]) for i in 1:600]

for j in 1:100
    Flux.train!(loss, params(m), dataset, opt, cb = throttle(evalcb, 1))
end

結果

最初ギザギザから思い出したように収束を始めるのは一緒っぽい

時間比で見るとおそろしく速い julia 側のデータが生配列ってことを差っ引いても速いと思います

まとめ

julia is GOD. あとはバグさえ少なければ…

Footnotes

1. 元記事の脚注 1. ↩