Re:
We sincerely appreciate your detailed and thoughtful feedback. In this response, we have addressed your primary concerns regarding our algorithms. We are happy to answer more if further clarification is required. Regarding your suggestions for the presentation, we have incorporated many of them, and the change is in red. (For some of the suggestions, like replacing sections in the main text and appendix, we plan to do; but at this moment, to avoid confusion for other reviewers, we didn't implement it)
**Weakness: My main concern with the paper is the following: The terminal reward $v_0$ for learning the drift $q$ is approximated with a neural network using a regression objective. Hence, the learned model could be arbitrarily bad in regions with no data which is also the cause for the overoptimization problem. It thus seems like the problem has just shifted to the initial value function estimation and thus not properly addressed.**
While we acknowledge your observation, we think it won't be a primary concern for the following reasons. In summary, we do not encounter the distribution shift issue in this context because error guarantees are required solely within the original data distribution. Specifically, the reasoning is as follows:
First, recall from Corollary 2, the optimal initial distribution is
$$
\exp(v^{ \star}\_0( \cdot)/ \alpha) \nu\_{ini}(\cdot)/C.
$$
Then, leveraging the relation from Lemma 1:
\begin{align*}
\exp\left(\frac{v^{\star}\_{0}(x)}{\alpha} \right)= \mathrm{E}\_{\mathbb{P}^{data}}\left [\exp \left (\frac{r(x_T)}{\alpha } \right)|x_0=x \right],
\end{align*}
our algorithm approximates the value function $v^{\star}$ in the optimal initial distribution with supervised learning:
\begin{align*}
\hat g = \mathrm{argmin}\_{g \in \mathcal{G}} \sum\_{i=1}^n \left ( \exp \left (\frac{r(x^i_T)}{\alpha } \right )- g(x^i_0) \right )^2, x^i_0 \sim \nu\_{ini}, x^i_T\sim \mathbb{P}^{data}(x_0),
\end{align*}
where $\mathcal{G}$ denotes a function class, such as neural networks. Under the well-specification assumption (a standard assumption as used in Wainright 2019), we achieve the following mean squared error (MSE) guarantee:
\begin{align*}
\mathrm{E}\_{x_0\sim \nu\_{ini}}\left [ \left (\hat g(x_0)- g^{\star}(x_0) \right)^2 \right ] = O\left (\frac{Cap(\mathcal{G})}{n} \right)
\end{align*}
where $Cap(\mathcal{G})$ represents the capacity of the function class $\mathcal{G}$ and $g^{\star}$ is $\exp(v_0(\cdot)/\alpha)$.
This result implies that we obtain the MSE guarantee within the original initial distribution $\nu\_{ini}$. Although this guarantee does not extend beyond the support of $\nu\_{ini}$, it is not an issue because we know that the support of the optimal initial distribution is restricted to $\nu\_{ini}$ (recalling Corollary 2), and in the second step of approximating the optimal initial distribution, in our algorithm, we consistently do sampling on the support of the original initial distribution $\nu\_{ini}$.
**Q. While inference time is already a bottleneck for diffusion models, the proposed approach requires simulating two SDEs.The approach requires solving two stochastic optimal control problems and an additional regression problem for obtaining the initial value function.**
While we ackowldge your opinion, we have an explanation that this won't be a concern in Appendix G.1. In particualr, the following paragraph demonstrates that learning the optimal initial distribution is not a practical bottleneck for our algorithm.
*Computational Cost: Cost of Learning Initial Distributions: Our overhead in learning the initial distribution is low. When learning a second diffusion chain, our goal is to learn $\exp(v^{\star}_0(x))\nu\_{\mathrm{ini}}(x)$. This distribution is much simpler and smoother compared to the target distribution $\exp(r(x)/\alpha)p\_{\mathrm{data}}(x)$ in the first diffusion chain. Therefore, we require fewer epochs to learn this distribution. For instance, in image generation, we use only $200$ reward queries to learn initial distributions while the main diffusion chain takes $12000$ queries for finetuning. The wall time of learning the initial distribution (i.e., the second diffusion chain) is $30-40$ minutes in image experiments, while the wall time of learning the second chain is roughly $1800$ minutes.*
**Weakness. The method requires a) gradient checkpointing and b) low-rank modules to train**
This technique is commonly employed when fine-tuning large models. Representative works for fine-tuning diffusion models (Black et al., 2023; Prabhudesai et al., 2023; Clark et al., 2023) also utilize these techniques, indicating that this is not a unique aspect of our work. Hence, we believe this is not a particular weakness of our methods. Additionally, we clarify that for smaller models, such as those used in experiments with biological diffusion models, employing this technique is not necessary.