Kolmogorov-Arnold Energy Models: Fast, Interpretable Generative Modeling

TL;DR

The paper introduces KAEM, a new generative model that combines efficiency, interpretability, and high-quality sampling by leveraging a novel latent structure and adaptive inference strategies.

Contribution

KAEM applies a univariate latent structure based on Kolmogorov-Arnold theorem, enabling exact inference and improved sampling efficiency in generative modeling.

Findings

KAEM achieves the best FID scores among latent-prior models on SVHN and CIFAR10.

KAEM allows sampling in a single forward pass.

KAEM provides an interpretable prior built from 1D densities.

Abstract

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 new opportunities for latent-space interpretability. Based on a novel adaptation of the Kolmogorov-Arnold Representation Theorem, KAEM imposes a univariate latent structure on the prior, enabling exact inference via the inverse transform method. With a low-dimensional latent space and appropriate inductive biases, importance sampling becomes a tractable, unbiased, and efficient posterior inference method. For settings where this fails, we propose a population-based strategy that decomposes the posterior into a sequence of annealed distributions, a new remedy for poor mixing in Energy-based Models. We compare KAEM against VAEs and the neural latent EBM architecture. KAEM attains the best Fr\'echet Inception Distance among latent-prior models on SVHN and CIFAR10, while sampling in a single forward pass and exposing an interpretable prior built from 1D densities.