Finding the Diagonal of Hessian wrt Input for Vector-Valued Functions

jdeo: <p>My goal is to compute the hessian wrt the input to the function / neural net, not wrt its parameters.</p>
<p>Suppose you have a vector-valued function in pytorch f:R^n → R^m.<br>
Computing its hessian should return a [M, N, N] tensor, and its diagonals would return a [M,N] matrix</p>
<p>Computing this whole hessian and then diagonalizing is too expensive for large functions like neural nets.</p>
<p>There also already exists a technique provided in the Pytorch docs to do something similar when dealing with a scalar-valued function (f:R^n → R, and its hessian diagonal would return a [N] vector) using hvp</p>
<pre><code class="lang-auto">def hvp(f, x, v):
    return jvp(grad(f), (x,), (v,))[1]
</code></pre>
<p>but using it directly on vector-valued functions returns an error.</p>
<p>I’ve tried coming up with my own implementation with hvp_vecfwd</p>
<pre><code class="lang-auto"># My attempt
def hvp_vecfwd(f, x, v):
    return jvp(jacrev(f), (x,), (v,))[1]
</code></pre>
<p>It does produce the correct shape, but using the cotangents</p>
<pre><code class="lang-auto"> torch.ones_like(x_single)
</code></pre>
<p>the values inside are not all correct even on simple vector functions (f_single_vector in code below). Is this problem with my choice of cotangents, my implementation, or both?</p>
<p>Any help or insights on this would be great!</p>
<pre><code class="lang-auto">import torch
from torch.func import vmap, jvp, vjp, grad, jacrev, jacfwd, hessian
from functools import partial

x_single = torch.tensor([1.0, 2.0, 3.0], requires_grad=False)
x_batch = torch.tensor([[1.0, 2.0, 3.0],
                        [4.0, 5.0, 6.0],
                        [7.0, 8.0, 9.0]], requires_grad=False)
layer = torch.nn.Linear(3, 5).requires_grad_(False)


# For Scalar Functions
# Provided by PyTorch Docs
def hvp(f, x, v):
    return jvp(grad(f), (x,), (v,))[1]

# For Vector Functions
# My attempt
def hvp_vecfwd(f, x, v):
    return jvp(jacrev(f), (x,), (v,))[1]

# Simple functions for debugging

# Batched f: [B, R^N] -&gt; [B, R]
def f_batch_scalar(x):
    return torch.sum(x*x, dim=-1)


# f: R^N -&gt; R^M
def f_single_vector(x):
    return torch.sigmoid(layer(x))


# Batched f: [B, R^N] -&gt; [B, R^M]
def f_batch_vector(x):
    return torch.sigmoid(layer(x))


print("Batched Full Hessian of Scalar Function")
# Batched f: [B, R^N] -&gt; R, d^2f/dx^2 = [B, N, N]
print(vmap(hessian(f_batch_scalar))(x_batch))

print("Batched Diagonal Hessian of Scalar Function")
# Batched f: [B, R^N] -&gt; R, diag[ H] = [B, N]
print(vmap(hvp, in_dims=(None, 0, 0))(f_batch_scalar, x_batch, torch.ones_like(x_batch)))

print("Non-Batched Full Hessian of Vector Function")
# f: R^N -&gt; R^M, d^2f/dx^2 = [M, N, N]
print(hessian(f_single_vector)(x_single))

print("Non-Batched Diagonal Hessian of Vector Function")
# f: R^N -&gt; R^M, diag[H] = [M, N]
# WRONG
print(hvp_vecfwd(f_single_vector, x_single,  torch.ones_like(x_single)))

print("Batched Full Hessian of Vector Function")
# Batched f: [B, R^N} -&gt; [B, R^M] d^2f/dx^2 = [B, M, N, N]
print(vmap(hessian(f_batch_vector))(x_batch))

print("Batched Diagonal Hessian of Vector Function")
# Batched f: [B, R^N} -&gt; [B, R^M], diag[H] = [B, M, N]
# WRONG
print(vmap(hvp_vecfwd, in_dims=(None, 0, 0))(f_batch_vector, x_batch, torch.ones_like(x_batch)))
</code></pre>

KFrank: <p>Hi jdeo!</p>
<aside class="quote no-group" data-username="jdeo" data-post="1" data-topic="201278" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/eada6e/48.png" class="avatar"> jdeo:</div>
<blockquote>
<p>Computing this whole hessian and then diagonalizing is too expensive for large functions like neural nets.</p>
</blockquote>
</aside>
<p>Why do you say this?  You say that you want the (diagonal of the) hessian with<br>
respect to your R^n <em>input,</em> so the size of the network doesn’t affect how large<br>
your <code>[N, N]</code> hessian matrix is.</p>
<p>Consider your gradient of shape <code>[N]</code>.  To compute the <code>[0, 0]</code> element of your<br>
hessian you compute the gradient of the <code>[0]</code> element of your gradient and<br>
keep just the <code>[0]</code> element of that result.  But (for any realistic network) you<br>
wouldn’t save any significant work by computing just the <code>[0]</code> element of<br>
that gradient.</p>
<p>When you compute the gradient of the <code>[0]</code> element of the first gradient,<br>
autograd chains together a series of jacobian-vector products for each layer.<br>
Roughly speaking, all of the terms get mixed together in each layer, so you<br>
<em>have to</em> perform the full jacobian-vector products for all of the intermediate<br>
layers even though you know that you will be discarding all but one element<br>
of the final gradient.</p>
<p>Just go ahead and let autograd compute the full <code>[N, N]</code> hessian, and then<br>
discard all but the diagonal elements.</p>
<aside class="quote no-group">
<blockquote>
<p>(f:R^n → R, and its hessian diagonal would return a [N] vector) using hvp</p>
</blockquote>
</aside>
<p>You can’t use a hessian-vector product to obtain the diagonal of the hessian.</p>
<p>Consider some <code>[N, N]</code> matrix.  There is no vector (which would have shape <code>[N]</code>)<br>
for which matrix times vector gives you the diagonal of the matrix.  That’s just<br>
not how matrix-vector multiplication works.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">print("Batched Full Hessian of Scalar Function")
...
print("Non-Batched Full Hessian of Vector Function")
...
print("Batched Full Hessian of Vector Function")
</code></pre>
</blockquote>
</aside>
<p>Leaving aside the issue of obtaining the diagonal of the hessian, are you<br>
asking more about how to do this for a vector-valued function, for a batch<br>
of independent inputs, or for both?</p>
<p>Best.</p>
<p>K. Frank</p>