Significant time difference of calculating Jacobian matrix using jacrev and oracle functions

tiantt: <p>Specifically, I found that using <code>jacrev</code> would be much slower than using the oracle Jacobian function:</p>
<pre><code class="lang-python">from torch.func import vmap, jacrev
import torch
import time

a = torch.rand(10000, 10000)

def f(x):
    return (x ** 2).sum(-1)

def df(x):
    return 2 * x


t0 = time.time()
b = df(a)
t1 = time.time()
c = vmap(jacrev(f))(a)
t2= time.time()

assert torch.allclose(b, c)
print(t1 - t0, t2 - t1)
</code></pre>
<p>result: 0.10568618774414062 0.9206998348236084</p>
<p>Given that oracle’s Jacobian is readily available in neural networks, I wonder why using <code>jacrev</code> is so much slower? Is there something wrong with me?</p>
<p>Of course, I can actually rewrite each layer of the neural network to obtain the value and Jacobian at the same time, but calculating the Hessian matrix is too troublesome. It would be great if <code>jacrev</code> could be faster.</p>

KFrank: <p>Hi Frank!</p>
<aside class="quote no-group quote-modified" data-username="tiantt" data-post="1" data-topic="184392" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tiantt/48/30110_2.png" class="avatar"> tiantt:</div>
<blockquote>
<p>I found that using <code>jacrev</code> would be much slower than using the oracle Jacobian function:</p>
<pre data-code-wrap="python"><code class="lang-python">def f(x):
    return (x ** 2).sum(-1)

def df(x):
    return 2 * x

b = df(a)
c = vmap(jacrev(f))(a)
</code></pre>
<p>result: 0.10568618774414062 0.9206998348236084</p>
</blockquote>
</aside>
<p>This result is to be expected.</p>
<p>In the first case, you perform one direct computation – no autograd, no<br>
computation graph, and hence none of the associated overhead.  Note,<br>
in the first case, you never compute <code>x ** 2</code>, so if you need that result,<br>
you would have to compute it separately.</p>
<p>In the second case, <code>jacrev()</code> performs the forward pass that computes<br>
<code>x ** 2</code> and then performs the backward pass that computes the gradient<br>
of the forward pass (which happens to be <code>2 * x</code>).  It does this with all the<br>
overhead of the autograd machinery.</p>
<p>If your real-world use case only has one “layer,” then using autograd might<br>
well be overkill – the real benefit of autograd is that it knows how to chain<br>
multiple layers together.</p>
<p>As an aside, the jacobian in your example is really just the simpler case of<br>
a gradient (because after applying <code>vmap()</code> you are performing <code>x ** 2</code> on<br>
a single row of <code>a</code> at a time, after which <code>.sum (-1)</code> leaves you with a single<br>
scalar value, and the “jacobian” of a scalar-valued function is its gradient).</p>
<p>You can verify this by replacing <code>vmap (jacrev (f)) (a)</code> in your example<br>
with <code>vmap (grad (f)) (a)</code>.</p>
<p>Here’s an example script similar to yours that shows the overhead of the<br>
autograd forward / backward pass when computing just the gradient,<br>
rather than the (scalar-function) jacobian:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

import time

_ = torch.manual_seed (2023)

a = torch.randn (100000000)

def f (x):
    return  x.sin()

def df (x):
    return  x.cos()

t0 = time.time()
b = df (a)
t1 = time.time()
c = torch.func.vmap (torch.func.grad (f))(a)
t2 = time.time()

print ('torch.allclose(b, c) =', torch.allclose(b, c))

print(t1 - t0, t2 - t1)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.0.0
torch.allclose(b, c) = True
0.07053017616271973 0.4128530025482178
</code></pre>
<p>Best.</p>
<p>K. Frank</p>