Kolmogorov-Arnold Energy Models: Fast, Interpretable Generative Modeling
For practitioners needing fast, interpretable generative models, KAEM offers a practical trade-off between efficiency and quality, though it is incremental over existing latent variable models.
KAEM introduces a generative model that combines fast, exact inference via inverse transform with competitive sample quality, achieving inference speeds orders of magnitude faster than diffusion models while matching VAE sample quality on standard vision datasets.
Generative models typically rely on either simple latent priors (e.g., Variational Autoencoders, VAEs), which are efficient but limited, or highly expressive iterative samplers (e.g., Diffusion and Energy-based Models), which are costly and opaque. We introduce the Kolmogorov-Arnold Energy Model (KAEM) to bridge this trade-off and provide a new avenue for latent-space interpretability. Based on a novel interpretation of the Kolmogorov-Arnold Representation Theorem, KAEM imposes a univariate latent structure that enables fast and exact inference via the inverse transform method. With a low-dimensional latent space and appropriate inductive biases, we show that importance sampling becomes a viable, unbiased, and highly efficient posterior inference method. For settings where importance sampling fails, we propose a population-based strategy that decomposes the posterior into a sequence of annealed distributions to improve mixing during sampling, a common pitfall in Energy-based Models. We present initial comparisons of KAEM against VAEs for standard vision datasets, demonstrating its potential for competitive sample quality, inference speed, and interpretability.