Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference

We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.

---

The authors claim "minimal computational overhead" for the dual optimization of Lagrange multipliers. However, empirical comparisons with baselines are only qualitatively sketched, not based on running time. It would be better to see the benefits under the same computational cost.

We thank the reviewer for the helpful suggestion and have accordingly added quantitative results on running time. Specifically, we evaluate two aspects:

(1) Overhead from dual optimization: For the dual optimization step, we employed a line search algorithm based on a bracketing method. In the diffusion-based sampling experiments, for example, computing $\lambda^*$ from 50k importance weights took approximately 0.014 seconds per optimization (averaged over 10k trials on an RTX 3060 GPU). Performing 150 such updates during training, the total overhead from dual optimization amounts to only 2.1 seconds, which is negligible relative to the overall training time.

(2) Impact of the replay buffer in fine-tuning: We also compared TR-LV, which uses a buffer, against AM (which does not use a buffer). We observed a significant improvement in efficiency when using a buffer: – Without a buffer, fine-tuning experiments averaged 1 day and 23 minutes. – With a buffer, the same experiments required only 11 hours and 9 minutes on average. This clearly demonstrates that incorporating a replay buffer can substantially reduce computational cost in practice.

We thank the reviewer for this helpful suggestion and will include the corresponding runtime comparisons in the revised manuscript.

---

Hyper-parameter guidance on the effect of $\varepsilon$ and buffer size is missing, which could critically affect algorithmic performance.

We appreciate the reviewer’s observation regarding the importance of hyperparameter guidance for both the trust-region bound $\varepsilon$ and the buffer size.

Trust-region bound $\varepsilon$: We discuss the role of $\varepsilon$ in Appendix E.3, “Influence of Trust Region Bounds”, where we highlight its connection to the variance of the importance weights and the resulting effective sample size. This theoretical relationship often provides an intuitive basis for selecting $\varepsilon$. In many of our experiments, we used $\varepsilon = 0.1$, which corresponds approximately to an effective sample size of 0.8 and showed good empirical performance.

Buffer size: For transition path sampling, we adopted buffer sizes consistent with prior work. In the fine-tuning experiment, we evaluated several buffer sizes and found that a size of 100 offered a good balance between performance and computational efficiency, as larger buffer sizes did not yield significant improvements.

In contrast, for diffusion-based sampling, we observed that buffer size has a substantial impact on performance. To explore this further, we conducted an additional experiment on the GMM task using the TR-LV loss, varying both the buffer size and dimensionality. We report the control L2 error in the table below:

The results show that larger buffer sizes consistently lead to lower control error, particularly as dimensionality increases.

We thank the reviewer for this valuable suggestion. We will incorporate these results and practical guidelines into the revised manuscript.

---

Could the authors explain why their proposed methods TR-SOCM and TR-LV are scalable to high dimensions (e.g., in Figure 2)?

Both, the SOCM and LV loss do not require differentiation through simulations which makes them more memory efficient and therefore more scalable. For the LV loss we refer to Proposition 5.7 in [1], which shows beneficial scaling properties of this loss in high dimensions. Further, note that the additional use of trust regions allows losses that require importance weighting - such as the SOCM loss - to scale to high-dimensions as well (see e.g. Fig. 11) since the trust region keeps the variance of the importance weights approximately constant (see Figure 10).

However, we agree that the statement about scalability in Figure 2 could potentially be misleading. We, therfore, changed it to "We observe that our trust region methods (TR-SOCM and TR-LV) are the only methods that perform well in high dimensions."

---

Have the authors tried shrinking $\varepsilon$, which is commonly used in optimization?

Shrinking the trust-region bound is indeed a promising idea that could potentially enhance performance. However, to maintain simplicity and avoid introducing additional hyperparameters, such as a shrinking schedule, we opted to use a fixed value for $\varepsilon$ in our current implementation. We acknowledge this as a valuable direction and highlight it as a potential extension for future work. At the same time, however, we want to highlight a tradeoff between a small $\varepsilon$ and a potentially long runtime: the smaller $\varepsilon$, the larger is the number of the overall iterations $I$, which might result in better convergence to the optimal control $u^*$, but at the same time typically increases the overall runtime.

---

How is the length of the sequence $I$ affect the algorithmic performance?

The length of the sequence is intimately connected with the choice of the trust-region bound $\varepsilon$: Small trust-region bounds require more iterations to approach the target and therefore larger $I$. Conversely, chosing large values of $\varepsilon$ requires fewer iterations and therefore a smaller $I$.

As such, $I$ represents a trade-off between computational cost and performance. To see this, we kindly refer the reviewer to the supplementary material (Figure 10), where the left subfigure shows consistent performance improvements with smaller $\varepsilon$ values. However, the middle subfigure highlights that these settings require more target evaluations and thus larger $I$ to reach convergence. (For reference, $(\beta_i)_i$ denotes the annealing sequence, where $1 - \beta_i = 0$ indicates that the target distribution has been reached.)

---

References

[1] Nüsken, N., & Richter, L. (2021). Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2(4), 48.

---

We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.

---

The approach appears sound and is applied across multiple domains. How sensitive is the method to domain-specific hyperparameters or design choices?

The domain-specific hyperparameters and design choices were largely adopted from the respective baselines. Specifically, for transition path sampling, we followed the settings used in [1]; for fine-tuning, we adopted the configuration from [2]; and for diffusion-based sampling, we partially followed the setup in [3]. We kindly refer the reviewer to Appendices E.1, F.1, and H.1 for further details on these design choices.

Regarding the trust-region bound $\varepsilon$, we find that our method is generally robust across different domains. For all diffusion-based sampling tasks, we consistently used $\varepsilon = 0.1$, which performed well without additional tuning. In the case of transition path sampling and fine-tuning, we observed that smaller trust-region bounds can yield further improvements. We attribute this robustness to the theoretical relationship between $\varepsilon$ and the variance of the importance weights, as discussed in Appendix E.3, “Influence of Trust Region Bounds”. This relationship holds independently of the specific application domain (see also our response to the final question for further details).

---

The experimental results demonstrate sample efficiency on Stable Diffusion, but the performance is similar to the baseline. Do the authors expect their method to outperform the baseline with more samples?

Our primary goal in the fine-tuning experiments was to demonstrate improved sample efficiency, i.e., achieving competitive performance with significantly fewer samples. As noted, our method matches the baseline’s performance while requiring substantially less computational effort.

We are currently uncertain whether additional samples would allow our method to consistently outperform the baseline in terms of final performance, particularly given the strength of the baseline itself. However, since we observe consistent improvements under equal computational budgets in other domains, we plan to further investigate this question in the context of fine-tuning experiments.

---

The proposed approach relies on the use of constraints to improve robustness. Could the authors provide an ablation or analysis showing how different values of the $\varepsilon$ parameters influence performance and convergence? Is there a clear trade-off? How to set the parameter?

We have added such an ablation in Appendix E.3, “Influence of Trust Region Bounds” in the context of diffusion-based sampling. Indeed, in that experiment there is a clear trade-off between computational cost and performance: Smaller trust-region bounds lead to improved performance as shown in the left plot of Figure 10, but take longer to approach the target (see middle plot of Figure 10). We agree that this trade-off deserves more visibility, and have therefore integrated the ablation study in the main part of the manuscript.

To guide the choice of $\varepsilon$, we have further analyzed its effect in Appendix E.3 by examining its relationship to the variance of the importance weights and the corresponding effective sample size. This connection often provides a practical and intuitive way to set $\varepsilon$. In many of our experiments, we used $\varepsilon = 0.1$, which corresponds approximately to an effective sample size of 0.8.

---

We hope that our response addresses the reviewer’s question, and we would be glad to provide further clarification if needed.

---

References

[1] Seong, Kiyoung, et al. "Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers. ICLR 2025

[2] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." ICLR 2025

[3] Domingo-Enrich, Carles, et al. "Stochastic optimal control matching." NeurIPS 2024

---

We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.

---

[...] a result like this holds for path measures as well and could be connected to your result in Proposition 2.2, or be used to arrive at the trust-region method from a different angle?

Indeed, Proposition 2.2 can also be seen as an instance of the variational characterization you mention with the key difference that $\pi_0$ is replaced with $\pi_i$. To make this more explicit, consider the following optimization objective, which shares the same critical points in $u$ as $\mathcal{L}\_{\mathrm{TR}}^{(i)}(u, \lambda_i)$ (Eq. (5)):

$$\mathcal{L}(u) = \tfrac{1}{1+\lambda_i}D\_{\text{KL}}\left(\mathbb{P}^u | \mathbb{Q} \right) + \tfrac{\lambda_i}{1+\lambda_i} D\_{\text{KL}}(\mathbb{P}^u | \mathbb{P}^{u_i} ).$$

Next, defining $\alpha_i = \tfrac{1}{1+\lambda_i}$, we arrive at

$$\mathcal{L}(u) = \alpha_i D\_{\text{KL}}\left(\mathbb{P}^u | \mathbb{Q} \right) + (1-\alpha_i) D\_{\text{KL}}(\mathbb{P}^u | \mathbb{P}^{u_i} ).$$

Solving $u_{i+1} = \text{argmin}\_u \mathcal{L}(u)$ gives

$$\frac{\mathrm d\mathbb{P}^{u_{i+1}}}{\mathrm d\mathbb{P}} \propto \left(\frac{\mathrm d\mathbb{P}^{u_{i}}}{\mathrm d\mathbb{P}}\right)^{1-\alpha_i}\left(\frac{\mathrm d\mathbb{Q}}{\mathrm d\mathbb{P}}\right)^{\alpha_i} \quad \iff \quad \frac{\mathrm d\mathbb{P}^{u_{i+1}}}{\mathrm d\mathbb{P}^{u_i}} \propto \left(\frac{\mathrm d\mathbb{Q}}{\mathrm d\mathbb{P}^{u_i}}\right)^{\alpha_i}$$

which is equivalent to Eq.(7) in Proposition 2.2.

We thank the reviewer for highlighting this connection. We will add additional details and references to the manuscript.

---

Proposition 2.3 reminds me of a similar result in [Chopin et al. 2024] [...]

We agree that the work of Chopin et al. (2024) is highly relevant to our approach. In particular, the connection between geometric annealing and entropic mirror descent is intriguing. We are currently investigating whether similar connections can be established in the context of path measures, which we believe could offer valuable theoretical insights.

---

This is minor, but in Algorithm 1 should be incremented on every iteration? And should $u_{i+1}$ be returned instead of $u_i$?

We thank the reviewer for pointing this out. We have corrected the indexing and incrementing accordingly in the revised manuscript.

---

The explanation of how one goes from the cross-entropy loss to the SOCM loss is very brief and relies on prior knowledge of adjoint-matching.

We agree with the reviewer that the SOCM loss would benefit from further clarification. Accordingly, we have expanded the explanation in the revised manuscript to improve clarity and provide additional intuition.

---

Once again, we sincerely thank the reviewer for the valuable feedback and would be happy to provide any further clarifications if needed.