Skip to content

* add gumbel sampling#94

Open
mattlST wants to merge 3 commits into
markusschmitt:masterfrom
mattlST:dev/stm58133/2026-03-24-gumbel-sampling
Open

* add gumbel sampling#94
mattlST wants to merge 3 commits into
markusschmitt:masterfrom
mattlST:dev/stm58133/2026-03-24-gumbel-sampling

Conversation

@mattlST
Copy link
Copy Markdown

@mattlST mattlST commented Mar 24, 2026

"Manually rebased" on master
Let's see if the tests work now, locally (cpu) all tests passed
Old pull request:

Hi Markus,

mostly add a sample technique without repetition, aka gumbel sampling.
The implementation is based on the article: https://arxiv.org/pdf/2408.07625v1
The algorithm is implemented for general local Hilbert space dimension, I mostly used it for peaked bosonic systems (Bose Hubbard).

The core addition is the gumbel_wrapper in the util subfolder, some remarks about the wrapper:

It require autoregressive networks which gives conditional normalized probabilities. Further, it requires a network with a step-wise calculation method as implemented in RWKV (written by J. Rigo, which is also added in the net subfolder)
It only runs on one device (one GPU) as in each step, the algorithm needs all conditional probabilities to estimate unique samples.
The number of samples must be smaller than the total Hilbert space dimension, as only unique samples should be returned. Hence, the wrapper throws an error if one chooses a larger sample size.
Patching is untested, but I think it should be fine ....
I removed the feature of conserved quantities (particle number conservation in bosonic systems) to be more compatible
Some minor changes are in sampler and nqs class to account for the gumbel sampling, i.e., nqs has an property "is_gumbel" and the sampler class has an additional if condition in the sampling method
Gumbel with symmetrization does not work! I think it is also not easy/possible to integrate symmetrized wave-functions in this way.
I wrote a test file "gumbel_test.py" inspired by the sample test file "sampling_test.py". It only takes autoregressive, because they are required for gumbel sampling. The tests check sample size, unique samples,, sampling frequencies and reproducibility for different local Hilbert space dims.

I hope this can be useful to some folks! Most likely, I forgot something important, so if there are any questions, problems, let me know!

Best,
Mathias

mattl added 3 commits March 24, 2026 22:11
- gumbel wrapper takes right now only RWKV, needs an autoregressive network with stepwise sampling implementation
- RWKV implementation based on Jonas Rigo
- some basis tests for gumbel sampling
- adapted sampler and vqs to support gumbel sampling
 - a comment about implementing setting MPI rank to 0 in gumbel_test (but probably deprecated soon)
    to identify which lDim is being tested
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant