Hello
I am trying to use Hutchinson’'s trick to compute an estimator of the trace of a hessian matrix.
To make this estimator precise, I need to run the same backward pass several times with different stochastic vectors v that are used for vector-jacobian product.
The code using a for loop would look like this:
def hutch_loop(loss, params, niter, v):
g = autograd.grad(loss, params, create_graph=True)
out = 0.0
for i in range(niter):
vsplit = [torch.randn_like(p) for p in params]
Hv = autograd.grad(g, params, vsplit, retain_graph=True)
out += group_product(Hv, vsplit)
return out/niter
def group_product(xs, ys):
return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)])
Ideally, I would like to be able to do this operation in batch mode, alleviating the need to compute as many backward passes than the number of samples i need to draw. Is there any way to do this efficiently?
Thanks!