Autograd Jacobian vs. Finite Difference Approximation

Dominik: <p>I just took a look at the implementation of <a href="https://github.com/pytorch/pytorch/blob/4cdbe5c4952ee6b7be24ff15b1d45d0fb593b13f/torch/autograd/functional.py#L366" rel="nofollow noopener"><code>torch.autograd.functional.jacobian</code></a> and as it appears it does the following (simplified for <em>single</em> input and output tensor):</p>
<pre><code>input.requires_grad = True
output = func(input)
jac_rows = []
for o in output:
    input.grad.zero_()
    o.backward(retain_graph=True)
    jac_rows.append(input.grad.data)
jac = torch.stack(jac_rows)
</code></pre>
<p>So if we have 1d-input and -output tensors of shape <code>(N,)</code> and <code>(M,)</code> respectively then this performs <code>M</code> backward passes through the graph. In order to estimate the Jacobian one could also use <a href="https://en.wikipedia.org/wiki/Finite_difference#Relation_with_derivatives" rel="nofollow noopener">Finite Difference Approximation</a> which requires <code>N</code> forward passes instead. One can also assume that the cost of a forward and backward pass are roughly the same.</p>
<p>Now the thing is, I’m using this <code>jacobian</code> function specifically because I want to use the <a href="https://en.wikipedia.org/wiki/Levenberg%E2%80%93Marquardt_algorithm" rel="nofollow noopener">Levenberg-Marquardt algorithm</a> to fit a small network (~ 100 parameters) which is more efficient than other gradient-based optimizers. The Levenberg-Marquardt algorithm requires that <code>M &gt;= N</code> and hence this backward Jacobian approach is less efficient than forward difference approximation (M vs. N steps). Now I’m wondering why this implementation has been chosen for <code>jacobian</code> and if there’s a more efficient way. After all there must be a reason to use <code>jacobian</code> rather than just one of the available optimizers. If the reason is to implement a custom optimizer using Jacobian then most likely M &gt; N and the implementation is inefficient (e.g. in one of my cases I had <code>N = 84, M = 576</code> so using Autograd Jacobian is about 7 times slower than Finite Difference Approximation).</p>
<p>Also it seems strange to compute the backward pass separately for each of the output tensor’s elements. This is like using an explicit <code>for</code> loop over the batch dimension for the forward pass (which we don’t do of course). The forward pass uses all parallelization capabilities for the batch dimension, so I’m wondering whether something similar can be used for the backward pass and Jacobian, i.e. performing all the computations for <code>output</code> at once.</p>

kiwiriver: <p>I have almost the same issue as you had. I have a model with input N=15 and output M=55, and I need to compute the Jacobian with respect to the input parameters (15).</p>
<p>Using torch.autograd.functional.jacobian is almost 8 times slower than the finite difference method.</p>
<p>I am curious whether autograd can be implemented in a way which is faster than finite difference.</p>

Connor_Stone: <p>You want to use autograd jacobian in the “forward-mode” in this case for the “strategy” argument. In the forward mode it will do N passes and so will be as fast as finite differences, except that it will be exact.</p>

soulitzer: <p>There’s also torch.func.jacfwd and torch.func.jacrev (these use vmap by default underneath to vectorize the gradient computations)</p>
<p><a href="https://pytorch.org/docs/stable/func.api.html#module-torch.func" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/func.api.html#module-torch.func</a></p>