Sampling Unique states with autoregressive models

[jkocklaeuner]

Hi,
first of all thanks for the amazing package, I really enjoy working with it!
I'm currently experimenting with autoregressive models and stumbled across a paper which uses a different sampling algorithm which seems more efficient to me, as it simply samples the unique states and the number of occurences given a sample size.
Since I often have a lot of redundant samples this would be really advantageous for me in the following steps, e.g. the local energy and the gradients are only computed once for each unique state and scaled afterwards. I would really like to implement a similar algorithm, since I ran into some out-of-memory errors for larger systems.
While the implementation of the sampling procedure alone seems manageable to me, this would probably break the compatibility with all other functionalities provided by Netket, e.g. the calculation of the gradients, QGT etc., since they handle sampled states differently.
So what I'm basically asking is:

1. Do you plan to integrate a similar sampling procedure to your toolbox ?
2. Would it lead to an speedup/memory saving or is a similar "contraction" to unique states already happening under the hood and I missed it?
3. Is there any obvious reason why this wouldn't work in Jax/Netket that I didn't get?

I just feel like this is a much more efficient way to exploit the normalization of the autoregressive models compared to the current ARDirectSampler and I would be very thankful if I would get your opinion on this :)
Thanks in advance!
Jannis

[gcarleo]

My understanding is that the approach you mention (which is very nice, I agree!) will generate occurencies of given samples, thus it will require using weighted averages of all observables (including the QGT, energy etc). This is a major difference wrt to the current implementation of these objects in netket. It's not impossible to do, but it is certainly a lot of work, and currently we don't have anything planned along those lines. The main inconvenience I see is that this approach will pay off only if the sampled nqs has a few dominating states (say, for chemistry in the HF basis and for weakly correlated molecules)... But of course if you are interested in implementing it, we can try to help you out. @wdphy16 might also comment more on this algorithm and its possible implementation

[PhilipVinc]

Hi @jkocklaeuner, thanks for the comment.

It's actually not so much work to implement weighted averages of observables, and I would be in favour of implementing support for it internally. We already do support it in one QGT implementation, as we need it to perform full-space calculations (I think its supported by QGTJacobianPyTree), and supporting it everywhere should not be particularly difficult.

From an API point of view, this approach would probably result in implementing a different kind of Variational State, maybe a MCWeightedState or something similar, that only works with particular kinds of samplers (in this case, a specific implementation of the Autoregressive sampler that returns weights).

About your questions:

1. It's not planned and we won't dedicate manpower to this, but as I said above we would welcome the internal changes necessary to make this work.
2. What Giuseppe says. It's very problem-specific.
3. No. It can be implemented and it would work.