Parallel Simulation for Log-concave Sampling and Score-based Diffusion Models

We thank the reviewer for their detailed and constructive feedback. We are encouraged by your recognition of the core contribution—the diagonal reordering of Picard iterations to achieve improved parallel iteration complexity, the soundness of the theoretical analysis, and the potential applicability of our approach in high-compute environments (e.g., with sufficient RAM and cores). We address the specific concerns and questions raised below:

"Are there any ways to empirically validate whether the approaches outlined here hold? If so, could they be added to the paper and compared against alternatives?"

We agree that empirical validation is important. While we are not aware of existing implementations of our diagonal update scheme, we believe it could be adapted to diffusion-based samplers by modifying their update scheduling. We will add a discussion on this potential direction in the revised version.

Regarding typos: We will carefully proofread the manuscript and correct all grammatical and typographical errors to improve the clarity and overall presentation in the revised version.

Thank you for your detailed and positive feedback. We greatly appreciate your recognition of our novel parallel sampling technique, the clarity of the paper’s structure, the soundness of the theoretical analysis, and the potential of our approach to substantially improve time complexity and contribute meaningfully to the field. Below, we address your specific suggestions and questions:

"...provide additional intuitive explanations and justifications for the validity of the claim $L_n^j\leq aL_{n−1}^j+bL_{n}^{j−1}$...selection of $J$, $N$, and $P$..."

We appreciate your thoughtful observation that this recurrence forms the core of our analysis. Recall we denote the truncation error at the $n$-th time slice and the $j$-th iteration as $L_n^j$ in the squared sense.

• For Problem (b): by Young’s inequality and the definition of the exponential scheme, we can bound the change of $L_n^j$ in one update as

with $a = O(e^{h/2})$ and $b = O(L^2 e^{h/2} h^2)$. Thus, by choosing the time step $h$ sufficiently small relative to the Lipschitz constant $L$, we ensure convergence along the Picard direction (indexed by $j$). Specifically, this choice guarantees that $b < 1$ and $a > 1$, with $a$ remaining bounded by a constant.

• For Problem (a): since the score itself is not Lipschitz smooth, an additional score estimation error term $c \delta^2$ appears in the one-step update of $L_n^j$. To ensure the total score estimation error remains bounded, it is necessary to have $a, b < 1$. Similarly, by Young’s inequality and the definition of the Euler–Maruyama scheme, we can bound the change of $L_n^j$ as

Thank you for your thoughtful comments and positive evaluation of our work, particularly your recognition of the sharp iteration complexity in terms of $\varepsilon$ and the rigor of our analysis. We also appreciate your conclusion that the work merits acceptance. Below, we address your questions and comments in detail:

"A more thorough discussion of condition number (kappa) dependence in both time and space would be appreciated, with reference to prior work."

Thank you for pointing this out. We will revise the paper to include a more explicit discussion of how the condition number $\kappa$ affects both the step size and the overall iteration complexity. In particular:

• For iteration complexity: Our method achieves an iteration complexity of $O(\kappa)$, which matches the state-of-the-art sequential query complexity established by [AC24]. In the parallel setting, we believe it is inherently difficult to break the $O(\kappa)$ barrier using the Picard method, as the length of each time slice must scale inversely with the smoothness in order to preserve non-expansiveness along the Picard direction.

• For space complexity: Our method builds on the regrouping of scheduled grids along the time direction, and such schedual was proposed by [VW19]. Since we update all grids simultaneously in the worst case, the space complexity is given by $d N_{\text{seq}} = d \kappa^2$, where $N_{\text{seq}} = \kappa^2$ is the number of time steps required in the sequential setting.

[AC24] Shifted Composition III: Local Error Framework for KL Divergence, Jason M. Altschuler, Sinho Chewi, 2024

[VW19] Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted Langevin algorithm: isoperimetry suffices, 2019

"How difficult is it to handle the underdamped Langevin dynamics? What are the barriers and why can the current analysis not immediately give results?""

Applying our method to underdamped Langevin dynamics, as in [ACV24], yields $\log(d/\varepsilon^2)$ iteration complexity and improved space complexity $d^{3/2}/\varepsilon^2$. However, due to the failure of the triangle inequality for KL divergence, our current analysis cannot establish optimal iteration complexity for KL convergence—our main goal—so we leave this as future work.

[ACV24] Fast parallel sampling under isoperimetry. Nima Anari, Sinho Chewi, Thuy-Duong Vuong, 2024

"Can the Lipschitz score estimation assumption be removed, similar to the work from Benton et al.?"

We believe not, at least for Picard-type methods, as the Lipschitz condition is crucial to ensure convergence along the Picard direction.

Thank you for your insightful questions and appreciation of our work, particularly recognizing the clarity of the writing and the effectiveness of the diagonal communication scheme for parallel sampling. Below, we address your questions and comments:

"...gradient mapping contraction...missing the strong-log-concavity assumption?"

You're absolutely right—the contraction property used in Lemma B.5 indeed requires the strong log-concavity of the target distribution. We will revise Theorem 4.2 to explicitly state this assumption and clarify its role in ensuring convergence under the diagonal update scheme. Additionally, we will update all related claims and discussions throughout the paper to consistently reflect this assumption.

"...by closely extending the work of Anari et al 2024... build on the proofs of Chen et al 2024... (understandable) closeness to prior work..."

While we adopt several standard tools in parts of our analysis—such as interpolation arguments (e.g., VW19 or Section 4.2 of Che23), and Girsanov-based comparisons in Appendix C.2 for the diffusion model—we emphasize that our work addresses analytical challenges not present in prior literature.

In particular, our diagonal update scheme introduces changing start points, which presents analytical challenges not encountered in prior work. As discussed in Section 3, this required new analysis to control error propagation along the Picard iteration and time direction. As Reviewer KnEh pointed, the recurrence

captures the core of our theoretical analysis. We appreciate this insight and will provide additional intuition for this recurrence, highlighting its importance and distinguishing our analytic contributions from previous work.

[VW19] Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices.

[Che23] Sinho Chewi. Log-concave sampling.

[CRYR24] Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity

"The L∞ score approximation... unnecessary notational burden... does not capture any realistic setting..."

We would like to clarify that the $L^\infty$ approximation follows the same formulation as in [ACV24]. As discussed in their paper, this form of approximation is natural in the context of discrete sampling, where approximate scores can be efficiently computed in $O(1)$ iterations using access to a weighted counting oracle or Laplace transform. Furthermore, in spin glass models, the score function can be approximated in the $L^\infty$ norm using approximate message passing as shown in [AM22].

Moreover, since the final accumulated error in our method depends on the score approximation error at each queried point with varying weights, a natural refinement is to define a weighted version of Assumption 5.1. Also, another alternative relaxation is to adopt a bias-variance decomposition, as in Assumption 4 of [BCESZ22], though it is still formulated in an $L^\infty$-style metric.

[ACV24] Fast parallel sampling under isoperimetry.

[BCESZ22] Towards a Theory of Non-Log-Concave Sampling: First-Order Stationarity Guarantees for Langevin Monte Carlo

[AM22] Sampling from the Sherrington-Kirkpatrick Gibbs measure via algorithmic stochastic localization

"The Lipschitz constant of the score... affects the stepsize $h$...Could the authors expand on the possible dimension dependence on the uniform Lipschitz assumption and its effects on their bound?"

We acknowledge that the uniform Lipschitz assumption in Problem (b) is indeed strong; however, it is a common simplification adopted in prior works [CRYR24], [CCLL22]. As discussed in [SBDD22], the Lipschitz constant must be large to capture multimodal distributions, reflecting a trade-off between model expressivity and training stability. In particular, in settings with well-separated modes, the score's Lipschitz constant can grow significantly, even becoming unbounded near the zero point [YFZS24].

To ensure convergence of the Picard iterations under these conditions, the product $L_s^2 e^{h_n} h_n$ must be sufficiently small. This requirement implies that the length of each time slice, $h_n$, should scale as $O(1/L_s^2)$. Consequently, the number of time slices becomes $N = O(L_s^2 \log d)$, leading to an overall iteration complexity of $O(L_s^2 \log d)$.

Furthermore, our regrouping of time grids provides a flexible framework that could help mitigate Lipschitz singularities in other diffusion models, such as E-TSDM, as studied in [YFZS24]. We will expand on this point in the revised version.

[CRYR24] Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity

[CCLL22] Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions

[YFZS24] Lipschitz Singularities in Diffusion Models

[SBDD22] Can Push-forward Generative Models Fit Multimodal Distributions?