Efficient Computation of the Product between the Squared Jacobian and a Vector

1

I am trying to compute the product between the squared Jacobian and a vector. Is the best way to edit the C++ jvp code and rebuild PyTorch from source, or is there a better approach? Thanks!

2

Hi! By square do you mean element-wise square or J @ J^T? (where J is the Jacobian matrix).

Also, do you want to do that for forward-mode autodiff? I’m asking that because you mention jvp rather than vjp.

3

Yes. I am looking to use the element-wise squared Jacobian. I need autodiff, but I do not know if I need forward vs. reverse autodiff.

4

I think it depends a bit on what you want to achieve. Internally, torch.func.vjp doesn’t necessarily materialize the Jacobian and make a product between the vector and the Jacobian (in many cases it’s possible to compute the result of the vector-Jacobian product without storing the full Jacobian in memory at any point). So it’s not easy to change it so that it returns the product between a vector and the element-wise squared Jacobian.

What you could do, however, is to compute the Jacobian yourself, square it, and multiply it by a vector. If you want to use the functional API, you can do that easily with torch alone, with for example jacrev. If you prefer to use the imperative API, I would suggest to use torchjd.autojac.jac from the library torchjd (disclaimer: I’m a maintainer of this library).

In any case, you’ll obtain a tuple of Jacobians, each with the same shape as the corresponding parameter. You then have to reshape them, concatenate, square, and multiply by a vector.

With functional API:

import torch
from torch.func import jacrev

weight = torch.tensor([[1., 2.], [3., 4.]])  # shape: [2, 2]
bias = torch.tensor([0.5, -0.5])  # shape: [2]
input_vec = torch.tensor([1., -1.])

def f(weight, bias):
    # Compute arbitrary quantities that are function of weight and bias
    y1 = weight @ input_vec + bias  # shape: [2]
    y2 = (weight ** 2).sum() + (bias ** 2).sum()  # shape: [] (scalar)
    return torch.cat([y1, y2.unsqueeze(0)])  # shape: [3]

jacobians = jacrev(f, argnums=(0, 1))(weight, bias)  # shapes: [3, 2, 2], [3, 2]
jacobian_matrices = tuple(j.flatten(1) for j in jacobians)  # shapes: [3, 4], [3, 2]
J = torch.concat(jacobian_matrices, dim=1)  # shape: [3, 6]
J_sq = J ** 2

vector = torch.tensor([1., 2., 3.])
result = vector @ J_sq
print(result)

=> tensor([ 13., 49., 110., 194., 4., 5.])

With torchjd, after a pip install torchjd, the full code would look something like that:

import torch

from torchjd.autojac import jac

weight = torch.tensor([[1., 2.], [3., 4.]], requires_grad=True)  # shape: [2, 2]
bias = torch.tensor([0.5, -0.5], requires_grad=True)  # shape: [2]
# Compute arbitrary quantities that are function of weight and bias
input_vec = torch.tensor([1., -1.])
y1 = weight @ input_vec + bias  # shape: [2]
y2 = (weight ** 2).sum() + (bias ** 2).sum()  # shape: [] (scalar)

jacobians = jac([y1, y2], [weight, bias])  # shapes: [3, 2, 2], [3, 2]
jacobian_matrices = tuple(j.flatten(1) for j in jacobians)  # shapes: [3, 4], [3, 2]
J = torch.concat(jacobian_matrices, dim=1)  # shape: [3, 6]
J_sq = J ** 2

vector = torch.tensor([1., 2., 3.])
result = vector @ J_sq
print(result)

=> tensor([ 13., 49., 110., 194., 4., 5.])

However, keep in mind that none of those two methods are maximally memory-efficient since you may be able to compute your desired result without ever storing the full Jacobian (which may be extremely heavy) in memory in theory. But I don’t know how to do that. It’s probably an open problem to do that efficiently, unless there’s a trick I didn’t think of.

I hope this helps!