This feels extremely close in nature to being a generalisation of discrete distribution networks.

Paper: https://arxiv.org/abs/2401.00036 Project page: https://discrete-distribution-networks.github.io/

Given that these were published at ICML 2025, at which the author was a top reviewer as per their own website https://alexiglad.github.io/, I wonder how influenced they were to pursue this avenue of the back of it. They do have a related works section in appendix E, but it somehow seems to miss this. Which is odd, as the paper made a splash at the time at the conference and even made it close to the top of hacker news due to the novelty.

Of course, DDNs are a fundamentally new architecture, whereas this is more a generaliseable training strategy, but there is a very similar core.