Provable Maximum Entropy Manifold Exploration via Diffusion Models

Riccardo De Santi, Marin Vlastelica, Ya-Ping Hsieh, Zebang Shen, Niao He, Andreas Krause

TL;DR: A scalable algorithmic principle for provable manifold exploration via diffusion model fine-tuning.

Abstract: Exploration is critical for solving real-world decision-making problems such as scientific discovery, where the objective is to generate truly novel designs rather than mimic existing data distributions. In this work, we address the challenge of leveraging the representational power of generative models for exploration without relying on explicit uncertainty quantification. We introduce a novel framework that casts exploration as entropy maximization over the approximate data manifold implicitly defined by a pre-trained diffusion model. Then, we present a novel principle for exploration based on density estimation, a problem well-known to be challenging in practice. To overcome this issue and render this method truly scalable, we leverage a fundamental connection between the entropy of the density induced by a diffusion model and its score function. Building on this, we develop an algorithm based on mirror descent that solves the exploration problem as sequential fine-tuning of a pre-trained diffusion model. We prove its convergence to the optimal exploratory diffusion model under realistic assumptions by leveraging recent understanding of mirror flows. Finally, we empirically evaluate our approach on both synthetic and high-dimensional text-to-image diffusion, demonstrating promising results.

Lay Summary:

Modern generative models, like diffusion models, excel at reproducing patterns from their training data, but they struggle to venture into truly novel regions of valid designs. For tasks such as discovering new molecules or materials, it’s crucial not just to mimic known examples but to systematically explore under-sampled, potentially groundbreaking areas of the design space.


We recast “exploration” as maximizing how uniformly samples cover the space (“manifold”) of valid designs learned by a pre-trained diffusion model. By iteratively adapting a pre-trained diffusion model in a specific manner, our method pushes the model toward one that spreads its samples across the space of valid structures.


This work outlines a practical, theory-backed approach for encouraging diffusion models to explore more broadly across their learned design spaces. While additional empirical validation is needed to assess its impact on real-world discovery tasks, this method could, for example, be applied to molecule or material design—potentially uncovering under-explored but valid regions that standard models might miss

Primary Area: General Machine Learning

Keywords: diffusion models, exploration, fine-tuning, maximum state entropy reinforcement learning

Submission Number: 11753