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もバックエンドに利用できますが, プリコンパイルに時間がかかるのでここでは省きました.
import Pkg; Pkg.add("DifferentiationInterface")
import Pkg; Pkg.add("ForwardDiff")
import Pkg; Pkg.add("Zygote")
使うときは, 次のようにバックエンドも一緒に読み込みます.
using DifferentiationInterface
import ForwardDiff
import Zygote
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
正しく計算できていますね. このように, 2つ目の引数で利用したいバックエンドを指定することによって, 簡単にバックエンドを切り替えることができます.
2階微分
上記と同じ関数の2回微分を計算します.
が計算できれば正解です.
julia> second_derivative(f, AutoForwardDiff(), 2.0)
48.0
julia> second_derivative(f, AutoZygote(), 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コミュニティでは全世界的に高速な自動微分パッケージEnzyme.jlへの移行が検討されており, 2024年12月の第146回計算科学コロキウム以降, 日本でもEnzyme採用の機運が高まっています. Enzymeは他のパッケージと比べると使い方がやや難しい気がしますが, このインターフェースなら簡単に利用できそうですね. 今回は省略しましたがAutoEnzyme()で利用できるはずです. ぜひお試しください.
参考文献
ノートブック
下記のリンクから, Google Colab上で実行できます.
Discussion