Computing vector-Jacobian and Jacobian-vector product efficiently

username: Tidan
cooked: <p>I need to compute both the vector-Jacobian product and the Jacobian-vector product at the same time, and then to backprop through both. I have the following code that I have tested and I believe it works correctly:</p>
<pre><code class="lang-auto">def vjp(f, x, v, create_graph=True):
    x = x.detach().requires_grad_()
    y = f(x)
    y.backward(v, create_graph=create_graph)
    return x.grad

def jvp(f, x, v, create_graph=True):
    g = lambda v: vjp(f, x, v, create_graph=True)
    return vjp(g, v, v, create_graph=create_graph)

def get_loss(f, x, v):
    vjp_val = vjp(f, x, v)
    jvp_val = jvp(f, x, v)

    return (vjp_val - jvp_val).norm(1)
</code></pre>
<p>It is however inefficient, as it effectively computes <code>f(x).backward(v)</code> twice. Hence I would like to rewrite it in such way that it only does so once. Here is my attempt:</p>
<pre><code class="lang-auto">def get_loss_fast(f, x, v):
    x = x.detach().requires_grad_()
    y = f(x)
    y.backward(v, create_graph=True)
    vjp_val = x.grad

    vjp_val.backward(v, create_graph=True)
    jvp_val = x.grad

    return (vjp_val - jvp_val).norm(1)
</code></pre>
<p>This code always returns zero. In fact, inside the <code>get_loss_fast</code> function, <code>vjp_val is jvp_val</code> is <code>True</code>, which means that the second <code>backward()</code> does not overwrite the output of the first one.</p>
<p>How can I compute this loss efficiently and correctly?</p>

username: albanD
cooked: <p>Hi,</p>
<p>First thing for such tasks, I would advise to use <code>autograd.grad()</code> instead of <code>.backward()</code> as <code>.backward()</code> might create reference cycles.</p>
<p>Some thing like this should work no?</p>
<pre><code class="lang-python">def get_loss_fast(f, x, v):
    x = x.detach().requires_grad_()
    v = v.detach().requires_grad_()
    y = f(x)
    grad_x = autograd.grad(y, x, v, create_graph=True)[0]
    vjp_val = grad_x

    jvp_val = autograd.grad(grad_x, v, v.detach())[0]

    return (vjp_val - jvp_val).norm(1)
</code></pre>
<p>note that because of the detach at the beginning, no gradient will flow back to the input x or v if you try to backward the result of that function. Is that something you want?</p>

username: Tidan
cooked: <p>Thank you, this seems to return correct results. One question though: why did you put <code>v.detach()</code> in the expression for <code>jvp_val</code>? I.e., why would something like this not work correctly?</p>
<pre><code class="lang-auto">def get_loss_fast(f, x, v):
    x = x.detach().requires_grad_()
    v = v.detach().requires_grad_()

    vjp_val = autograd.grad(f(x), x, v, create_graph=True)[0]
    jvp_val = autograd.grad(vjp_val, v, v)[0]  # instead of vjp_val, v, v.detach()

    return (vjp_val - jvp_val).norm(1)
</code></pre>
<blockquote>
<p>Because of the detach at the beginning, no gradient will flow back to the input x or v if you try to backward the result of that function. Is that something you want?</p>
</blockquote>
<p>Yes, it is exactly as I intend.</p>

username: albanD
cooked: <p>Just because we don’t want to create the graph for that backward, so I was making the grad_outputs not require gradients. But I guess that is not strictly necessary as we do not set <code>create_graph</code> <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

username: Tidan
cooked: <p>Since I want to backprop through the resulting loss, I (think I) need to set <code>create_graph</code> in the second call to <code>grad()</code>. Otherwise, I am getting the following:</p>
<pre><code class="lang-auto"># loss = get_loss_fast(f, x, v)
# loss.backward()

RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.
</code></pre>
<p>where <code>f</code> is an <code>nn.Linear</code>, and <code>x</code> and <code>v</code> are two random vectors. When I specify <code>create_graph=True</code>, there is no such error.</p>
<p>In any case, as far as I understand, the performance penalty for either doing a <code>.detach()</code> or doing a redundant gradient computation is negligible, so there is little point in benchmarking whether this <code>.detach()</code> is useful or not.</p>
<p>A working version of the code for posterity:</p>
<pre><code class="lang-auto">def get_loss_fast(f, x, v):
    x = x.detach().requires_grad_()
    v = v.detach().requires_grad_()

    vjp_val = autograd.grad(f(x), x, v, create_graph=True)[0]
    jvp_val = autograd.grad(vjp_val, v, v.detach(), create_graph=True)[0]

    return (vjp_val - jvp_val).norm(1)
</code></pre>

username: anantguptadbl
cooked: <p>All the codes are giving me a result of 0</p>
<pre><code class="lang-auto">x=torch.FloatTensor([1,2,3])
v=torch.FloatTensor([1,1,1])
x = x.detach().requires_grad_()
v = v.detach().requires_grad_()
y = x**3 - 6*x
grad_x = torch.autograd.grad(y, x, v, create_graph=True)[0]
print("Grad_x is {0}".format(grad_x))
vjp_val = grad_x
jvp_val = torch.autograd.grad(grad_x, v, v.detach(),create_graph=True)[0]
print("jvp_val is {0}".format(jvp_val))
print((vjp_val - jvp_val).norm(1))
</code></pre>
<p><strong>Results</strong></p>
<pre><code class="lang-auto">tensor([ 6., 12., 18.], grad_fn=&lt;CloneBackward&gt;)
Grad_x is tensor([-3.,  6., 21.], grad_fn=&lt;AddBackward0&gt;)
jvp_val is tensor([-3.,  6., 21.], grad_fn=&lt;AddBackward0&gt;)
tensor(0., grad_fn=&lt;NormBackward0&gt;)
</code></pre>

username: albanD
cooked: <p>Well if you want gradients, you might want to double check that.<br>
In particular, you will need to remove the detach at the beginning and make sure they require gradients.<br>
Also remove the detach in the last grad and add a create_grad=True</p>
<p>The following should work:</p>
<pre><code class="lang-python">import torch
from torch import autograd

def f(x):
    return 2 + x**4 + x[0] ** 5

def get_loss_fast(x, v):
    y = f(x)
    grad_x = autograd.grad(y, x, v, create_graph=True)[0]
    vjp_val = grad_x

    jvp_val = autograd.grad(grad_x, v, v, create_graph=True)[0]

    return (vjp_val - jvp_val).norm(1)

inp = (torch.rand(2, dtype=torch.double, requires_grad=True), torch.rand(2, dtype=torch.double, requires_grad=True))
print(autograd.gradcheck(get_loss_fast, inp))

get_loss_fast(*inp).backward()
print(inp[0].grad)
print(inp[1].grad)
</code></pre>

username: Tidan
cooked: <p>In my particular case, I only need the gradients for <code>f</code>, and not for <code>x</code> or <code>v</code>. Thanks for pointing this out though!</p>

username: Tidan
cooked: <p>Here is a working demo on Colab: <a href="https://colab.research.google.com/drive/16RO6uW1L8RS8neC3NUs5wQxfObwHP0i6" rel="nofollow noopener">https://colab.research.google.com/drive/16RO6uW1L8RS8neC3NUs5wQxfObwHP0i6</a></p>

username: anantguptadbl
cooked: <p>Thank you <a class="mention" href="/u/tidan">@Tidan</a><br>
Can you also tell me the mathematical background behind calculating the  (vjp_val - jvp_val).norm(1) as loss. Somehow, I am not able to get my head around this</p>
<p>An external link etc would also suffice</p>