Implementing MinSR algorithm

[chrisrothUT]

This paper was very recently posted on Arxiv which seems like a major breakthrough for stochastic reconfiguration (SR) for NQS. The main idea is that the stochastic estimate of the Quantum Fisher Information can be described by a matrix [n_samples, n_samples] instead of [n_params, n_params]. This allows you to do exact SR with matrix inversion so long as you can fit [n_samples, n_samples] into memory.

I'm happy to lead the charge in implementing this in NetKet. Right now I foresee a technical problem in computing (O O^+) where we cannot fit O into memory. Ideally we'd like to write something like:

def sample_fisher():
  sample_fisher = 0
  for layer in range(final_layer, first_layer):
    grad_psi_hid = backpropogate(grad_psi_hid)
    grad_psi_W = grad(grad_psi_hid)
    sample_fisher = sample_fisher+ matmul(grad_psi_W,g grad_psi_W.conj.T)

Does anyone have any idea how to do this kind of thing in Jax?

[PhilipVinc]

I already implemented this more than a year ago...

You want to implement the neural tangent kernel (what they call T) and then you can pull off a few tricks.

I'll try to find my code...

[PhilipVinc]

I think that the library you linked is for the infinite width limit (where you don't use a jacobian because... well, it would be infinitely large, but can compute its action under some approximations).

The NTK is literally the gram transpose of the QGT.
You could implement a Jacobian-based version by copy-pasting the QGTJacobianXXX code and simply transposing the two jacobian vector multiplications

[chrisrothUT]

Yeah the NTK function is very straightforward to implement. But I want to make sure the version plays nicely with the GPUs.

As for the algorithm I'll just put it in nk.experimental for the moment and we can think about how it fits into the main library. The issue is that it's preconditioning over the sample dimension so it doesn't fit cleanly into the current interface.

[PhilipVinc]

Ah, here I found it. I had hidden it away in a branch named furbizia about an year ago
https://github.com/netket/netket/tree/pv/furbizia

It's similar to what you started prototyping already.
I think the chunking of the NTK should be somewhat equivalent to what NTKOnTheFly with default chunking.

[chrisrothUT]

Thanks for the link. I actually don't see an NTKOnTheFly on that branch but it looks like this is an on the fly implementation where the jacobian or NTK is never fully instantiated unless you call _to_dense. Or am I mistaken? This is probably the complementary approach that I haven't coded so it's useful.

[PhilipVinc]

No I never implemented it.
Sorry, I was probably unclear.
What I meant is that NTKOnTheFly with chunking would behave as your double loop of computing one sub-part of jacobian.
Or the opposite. I think that your implementation is somewhat equivalent to what NTKOnTheFly+chunking would do.