sharding,chunking for netket_fidelity

[jangwon00]

hi = nk.hilbert.Spin(0.5, 4)
sampler = nk.sampler.MetropolisLocal(hilbert=hi, n_chains_per_rank=10)
model = nk.models.RBM(alpha=1, param_dtype=complex, use_visible_bias=False)
vs = nk.vqs.MCState(sampler=sampler, model=model, n_samples=10,chunk_size=5)
vs_target = nk.vqs.MCState(sampler=sampler, model=model, n_samples=10,chunk_size=5)

optimizer = nk.optimizer.Adam()
driver = nkf.driver.InfidelityOptimizer(vs_target, optimizer, variational_state=vs)

log = nk.logging.RuntimeLog()
driver.run(300, out=log)

##############################################################################
thank you for these remarkable contribution for NQS.

I'm trying to use fidelity or pTVMC for larger models, but i have noticed that i need to utilize either sharding or chunking. But as far as i know, nkf doesn't support it. then, it would be tricky to customize it? or do you by any chance know some simple way to apply chunking or sharding for nkf?

[PhilipVinc]

Hi @jangwon00 .

Indeed, sharding and chunking are not implemented.
Implementing chunking should be easy, by specifying the argument chunk size to nk.jax.expect used in the internal implementation of the infidelity operator (see https://netket.readthedocs.io/en/latest/api/_generated/jax/netket.jax.expect.html#netket.jax.expect ).

However, if I recall correctly, this won't work with sharding and will give several uncomprehensible jax errors.

--

Let me point out however a very, very important point.
In a recent paper arXiv: 2410.10720 by @lgravina1997 we learnt that the gradient estimator used by netket_fidelity, that we had discussed in the Quantum paper by Sinibaldi is largely suboptimal.

See for example Figure 2 from that paper,

where we are showing that the gradient estimator of netket_fidelity (cv-smc) has a SNR ratio 1-2 orders of magnitude lower than the smc/cmc estimator, which means that cv-smc require a factor of 100-10'000 more samples than other estimators.

Essentially this is telling you that you should compute the gradient without CV, but compute the infidelity itself with CV.
This is not implemented in netket fidelity.

This is even more striking if you look at the stability of natural gradient descent (figure 3),

which shows that cv-smc is much less stable (black in the plot above means that the simulations diverge 100% of the time) than the cmc estimator of the gradient.

--

In short: be careful if you us netket fidelity, as it is very hard to use in a stable way.
We are cleaning up the accompanying code for that paper. We are almost done, I think that @lgravina1997 will push the infidelity optimisation code in the next few days. You can also reach out to him for help.