# Improve speed of vector-jacobian product

**URL:** https://discourse.julialang.org/t/improve-speed-of-vector-jacobian-product/106312
**Category:** Performance
**Tags:** reversediff, autodiff
**Created:** [November 16, 2023, 11:51am UTC](https://discourse.julialang.org/t/improve-speed-of-vector-jacobian-product/106312 "2023-11-16T11:51:24Z")
**Posts on this page:** 3
**Page:** 1

<div class="post-metadata">

### Author: ![rveltz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rveltz/32/2707_2.png) [@rveltz](https://discourse.julialang.org/u/rveltz)
#### Post date: [November 16, 2023, 11:51am UTC](https://discourse.julialang.org/t/improve-speed-of-vector-jacobian-product/106312/1 "2023-11-16T11:51:24Z")

</div>

Hi,

I would like to improve some part of a code which relies on vector-jacobian product (vjp), ie computing the action of the adjoint of an operator. On my computer, the computation of vjp is ~100x slower than the evaluation of the linear operator itself.

I would be grateful if some of you can come up with better ideas 😃

Thanks a lot.

Here is a very simplified example of my problem

```julia
using Revise, Pkg, LinearAlgebra
cd(@ __DIR__ )
pkg"activate 2d"
using Revise, Parameters, ForwardDiff

function init(vmin, vmax, Nv)
    v = LinRange(vmin, vmax, Nv) |> collect
    Dv = v[2]-v[1]
    fv = @. exp(v)
    V = @. v^2
    return (;vmin, vmax, Nv, v, Dv, fv, V)
end

function inner(i, divergence, Nv, invDv, g, V)
    # Values
    gi = g[i]
    gim = i > 1 ? g[i-1] : zero(gi)
    gip = (i+1) <= Nv ? g[i+1] : zero(gi)
    # Upwind fluxes + null flux boundary conditions
    Fi = i > 1 ? (V[i] > 0 ? V[i] * gim : V[i] * gi) : zero(gi)
    Fip = i < Nv ? (V[i+1] > 0 ? V[i+1]* gi : V[i+1] * gip) : zero(gi)
    # Divergence
    divergence[i] = invDv * (Fip - Fi)
end

function rhs1d!(result, g, p)
    @unpack vmin, vmax, Dv, fv, Nv, V = p
    invDv = 1 / Dv

    Vtmp = V .+ (dot(g, fv) * Dv)
    for i=1:Nv
        inner(i, result, Nv, invDv, g, Vtmp)
    end
    result
end
rhs1d(g, p) = rhs1d!(similar(g), g, p)

mesh = init(-1, 3, 1000)

g0 = @. exp(-mesh.v^2/(2*0.1^2))
g0 ./= sum(g0)*mesh.Dv
# plot(mesh.v,g0)

@time rhs1d(g0, mesh);

import AbstractDifferentiation as AD
using ReverseDiff

vjp(x,p,dx) = AD.pullback_function(AD.ReverseDiffBackend(), z->rhs1d(z, p), x)(dx)[1]
@btime vjp($g0,$mesh,$g0); # 530.292 μs (33040 allocations: 1.46 MiB)
@btime rhs1d($g0, $mesh); # 2.338 μs (2 allocations: 15.88 KiB)

```

---

<div class="post-metadata">

### Author: ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [November 16, 2023, 12:36pm UTC](https://discourse.julialang.org/t/improve-speed-of-vector-jacobian-product/106312/2 "2023-11-16T12:36:20Z")

</div>

> [@rveltz](#):
>
> ```julia
> for i=1:Nv
> inner(i, result, Nv, invDv, g, Vtmp)
> end
> 
> ```

Reverse-mode differentiation may not understand that each iteration of this loop updates a different element of `result` independently — I’m guessing that it makes a copy of the `result` array for every iteration of the loop, before running the loop backwards in order to backpropagate the vjp.

Moreover, you are effectively doing a sparse matrix–vector multiplication in which the elements of the matrix depend on your parameters, and AD tools often struggle with backpropagating through a sparse-matrix construction for reasons I explained in this thread: [Zygote.jl: How to get the gradient of sparse matrix - #6 by stevengj](https://discourse.julialang.org/t/zygote-jl-how-to-get-the-gradient-of-sparse-matrix/59067/6)

I’ve typically found that you need to write a manual vJp for some step(s) in most reasonable-scale scientific problems, unless you are using a package like [DiffEqFlux.jl](https://github.com/SciML/DiffEqFlux.jl) that has done that for you. AD tools are more reliable for cookie-cutter ML-style problems where you are plugging together large components that it already knows about, and are only fiddling around the edges with small scalar functions or code written in a functional/non-mutating style (e.g. composing vectorized primitives).

---

<div class="post-metadata">

### Author: ![rveltz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rveltz/32/2707_2.png) [@rveltz](https://discourse.julialang.org/u/rveltz)
#### Post date: [November 16, 2023, 6:31pm UTC](https://discourse.julialang.org/t/improve-speed-of-vector-jacobian-product/106312/3 "2023-11-16T18:31:59Z")

</div>

> [@stevengj](#):
>
> Moreover, you are effectively doing a sparse matrix–vector multiplication in which the elements of the matrix depend on your parameters, and AD tools often struggle with backpropagating through a sparse-matrix construction for reasons I explained in this thre

thanks for this, it is worrying indeed. In my case I use `ReverseDiff` but I imagine that your remark remains.

> I’ve typically found that you need to write a manual vJp for some step(s) in most reasonable-scale scientific problems,

ah ! it is a nice suggestion. I will try to decompose the problem and use AD only where necessary.