Juliaの自動微分でラプラシアンを計算する

Julia言語の自動微分パッケージの一つであるForwardDiff.jlとZygote.jlを使って勾配やラプラシアンを計算します. 2025年1月5日現在, Enzyme.jlにはヘッシアンがないため候補から外しました. ForwardDiff.jlとZygote.jlにはヘッシアンが実装されているためトレースを取ればラプラシアンが簡単に計算できます. ヘッシアンの対角成分以外は不要なので間違いなく高速化できますが, 実装の手間がかかるので速度は諦めます. 永井さんによるとZygote.jlにはdiaghessianが用意されているとのことです. ForwardDiff.jlにはそのうちラプラシアンが実装されるかもしれません.

勾配とラプラシアン

勾配は gradient, ラプラシアンは hessian のトレースを取ることで計算できます.

ForwardDiff.jlによる勾配とラプラシアンの例

import LinearAlgebra
import ForwardDiff
f(x) = x[1]^3 + 2*x[2]^4 + 3*x[2]^5
∇f(x) = ForwardDiff.gradient(f, x)
∇²f(x) = LinearAlgebra.tr(ForwardDiff.hessian(f, x))
@show ∇f([1.0, 1.0, 1.0])
@show ∇²f([1.0, 1.0, 1.0])
# ∇f([1.0, 1.0, 1.0]) = [3.0, 23.0, 0.0]
# ∇²f([1.0, 1.0, 1.0]) = 90.0

Zygote.jlの場合は diaghessian の和としても計算できます.

Zygote.jlによる勾配とラプラシアンの例

import LinearAlgebra
import Zygote
f(x) = x[1]^3 + 2*x[2]^4 + 3*x[2]^5
∇f(x) = first(Zygote.gradient(f, x))
# ∇²f(x) = LinearAlgebra.tr(Zygote.hessian(f, x))
∇²f(x) = sum(first(Zygote.diaghessian(f, [1.0, 1.0, 1.0])))
@show ∇f([1.0, 1.0, 1.0])
@show ∇²f([1.0, 1.0, 1.0])
# ∇f([1.0, 1.0, 1.0]) = [3.0, 23.0, 0.0]
# ∇²f([1.0, 1.0, 1.0]) = 90.0

速度の比較

次の4つの関数でのラプラシアンの計算速度を比較します.

• Zygote.hessian
• Zygote.hessian_reverse
• Zygote.diaghessian
• ForwardDiff.hessian

Zygote.jlの diaghessian を使ったものが最速でした.

速度の比較

import BenchmarkTools
import ForwardDiff
import Zygote
f(x) = x[1]^3 + 2*x[2]^4 + 3*x[2]^5
BenchmarkTools.@btime LinearAlgebra.tr(Zygote.hessian(f, [1.0, 1.0, 1.0]))
BenchmarkTools.@btime LinearAlgebra.tr(Zygote.hessian_reverse(f, [1.0, 1.0, 1.0]))
BenchmarkTools.@btime sum(first(Zygote.diaghessian(f, [1.0, 1.0, 1.0])))
BenchmarkTools.@btime LinearAlgebra.tr(ForwardDiff.hessian(f, [1.0, 1.0, 1.0]))
# 309.426 ns (15 allocations: 832 bytes)
# 393.900 μs (2939 allocations: 247.36 KiB)
# 271.617 ns (14 allocations: 656 bytes)
# 472.959 ns (12 allocations: 1.41 KiB)

検算

ガウス関数の勾配とラプラシアンを計算し, 下記の解析的に導いた式と一致するか確認します. 勾配のx方向の成分は次のように計算できます. y,z成分についても同様です.

$$∂∂xexp⁡(−ar2)=∂(−ar2)∂x∂∂(−ar2)exp⁡(−ar2)=∂(−ar2)∂xexp⁡(−ar2)=∂(−a[x2+y2+z2])∂xexp⁡(−ar2)=−2axexp⁡(−ar2)‾ \begin{aligned} \frac{\partial}{\partial x} \exp(-a r^2) &= \frac{\partial (-a r^2)}{\partial x} \frac{\partial}{\partial (-a r^2)} \exp(-a r^2) \\ &= \frac{\partial (-a r^2)}{\partial x} \exp(-a r^2) \\ &= \frac{\partial (-a [x^2 + y^2 + z^2])}{\partial x} \exp(-a r^2) \\ &= \underline{-2ax \exp(-a r^2)} \end{aligned}$$

ガウス関数の勾配

import Zygote
f(x) = exp(-(x[1]^2 + x[2]^2 + x[3]^2))
∇f(x) = first(Zygote.gradient(f, x))
∇f([1.0, 1.0, 1.0])
# 3-element Vector{Float64}:
#   -0.09957413673572789
#   -0.09957413673572789
#   -0.09957413673572789

解析的な勾配

∇f_analytical(x; a=1.0) = [
    -2*a*x[1]*exp(-a*(x[1]^2 + x[2]^2 + x[3]^2))
    -2*a*x[2]*exp(-a*(x[1]^2 + x[2]^2 + x[3]^2))
    -2*a*x[3]*exp(-a*(x[1]^2 + x[2]^2 + x[3]^2))
]
∇f_analytical([1.0, 1.0, 1.0])
# 3-element Vector{Float64}:
#  -0.09957413673572789
#  -0.09957413673572789
#  -0.09957413673572789

ガウス関数に対するラプラシアンは次のように計算できます.

$$∇2exp⁡(−ar2)=[∂2∂x2+∂2∂y2+∂2∂z2]exp⁡(−ar2)=[(−2a+4a2x2)+(−2a+4a2y2)+(−2a+4a2z2)]exp⁡(−ar2)=[−6a+4a2r2]exp⁡(−ar2)‾ \begin{aligned} \nabla^2 \exp(-a r^2) &= \left[ \frac{\partial^2}{\partial x ^2} + \frac{\partial^2}{\partial y ^2} + \frac{\partial^2}{\partial z ^2} \right] \exp(-a r^2) \\ &= \left[ (-2a + 4a^2 x^2) + (-2a + 4a^2 y^2) + (-2a + 4a^2 z^2) \right] \exp(-a r^2) \\ &= \underline{\left[ -6a + 4a^2 r^2 \right] \exp(-a r^2)} \end{aligned}$$

ガウス関数のラプラシアン

import Zygote
f(x) = exp(-(x[1]^2 + x[2]^2 + x[3]^2))
∇²f(x) = sum(first(Zygote.diaghessian(f, [1.0, 1.0, 1.0])))
∇²f([1.0, 1.0, 1.0])
# 0.29872241020718365

# 解析的なラプラシアン
∇²f_analytical(x; a=1.0) = (-6*a + 4*a^2*(x[1]^2 + x[2]^2 + x[3]^2)) * exp(-a * (x[1]^2 + x[2]^2 + x[3]^2))
∇²f_analytical([1.0, 1.0, 1.0])
# 0.29872241020718365

まとめ

ForwardDiff.jlとZygote.jlそれぞれで勾配とラプラシアンが計算できました. 解析的に計算したガウス関数の勾配とラプラシアンに一致する結果が得られています.

バージョン情報

バージョン情報

Julia v1.11.2
ForwardDiff v0.10.38
Zygote v0.6.75

https://gist.github.com/ohno/85874ac32d60cbd89938be3356d94017