Trying to understand the Hutchinson diagonal Hessian approximation

Question

I am reading about his paper [1] and I have an implementation taken from here. At some point of the code the diagonal of the Hessian matrix is approximated by a function set_hessian you can find below. In the end of set_hessian(), it is mentioned that # approximate the expected values of z*(H@z). However, when I print p.hess I get

tensor([[[[ 2.3836e+01,  1.4929e+01,  4.1799e+00],
          [-1.6726e+01,  6.3954e+00, -5.1418e+00],
          [ 2.2580e+01, -1.1916e+01, -2.5049e+00]],
         [[-1.8261e+01,  8.7626e+00,  1.8244e+00],
          [-1.0819e+01, -2.9184e-01,  1.1601e+01],
          [-1.6267e+01,  5.6232e+00,  3.4282e+00]],
         ....
         [[-3.1088e+01,  4.3013e+01, -4.2021e+01],
          [ 1.5338e+01, -2.9806e+01, -3.0049e+01],
          [-9.8979e+00, -2.2835e+00, -6.0549e+00]]]], device='cuda:0')

How p.hess is considered a diagonal approximation of the Hessian? The reason I am trying to understand this structure is because I want get the smallest eigenvalue, the inverse of the diagonal matrix, and the product between the Hessian and the gradient which is a vector. We know that the smallest eigenvalue of a diagonal matrix is the smallest element of the diagonal, while the inverse of a diagonal matrix can be computed by inverting the elements of the diagonal. Could you please someone cast some light on the structure of p.hess?

    @torch.no_grad()
    def set_hessian(self):
        """
        Computes the Hutchinson approximation of the hessian trace and accumulates it for each trainable parameter.
        """

        params = []
        for p in filter(lambda p: p.grad is not None, self.get_params()):
            if self.state[p]["hessian step"] % self.update_each == 0:  # compute the trace only each `update_each` step
                params.append(p)
            self.state[p]["hessian step"] += 1

        if len(params) == 0:
            return

        if self.generator.device != params[0].device:  # hackish way of casting the generator to the right device
            self.generator = torch.Generator(params[0].device).manual_seed(2147483647)

        grads = [p.grad for p in params]

        for i in range(self.n_samples):
            zs = [torch.randint(0, 2, p.size(), generator=self.generator, device=p.device,
                                dtype=torch.float32) * 2.0 - 1.0 for p in params]  # Rademacher distribution {-1.0, 1.0}
            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, only_inputs=True,
                                       retain_graph=i < self.n_samples - 1)
            for h_z, z, p in zip(h_zs, zs, params):
                p.hess += h_z * z / self.n_samples  # approximate the expected values of z*(H@z)

[1] ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning

Answer 1

Hutchinson gives to you a approximation of Trace of hessian matrix, not a diagonal of Hessian matrix.

Answer 2

One of the subquestions here was "the product between the Hessian and [...] a vector". This can be computed exactly and efficiently, using the "Hessian-vector product" method. No need for Hutchinson or approximations.