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?