Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion

URL for Additional Figures: https://imgur.com/a/3UiYDVF

$**[A1]$H_\theta$estimates the vector field and$\gamma_\phi$performs inference-time planning**$

We clarify the misunderstanding here. As shown in Additional Figure 2, the roles of $\theta$ and $\phi$ differ clearly. $H_{\theta}$ estimates the vector field by predicting the endpoint $x_0$ regardless of trajectory curvature. Adaptive coefficients ($\gamma_{\phi}$) determine nonlinear trajectory planning by adjusting velocity directions and step sizes based on predictions from $H_{\theta}$. Thus, the capability for nonlinear trajectory planning depends entirely on $\gamma_{\phi}$, which adaptively guides the inference until endpoints are reached.

For clarity, we added Additional Figure 1 to illustrate MAC enabling distinct, adaptively curved trajectories for inference planning. Our core contribution, MAC, provides inference-time planning, meaning optimal inference plans are computed and optimized offline via simulation and then directly deployed during inference without additional computations.

$**[A2] Scalability Experiments on ImageNet**$

We agree with the reviewer that experiments on ImageNet are essential for demonstrating scalability and practical value. However, due to substantial resource requirements (32 A100 GPUs for 2 weeks needed for pre-training $EDM_\gamma$ on ImageNet-64), conducting large-scale experiments within the rebuttal period is currently impractical.

Nevertheless, we emphasize our method’s scalability, supported by the following points:

As in current SoTA methods on ImageNet-256 ([5], [6], [7], [8]), diffusion models’ scalability primarily relies on latent diffusion paradigms [4]. Our ongoing research explicitly integrates MAC with LDM, discrete diffusion frameworks [7], and IMM [8]. These frameworks substantially reduce computational demands, making ImageNet-scale experiments feasible by leveraging available pre-trained models.

Specifically, given that discrete diffusion frameworks naturally use highly multidimensional interpolated values for pre-training, optimizing MAC on discrete diffusion frameworks enables large-scale experiments using existing pre-trained models. Additionally, combining MAC with advanced models like IMM [8] can achieve high performance on large-scale datasets by using existing pre-trained models (explained in our response to Reviewer TFVT [A3]).

MAC can be readily integrated with these frameworks, and we are currently preparing dedicated large-scale follow-up studies.

For revision, we propose including experiments on ImageNet-64 using our existing setup with available $EDM_\alpha$ models. We believe this sufficiently demonstrates our method’s scalability within the scope of the original submission.

$**[A3] Essential References Discussion**$

Thank you for identifying the missing literature ([1], [2], [3]). Briefly, the major conceptual difference is that prior works ([1], [2], [3]) optimize vector fields by predefined trajectory properties (e.g., straightness or minimal curvature) guided by optimal transport theory to reduce numerical errors. In contrast, our method does not enforce predefined inference trajectory optimality; instead, we adaptively discover optimal nonlinear trajectories by final transportation quality through simulation-based inference-time planning. This allows greater flexibility and potentially superior performance.

$**[A4] Notation Clarifications**$

$G_{\theta, \phi}(\tau, x_T) = x_{\text{est}, \theta, \phi}$ (Details in Appendix C).

$\hat{x}$ values are predictions from $H_{\theta}(t, x(t))$ for terminal points $x_0 \sim \rho_0$ and $x_1 \sim \rho_1$.

$x_T$ denotes the actual initial sample point, while $z$ is sampled from the same distribution but unused during inference.

We sincerely appreciate your thoughtful consideration and hope these clarifications effectively address your concerns.

[1] Minimizing Trajectory Curvature of ODE-based Generative Models, ICML, 2023

[2] Constant Acceleration Flow, NeurIPS, 2024

[3] Simple ReFlow: Improved Techniques for Fast Flow Models, ICLR, 2025

[4] High-resolution Image Synthesis with Latent Diffusion Models, CVPR, 2021

[5] Scalable Diffusion Models with Transformers, ICCV, 2023

[6] SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers, ECCV, 2024

[7] [Mask] is All You Need, arXiv, 2025

[8] Inductive Moment Matching, arXiv, 2025

URL for Additional Figures: https://imgur.com/a/3UiYDVF

$**[A1] MAC’s core value lies in inference-time planning, orthogonal to vector field tuning methods like CTM and CD+GAN**$

CTM and CD+GAN optimize trajectories by adjusting the vector field parameter $\theta$. In contrast, our method (MAC) performs inference-time planning, meaning the optimal inference plans are computed and optimized offline via simulation and then directly deployed during inference without additional computations.

As shown in Additional Figures 1 and 2, without MAC, trajectory planning’s search scope is limited to linear paths, restricting performance gains. MAC enables adaptive, nonlinear, dimension-wise trajectory planning, significantly expanding optimization flexibility.

CTM and CD+GAN do not employ actual inference-time simulation feedback. Our method dynamically optimizes trajectories and timestep plans using simulations identical to inference, making it uniquely dynamic.

MAC acts as a final-stage inference planning strategy, enhancing performance beyond what is achievable by optimizing the vector field alone. This means that MAC can be integrated with existing frameworks such as CTM and other diffusion or distillation methods, as their methodologies attempt to solve different problems.

$**[A2] MAC Reduces Hyperparameter Engineering Costs**$

The modest gains observed with EDM result from EDM’s already highly optimized configuration, leaving limited margin for further improvement. (This is connected to the reason why we use $EDM_\gamma$ and coefficient labeling, explained below in [A3].) However, Sections 4.1 and 4.2 demonstrate significant improvements (~10x NFE efficiency) in less optimized frameworks (DDPM, FM, SI) by using MAC.

MAC alleviates the extensive manual tuning burden associated with framework-specific hyperparameters, like in EDM. Although MAC requires tuning of hypothesis space parameters ($M, s, q$, LPF configs), these parameters are framework-independent. Hence, once identified, these parameters generalize well across various flow and diffusion models, providing considerable performance benefits without framework-specific engineering, as shown in our experiments.

Regarding baseline comparison with similar parameterization: As demonstrated in Table 6, directly incorporating pre-training parameterization into EDM ($EDM_\gamma$) negatively impacts performance. Using this as a teacher network will degrade CTM performance unless combined with MAC for inference-time planning. MAC provides clear advantages specifically when used for adaptive inference-time planning, not as a general pre-training strategy.

$**[A3] Practicality via Existing Pre-trained Models**$

Our method indeed requires higher computational costs compared to CTM (approximately 201 Mimg vs. 25 Mimg) when using the specific $EDM_\gamma$ pre-training setup, chosen deliberately to push the limits and achieve SoTA performance on CIFAR-10. However, we emphasize that using $EDM_\gamma$ and coefficient labeling is optional. MAC can practically and effectively leverage existing pre-trained models (such as $EDM_\alpha$), significantly reducing computational overhead without substantially compromising performance.

Using $\gamma$ increases the probability of multidimensional interpolated values $x(t)$, but even using $\alpha$ inherently yields multidimensional interpolations, as Gaussian noise is independently sampled across dimensions. Specifically, given $x_1 \sim \rho_1 = \mathcal{N}(0, I_d) \in \mathbb{R}^d$, each dimension $x_{1,i}$ is independently drawn from $\mathcal{N}(0,1)$, and thus even a linear interpolation $x(t) = t \cdot x_1 + (1-t) \cdot x_0$ naturally results in dimension-wise distinct noise contributions.

Empirical evidence from our experiments supports this claim. Section 4.1 demonstrates effective nonlinear trajectory planning in $OT-SI_{MAC}$ using models pre-trained with $\alpha$. Additionally, Figure 6 shows the trained $\phi$‘s distribution significantly differs from the pre-training distribution, indicating MAC’s effectiveness is not critically dependent on pre-training specifics.

Thus, for large-scale practical scenarios where training from scratch is costly, MAC can effectively utilize existing models (e.g., $EDM_\alpha$), ensuring practicality and computational efficiency.

$**[A4] Scalability (ImageNet)**$

Please refer to our response to Reviewer jRK3’s comment [A2].

We sincerely appreciate your thoughtful consideration and hope these clarifications effectively address your concerns.