Provable Maximum Entropy Manifold Exploration via Diffusion Models

We thank the Reviewer for appreciating our work and asking interesting questions. In the following, we address several important points mentioned within the review that can hopefully let the Reviewer appreciate more the content of this work.

Asymptotic convergence guarantees We thank the Reviewer for raising these concerns and would like to provide a clarification: We can in fact establish convergence guarantees under three different assumptions, ordered by increasing generality:

• Perfect fine-tuning: If fine-tuning is exact, our algorithm terminates in a single step (see Section 5).
• Unbiased noise oracle: When the noise oracle is unbiased (i.e., $b_k = 0$ in Eq. (16)), standard mirror descent analysis yields a convergence rate of $\mathcal{O}(k^{-1/2})$ (see [3]).
• General bias term: If the noise oracle in Eq. (16) includes a general bias term, then under the arbitrary slow decay assumption in Eq. (18), a polynomial-time guarantee is no longer feasible [3]. In this case, stochastic approximation techniques are required, and the best achievable rate is $\tilde{\mathcal{O}}((\log\log k)^{-1})$, which follows from our proof. We chose to present an asymptotic result rather than explicitly stating this rate, as the difference is negligible and in this case it is conventional to present the asymptotic result [3,4].

Among these, the third setting is the most practical, which is why we focused on it in Section 7. We will incorporate these clarifications into the revision.

Empirical results clarification Within the Experimental Evaluation section, we evaluate qualitatively (i.e., visually) and quantitatively (i.e., via the metrics within Fig. 2-d and Table 1) the performance of models $\pi_k$ obtained by fine-tuning a pre-trained model $\pi^{pre}$ for $k$ iterations of S-MEME. In the case of text-to-image experiments, Fig. 3 shows on the top row a set of images obtained by sampling the pre-trained model $\pi^{pre}$, which in this case corresponds to Stable Diffusion 1.5 [1] trained on the LAION-5B dataset [2], while the bottom row shows images sampled via $\pi_3$, which is the diffusion model obtained after $3$ iterations of our algorithm. Qualitatively, one can visually notice that the images within the top row appear mostly gray and similar to each other, while the images from the bottom row show more diversity both among them and compared with the ones above, while preserving semantic meaning. Quantitatively, ideally we would want to estimate the entropy of the marginal density induced by the fine-tuned model, but this is hard in practice, as explained within Sec. 8.2. As a consequence, within Table 1 we show several proxy metrics to give numerical insights about the performances of the fine-tuned models. Within this table, the label 'S-MEME 1' refers to the model obtained after one single iteration of S-MEME, 'S-MEME 2' after two iterations, and 'S-MEME 3' after three iterations. Crucially, we show that the FID and cross-entropy scores, which aim to capture the degree of diversity of the marginal density induced by a model $\pi_k$ from the pre-trained model, increases over the iterations of S-MEME, meanwhile the CLIP score, which assesses the naturalness, or image quality, of the generated samples, is kept high across increasing iterations of S-MEME. Clearly, the trade-off between surprise maximization (i.e. diversity from pre-trained model) and naturalness due to regularization with the pre-trained model, can be chosen arbitrarily by tuning the $\{\alpha_k\}$ parameters of the algorithm, which manage this trade-off as shown In Eq. (9). The experimental results in Sec. 8 aim to evaluate a specific parameter choice to show practical relevance of the proposed scheme.

We thank the Reviewer for these question sand hope that the explanations given can help the Reviewer better appreciate our work. We will update a revised version of the paper with better clarification of these aspects.

References

[1] Rombach et al., High-resolution image synthesis with latent diffusion models. CVPR 2022.

[2] Schuhmann et al., An open large-scale dataset for training next generation image-text models. NeurIPS 2022.

[3] Karimi et al., Sinkhorn flow as mirror flow: A continuous-time framework for generalizing the sinkhorn algorithm. AISTATS 2024.

[4] Borkar et al., The ODE method for convergence of stochastic approximation and reinforcement learning. SIAM Journal on Control and Optimization 2000.

We thank the Reviewer for reading our work. In the following, we address several fundamental points and questions mentioned within the review that can hopefully let the Reviewer appreciate more the content of this work.

Motivation of the work In the following, we aim to make clear the main motivation of this work. Pre-trained generative models can generate plausible objects of a certain data type, e.g. valid images or molecules. Nonetheless, the probability of sampling rare, novel, objects will be very low. This is because the model has been trained on a dataset of existing objects to approximate the data distribution. Therefore rare objects in the data will be rare to sample. In this work, we adapt the generative model so that its induced distribution is not proportional to that of existing data, but is shifted towards lower-probability regions where novel designs can be sampled, while preserving plausibility (i.e., staying within the approximate data manifold). To the best of our knowledge, this is a fundamental (open) problem for discovery via generative models.

Question 1 Yes, it is. In fact, if $\Omega_{pre}$ would be a known set representable in a computer one could simply define a uniform distribution on such a set and then sample from it. But what if $\Omega_{pre}$ is an unknown and possibly very complex set, e.g., space of valid molecules or images, that cannot be represented in a computer if not implicitly via a generative model? How can we explore such a set? This work tackles exactly this problem, which is both fundamental and non-trivial. We thank the Reviewer for asking this question as we believe it is of central importance, and we will make sure to further clarify this point in a revised version of the work.

Question 2 Equation (7) describes an optimization problem over the set of densities supported over a set $\Omega_{pre}$ represented only implicitly by a pre-trained diffusion model. This fully captures the manifold exploration problem. On the other hand, Eq. (9) is an unconstrained KL-regularized optimization problem of a reward function obtained by linearizing the entropy, as explained in details within Sec. 4. As explained in Sec. 6, a concave function like entropy can be maximized via a Mirror Descent scheme [1], which relies on a sequence of simpler optimization problems. In this context, Eq. (7) represents the concave optimization problem, while Eq. (9) captures each simpler optimization problem. Crucially, via this viewpoint, we can use scalable methods to solve each smaller problem to solve the complex exploration problem. We will make sure to make this more explicit within an updated version of the work.

Question 3 Within this work, $s^{pre}$ is the score of a pre-trained diffusion model that one wishes to fine-tune by optimizing a regularized measure of surprise (Eq. (9)).This quantity has to be estimated with respect to the previous model, which at the first iteration of S-MEME corresponds to the pre-trained model with score $s^{pre}$. In practice, as mentioned within Sec. 4.3, the score of the fine-tuned model $s^\pi$ can be initialized as $s^\pi = s^{pre}$. After this step, there is no difference in using $s^\pi$ or $s^{pre}$ to estimate the induced marginal density as these score networks are equal. We will make sure to make this more clear in order to prevent any doubt.

Question 4 Reinforcement learning spans several problems, including pure-exploration ones, e.g., [2,3], where the goal is to explore a certain space, typically via maximization of an intrinsic reward, which in the case of this work would correspond to surpise (i.e., the entropy first variation) in Eq. (9) and (10). Pure exploration problems often are relevant on their own. In this case, surprise maximization is clearly relevant for discovery applications where it can make it possible to sample surprising designs, or, as mentioned by Reviewer WBXm, to de-bias a given pre-trained model by inducing a more balanced distribution than the data distribution. Nonetheless, ideas from pure exploration can be used for exploration schemes in exploration-exploitation settings where there is an unknown quantity (e.g., a reward function) to be learned and optimized. In this case, the exploration principle introduced within Eq. (9), which makes it possible to scalably maximize a measure of surprise, could be used as a way to regularize an exploratory scheme for reward-learning in black-box optimization settings, e.g., [4].

References

[1] Nemirovski et al., Problem complexity and method efficiency in optimization, 1983.

[2] Hazan et al., Provably efficient maximum entropy exploration. ICML 2019.

[3] Mutny et al., Active exploration via experiment design in markov chains. AISTATS 2023.

[4] Uehara et al., Feedback Efficient Online Fine-Tuning of Diffusion Models. ICML 2024.

We thank the Reviewer for the interesting questions. In the following, we address several points that can hopefully let the Reviewer appreciate more this paper, and which will be included in a revised version.

Uncertainty quantification We agree, our method does not employ explicit uncertainty quantification, but arguably relies on a form of implicit uncertainty quantification, although not the classic Bayesian one.

Toy experiment

• The y-range of plot 2.a should be the same as plots 2.b and 2.c.
• Since Fig. 2.b and 2.c are histogram plots, for a smooth cmap, points with low probability have a similar color to points with zero probability. We will try to overcome this issue by testing non-smooth cmaps.
• The original density is in Fig. 2.a, namely two uniform distributions.
• Pre-training was performed by standard denoising score-matching and uniform samples (namely 10K) from the two distributions in Fig. 2.a.

Text-to-Image experiment

• The references to Fig. 9 in Sec. 8 should be references to Fig. 3.
• We agree with the Reviewer that the method presented in this work is relevant for molecular spaces. Nonetheless, applications in scientific discovery go beyond the scope of this paper, which captures mathematically a novel and relevant problem, proposes fundamental algorithmic advances, and provides a novel type of theoretical analysis for diffusion fine-tuning. We believe that exploration in the visual domain is relevant for multiple applications, and particularly amenable for evaluation as it does not require specific background knowledge.
• Since the method proposed directly maximizes a gradient of entropy, one can show that for $\alpha \geq 1$ (in Eq. 9), the entropy value increases monotonically. One could attempt at increasing entropy by reducing guidance strength, but this would result in being closer to the unconditional distribution thus changing the support significantly. On the contrary, our algorithm (provably) explores within the support of the conditional distribution via entropy maximization, which is substantially different.

Computational complexity and hyperparameter selection Under exact fine-tuning one iteration is sufficient (see Sec. 5). In practice, we show in Sec. 8 that very few iterations (e.g., $K=3$) can lead to useful fine-tuned models. Therefore, the computational cost aligns with typical fine-tuning, e.g. [3]. The parameter $N$ is significantly problem dependent, and as in the case of Adjoint Matching [3], it requires experimental tuning.

Question 1 It is common to express the reverse process of a diffusion process via a SDE parametrized by a parameter $\eta \in [0,1]$ (e.g., in [1]). The processes obtainable for any valid $\eta$ induce the same marginal distributions.

Question 2

• The truncation in Prop. 1 is merely a technical aspect to build a theoretical formulation that matches the real sampling process, which always gives finite values. The bounding interval can be any as long as it is finite. Therefore it can fully capture any real scenario.
• Diffusion notation uses $0$ for the time-step corresponding to data (e.g., [1]), while fine-tuning works denote by $0$ the noise level (see, e.g., [2]). This is formally not fully correct, since to compute closed-form the noise distribution one has to take the limit to infinity. A recent rigorous solution (see, e.g., [1]) is the one used in our paper. Nonetheless, we agree with the Reviewer that the sign flip might create confusion and have already added a drawing that clarifies this in a new version.

Question 3 Sec. 4 shows that given exact fine-tuning, the algorithm convergences in $1$ iteration with $\alpha_1 = 1$. It should be set to lower values to prioritize the surprise term in Eq. 9, and to higher values to prioritize regularization, as discussed in Sec. 4.4. When considering apx. fine-tuning oracles, Theorem 7.1 gives necessary conditions on $\alpha_k = 1/\gamma_k$ for the guarantees to hold. In practice, $\alpha_k$ should be experimentally tuned for the specific application.

Question 4 Ass. 7.1 could be difficult to verify in some cases. However, it can be relaxed as $supp(p_T^{\pi_k}) \subset \tilde{\Omega} \text{ for all } k,$ and $supp(p_j^{\pi_k}) = \tilde{\Omega}$ for some $j.$ Using the same analysis, we can show that the algorithm can solve the exploration problem on $\tilde{\Omega}$. In particular, if $\tilde{\Omega}$ approximates the true support, as in Fig. 2, our analysis guarantees exploration of the approximate support.

References

[1] Zhao et al., Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning.

[2] Uehara et al., Feedback Efficient Online Fine-Tuning of Diffusion Models. ICML 2024.

[3] Domingo-Enrich et al., Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. ICLR 2025.

We thank the Reviewer for recognizing our work as very well written, theoretically solid, and practical. In the following, we address several points and questions mentioned within the review.

Inception Score (IS) and Vendi Score (VS).

IS: To the best of our understanding, the Fréchet Inception Distance (FID) has been introduced as a way to mitigate the inability of the IS score to recognize intra-class mode collapse [1], which is related to what we are aiming to measure in our setting. In particular, in our setup with images, the generated points would all be part of one class. As a consequence, IS might not be a better evaluation measure than FID.

VS: We agree with the Reviewer that this measure of diversity might be relevant in this context. Nonetheless, choosing a relevant kernel in practice might be non-trivial, and if the feature space is obtained via a non-linear map, which is typically the case, it is unclear how well this measure aligns with entropy. In any case, we are grateful for this suggestion and will certainly explore it in this context.

Relevance to de-bias pre-trained generative models Although not strictly related to the motivation presented in this work, we agree that the presented method is very relevant to de-bias generative models, and further research could focus on its potential impact on this important problem.

Question 1. The presented method can run using any state-of-the-art (linear) diffusion model fine-tuning scheme, such as Adjoint Matching [2] or alternative methods, e.g., [3]. To the best of our knowledge, these methods are very scalable and they have been shown to successfully fine-tune large-scale pre-trained models for images [2,3], molecules [3], and proteins [3], among others. Crucially, due to Eq. (12) within Sec. 4.3, our algorithm becomes effectively as scalable as the state-of-the-art methods mentioned above, which have already proved successful in relevant real-world applications with larger pre-trained models [2,3].

Question 2. We can establish convergence guarantees under three different assumptions, ordered by increasing generality:

• Perfect fine-tuning: If fine-tuning is exact, our algorithm terminates in a single step (see Section 5).

• Unbiased noise oracle: When the noise oracle is unbiased (i.e., $b_k = 0$ in (16)), standard mirror descent analysis yields a convergence rate of $O(k^{-1/2})$ (see [7]).

• General bias term: If the noise oracle in Eq. (16) includes a general bias term, then under the arbitrary slow decay assumption in Eq. (18), a polynomial-time guarantee is no longer feasible [6]. In this case, stochastic approximation techniques are required, and the best achievable rate is $\tilde{\mathcal{O}}((\log\log k)^{-1})$, which follows from our proof. We chose to present an asymptotic result rather than explicitly stating this rate, as the difference is negligible and in this case it is conventional to present the asymptotic result [6,7].

Among these, the third setting is the most practical, which is why we focused on it in Section 7. Notably, our results show that convergence does not require running $k \to \infty$ even with the most general assumptions (i.e., case 3 above).

Question 3. The presented version of S-MEME relies on optimal control methods as discussed in Sec. 4.2. Nonetheless, it is standard to refer to this recent problem of optimization over the space of admissible state distribution to maximize a known functional as Convex RL [4] or General Utilities RL [5]. This is the main reason why we use the RL notation. Moreover, to the best of our understanding, nothing prevents to replace the control oracle currently used with classic MDP planning methods used in RL. In conclusion, we agree with the Reviewer that the choice of introducing the problem with RL notation is mostly a notational convention.

Ultimately, we want to thank the Reviewer for reporting a typo, which we have corrected, and mentioning several relevant related references of which we were not aware.

References

[1] Lucic et al., Are GANs Created Equal? A Large-Scale Study. NeurIPS 2018.

[2] Domingo-Enrich et al., Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. ICLR 2025.

[3] Uehara et al., Feedback Efficient Online Fine-Tuning of Diffusion Models. ICML 2024.

[4] Mutti et al., Challenging common assumptions in convex reinforcement learning. NeurIPS 2022.

[5] Zhang et al., Variational policy gradient method for reinforcement learning with general utilities. NeurIPS 2020.

[6] Karimi et al., Sinkhorn flow as mirror flow: A continuous-time framework for generalizing the sinkhorn algorithm, AISTATS 2024.

[7] Borkar et al., The ODE method for convergence of stochastic approximation and reinforcement learning. SIAM Journal on Control and Optimization 2000.