Kolmogorov-Arnold Energy Models: Fast and Interpretable Generative Modeling
This work addresses efficiency and interpretability issues in generative modeling for researchers and practitioners, though it appears incremental as it builds on existing energy-based model frameworks.
The paper tackles the challenges of slow sampling and poor interpretability in energy-based generative models by proposing the Kolmogorov-Arnold Energy Model (KAEM), which enables fast and exact inference via the inverse transform method and improves sampling efficiency with importance sampling or population-based Langevin Monte Carlo.
Learning an energy-based model (EBM) in the latent space of a top-down generative model offers a powerful framework for generation across many data modalities. However, it remains unclear how its interpretability can be used to guide model design, improve generative quality, and reduce training time. Moreover, the reliance on Langevin Monte Carlo (LMC) sampling presents challenges in efficiency and sampling multimodal latent distributions. We propose a novel adaptation of the Kolmogorov-Arnold representation theorem for generative modeling and introduce the Kolmogorov-Arnold Energy Model (KAEM) to take advantage of structural and inductive biases. By constraining the prior to univariate relationships, KAEM enables fast and exact inference via the inverse transform method. With the low dimensionality of the latent space and suitable inductive biases encoded, we demonstrate that importance sampling (IS) becomes a viable, unbiased, and highly efficient posterior sampler. For domains where IS fails, we introduce a strategy based on population-based LMC, decomposing the posterior into a sequence of annealed distributions to improve LMC mixing. KAEM balances common generative modeling trade-offs, offering fast inference, interpretability, and stable training, while being naturally suited to Zettascale Computing hardware.