CODEWITHSUNDEEP
pythonpytorchgradient-descentderivativeautograd
Ray Bern
Ray Bern

Reputation: 135

PyTorch function involving softmax and log2 second derivative is always 0

I'm trying to compute the second derivatives (Hessian) of a function t with respect to a tensor a using PyTorch. Below is the code I initially wrote:

import torch

torch.manual_seed(0)
a = torch.randint(0, 10, (10,), dtype=float, requires_grad=True)
b, *_ = a.sort(descending=True)
c = (b.unsqueeze(0) - a.unsqueeze(1)).abs().neg()
d = c.softmax(0).matmul(torch.arange(c.size(0), dtype=c.dtype))
e = torch.randint(0, 3, (10,), dtype=float)
t = torch.sum(e * torch.log2(d + 1))

grad, *_ = torch.autograd.grad(t, a, create_graph=True)
hess, *_ = torch.autograd.grad(grad, a, torch.ones_like(a))
print(hess)

I'm expecting the code to compute the component-wise second derivatives of t with respect to a. However, the Hessian vector returned consists entirely of zeros. This result is puzzling to me because the function t involves a softmax operation and logarithms, which I would expect to yield non-zero second derivatives.

To investigate further, I attempted a different approach:

import torch

torch.manual_seed(0)
a = torch.randint(0, 10, (10,), dtype=float, requires_grad=True)
b, *_ = a.sort(descending=True)
c = (b.unsqueeze(0) - a.unsqueeze(1)).abs().neg()
d = c.softmax(0).matmul(torch.arange(c.size(0), dtype=c.dtype))
e = torch.randint(0, 3, (10,), dtype=float)
t = torch.sum(e * torch.log2(d + 1))

t.backward(create_graph=True)
grad = a.grad.clone()
hess = torch.zeros_like(a)
for i in range(len(a)):
    a.grad.zero_()
    grad[i].backward(retain_graph=True)
    hess[i] = a.grad[i]
print(hess)

In this second attempt, while the gradient (grad) matches what I obtained in the first method, the Hessian (hess) does not. Clearly, there's a difference between the two approaches, but I'm not sure what it is.

Questions:

  1. Why does the first approach return a Hessian of all zeros?
  2. Is either approach correct for computing the second derivatives in this context?
  3. If not, what is the correct method to compute the Hessian?

Any insights or explanations would be greatly appreciated. Thank you!

Upvotes: 1

Views: 46

Answers (1)

Karl
Karl

Reputation: 5373

  1. Why does the first approach return a Hessian of all zeros?

This has to do with the softmax operation. When you compute torch.autograd.grad(grad, a, torch.ones_like(a)), you are essentially computing torch.autograd.grad(grad.sum(), a). If you compute grad.sum(), you will find the output is always zero (or near zero due to numeric issues). This is because the sum of gradients through a softmax is always zero.

  1. Is either approach correct for computing the second derivatives in this context?

The second approach is correct, but likely slow due to looping at the python level. It works because you are backproping from individual elements of grad rather than the sum of grad, so you don't have the zero issue.

  1. If not, what is the correct method to compute the Hessian?

It's probably more efficient to use pytorch methods to compute the full hessian and take the diagonal. This is what your second method does, only we move the operations to a lower level.

def my_func(a):
    b, *_ = a.sort(descending=True)
    c = (b.unsqueeze(0) - a.unsqueeze(1)).abs().neg()
    d = c.softmax(0).matmul(torch.arange(c.size(0), dtype=c.dtype))
    e = torch.randint(0, 3, (10,), dtype=float)
    t = torch.sum(e * torch.log2(d + 1))
    return t

torch.manual_seed(0)
a = torch.randint(0, 10, (10,), dtype=float, requires_grad=True)
hessian = torch.autograd.functional.hessian(my_func, a)
hess_diag = torch.diag(hessian)

Upvotes: 0

Related Questions