Juliaの自動微分ならDifferentiationInterface.jlがおすすめ
Juliaには多数の自動微分パッケージがあり, それぞれ使い方が異なります. ここでは, これらを共通のインターフェースから利用できるDifferentiationInterface.jlを紹介します.
はじめに
2025年7月に開催されたJuliaHEP 2025 Workshopに参加したところ, 自動微分に詳しい方に上記パッケージを紹介してもらえました. バックエンド一覧をによると, 対応している自動微分パッケージは下記の通りです.
- ChainRules
- Diffractor
- Enzyme
- FastDifferentiation
- FiniteDiff
- FiniteDifferences
- ForwardDiff
- GTPSA
- Mooncake
- PolyesterForwardDiff
- ReverseDiff
- Symbolics
- Tracker
- Zygote
これらのパッケージを同じインターフェースから利用できるようになり, 非常に便利なので共有します.
インストール
DifferentiationInterface.jlを利用するときは, バックエンドとして利用する自動微分パッケージも自分でインストールする必要があります. ここではForwardDiff.jlとZygote.jlをバックエンドとして使用するのでインストールしておきます. もちろん多変数の場合で最速の自動微分パッケージEnzyme.jlもバックエンドに利用できます. Google Colabの場合, Pkg.add("Enzyme")は998秒, import Enzymeは42秒かかったのでコーヒーでも淹れて待ちましょう.
import Pkg; Pkg.add("DifferentiationInterface")
import Pkg; Pkg.add("ForwardDiff")
import Pkg; Pkg.add("Zygote")
import Pkg; Pkg.add("Enzyme")
使うときは, 次のようにバックエンドも一緒に読み込みます.
using DifferentiationInterface
import ForwardDiff
import Zygote
import Enzyme
1階微分
試しに
が計算できれば正解です. 実際にderivativeという関数で微分できることを確かめてみましょう.
julia> f(x) = x^4
f (generic function with 1 method)
julia> derivative(f, AutoForwardDiff(), 2.0)
32.0
julia> derivative(f, AutoZygote(), 2.0)
32.0
julia> derivative(f, AutoEnzyme(), 2.0)
32.0
正しく計算できていますね. このように, 2つ目の引数で利用したいバックエンドを指定することによって, 簡単にバックエンドを切り替えることができます.
2階微分
上記と同じ関数の2回微分を計算します.
が計算できれば正解です.
julia> f(x) = x^4
f (generic function with 1 method)
julia> second_derivative(f, AutoForwardDiff(), 2.0)
48.0
julia> second_derivative(f, AutoZygote(), 2.0)
48.0
julia> second_derivative(f, AutoEnzyme(), 2.0)
48.0
正しく計算できていますね. 2階微分の利用方法はドキュメントに書かれていないことが多いので, ここまで簡単に計算できることには感動しました.
勾配
勾配も次のように計算できます.
julia> f(x) = x[1] + x[2]^2
f (generic function with 1 method)
julia> gradient(f, AutoForwardDiff(), [1.0, 1.0])
2-element Vector{Float64}:
1.0
2.0
julia> gradient(f, AutoZygote(), [1.0, 1.0])
2-element Vector{Float64}:
1.0
2.0
julia> gradient(f, AutoEnzyme(), [1.0, 1.0])
2-element Vector{Float64}:
1.0
2.0
どのバックエンドでも同じ結果が得られていますね.
ヘッシアン
ヘッシアンも次のように計算できます.
julia> f(x) = x[1] + x[2]^2
f (generic function with 1 method)
julia> hessian(f, AutoForwardDiff(), [1.0, 1.0])
2×2 Matrix{Float64}:
0.0 0.0
0.0 2.0
julia> hessian(f, AutoZygote(), [1.0, 1.0])
2×2 Matrix{Float64}:
0.0 0.0
0.0 2.0
julia> hessian(f, AutoEnzyme(), [1.0, 1.0])
2×2 Matrix{Float64}:
0.0 0.0
0.0 2.0
どのバックエンドでも同じ結果が得られていますね.
高速化について
こちらのページで紹介されているように, 途中計算をキャッシュしておき, 利用時に引数として一緒に渡すことで高速化できる場合があります. また, !が付くインプレース版の関数には事前にメモリを確保して渡すので, メモリアロケーションの分だけコストが下がります. 勾配を例に通常版, キャッシュ版, キャッシュ＋インプレース版の使い方と速度を比較してみました.
import Pkg; Pkg.add("BenchmarkTools")
using BenchmarkTools
f(x) = x[1] + x[2]^2
x = [1.0, 1.0]
backend = AutoEnzyme()
prep = prepare_gradient(f, backend, x)
value = [0.0, 0.0]
gradient(f, backend, x) |> println # 通常版
gradient(f, prep, backend, x) |> println # キャッシュ版
gradient!(f, value, backend, x) |> println # インプレース版
gradient!(f, value, prep, backend, x) |> println # キャッシュ＋インプレース版
# [1.0, 2.0]
# [1.0, 2.0]
# [1.0, 2.0]
# [1.0, 2.0]
@btime gradient($f, $backend, $x)
@btime gradient($f, $prep, $backend, $x)
@btime gradient!($f, $value, $backend, $x)
@btime gradient!($f, $value, $prep, $backend, $x)
# 48.437 ns (2 allocations: 80 bytes)
# 51.485 ns (2 allocations: 80 bytes)
# 27.112 ns (0 allocations: 0 bytes)
# 23.077 ns (0 allocations: 0 bytes)
prep = prepare_hessian(f, backend, x)
value = zeros(2,2)
hessian(f, backend, x) |> println # 通常版
hessian(f, prep, backend, x) |> println # キャッシュ版
hessian!(f, value, backend, x) |> println # インプレース版
hessian!(f, value, prep, backend, x) |> println # キャッシュ＋インプレース版
# [0.0 0.0; 0.0 2.0]
# [0.0 0.0; 0.0 2.0]
# [0.0 0.0; 0.0 2.0]
# [0.0 0.0; 0.0 2.0]
@btime hessian($f, $backend, $x)
@btime hessian($f, $prep, $backend, $x)
@btime hessian!($f, $value, $backend, $x)
@btime hessian!($f, $value, $prep, $backend, $x)
# 3.663 μs (37 allocations: 1.55 KiB)
# 174.773 ns (8 allocations: 352 bytes)
# 3.538 μs (29 allocations: 1.20 KiB)
# 93.782 ns (0 allocations: 0 bytes)
Enzymeではキャッシュとインプレースの効果が顕著です. 勾配とヘッシアンの両方において, キャッシュ＋インプレース版ではメモリアロケーションをゼロに抑えて大幅に高速化することができました. なお, 効率化の度合いはバックエンドによって大きく異なり, Zygoteでは逆効果になったりすることもありました.
おわりに
Juliaコミュニティでは全世界的に高速な自動微分パッケージEnzyme.jlへの移行が検討されており, 2024年12月の第146回計算科学コロキウム以降, 日本でもEnzyme採用の機運が高まっています. Enzymeは他のパッケージと比べると使い方がやや難しい気がしますが, このインターフェースなら簡単に利用できました. ぜひお試しください.
参考文献
ノートブック
下記のリンクから, Google Colab上で実行できます.
関連記事
Discussion