Sparsity
We present sparsity handling with DifferentiationInterface.jl.
using BenchmarkTools
using DifferentiationInterface
import ForwardDiff, Zygote
Sparse AD is very useful when Jacobian or Hessian matrices have a lot of zeros. So let us write functions that satisfy this property.
f_sparse_vector(x::AbstractVector) = diff(x .^ 2) + diff(reverse(x .^ 2))
f_sparse_scalar(x::AbstractVector) = sum(f_sparse_vector(x) .^ 2)
Let's also pick a random test vector.
x = float.(1:8);
8-element Vector{Float64}:
1.0
2.0
3.0
4.0
5.0
6.0
7.0
8.0
Dense backends
When we use the jacobian
or hessian
operator with a dense backend, we get a dense matrix with plenty of zeros.
dense_first_order_backend = AutoForwardDiff()
J_dense = jacobian(f_sparse_vector, dense_first_order_backend, x)
7×8 Matrix{Float64}:
-2.0 4.0 0.0 0.0 0.0 0.0 14.0 -16.0
0.0 -4.0 6.0 0.0 0.0 12.0 -14.0 0.0
0.0 0.0 -6.0 8.0 10.0 -12.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 6.0 -8.0 -10.0 12.0 0.0 0.0
0.0 4.0 -6.0 0.0 0.0 -12.0 14.0 0.0
2.0 -4.0 0.0 0.0 0.0 0.0 -14.0 16.0
dense_second_order_backend = SecondOrder(AutoForwardDiff(), AutoZygote())
H_dense = hessian(f_sparse_scalar, dense_second_order_backend, x)
8×8 Matrix{Float64}:
112.0 -32.0 0.0 0.0 0.0 0.0 -112.0 128.0
-32.0 96.0 -96.0 0.0 0.0 -192.0 448.0 -256.0
0.0 -96.0 256.0 -192.0 -240.0 576.0 -336.0 0.0
0.0 0.0 -192.0 224.0 320.0 -384.0 0.0 0.0
0.0 0.0 -240.0 320.0 368.0 -480.0 0.0 0.0
0.0 -192.0 576.0 -384.0 -480.0 1120.0 -672.0 0.0
-112.0 448.0 -336.0 0.0 0.0 -672.0 1536.0 -896.0
128.0 -256.0 0.0 0.0 0.0 0.0 -896.0 1120.0
The results are correct but the procedure is very slow. By using a sparse backend, we can get the runtime to increase with the number of nonzero elements, instead of the total number of elements.
Sparse backends
Recipe to create a sparse backend: combine a dense backend, a sparsity detector and a compatible coloring algorithm inside AutoSparse
. The following are reasonable defaults:
using SparseConnectivityTracer: TracerSparsityDetector
using SparseMatrixColorings: GreedyColoringAlgorithm
sparse_first_order_backend = AutoSparse(
AutoForwardDiff();
sparsity_detector=TracerSparsityDetector(),
coloring_algorithm=GreedyColoringAlgorithm(),
)
sparse_second_order_backend = AutoSparse(
SecondOrder(AutoForwardDiff(), AutoZygote());
sparsity_detector=TracerSparsityDetector(),
coloring_algorithm=GreedyColoringAlgorithm(),
)
Now the resulting matrices are sparse:
jacobian(f_sparse_vector, sparse_first_order_backend, x)
7×8 SparseArrays.SparseMatrixCSC{Float64, Int64} with 26 stored entries:
-2.0 4.0 ⋅ ⋅ ⋅ ⋅ 14.0 -16.0
⋅ -4.0 6.0 ⋅ ⋅ 12.0 -14.0 ⋅
⋅ ⋅ -6.0 8.0 10.0 -12.0 ⋅ ⋅
⋅ ⋅ ⋅ 0.0 0.0 ⋅ ⋅ ⋅
⋅ ⋅ 6.0 -8.0 -10.0 12.0 ⋅ ⋅
⋅ 4.0 -6.0 ⋅ ⋅ -12.0 14.0 ⋅
2.0 -4.0 ⋅ ⋅ ⋅ ⋅ -14.0 16.0
hessian(f_sparse_scalar, sparse_second_order_backend, x)
8×8 SparseArrays.SparseMatrixCSC{Float64, Int64} with 40 stored entries:
112.0 -32.0 ⋅ ⋅ ⋅ ⋅ -112.0 128.0
-32.0 96.0 -96.0 ⋅ ⋅ -192.0 448.0 -256.0
⋅ -96.0 256.0 -192.0 -240.0 576.0 -336.0 ⋅
⋅ ⋅ -192.0 224.0 320.0 -384.0 ⋅ ⋅
⋅ ⋅ -240.0 320.0 368.0 -480.0 ⋅ ⋅
⋅ -192.0 576.0 -384.0 -480.0 1120.0 -672.0 ⋅
-112.0 448.0 -336.0 ⋅ ⋅ -672.0 1536.0 -896.0
128.0 -256.0 ⋅ ⋅ ⋅ ⋅ -896.0 1120.0
Sparse preparation
In the examples above, we didn't use preparation. Sparse preparation is more costly than dense preparation, but it is even more essential. Indeed, once preparation is done, sparse differentiation is much faster than dense differentiation, because it makes fewer calls to the underlying function. The speedup becomes very visible in large dimensions.
n = 1000
jac_extras_dense = prepare_jacobian(f_sparse_vector, dense_first_order_backend, zeros(n))
jac_extras_sparse = prepare_jacobian(f_sparse_vector, sparse_first_order_backend, zeros(n))
@benchmark jacobian($f_sparse_vector, $jac_extras_dense, $dense_first_order_backend, $(randn(n)))
BenchmarkTools.Trial: 840 samples with 1 evaluation.
Range (min … max): 4.866 ms … 25.719 ms ┊ GC (min … max): 0.00% … 3.92%
Time (median): 5.462 ms ┊ GC (median): 12.39%
Time (mean ± σ): 5.947 ms ± 1.750 ms ┊ GC (mean ± σ): 12.54% ± 3.17%
▁█▅▃
▄▁████▇▇▅▁▆▄▇▆▅▄▅▄▅▄▄▄▄▁▄▄▄▄▄▄▄▄▁▄▄▄▄▁▄▁▁▄▄▆▆▄▄▄▄▄▁▁▄▁▄▁▄▅ ▇
4.87 ms Histogram: log(frequency) by time 14.2 ms <
Memory estimate: 57.62 MiB, allocs estimate: 1011.
@benchmark jacobian($f_sparse_vector, $jac_extras_sparse, $sparse_first_order_backend, $(randn(n)))
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Range (min … max): 25.738 μs … 2.480 ms ┊ GC (min … max): 0.00% … 95.53%
Time (median): 33.242 μs ┊ GC (median): 0.00%
Time (mean ± σ): 41.056 μs ± 80.365 μs ┊ GC (mean ± σ): 12.64% ± 6.54%
▅██▄▂▁ ▂
▇███████▇▇▆▆▆▅▄▄▃▁▁▁▁▁▁▃▁▄▄▁▁▁▁▁▁▁▁▁▁▃▁▁▁▁▁▃▁▁▁▁▁▁▃▁▁▃▁▁▅██ █
25.7 μs Histogram: log(frequency) by time 172 μs <
Memory estimate: 399.66 KiB, allocs estimate: 26.