Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Authors: Dohoon Lee
Jaehyun Park
Hyunwoo J. Kim
Kyogu Lee
Venue: ICML 2025 poster (2025)
Keywords: Multidimensional Adaptive Coefficient
Inference Trajectory Optimization
Flow
Diffusion
Adversarial Optimization
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We introduce MAC, a plug-in for flow and diffusion models that enables multidimensionality and adaptability in inference trajectories, improving generation quality and efficiency.

Abstract: Flow and diffusion models have demonstrated strong performance and training stability across various tasks but lack two critical properties of simulation-based methods: freedom of dimensionality and adaptability to different inference trajectories. To address this limitation, we propose the Multidimensional Adaptive Coefficient (MAC), a plug-in module for flow and diffusion models that extends conventional unidimensional coefficients to multidimensional ones and enables inference trajectory-wise adaptation. MAC is trained via simulation-based feedback through adversarial refinement. Empirical results across diverse frameworks and datasets demonstrate that MAC enhances generative quality with high training efficiency. Consequently, our work offers a new perspective on inference trajectory optimality, encouraging future research to move beyond vector field design and to leverage training-efficient, simulation-based optimization.

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Dohoon Lee Jaehyun Park
Hyunwoo J. Kim
Kyogu Lee
24
Abstract
Flow and diffusion models have demonstrated strong performance and training stability across various tasks but lack two critical properties of simulation-bart methods: freedom of dimension- ality and adaptability to different inference trajec- tories. To address this limitation, we propose the Multidimensional Adaptive Coefficient (MAC), plug-in module for flow and diffusion models that extends conventional unidimensional coeffi- cients to multidimensional ones and enables infer- ence trajectory-wii wise adaptation. MAC is trained via simulation-bal based feedback through adversar- ial refinement. Empirical results across diverse frameworks and datasets demonstrate that MAC enhances generative quality with high training ef- ficiency. Consequently, our work offers new perspective on inference trajectory optimality, en- couraging future research to move beyond vec- tor field design and to leverage training-efficient simulation- based optimization.
(a)
(b)
(c) To (MAC)
Figure 1: Comparison of and 7ø- As shown, em- ploying (MAC) expands the search space for trajectory optimization by enabling adaptive, curved trajectories and time steps that are distinct for each sample.
erties, they exhibit limitations in quality and training effi- ciency compared to flow and diffusion. We integrate these two properties of simulation-basee methods into flow and diffusion models to enhance performance while preserving training efficiency, effectively combining their advantages.
1. Introduction
Compared simulation-based methods like NeuralODE (Chen et al., 2018), flow and diffusion modeling (Sohl- Dickstein et al., 2015; Lipman et al., 2023) demonstrates remarkable performance and training stability across various tasks and has become a standard approach for generation tasks. However, trades off two critical properties that simulation-ball based methods offer, which could enhance the quality of transportation: freedom of dimensionality and adaptability with respect different inference trajectories. While iimulation-based methodologies possess these prop-
To achieve this, we introduce the Multidimensional Adap- tive Coefficient (MAC). As described by Albergo et al. (2023), the trajectory with and diffusion for t E [0, can be written as x(t) oo (1)zn + (1)x)),, Np. 21 E Rd, where, convention- ally, the coefficients (n a1 (t) E R are unidimensional. We extend this by introducing Multidimensional Coeffi- cient 7o()), 71 (1) E Rddd which allows for different time scheduling across all data dimensions. Based this, introduce MAC nss xS.)). parameterized by Φ. to adapt to different inference trajectories x².3 To optimize MAC, we use to construct flow- or diffusion-baas differential equation solver Ge.s based flow diffusion model Ho. and ad adversarially optimize 0 and & using discriminator D By optimizing MAC, we achieve full multidimensionality and adaptability in the inference trajectory of flow and diffu- sion models. This can be applied to any type of framework as a plug-in module enhance inference performance.
Po and
in flow
'Music and Audio Research Group, Department of Intelli- gence and Information, Scoul National University, Scoul, Korea Interdisciplinary Program in Artificial Intelligence, Seoul Na- tional University, Seoul, Korea "School of Computing, KAIST, Daejeon, Korea "Attificial Intelligence Institute, Seoul National University, Scoul, Korea. Correspondence to: Dohoon Lee <do- hoonlee.researché gmail.com>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusson
3o (4)
Io (44)
20.8
00,0.6
(a) Inference without MAC
(b) Inference with MAC
Figure 2: Comparison of inference trajectory optimization with and without MAC, given the same vector field estimated by Hg, indicated by gray arrows in the background. (a) Without MAC, inference trajectories strictly follow the directions defined by the vector field Ho. allowing adjustments only in step sizes, which must remain consistent across all samples. (b) With MAC, the search space (shaded in blue) expands, enabling flexible adjustments in both trajectory directions and step sizes, adaptively optimized for each individual sample. Thus, even when a suboptimal vector field is given, MAC effectively corrects errors by identifying optimal inference plans that maximize final transportation quality.
For a given vector field defined by Ho, optimizing MAC corresponds to inference trajectory optimization. In par- ticular. optimal inference plans computed and refined offline via simulation, and subsequently deployed during inference without incurring additional optimization costs. Before MAC, it was challenging to induce significant per- formance gains through inference trajectory optimization alone, due to the absence of multidimensionality and adapt- ability. In contrast, as shown in Figure MAC substantially expands the search space for optimization by introducing dimension-wise adaptive control, thereby enhancing the effectiveness of inference trajectory optimization.
transportation. Our approach addresses this limitation by introducing training-efficient simulation, thereby making simulation-baase based inference trajectory optimization practi- cally feasible. In summary, our main contributions are:
We introduce the concept of Multidimensioonl Adap- tive Coefficient (MAC) in flow and diffusion, laying the foundation for achieving both freedom of dimension- ality and adaptability in inference trajectories, easily implemented as a plug-in in module for any framework.
2.
We propose a highly efficient training strategy for in- ference trajectory optimization by leveraging MAC in combination with adversarial training.
Our experiments span various frameworks (DDPM (Ho et al., 2020), FM (Lipman et al., 2023), EDM (Karras et al., 2022), SI (Albergo et al., 2023)) across multiple datasets (CIFAR-10 10 (Krizhevsky Hinton, 2009), FFHQ (Karras et al., 2019), AFHQ (Choi al., 2020), and ImageNet (Deng et al., 2009)). Notably, inference using MAC yields consistent performance improvements, achieving state-of- the-art results in CIFAR-10 10 conditional generation. These results are obtained with high training efficiency, suggesting that we successfully integrate the advantages of simulation- based and simulation-free dynamics.
We suggest exploring the notion of trajectory optimal- ity beyond predefined properties, emphasizing evalua- tion based on the final transportation quality.
2. Related Works
Trajectory Optimizations in Flow and Diffusion Vari- ous trajectory optimization approaches predefine optimal- ity before training using simulation-fri -ree objectives. Meth- ods such Liu et al. (2023) and Tong et al. (2024) define straightness the optimality criterion and optimize trajec- tories by maintaining consistency between (xo. 2i) when training flow and diffusion models, aligning with an Optimal Transport (OT) perspective. Other approaches, including Lee et al. (2023); Park et al. (2024); Kim et al. (2025), also adopt the OT perspective as a predefined criterion for tuning the vector field. Singhal et al. (2023) defines optimality through fixed sequence of diffusion steps aimed at reduc- ing inference complexity. Bartosh et al. (2024) introduces a neural flow model that implicitly sets trajectory optimality within the diffusion process. Kapušniak et al. (2024) op- timizes trajectories by defining data-dependent geodesics. There are also methods that refine trajectories after training.
Consequently, we propose exploring the optimality of in- ference trajectories through training-efficient simulation in flow- and diffusion-based generative models. Prior ap- proaches primarily focused on shaping the vector field, often relying on predefined, simulation-fis notions of optimality, such as straightness (Liu et al., 2023; Tong et al., 2024). However, the inference trajectory is governed not solely by the vector field, but by the interplay between the vector field and the inference plan. From this perspective, simulation- free vector field training deviates from true transportation optimality, which should be defined by the final quality of
2

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusson
such as Albergo al. (2024), which defines optimality by minimizing trajectory length under the Wasserstein-2 metric, focusing on a shortest-distance criterion.
directly model the vector field (t)() vg (1))
with the flow model Ho. All these methods achieve genera- tive modeling by numerically solving an ODE or SDE using the predicted vector field. We use DDPM (Ho et al., 2020), FM (Lipman et al., 2023), EDM (Karras et al., 2022), and SI (Albergo et al., 2023), as summarized in Appendix A.
Despite these diverse perspectives on trajectory optimality, there are two significant differences between these methods and ours. First, we assess optimality solely based on the final transportation quality through simulation, which is crucial factor in generative modeling. Second, none of the existing methods achieve full flexibility in inference trajectory design on two fronts: multidimensionality and adaptability with respect to different inference trajectories.
3.1. Multidimennsonall Adaptive Coefficient
To introduce multidimensionality into the coefficient, we first define the mulridimenssonal coefficient, 7() ho(t), (t)] E Rdxdx2, which generalizes the unidimen- sional coefficient by extending it to higher dimensions.
Few-Step Generation via Distillation and Fine- Tunn-T Existing approaches can be broadly categorized into non- adversarial and adversarial methods. Non-adversarial meth- ods, such as Geng et al. (2023); Berthelot et al. (2023); Yin et al. (2024); Song & Dhariwal (2024), focus on I-step dis- tillation techniques without adversarial objectives. These approaches aim for few-step generation by leveraging distri- butional losses and equilibrium models, effectively distilling the diffusion process without involving discriminator. Con- versely, adversarial approaches such as Wang et al. (2023); Lu et al. (2023); Luo et al. (2023); Xu et al. (2024), integrate a discriminator within diffusion models via a GAN-based framework. Kim et al. (2024) also generalizes Song et al. (2023) to enable efficient sampling using a discriminator.
Definition 3.1. (Multidimensional Coefficient) Given a trajectory defined by xx1) - (o(()) + m (4)x) where Zo, 21 E Rd, the multidimensional coefficient is defined as:
(t) = ho)(), n (t) : [0,7] -> Rddxxx
subject to the following conditions: 70(0) = I, (0) = Oaxd, (T() = TI, and r(t) E (²O Rdxddxx2))
7(t) E C* ([0, T). Rdxdx2) indicates that is continuously differentiable to the first order with respect to t on the in- terval [0, 7j. The boundary conditions ensure that (1() be- comes Zo and zT fort = 0 and t = T', respectively, which is necessary for transportation. The unidimensional coefficient a can be regarded as special case of the multidimensional coefficient y. where 7 is scalar matrix for given t.
While our method shares similarities in achieving few-step generation by leveraging a pre-trained diffusion model, key difference is that prior works optimize only the vector field, omitting the inference plan. In contrast, our method jointly optimizes both the vector field and the inference plan, thereby enabling inference trajectory optimization.
For adaptability with respect to different inference trajecto- ries, we define multidimensional coefficient 7g- parameter- ized by S. allowing it to adapt to different inference trajec- tories (t) over the inference times т = (fo...... :::):
Definition 3.2. (Multidimensional Adaptive Coeffi- cient (MAC)) Let S = {t().....
3. Methodology
S [0,T]
be an arbitrary set of inference times, and define the corresponding inference trajectory as x (t"))) ...., ((0,))))..,,,,((( (t(()) E Rdxe Then, the MAC n.oll, xx, is defined as: (1, x = -10,,0 (l, x)
We consider the task of transporting between two distribu- tions xo ~ Po and x1 ~ Pl, where To; .1 E Rd. Following Albergo et al. (2023), for te E [0, T). the trajectory x(t) is defined as:
(t,,, x 10. T x dxe + Rdxdx2
(t) = oo(()ao + +n v(1,x(1)= = i(),
(1)
where (t) = [ao(t), 01 (t)) E 2 represents the unidimen- sional coefficients, and (!) denotes the derivative of (t))
and is parameterized by d. Boundary conditions and smooth- ness follow Definition 3.1.
with respect to t. The diffusion model Ha estimates the
Efficient Computation of MAC We adopt a diagonal matrix for Ye. which significantly reduces the output di- mensionality and the size of the neural network for MAC, while retaining sufficient capacity to model non-linear tra- jectories. Accordingly, consider >(t) E Rd instead of 7()) dxxx2 (details in Appendix B). Furthermore, we evaluate the d-parameterized part of to only once at t = T, and reuse the resulting function throughout the entire infer- ence schedule т. This allows us to compute % over the full schedule with a single forward pass through the denetwork.
vector field v as follows:
x()) = cia (o) + de (1)
(2)
Ho(t, = L1.e.
For example. Song et al. (2021) predict the score value V, log t() = a to obtain the vector field.
There are also flow-based methods, such as Lipman et al. (2023), that do not explicitly target $o.0 or 01,0, but instead
3

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
3.2. Inference Trajectory Optimization
Consider a flow- and diffusion-basee differential equation solver & with fixed configurations (c.g., NFE, integration schemes), where f represents the vector field from the pre- trained flow or diffusion model parameters, and ø is the parameter for MAC. We propose two optimization problems. The first is:
(a)
(b)
6 arg arg min D (po: Po.9,0)
(3)
where Po.0. denotes the generated distribution from Ga and is a divergence metric. Solving Equation 3 corre- sponds to inference trajectory optimization with a fixed vector field. This setting allows us isolate the impact of optimizing MAC and clearly demonstrate its effectiveness.
Aite LS 14
to
(c)
(d)
To achieve higher transportation performance, we optimize the full inference trajectory by jointly optimizing the vector field A and the inference plan di, as follows:
Figure Crude coefficients: (a) Oscillatory behavior in t due high-frequency components; (b) High adjacent pixel
differences. Refined coefficients: (c) Constrained multidi- mensionality for larger i in pre-training; (d) Unconstrained multidimensionality for adversarial training.
8" o² arg min D (po. Po.0,)) 0,0
(4)
high-frequency components in & and across different dimen- sions. Given the vast size of the coefficient set, it is crucial to exclude such crude coefficients using appropriate con- straints and define a well-designed hypothesis set for % to explore, Second, the computation of % via the parameter & should require low NFE. Lastly, for pre-training the flow or diffusion models, the cost of sampling random ~ I'm should be low. Considering these factors, we choose to model the weights of ainusoidals-akin to Fourier expan- sion- by parameter o, as in Albergo et al. (2024). Our chosen design is:
For low-dimensional datasets, the divergence terms above can be computed directly using the Wasserstein distance. For high- sigh-dimensional datasets, following Goodfellow al. (2014), minimize Equations 3 and by using the follow- ing min-max problem:
min max Ezo [log (ro)) Exr [log (1 D& (T,))))) 8.0
(5)
where De represents the discriminator, and ØT pr is the starting point of the inference trajectory. As shown, 0 and ø in Go.ş aim to deceive Da, Based on Equation 5, we use the simulation-based adversarial objective exclusively for o (details in Section 3.2.3). Optionally, before trajectory optimization, we pre-traee pre-train the flow and diffusion model He using various 7 sampled from a well-designed hypothesis set L' (details in Section 3.2.2).
(T(), [0, x Rd - Rdx2 where %o(t, = ho, xT). (, oT),.
TB
(t, vT) = Fo (t), us (xT)) (1, TT) = Fi (t), ws (xT)), & E P
where P represents the parameter set for ø. The function F constructs the weighted sinusoidal expansion using basis functions bm(()) and their weights was defined as:
3.2.1. HYPOTHESIS SET FOR MAC
We define the multidimensional coefficient set 1 as:
(8) ws (xT) LPF tanh (NNa(IT)), SER, (9)
bm (1) sin sin (Fm((/T))1)) ER, GER,
7(1) [0, T] dx2 where
7(1 = (1)) 7(1) E Co ([0, Rd2)) (0) = 1a, (0) 02, 7 (T) = Tla
where wa(xx) Rxx represents the multidimensional weights for the sinusoidals, modeled by the neural network NNg- LPF refers to low-pass filtering, implemented as con- volution with Gaussian kernel, applied across different d dimensions to exclude high frequencies. This parameteriza- tion always satisfies DA S P. For image generation, we em- ploy a U-Net Ug for NN2, and condition it as Us(2r, c) for conditional generation. Details are provided in Appendix
(6)
To design the hypothesis set T C I for exploring Ya. we consider three main properties. First, the hypothesis set should be broad enough to include the optimal coefficient To- while avoiding unnecessary complexity to reduce the burden on the flow and diffusion model. As shown in Fig- ure some multidimensional coefficients exhibit excessive

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusson
This parameterization satisfies all three properties men- tioned above. First, high-frequency components in t and d can be excluded by controlling hyperparameters such as s and LPF configurations. Second, Y can be com- puted with NFE using NNg- Lastly, because the tanh in Equation 9 bounds the output of NN₂ to (-1, we can efficiently sample a random multidimensional coefficient nt,, u) ~ I' by directly sampling sinusoidal weights w from uniform distribution rather than evaluating NN₂ dur- ing pre-training. Specifically, w(u)- = s LPF o u, u U(-1,1) Rdxdxx2
tionally prohibitive, MAC remains compatible with existing pre-trained models, ensuring practical applicability.
3.2.3. ADVERSARIAL. OPTIMIZATION
For Equation 5, we use the hinge loss (Lim & Ye, 2017) with the StyleGAN-XL (Sauer et al., 2022) discriminator for D2, as used in Kim et al. (2024). Details for the flow- and diffusion-basee provided in Appendix D. We use separate loss functions for @ and Φ, applying the simulation- -aased objective exclusively to Φ as follows:
differential equation solver Go.o are
3.2.2. OPTIONAL PREE- TRAINING FOR Ho
L+ = -E (Go( xT)]]
f Ezo [ml (0, (0,1-D((zx))
(10)
To ensure that Ha can better accommodate y, we introduce a pre-training procedure. The loss functions for pre-trainn vre-tranning Ho follow the original flow and diffusion losses, except that y is used instead of a, and coefficient conditioning is incor- porated. For example, in the case of DDPM, the loss func- tion is given by fg ,1 Ho (t, (t), 7(, u)) mol12. Detailed versions of the loss functions for each flow and diffusion framework are provided in Appendix E.
+ Ex, |max(0,1 + De(Go,s(T,
where L. and Lup indicate that gradients are computed with respect to & and uh respectively. Through these loss func- tions, e is optimized to determine better coefficients via simulation through Gø.6. For 0 we use the adversarial loss defined as:
Lo - 10 (De (Holl, (0(1,
Coefficient Conditioning When adopting existing pre- trained models for Hg. we apply (t) dur- ing trajectory optimization, where To (t,, xTy) E R denotes the scalar average of (1, Zr)) over all dimensions. How- ever, to allow Y to receive precise gradients through Ho. we incorporate To (t, .T) directly into the model during infer- ence trajectory optimization by using Ht x(t), (t, xr ))- To support this without modifying the structure of Ho, we concatenate (f, ÆT) with xt) along the channel axis as input to Hg. We also pre-train Ha with coefficient condi- tioning in the form He (t, (t))) (t, u)), consistent with its usage during inference trajectory optimization.
(11)
xxt) to,a(t, 2) @ zo + (,, 2) O I1,
where ~ pr (sampled like IT, used only to train Ho), To ~ Po. ~ pl, t 1~ and (t( x()) is used when coefficient conditioning is disabled. Given that Hs approximates Zo. it can be adversarially optimized with De Since Ha only needs to handle coefficients from Yo rather than the full hypothesis set T,, it is trained exclusively using Ye. thereby reducing its computational burden, as 76 is significantly smaller than i The final loss terms are:
Dasty = Ls + Lap: Cady 8.000 = Lo + Lis + Leg-
(12)
Ta for Pre- Trainingn Given that a large r of coefficients for pre-training can burden Hg and potentially degrade per- formance, we use a smaller A during pre- pre-training and fully enable multidimensionality across & during adversarial train- ing, shown in Figure 3 (c) and (d). Specifically, we use y with high multidimensionality near & = O and low mul- tidimensionality at larger by configuring the LPF during the pre-trate wre-training of Hg. For adversarial training, we fully enable multidimensionality for Yo across the entire range of t. Details are provided in Appendix C.
(13)
Here, Equation 12 corresponds to Equation 3, and Equa- tion 13 corresponds to Equation 4. The procedure corre- sponding to Equation 13 is summarized in Algorithm
Algorithm 1 7-ee Jy-Pre-Training and Adversarial Optimization
Input: Po: Pr, Ho, To: Go.o; Da; T,7 if p-rettraiiig then repeat
Sample 20 ~ pa: 21 ~ pI, u a~ Ul 1(,,,,,,, INT 21. (-))+=))))) (L),
Optionality of Pre-Training We emphasize that this pre- training stage is optional, as MAC can effectively leverage existing models pre-trained with a without substantially compromising performance. While using 7 for pre-training increases the probability of generating multidimensional interpolated values s(). multidimensional interpolation in- herently arises even with a standard coefficient a, since Ii ~ PI = (0, E Rd implies that each dimension 2 is independently drawn from (o, 1). Consequently, for large-scale datasets where training from scratch is computa-
Compute Cge, Update 0
until @ converges
end if
repeat
Sample uo ~ po; II ~ Pl, 2T ~ pT, 2 PT, L INT
Holl, ((44)
with 7-pretraining w/o p-petraiiing
$0.0.6 +
Ho (t,
(t.)(tx) + Go,
Compute Cady Lo...u: Update 0, d.
until 0, á. v converge
5

---

Multidimensional Adaptive Coefficient for Inferenee Trajectory Optimization in Flow and Diffusson
Table 1: W2 distance (4) for 2-dimensional transportation results. The best performance is highlighted.
Gaussian 8 Gaussians
Gaussian to Moons
Gaussians to Moons
Moons 8 Gaussians
Method
\
NFF
5
10
5
10
5
10
5
10
SIa
0.763
0.673
0.882=aass
0.643
0.981
+0.12
0.649
1.271
0.998
SI, opt Φ
0.72 10.082
0.452
0.682
0.359
100098
0.924
10235
0.31
10051
0.908
10.109
0.5001.0m
ST
0.457
0.440
0.245_0.023
0.217=a.0
0.321
10004
0.318
200ss
0.48822.0
0.492
SI + opt
0.399 10.017
0.415 10.016
0.230
0.188 10.006 0.258
1:0015
0.22 1 10014
0.42 1 100012
0.407 100001
Table 2: FID (4) of models pre-trained with a and y.
CIFAR-10
ImageNet-32
Method \NFE
100
150
200
100
150
200
SIa
4.75
4.51
4.30
8.08
7.79
7.63
SI,
3.98
3.74
3.63
6.33
6.21
6.20
FM,
4.52
4.23
4.07
7.78
7.53
7.38
FM, 3.59
3.42
3.42
6.18
6.03
6.01
DDPM,
6.64
4.84
4.10
8.13
7.40
7.14
DDPM, 4.73
4.11
3.83
6.84
6.51
6.42
(a) Sla
(b) SI, + opt ạ
Table 3: FID after adversarial optimization for varying NFE on CIFAR-10.
Method NFE
4 (+)
6 (+)
8 (+)
10 (+)
SI, + adv & FM, + adv di
20.59
6.62
4.85
4.14
16.42
8.17
6.56
6.13
DDPM,, + adv Φ
72.64
20.13
13.72
10.04
model to learn a straight trajectory by pairing xo and æ1 as an OT solution within a minibatch during pre-training, where optimality for the inference trajectory is defined as straight. Details are provided in Appendix E.
(c) STOT
(d) SI + opt Φ
Figure 4: Comparison of inference trajectories from Gaus-
sians to Moons.
As shown in Table 1, using MAC consistently achieves the best results, even for models trained with minibatch-OT.
4. Experiments
This suggests that a straight trajectory is not always optimal, even in OT-trained models, and that MAC can adaptively discover better inference trajectories to correct errors. Fig- ure 4 further illustrates how MAC adjusts the trajectory direction to optimize transportation, resulting in a path that is not straight. A comparison of (c) with (d) reveals dis- tinct piecewise linear trajectory in (d), indicating that non- straight trajectory achieves superior performance. These results empirically suggest that the optimality of the infer- ence trajectory should be defined and explored in terms of transportation quality, rather than by a predefined property.
4.1. 2-Dimensional Transportation: Optimizing Only
We conduct experiments on 2-dimensional synthetic datasets used by Tong et al. (2024) with SI (Albergo et al., 2023). To validate the benefits of using MAC Yo- train & exclusively while freezing 0, as in Equation 3, en- suring a fair comparison with baseline methods. We use the Wasserstein loss function L.- = W2 (2o, GT, 2T)) instead of discriminator. To further eval- uate whether MAC improves transportation even when opti- mality is defined as straight tra- jectory, we test additional con- figurations of minibatch pairing for pre-training (xo, x ) random pairing and OT pairing (Tong et al., 2024). The minibatch- OT method encourages the flow and diffusion
conductexperrmentt
Table 4: FID after adversarial optimization with different conditionings of (t, -) on CIFAR- 10 (10 (+) NFE).
Figure 5: FT PT (black) and zo Po (blue).
Method \ Conditioning
1
PT
FT PT
SI, + ady Φ FM, + adv Φ
7.84
6.48
4.14
9.20
9,06
6.13
DDPM, + adv d
26.09
23.31
10.04
6

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Table 5: Performance comparisons on CIFAR-10, FFHQ, AFHQ, and ImageNet. indicates metrics computed by us.
ImageNNtt-66 cond.
Model
NFE
FTD (J) DDDINNN2 (L)
CIFAR-I0 uncond. CIFAR-I0 cond.
Model
NFE
FID (U)
FID (L)
GAN
StyleGAN XL. (Sauer el 2022)
2.09
GAN
StyleGAN-Ade Ada (Karras et al., 2020)
2.92
2.42
Diffusion Model
StyleGAN XL (Saver al., 2022)
1.85
ADM, (Dhariwal & Nichol, 2021)
250
2.07
EDMg (Karras et al,, 2022)
79
2.44
Diffusion Model
DDPM,, (IIo et al., 2020)
1000
3.17
Consistency Model (Distillation)
Score SDE, (Song et al., 2021)
2000
2.20
CDa (Song et al., 2023)
6.20
EDM, (Karaa al., 2022)
35
1.98
1.79
2
4.70
CTM, + adv (Kim et al., 2024)
1.92
160.8"
Rectified Flow (Distillation)
1.73
157.7"
2 RRctiiieed Flow. (Liu et al., 2023)
4.85
3.02"
195.0"
2-Rectified Flow + (l.ce et al., 2024)
205.1"
2
2.40
Inference Trajectory Optimization EDM,, + adv O. (Ours)
Consistency Model (Distillation) CD2 (Song et al., 2023)
5 (+) 1488
70.2"
3,55
CD. + adv 8 (Lu et al. 2023) CTM., (Kim et al.. 2024)
5.19
2,65
Model
NFE
FFHQ-64
(L)
AFHQ-64
(I)
CTM, + adv 0 (Kim et al., 2024)
1.98
1.73
1.87
1.63
Diffusion Model
1,86"
1.98"
HDM,, (Karras al., 2022)
2.39
1.96
1.93"
2.04
Rectilied Flow (Distillation)
Inference Trajectory Optimization EDM., + adv f, o (Ours)
2 Rectifie Flow.,++ ++ (Lece et al., 2024)
5.21
4.11
5
1.69
1.7
4.26
3.12
Inference Trajectory Optimization HDM, adv 6, o (Ours)
F1D
RMS,, + Afc. e
2,04"
18391, d.
Table 6: FID from the ablation study on CIFAR-10. Experiments used 500k images, compared to 1500k in Table 5.
Unconditional
Conditional
5mmg
100
201
300
100
(a) CIFAR- 10 uncond.
Configuration NFF
5 (Euler)
35 (Heun)
5 (Euler)
35 (Heun)
EDM, Ads.
EDM.
68.73
1.97
48.76
1.79
EDM-
69.58
2.08
48.53
1.81
EDM, + adv Φ (no multi.)
33.55
25.56
EDM, + adv
18.67
7.77
FDM,, adv
2.28
2.14
king
100
200
300
4041
EDM, + adv 0, Φ
1.81
1.42
(b) CIFAR-10 cond.
Figure 6: EDM, + adv 0 and 0, d.
4.2. Image Generation: Optimizing Only 6
Adversarial Training We achieve approximately 10x NFE efficiency across all frameworks. Specifically, in SI, we obtain FID scores of 4.14 and 7.06 on CIFAR-10 and ImageNet-32, respectively, using 10 (+) NFE (where (+) accounts for the computation of %/o given that the size of Us is smaller than that of Hg). These results demonstrate the compatibility of MAC across various frameworks. As shown in Table 3, MAC can also be optimized for different sampling configurations. Additionally, we validate the im- pact of using the starting point .T as a condition for MAC. As shown in Table 4, incorporating JT as an input to % consistently improves performance across all frameworks, providing empirical evidence of MAC's adaptability to zT-
To isolate the benefits of using MAC in image generation, we conduct experiments that optimize only é. We employ DDPM (Ho et al. 2020), FM (Lipman al., 2023), and SI (Albergo et al., 2023) on the CIFAR-10 (Krizhevsky & Hinton, 2009) and ImageNet-32 (Deng et al., 2009) datasets. Detailed training configurations and additional results are presented in Appendix E.2.
Pre-Training Before adversarial optimization, we com- pare the Fréchet Inception Distance (FID) (Heusel et al., 2017) of models trained with a and 7. Interestingly, as shown in Table 2, models trained with y achieve the lowest FID across all frameworks and NFE compared to models trained with a. These results suggest that training flow and diffusion models with y ~ F and coefficient labeling can improve the innate performance of the model, despite the increased complexity compared to a.
4.3. Image Generation: Optimizing 0 and
We optimize 0 and 6 using Equation 13 on the CIFAR- 10 (Krizhevsky & Hinton, 2009), FFHQ-64 (Karras et al., 2019), AFHQv2-64 (Choi et al., 2020), and ImageNet-64
7

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
FID
F, F, F
Automobike
35
F. T, F
Horse
Ship
Cal Dog
F. E. T
Deer
30
PP, T. T
Airplane
IT, T, T1
Frog
25
Bird
Truck
20
Iter.
1000
2000
3000
(a) CIFAR-10 10 uncond.
FID
FF, F, P]
(a) CIFAR- 10 uncond.
(b) CIFAR- 10 cond.
35
F"
F, F. T
30
F, T
[T.
25
20
Iter.
1000
2000
3000
(b) CIFAR- 10 cond.
Figure 7: Training with different multidimensionalitiss. T/F indicate whether multidimensionality applied along the channel, height, and width axes, respectively.
(c) FFHQ.
(d) AFHQ.
Figure 8: T-SNE visualization of various coefficients. Yel-
(Deng et al., 2009) datasets with EDM (Karras et al., 2022). As mentioned in Section 3.2.2, since pre-training Ho with is optional for practicality, we use the existing pre-trained model FDM, for ImageNet. We measure FID and FDDINOv2 (Oquab et al., 2024). Details are provided in Appendix E.3.
low: w = (corresponding to a); Purple: w su;
Green: w s sLPF (used for pre-training); Blue:
" = s LPF o (an (trained).
Impact of Multidimennionalltty of MAC To examine how the multidimensionality of MAC influences perfor- mance, we train & with different levels of multidimension- ality by averaging tanh(Us (xr)) in Equation 9 across spe- cific axes. For instance, to retain multidimensionality solely in the height axis ([F, T. FJ), we use the same Wy across the channel and width axes by averaging along those di- mensions. As shown in Figure 7, incorporating more axes consistently improves performance, indicating that MAC's multidimensionality positively impacts generation quality.
Impact of Trajectory Optimization As shown Table 5, our approach generates high-quality samples across various datasets with only 5 (+) NFE, achieving a state-of-the the-thh-rrt result (FID 1.37) CIFAR-10 conditional generation. For fair comparison in terms of NFE, we select CTM (Kim et al., 2024) due to its popularity, strong performance, and use of the same model architecture based on EDM and adversarial training with StyleGAN-XI. discriminator. We then compute FID using 5 and 6 NFE on CIFAR- 10 and ImageNet-64. As shown in Table 5, increasing the NFE of CTM does not significantly reduce the FIDD, and in some cases, the FID even increases. This indicates that MAC provides additional performance gains that cannot be achieved by distillation or vector field tuning methods alone, even with increased NFE.
Analysis of Trained MAC To analyze how MAC is trained, we plot t-SNE embeddings of four different si- nusoidal weights w across multiple datasets, as shown in Figure 8. Notably, the trained we diverges from weights randomly sampled from the predefined hypothesis set and is far from a. This suggests that during adversarial train- ing, Yo adaptively identifies optimal coefficients without heavily depending on the pre-trained distribution of y Th- Interestingly, Yo exhibits sparser distribution in CIFAR- 10 conditional generation than in the unconditional setting, with nearly identical ws values for samples sharing the same label condition c. This indicates that the optimality of the coefficient depends more on the label condition e than on the starting point Ir of the differential equation.
To empirically validate the benefit of MAC, we conduct ablation studies by training either & individually. As presented in Table 6, the results reveal that jointly training 0 and yields the best performance. Figure 6 further illus- trates that FID decreases more rapidly during joint training compared to training alone. These findings suggest that performance improvements stem not only from the adver- sarial training of Hg. but also from the combined training of both Hg and 7a. demonstrating the benefits of MAC.
8

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Training Efficiency Our method demonstrates notable training efficiency despite incorporating simulation-based training, As shown in Table 7, the number of train- ing images required by our approach is significantly lower across all approxis-- 20 to 2000 times fewer than distillation-based methods such as CTM. Training times are 10. 2, 6, and hours for CIFAR-10, FFHQ, AFHQ, and ImageNet, respectively, which are also comparatively low. The primary AFHQ
Impact Statement
Our research introduces the Multidimensional Adaptive Co- efficient (MAC), enhancing the generative quality and ef- ficiency of flow and diffusion models. The broader impli- cations of this advancement are twofold. On the positive side, MAC improves the accessibility and practicality of generative models, potentially benefiting diverse applica- tions such as medical imaging, drug discovery, and creative industries by reducing computational costs and enhancing output quality. However, improved generative technologies also carry potential risks, including the misuse in creating deepfakes or synthetic biological threats. It is essential for future work to develop robust safeguards and detection methods to mitigate such harmful applications. Continued interdisciplinary efforts in governance, ethics, and technical controls will be critical to ensuring the beneficial integration of our contributions into society.
Table kimg (L) for adversarial training.
Dataset
CTMa + adv 0 Ours
CIFAR-10
25.6k
1.5k
FFH
38.4k
1.0k
51.2k
1.0k
cost of simulation ImageNet
61.4k
30
dynamics arises from VRAM requirements. However, this remains practically feasible, as strong performance can be achieved with just 5 (+) NFE. These results not only highlight the effectiveness of simulation- but also showcase the strength of combining simulation-free and simulation-bast based methodologies- -eeerragggg the advantages of each while mitigating their limitations.
optimality,
Acknowledgement
This work was supported by Institute of Information com- munications Technology Planning & Evaluation (IITP) grant funded by the Korea governmen((MSIT) [NO.RS-2021- I1211343, 1343, Artificial Intelligence Graduate School Program (Seoul National University)]
5. Conclusions
We have introduced the Multidimensional Adaptive Coef- ficient (MAC), a plug-in module for flow- and diffusion- based generative models that enables both freedom of di- mensionality and adaptability to different inference trajecto- pros-propesties previously found only in simulation-basee generative modeling. Our method improves generative per- formance across various datasets and frameworks with high training efficiency. These results highlight the potential of in- ference trajectory optimization with MAC via simulation a direction that has been underexplored compared to vector field design. We suggest broadening the notion of trajectory optimality beyond predefined criteria such as straightness, toward a more general view based on final transportation quality, We hope our work encourages further exploration.
References
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.0879) 2023.
Albergo, M. S., Boffi, N. M., Lindsey, M., and Vanden- Eijnden, E. Multimarginal generative modeling with stochastic interpolants. In The Twelfth International Con- ference on Learning Representations, 2024.
Bartosh, G., Vetrov, D., and Naesseth, C. A. Neural flow dif- fusion models: Learnable forward process for improved diffusion modelling. In Globerson, A., Mackey, L., Bel- grave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Sys- tems, volume 37, pp. 73952-73985. Curran Associates, Inc., 2024.
6. Future Directions
First, the current implementation Th is tied to a fixed NFE during trajectory optimization. Future work could explore adaptive control of NFE based on inference trajectories. Second, while To currently conditions only on ær, future extensions could, with appropriate design, condition on tra- jectories x over diverse inference times and incorporate other contextual signals. Third, optimizing MAC under al- ternative objectives beyond distributional divergence may broaden its applicability to domains where task-specific cri- teria are more relevant. Finally, interpreting MAC selection as a form of trajectory-level policy learning may inspire novel approaches, particularly when combined with ideas from control theory or reinforcement learning.
Berthelot, D., Autef, A., Lin, J., Yap, D. A., Zhai, S., Hu, S., Zheng, D., Talbott, W., and Gu, E. Tract: Denoising diffusion models with transitive closure -istiiiiillattin. arXiv preprint arXiv:2303.00 04248, 2023.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duve- naud, D. K. Neural ordinary differential equations. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Choi, Y., Uh Y., Jo, and Ha, J.-W. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Karras, T., Aittala, M., Hellsten, J., Laine, S., Lehtinen, J., and Aila, T. Training generative adversarial networks with limited data. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12104- 12114. Curran Associates, Inc., 2020.
Deng, J. Dong, W., Socher, R., L.-J., K., and Fei-Fei, L.. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion- based generative models. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), ), Advances in Neural Informa- tion Processing Systems, volume 35, pp. 26565-26577. Curran Associates, Inc., 2022.
Dhariwal, and Nichol, A. Diffusion models beat gans image synthesis. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Ad- vances in Neural Information Processing Systems, vol- ume 34, pp. 8780-8794. Curran Associates, Inc., 2021.
Karras, T., Aittala, M., Lehtinen, J., Hellsten, J., Aila, T., and Laine, S. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 24174-24184, June 2024.
Geng, Z., Pokle, A., and Kolter, J. Z. One-step diffusion distillation via deep equilibrium models. In Oh, A., Nau- mann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems, volume 36, pp. 41914 41931. Curran Associates, Inc., 2023.
Kim, B., Hsieh, Y.GG,, Klein, M., marco cuturi, Ye, J. C., Kawar, B., and Thortoon, J. Simple reflow: Improved techniques for fast flow models. In The Thirteenth Inter- national Conference on Learning Representations, 2025.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., and Weinberger, K. (eds.), Advances in Neural Information Processing Systems, vol- ume 27, Curran Associates, Inc., 2014.
Kim, D., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., Ue- saka, T., He, Y., Mitsufuji, Y., and Ermon, S. Consistency trajectory models: Learning probability flow ODE trajec- tory of diffusion. In The Twelfth International Conference on Learning Representations, 2024.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by two time-scale update rule converge to local nash equilibrium. In Guyon, I,, Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vish- wanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30, Curran As- sociates, Inc 2017.
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Lee, S., Kim, B., and Ye, J. C. Minimizing trajectory curvature of bDEed generative models. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, (eds.), Proceedings of the 40th Interna- tional Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 18957- 18973. PMLR, 23-29 Jul 2023.
Ho. J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840- 6851. Curran Associates, Inc., 2020.
Kapuśniak, K., Potaptchik, P., Reu, T., Zhang, L., Tong, A., Bronstein, M., Bose, A. J., and Di Giovanni, Metric flow matching for smooth interpolations on the data mani- fold. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 135011-13500 Curran Associates, Inc., 2024.
Lee, S., Lin, Z., and Fanti, G. Improving the training of rectified flows. In Globerson, A., Mackey, L., Belgrave, D.,, Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, vol- ume 37, pp. 00082-63199, Curran Associates, Inc., 2024.
Lim, J. H. and Ye, J. C. Geometric gan, 2017.
Karras, T., Laine, S., and Aila, T. A style-based genera- tor architecture for generative adversarial networks. In Proceedings of the JEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), June 2019.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and
Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
10

---

Multidimensional Adaptive Coefficient for Inferenee Trajectory Optimization in Flow and Diffusion
Liu, X., Gong, C.,, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations. 2023.
Proceedings of the 32nd International Conference on Ma- chine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2256-2265, Lille, France, 07-09 Jul 2015. PMLR.
Lu, H., Lu, Y., Jiang, D., Szabados, S. R., Sun, S., and Yu, Y. CM-GAN: Stabilizing GAN training with consis- tency models. In ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
Song. Y. and Dhariwal, P. Improved techniques for train- ing consistency models. In The Twelfth International Conference on Learning Representations, 2024.
Song, Y., Sohl-Dickstein, J., Kingma, P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representationss 2021.
Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z.. and Zhang, Z. Diff-instruct: A universal approach for transferring knowledge from pre-trained diffusion models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- Curran
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Con- sistency models. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, (eds.), Pro- ceedings of the 40th International Conference Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 3221 1-2255. PMLR, 23-29 Jul 2023.
cessing Systems, volume 36, pp. Associates, Inc., 2023.
Oquab, M., Darcet, T., Moutakanni, T., Vo, H. V., Szafraniec, M., Khalidov, V., Fernandez, P., HAZIZA, D., Massa, F., El-Nouby, A., Assran, M., Ballas, N., Galuba, W., Howes, R., Huang, P.-Y., Li, S.-W., Misra, I. Rabbat, M., Sharma, V., Synnaeve, G., Xu, H., Jegou, H., Mairal, J., Labatut, P., Joulin, A.. and Bojanowski, P. DINOv2: Learning robust visual features without su- pervision. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. Featured Certification.
Tong, A., FATRAS, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improv- ing and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. Expert Cer- tification.
Park, D., Lee, S., Kim, S., Lee, T., Hong, Y., and Kim, H. Constant acceleration flow. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Pro- cessing Systems, volume 37, pp. 90030-90060. Curran Associates, Inc., 2024.
Wang. Z., Zheng, H., He, P., Chen. W., and Zhou, M. Diffusion-GAN: Training GANs with diffusion. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Xu, Y., Zhao, Y., Xiao, Z., and Hou, T. Ufogen: You forward once large scale text-to-image generation via diffusion gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), pp. 8 8196-8206, June 2024.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted intervention-MICCAI 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234-241. Springer, 2015.
Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F,, Freeman, W. T., and Park, T. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6613-6622 2024.
Sauer, A., Schwarz, K., and Geiger, A. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 Conference Proceedings, SIGGRAPH '22, New York, NY, USA, 2022. Association for Computing Ma- chinery. ISBN 9781450393379. doi: 10. 10.1145/352823. 3530738.
Singhal, R., Goldstein, M., and Ranganath, R. Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions. In The Eleventh International Conference on Learning Representations, 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. Bach, F. and Blei, D. (eds.),

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
Outline of Appendix
Appendix provides an overview of flow and diffusion frameworks, Appendix B provides justification for the use of diagonal matrices in the design of Yg. Appendix C describes the hypothesis set for MAC. Appendix D presents details of the differential equation solver Gg.ø. Appendix E outlines experimental setups, including datasets, training configurations, and additional results. Appendix F explains evaluation metrics, and Appendix G provides visualizations of generated samples and inference trajectories.
A. Flow and Diffusion Frameworks
During inference, all frameworks transport samples between pr and Po. either in the forward or reverse direction.
Denoising Diffusion Probabilistic Models (DDPM) DDPM (Ho et al., 2020) is one of the most widely used diffusion- based generative frameworks. Given zn ~ Po = No I) and æ1 ~ Pr as data, the VP diffusion coefficient is:
[V 1 b1-1
(t) = [au((t), a1)))
be = - Sá c(s)ds c(s) = Cuin + s(Cmax - Cmin): T=1, (14)
where Cmin 0.1 and Cmax = 20. DDPM minimizes the loss La Et, [|H)(t, xo zolß] where t ~ L((0,)
Flow Matching (FM) FM (Lipman et al., 2023) introduces simple yet effective objective for generation. Given To ~ [(,, and 21 ~ data, it employs the following coefficient:
a(t) [co an) [1- (1- mmin t), T= 1.
(15)
FM directly models using the objective Le Er,zo.. [1|H)(t, (x1 (1- - )=)) where t ((0,
Elucidating Diffusion Model (EDM) EDM (Karras et al., 2022) refines and stabilizes diffusion model training. Given To ~> Po as data and œ P1 = (0, I), the coefficient is defined as:
(() = [1)((), on (4) = [1,0, 4. T 80.
(16)
The loss function is:
La E2,20.21 A(1)cou(1)2 Fo ((,
Cout (1) (wo csap (1)x((())
(17)
2
where:
Važ(t) Tinn Cout (t) (1(t) Odaa
1
Cin(t)
Cakip(t) af(() + 2
data
Vohun + ai())
(18)
at + J2
Cuaise (1)
1 Int, X(t) (ai(1) -data)"
with sampled from In(t) N.-1,2,1.29) and Odata 0.5. EDM models Ho x()) Coat (o)Fa (4), Gin (()x(()) = čn,a- During inference, EDM transports from pr (0 72/) to po.
(sai(t)):(t)
+
Stochastic Interpolant (SI) SI (Albergo al., 2023) facilitates transportation between arbitrary distributions Pu and P1.
The conventional coefficient design is:
(+t) (oo(t), t, T=
(19)
SI models (l(1,x((1) = [oo (t), = [ $1.)). The loss function is:
Cx.0 1(" Ero.er [1
2zz H (t,.(t,x((()) dt, = t~ 1).
(20)
12

---

Multidimensional Adaptive Coefficient for Inferenee Trajectory Optimization in Flow and Diffusion
B. Justification for Diagonal Matrices in MAC
We argue that using diagonal matrices is sufficient to represent non-linear trajectories for inference trajectory optimization. Theoretically, the expressible space of the interpolated value x(t) remains rr() E R for both the Hadamard product and matrix-vector multiplication. While the primary benefit of full matrix multiplication lies in its ability to model non-linear trajectories by incorporating cross-dieeniiona effects, contend that diagonal matrix is already adequate for this purpose-pree provided that the U-Net U/. used to parameterize Tp. possesses enough expressive power to generate diverse non-linear trajectories.
Moreover, using diagonal matrices for the output of MAC significantly reduces model size compared to using full matrices. One alternative is to use low-rank- tat-rran--lu--dagoon decomposition, Y D + UVT, where D y is a diagonal matrix, and U. V E TR with r being typically small rank (e.g.. r or 8). However, even with a small r. this approach would still require large number of output channels from Ug. significantly increasing model complexity. Specifically, Up would require output channel size of C x (2r + 1) M x 2. where C' is the image channel size and M is the number of sinusoidal functions. For instance, with r 4, this results in 3 x 9 x 5 x 2 = 270 output channels for Up when NFF 5. compared to only C' x M x 2 30 channels in the diagonal case.
In conclusion, we believe that using a diagonal matrix, in conjunction with a sufficiently expressive U-Net Us, provides ample capacity to model non-linear trajectories while keeping the model size efficient and manageable.
C. Hypothesis Set Design for MAC
We design the hypothesis set and parameterization for Ta as follows:
(1.. 2T)] [0, TJ x Rd Rdxa
where
ft
(t, Fo(bm (t), (o(T)) T fo XT) + 96 (T,
T
(21)
9 (t,
(1, XT) (bm (x(T T [T(1, + go (1, IT)
This parameterization can vary depending the flow and diffusion framework. For example, we use Ya described above for SI, but set (t, aT) la for EDM to align with its original formulation. The functions So and Pø are defined as:
M
2
M
2
folt, = T
+
Σ
w a (mot)mm())))
go(t,
t + T
Σ m=1
u mi
(0 (t)
(22)
m=l
1/9
bm (f) = sin
am
wo(xr) = LPF o tanh (NN)(IT)), ER,
E R,
s SER,
(23)
T
where M and q are hyperparameters for the sinusoidal basis, and the scaling factor s controls the output range of ws (xT) (-8,8). When s 0.0, y reduces a. For image generation, we use a U-Net Uo NNę. The outputs UĄ(TT) (UnolZT), U (xT) E XxCXXXXXM* where C, H, and W denote the channel, height, and width of the image. (1T(())) and U (xT) each have channel shape of C' x M, structured along the channel axis.
Design Choice of Sinusoidals
The inference time schedule of EDM is defined as:
('min* 7 /maxx +))² mmin 0.002, 4max 80, q q= 7.
4 =
/maxx
+ N
(24)
As illustrated in Figure 9, (t) sin(Tm(t/TT)//) effectively covers the entire EDM time schedule, whereas (t, (t) sin(Tm((/T)) has minimal magnitude when L < 1. Since wa is constrained to (-s, the choice of sinusoidal basis significantly influences the controllability of ys during simulation. Accordingly, we use q = 1 for DDPM, FM, and SI, and = 7 for EDM.
In addition, to avoid aliasing in by (t), we follow the Nyquist criterion, requiring the sampling rate f. = N 2 2f = M. Based on this, we set M- = N during adversarial training to ensure sufficient frequency resolution.
13

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
O
0.1
0.01
f
10
1
0.1
0.01
(a) sin sin(xm(t//T)
(b) sin(Tm((/T))//)
Figure 9: Comparison bm (t) for M = 10. The dotted line represents the EDM inference time schedule for N = 10.
Low-Pass Filtering (LPF) For LPF, we apply 2D convolution with a Gaussian kernel, where the kernel size is 20xresolution - and the standard deviation of the Gaussian is 1.0xresolution Here, resolution refers to the image height width. To mitigate boundary effects introduced by LPF, we apply zero-padding of kemel size+1 on all sides of Ug 's input and crop the borders after convolution restore the original shape. To preserve consistent scaling, compute the minimum and maximum values before LPF for each batch and rescale the post-LPF output to match the original range. LPF is applied only for image generation experiments.
Hypothesis Set for Pre-training and Adversarial Training Since using a large hypothesis set of coefficients during pre-training may overburden Hg and degrade performance, we adopt a smaller hypothesis set for pre-training and enable full multidimensionality across L only during adversarial training. For Section 4.2, LPF used exclusively during pre- training and omitted during adversarial training. For Section 4.3, we configure the convolution group size for I.PF during pre-training I, yielding an LPF output shape of [B, I, resolution, resolution]. This enforces high muliidimensionality for small t and low multidimensionality for large t. In contrast, during adversarial training, we convolution group size of [B, 2 x 3 M, resolution, resolution), producing an LPF output shape of [B, 2 x x M, resolution, resolution].
D. Details of the Differential Equation Solver
The displacement of the trajectory x((++1) z(4),, parameterized by vo.o. expressed as:
Al, Ug (1i, (14), xT) Al (li, aa Ad, in (4i, oo 2 $1,0
(25)
A Yo (t,, 2T) Co.0 An, xT) O $1
where the time displacement At, t++ - t is derived from the inference time schedule T = fto....,ty. tw with to IN = 0. The corresponding displacements for y are: (t, FT) o
70. 2T), At, 2T) (t. zT) - noo (t2, TT))
(26)
This approach helps reduce numerical errors when solving differential equations under curved y. Accordingly, the trajectory displacement for EDM is given by:
(71,0(i, nolti, xT)
2 (x(t:) Ho(t., (o with yettr--taaiinng
At, (t,
(27)
Anoltt
(a(t;) HT)
wlo pree-trinng
For SI:
At 0o (ti, (t)), Ato,o(t, 2T) 20,9 + An, xT) @ $1,0, Holli,2(((4)) 70(4,
(28)
with r-pre-training w/o y-pre-ttrainnng
Co,o, $1.0]
(29)
He(Fs (t:, 2T), xt:))
The differential equation solver Gg, using Euler discretization, is then defined as:
N
CT, ) = (t,((NN) = 2T +
Σ
Al. vo (li, I0. (0
(30)
1=0
(t.((t+11) +- (t.((ts) (t+)+ + At; vo. (t., 2 (ts),
14

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
E. Details for Experiments
All experiments are conducted on NVIDIA RTX 3080Ti, RTX 4090, and RTX A6000 GPUs. Euler discretization used for all inference procedures.
E.1. 2-Dimensional Transportation: Optimizing Only e
E.1.1. P-PRE-T --E------EINIIEEEE TRAINING
We follow the implementation from Tong et al. (2024), employing a multilayer perceptron (MLP) consisting of four linear layers with 64 hidden units and SiLU activation functions. The Stochastic Interpolant (SI) model is trained with batch size of 256 for 20,000 iterations. The loss function is defined
(o) = Feoiiii [1 w))² 2nd
7(1 di, == 0,1, t t~14 1), 1) E RxxMx, (31)
E.1.2. OPTIMIZATION OF Φ
To train "e. we use a batch size of 1024 and run for 2000 iterations with 8 = 0.1. Each configuration is trained with three different random seeds, and we report the mean and standard deviation of the Wasserstein distance.
E.2. Image Generation: Optimizing Only
E. TRAINING CONFIGURATIONS
Table 8: U-Net configurations for Ho-
Configuration
CIFAR-I0
ImageNet-32
Channels
128
256
Depth
2
3
Channel multipliers
1,2,2,
1,2,2,
Attention heads
4
4
Head channel size
64
64
Attention resolution
16
16
Dropout
0.1
0.1
We use the U-Net architecture from Dhariwal & Nichol (2021) for Hg and the U-Net from Ronneberger et al. (2015) for Ug. Configuration details for Ho given in Table 8. For Yo. use channel progression of [256, 512, 1024, 2048). The discriminator Da consists of four convolutional layers with 1024 channels, batch normalization, leaky ReL.U activations, and final sigmoid output.
Table 9: Hyperparameters for pre-training and adversarial training.
Hyperparameter
CIFAR-10
ImageNet-32
Pre-training
Adversarial training
Pre-training
Adversarial training
Batch size
128
16
512
15
GPUs
1
4
Iterations
400k
200k
250k
200k
Peak LR
2e-4
2e-4
2e-4
2e-4
LR Scheduler
Poly decay 5k
Poly decay
Poly decay 5k
Poly decay
Warmup steps
5k
5k
Warmup steps for De
20k
20k
We use the Adam optimizer with 3 = 0.9, B₂ = 0.999, weight decay of 0.0, and e = x 10 along with polynomial
decay learning rate scheduling. An exponential moving average (EMA) with a decay rate of 0.999 is applied during all training phases. For adversarial training, we use the vamilla GAN loss from Equation 5 instead of hinge loss for simplicity. FID evaluated every 10,000 steps, and we report the lowest value obtained. Hyperparameters are summarized in Table
15

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
E.2.2. EXPERIMENTS FOR HYPOTHESIS SET HYPERPARAMETER TUNING
Table 10: FID comparison between a and y under various configurations (c.g., 8, LPF) across different NFE.
CIFAR-10 10
ImageNet-32
Method \ NFE
10
100
150
200
10
100
150
200
SIa
14.43
4.75
4.51
4.30
17.72
8.08
7.79
7.63
SI (s=00000)
14.59
3.98
3.74
3.63
17.41
6.33
6.21
6.20
ST. (s=0.1
15.44
3.77
3.68
3.75
17.86
6.63
6.47
6.44
FM,
13.70
4.52
4.23
4.07
16.92
7.78
7.53
7.38
(s=0000))
13.81
3.59
3.42
3.42
16.85
6.18
6.03
6.01
(e=00, LPF)
15.13
3.64
3.57
3.64
17.52
6.40
6.27
6.31
DDPM,
98.47
6.64
4.84
4.10
111.54
8.13
7.40
7.14
DDPM (e=0.005)
74.44
3.77
5.96
7.84
139.69
7.67
12.37
11.70
DDPM DDPM,( (1=00055,LL)))
72.23
4.73
4.11
3.83
135.48
6.84
6.51
6.42
DDPM (a=0.11 LPF)
71.80
4.46
6.32
12.60
142.99
6.70
8.69
10.91
Table 11: FID for different NFE and Gaussian kernel a in LPF, evaluated on CIFAR-10 with SI (1=00.LL.P)))
NFE
10
20
30
40
50
100)
150
200
0.1
14.89
8.05
6.49
5.45
6.53
9.59
10.67
11.20
1.0
14.92
8.47
5.32
4.45
4.68
6.06
7.01
7.50
2.0
16.25
9.56
7.56
6.06
4.72
3.77
3.95
4.17
4.0
15.44
9.13
7.39
6.36
5.59
3.77
3.68
3.75
Table 12: FID after adversarial training with 10 NFE
CIFAR-10 for varying values of M.
Method .M
15
20
25
30
(s=00 LPF)
6.89
4.14
4.42
5.32
6.11
5.74
FM LPF)
5.93
6.13
6.70
6.18
5.97
6.42
(s LPF)
10.15
10.04
9.60
9.04
8.94
9.19
Table 13: FID) after adversarial training of SI under various configurations on CIFAR-10 with 10 NFE.
Method M
5
10
15
20
0.0
10.20
9.75
11.30
11.53
0.005
6.60
4.79
4.45
5.28
0.005, LPF
7.37
4.42
4.26
5.31
0.1, LPF
7.21
4.14
5.59
5.32
We conduct extensive experiments to tune hyperparameters for the hypothesis set. Tables 10 and 11 show the effects of different coefficients, LPF configurations, and other settings during pre-train pre-training. Tables 12 and 13 present results from adversarial training under various configurations. These results help identify optimal hyperparameter settings that enhance model performance both pre-training and adversarial training stages.
16

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
E.3. Image Generation: Optimizing / and
E.3.1. (-PRE-TRAININ
Table 14: Hyperparameters used for pre-training EDM.
Hyperparameter
CIFAR-10
FFHQ & AFHQ
Number of GPUs
8
Duration (Mimg)
200
200
Minibatch size
512
256
Learning rate
le-3
2c-4
LR ramp-up (Mimg)
10
10
EMA half-life (Ming)
0.5
0.5
Dropout probability
13%
5% (FFHQ) / 25% (AFHQ)
Channel multiplier
128
128
Channels per resolution
2-2-2
1-2-2-2
Augment probability
12%
15%
M
10
Low- pass filtering
True 0.05
True
s
0.05
We adopt the codebase from Karras et al. (2022) and follow the EDM configuration, replacing a with y and applying coefficient conditioning. The pre-training loss function is:
Ce Ee,
A(t, u)couutt, u)² Fo (Cooise (t), cin u)x((),, Cooett(t, u))
Cour (t, u) (xo Cakip(t, w)x(t))
(32)
where:
Coultt, u) 7/2 (t. u) Odata
Cin(t, u) = V7H u) + o
V + Tit Cakip(t, 21)
data
(t, u)
data
(33)
Croise (1) = Int, Cover((,, = (t, A(e,u) = 7f(t, u) + o2 oduaa.
(1, daaaa) )2
with In(1) [ 1.22), ((-1,1) RdxMxx and ddata
0.5. For coefficient conditioning (Section 3.2.2),
concatenate [Cin (t, ut)((t), Cooeff (t, u)] along the channel axis input to Fg. We use the Adam optimizer with B1 0.9. $2 = 0.999, and = 1e-8.
E.3.2. ADVERSARIAL TRAINING
Table 15: Hyperparameters used for adversarial training.
Hyperparameter
CIFAR-10
FFHQ AFHQ
ImageNet
Number of GPUs
Training duration for 7h (kimg)
1500
1000
30
Minibatch size for Ui
512
256
256
Minibatch size for %
128
64
32
Minibatch size for De Learning rate for Up Learning rate for Yo Learning rate for D.
128
64
512
le-5
le-5
le-5
lc-4
Ic-4
le-4
le-3
le-3
le-3
EMA half-life (kimg)
10
10
10
Low-pass filtering
True
True
True
0.05
0.05
0.05
17

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
For Ug. we adopt U-Net architecture based on Song et al. (2021) with the following configuration: 256 channels, channel multipliers [1, 4]. dimensionality multiplier of 4 blocks, and attention resolution of 16. The embedding layer for & is disabled. We disable dropout in both He and To to ensure deterministic behavior of Ga.ø. We use the Adam optimizer with 31 = 0.0, 32 0.99, and é c 1e-8. The sampling time t is drawn from In(t) (-122, 1.22) and is quantized according to the inference time schedule r. For ablation studies, each configuration is trained for 500 kimg, corresponding to approximately 4000 iterations. When training only & (Table 6, Figure 7), LPF is not applied.
E.3.3. OPTIMIZED INFERENCE TRAJECTORIES DEVIATE FROM LINEARITY
CIFAR-10 uncond.
CIFAR 10 cond.
40
FFHQ
AFHQQ2
20
t
10
0.1
0.01
Figure 10: L2 norm between the linear trajectory Min (1) and the optimized trajectory (t) over t.
To confirm that the optimized inference trajectories in the image generation experiment are not linear lines, we compute the L2 norm between linear trajectories Tim (t) f. + (1 4) (t. ) and the optimized trajectories (t)) As shown in Figure 10, the optimized inference trajectories clearly deviate from the linear path. This result demonstrates that our method effectively discovers superior, non-linear trajectories in nigh-dimensional spaces, leading to enhanced performance.
F. Metrics Calculation
For 10, AFHQ, and FFHQ experiments, we follow the evaluation protocol and code provided by Karras et al. (2022) to compute the Fréchet Inception Distance (FID). For ImageNet-64 experiments, we use the ImageNet dataset from EDM (Karras et al., 2022) as the reference and calculate both FID and FDDINOv2 using the code from EDM2 (Karras et al., 2024). All metrics are computed over 50,000 generated samples. Each experiment is run three times with different random seeds, and we report the minimum FID and DDNOOv values observed.
G. Generated Samples
Figure 11: Inference trajectories from EDM, + adv 0, Φ for conditional generation on CIFAR- 10.
18

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
FID 1.98
NFE 35
FID 1.69
NFE 5(+)
CIFAR-I 10 unconditional.
FID 1.79
NFE 35
FID 1.37
NFE 5(+)
CIFAR-I0 10 conditional.
FID 2.39
NFE 79
FID 2.27
NFE 5(+)
AFHQ-64.
FID 1.96
NFE 35
FID 2.04
NFE 5(+)
FFHQ-64
Figure 12: Generated samples CIFAR-10, FFHQ-64, and AFHQ-64 using EDM,, (left) and EDM, + adv 0,6 (right).
19

---

Multidimensional Adaptive Coefficient for Inference Trajectory Optimization in Flow and Diffusion
FID 1.48
FDDINOV2 70.2
NFE 5(+)
Figure 13: Generated samples ImageNet-64 using EDMg adv 0.
20


Reviews:
[{"review_id": "tXnPRn4xM7", "reviewer": "Reviewer_MR8r", "summary": "This paper introduces a new way to handle the interpolation between noise and data in diffusion and flow-based generative models. Unlike standard approaches that use an interpolation with a uniform scale across the entire image (like Rectified Flows, DDPM, and IDDPM), the authors propose extending this interpolation to a 2D space, matching the dimensions of the image and noise. This 2D interpolation is learned using a coefficient model, offering potentially more flexibility in the interpolation process. The training procedure involves three main components: a standard diffusion model trained with a 2D sampling of the time variable, a discriminator, and the coefficient model, all optimized jointly in the second stage to find better interpolation weights.\n\n## update after rebuttal\n\nAs I didn't receive any rebuttal from the author I'll remain my current rating unchanged. \n\nAfter reviewing the author's rebuttal with other reviewers, I believe most of my concerns remain unaddressed, and the paper, especially the experiment section, still appears incomplete and not yet ready from my perspective.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "The central claim is that the proposed 2D learned interpolation offers more flexibility and can lead to improved generative models. The paper presents experimental results on CIFAR-10 to support this claim. However, the evidence might be considered somewhat limited due to the restriction to a single, relatively low-resolution dataset. The lack of thorough ablation studies on the design choices of the coefficient parameterization also weakens the evidence for the specific design being optimal.", "methods_and_evaluation_criteria": "The proposed method, involving the joint training of a diffusion/flow model, a discriminator, and a coefficient model with a 2D time sampling, presents an interesting alternative to standard interpolation techniques. The evaluation seems to primarily rely on FID scores obtained on the CIFAR-10 dataset. While FID is a common metric for generative model evaluation, the limited scope of the experiments raises questions about the generalizability of the findings.", "theoretical_claims": "This paper doesn't include new proofs or theoretical claims", "experimental_designs_or_analyses": "The experimental design and analyses raise several concerns:\n\n- Limited Dataset: The evaluation is limited to CIFAR-10 (32x32 images), which is insufficient to demonstrate the effectiveness of the proposed method on more complex and higher-resolution datasets. I suspect that the increased dimensionality of the 2D time sampling might lead to optimization difficulties as the image size grows.\n- Lack of Ablation: The absence of ablation studies on the design choices for the coefficient parameterization is a significant oversight. It's quite hand-crafted in terms of the specific parameterization in equation 8, and whether other designs might yield better results.\n- Inconsistent Results: There are discrepancies in the reported performance across different tables (e.g., Table 2, 3, and 4 compared to Table 5), and it's unclear if different models or settings were used. The comparison in Tables 2-4 is limited to NFE=10, which is unusually low for standard diffusion and flow models to achieve good performance.\n- Missing Baselines: Table 5 lacks many baseline results, making it difficult to properly assess the improvement offered by the proposed method. Furthermore, the comparison only includes distillation methods and inexplicably omits 1-Rectified Flow (Flow Matching), which seems like a relevant baseline.", "relation_to_broader_scientific_literature": "Previous work like rectified flow (flow matching), DDPM or other variants mainly focus on using a predefined and fixed intepolation. This work extend it using a learned intepolation and expand it to 2D", "essential_references_not_discussed": "The paper include the related references.", "other_strengths_and_weaknesses": "Strengths: The paper proposes a novel idea of using a learned 2D interpolation for diffusion and flow models, which could potentially offer more flexibility than standard interpolation.\n\nWeaknesses: Please see the questions below", "other_comments_or_suggestions": "Please see the questions part.", "questions_for_authors": "1. In the introduction, you mention a connection to Neural ODEs. Could you please elaborate on the nature of this relationship and how your work builds upon or relates to concepts from Neural ODEs?\n2. Why was the specific parameterization for the coefficient model chosen as described in equation 8? Were other parameterizations explored, and if so, what were the results? An ablation study on this design choice would be beneficial.\n3. The experimental evaluation is currently limited to CIFAR-10. Do you have plans to extend the evaluation to higher-resolution datasets like Imagenet 64x64 or include more comprehensive comparisons on datasets like FFHQ or AFHQ? Given the 2D time sampling, are there potential scalability challenges with increasing image size? \n4. In equation 5, x_T is used but not defined. Similarly, in Algorithm 1, ρ_T is used without definition. What are these variables, and what is the difference between ρ_1 and ρ_T?\n5. From Section 3.2.2 it seems that the method is currently applied to diffusion models predicting noise. Could you clarify if this approach can be extended to other types of diffusion and flow models, such as those predicting velocities?\n6. Tables 2, 3, and 4 present results with NFE=10, which is unusually low for standard DDPM and Flow Matching methods. The performance also differs from Table 5. Were different models or training settings used in these tables? Please clarify these discrepancies.\n7. Table 5 is missing many baseline results, including 1-Rectified Flow (Flow Matching). Could you provide these missing results for a more comprehensive comparison? \n8. What is the primary motivation for using adversarial training in your method? Could you compare the performance of your method with and without the adversarial loss to isolate the impact of the proposed interpolation technique?  It's unclear whether the performance gains are solely due to the proposed interpolation method or if the adversarial training plays a significant role, especially in terms of FID improvement.\n9. In Table 6, the FID scores for the first two rows (presumably vanilla EDM) are significantly higher than the EDM results in Table 5. Could you explain this difference? I assume the first row is vanilla EDM so it should match table 5 row 3 in diffusion models?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "stAyagdqgp", "reviewer": "Reviewer_TFVT", "summary": "This paper introduces Multidimensional Adaptive Coefficients (MAC) for flow and diffusion models, allowing coefficients to vary across dimensions and adapt to different starting points. The two-stage approach combines pre-training with multidimensional coefficients and adversarial refinement. Experiments across multiple frameworks and datasets show improved performance, including state-of-the-art results on CIFAR-10 conditional generation.\n\n## update after rebuttal\nThe author resolved my confusion. However, considering that reviewer MR8r's review was not addressed, I'm unsure what happened. Finally, considering the contribution of the current work, my current rating remains unchanged.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "URL for Additional Figures: https://imgur.com/a/3UiYDVF\n\n$\\textbf{[A1] MAC’s core value lies in inference-time planning, orthogonal to vector field tuning methods like CTM and CD+GAN}$\n\nCTM 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.\n\nAs 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.\n\nCTM 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.\n\nMAC 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.\n\n$\\textbf{[A2] MAC Reduces Hyperparameter Engineering Costs}$\n\nThe 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.\n\nMAC 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.\n\nRegarding 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.\n\n$\\textbf{[A3] Practicality via Existing Pre-trained Models}$\n\nOur 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.\n\nUsing $\\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.\n\nEmpirical 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.\n\nThus, 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.\n\n$\\textbf{[A4] Scalability (ImageNet)}$\n\nPlease refer to our response to Reviewer jRK3’s comment [A2].\n\nWe sincerely appreciate your thoughtful consideration and hope these clarifications effectively address your concerns.", "claims_and_evidence": "The performance improvement claims are generally supported by comprehensive experiments across multiple frameworks and datasets.\n \nHowever, the \"training efficiency\" claim is problematic since the method requires training models from scratch rather than fine-tuning existing pre-trained models, which is a significant practical limitation.", "methods_and_evaluation_criteria": "The evaluation metrics and benchmark datasets are appropriate for the task.\n\nHowever, there is a concern with the methodology. The hypothesis space design for MAC introduces many additional hyperparameters (M, s, q, LPF configuration). The paper lacks analysis of how baseline methods might perform if augmented with similar parameterization advantages. For example, it would be valuable to see how CTM would perform if pre-trained with comparable additional parameters before distillation.", "theoretical_claims": "There is no theoretical claims in this paper.", "experimental_designs_or_analyses": "The experimental design is overall sound, with comprehensive testing across multiple frameworks and datasets. \n\nThe ablation study in Table 6 shows that adversarial training provides most of the performance gain. Since adversarial training is not a novel contribution of this paper, this suggests that MAC's specific contribution to performance improvement is relatively limited.", "relation_to_broader_scientific_literature": "The paper positions itself as novel by optimizing trajectory quality through simulation, but this approach is conceptually similar to:\n\n1. Kim et al. (2024) (CTM) which also uses adversarial refinement on diffusion trajectories\n2. Lu et al. (2024) (CD + GAN) which likewise employs adversarial refinement\n\nThe primary difference, multidimensional coefficients, appears to provide only incremental benefits over standard adversarial refinement of model parameters.", "essential_references_not_discussed": "I find that the existing citations sufficiently cover the related literature, though the novelty of the proposed approach remains limited.", "other_strengths_and_weaknesses": "**Strengths**:\n\n1. Comprehensive experiments across multiple frameworks demonstrate the method's versatility\n\n2. Good visualizations of learned trajectories help explain the method's operation\n\n3. The approach achieves state-of-the-art results on CIFAR-10 conditional generation\n\n**Weaknesses**:\n\n1. The requirement to retrain models from scratch makes the method impractical for large-scale applications\n\n2. The performance gains attributable specifically to MAC (versus general adversarial refinement) appear modest\n\n3. The method introduces significant additional complexity through MAC hyperparameters", "other_comments_or_suggestions": "The paper would benefit from providing a more detailed analysis of computational overhead. Additionally, testing on higher-resolution images would help demonstrate the method's scalability to more complex generation tasks.", "questions_for_authors": "When considering the total computation (pre-training + adversarial finetuning), how does your method compare to these alternatives?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xd4qWFs0Ot", "reviewer": "Reviewer_jRK3", "summary": "This paper proposes to optimize model parameters and forward process interpolation coefficients with respect to a simulation-based adversarial loss. The training process consists of two stages. In the first stage, the flow model is trained w.r.t. randomly sampled multi-dimensional interpolation schemes via flow matching. To constrain the hypothesis space of interpolation schemes, interpolation coefficients are parameterized as a sum of weighted sinusoidals. In the second stage, the flow model and interpolation coefficients are optimized to minimize a simulation-based adversarial loss. Here, the weights for sinusoidals are parameterized as neural net outputs conditioned on $x_T$, providing an additional level of adaptivity. The authors demonstrate that the proposed method achieves competitive FID scores on CIFAR-10, ImageNet $32 \\times 32$, FFHQ $64 \\times 64$, and AFHQ $64 \\times 64$.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "URL for Additional Figures: https://imgur.com/a/3UiYDVF\n\n$\\textbf{[A1] $H_\\theta$ estimates the vector field and $\\gamma_\\phi$ performs inference-time planning}$\n\nWe 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.\n\nFor 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.\n\n$\\textbf{[A2] Scalability Experiments on ImageNet}$\n\nWe 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.\n\nNevertheless, we emphasize our method’s scalability, supported by the following points:\n\nAs 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.\n\nSpecifically, 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]).\n\nMAC can be readily integrated with these frameworks, and we are currently preparing dedicated large-scale follow-up studies.\n\nFor 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.\n\n$\\textbf{[A3] Essential References Discussion}$\n\nThank 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.\n\n$\\textbf{[A4] Notation Clarifications}$\n\n$G_{\\theta, \\phi}(\\tau, x_T) = x_{\\text{est}, \\theta, \\phi}$ (Details in Appendix C).\n\n$\\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$.\n\n$x_T$ denotes the actual initial sample point, while $z$ is sampled from the same distribution but unused during inference.\n\nWe sincerely appreciate your thoughtful consideration and hope these clarifications effectively address your concerns.\n\n[1] Minimizing Trajectory Curvature of ODE-based Generative Models, ICML, 2023\n\n[2] Constant Acceleration Flow, NeurIPS, 2024\n\n[3] Simple ReFlow: Improved Techniques for Fast Flow Models, ICLR, 2025\n\n[4] High-resolution Image Synthesis with Latent Diffusion Models, CVPR, 2021\n\n[5] Scalable Diffusion Models with Transformers, ICCV, 2023\n\n[6] SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers, ECCV, 2024\n\n[7] [Mask] is All You Need, arXiv, 2025\n\n[8] Inductive Moment Matching, arXiv, 2025", "claims_and_evidence": "- **(Left column, lines 260-262) \"Given that $H_\\theta(t, x(t)) \\approx x_0$, $H_\\theta$ can be adversarially refined by optimizing it with the discriminator $D_\\psi$.\"**\n\nThis claim holds only when generative paths are sufficiently straight. However, as cited below, the authors also claim that the benefit of using MAC arises from its ability to discover better non-linear trajectories.\n\n> (Left column, lines 300-305) \"This suggests that a straight trajectory is not always optimal, even in OT-trained models, and MAC can adaptively discover better trajectories to correct errors that arise during transportation. Figure 3 further illustrates how MAC adjusts the trajectory direction to optimize transportation, resulting in a path that is not straight.\"\n\nHence, the authors' assumption $H_\\theta(t, x(t)) \\approx x_0$ that justifies the usage of adversarial learning on velocity $H_\\theta$ output instead of generator $G_{\\theta,\\phi}$ output seems to be at odds with the benefits of using MAC.", "methods_and_evaluation_criteria": "CIFAR-10, FFHQ-64, AFHQ-64 are standard datasets / FID and IS are standard metrics for evaluating the performance of generative models. However, based on previous works such as [1,2,3], I believe results on at least ImageNet $64 \\times 64$ are necessary to demonstrate the practicality and scalability of the proposed method.\n\n[1] Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion, ICLR, 2024\n\n[2] Constant Acceleration Flow, NeurIPS, 2024\n\n[3] Simple ReFlow: Improved Techniques for Fast Flow Models, ICLR, 2025", "theoretical_claims": "There are no new theoretical claims.", "experimental_designs_or_analyses": "I did not find any issues with experimental designs or analysis.", "relation_to_broader_scientific_literature": "This paper proposes to adversarially optimize forward process interpolation coefficients for better generative modeling.", "essential_references_not_discussed": "This paper is missing a discussion of [1,2,3]. [1] also optimizes the forward process to learn faster flows. [2,3] also incorporates adversarial learning into flow matching for improved generative modeling.\n\n[1] Minimizing Trajectory Curvature of ODE-based Generative Models, ICML, 2023\n\n[2] Constant Acceleration Flow, NeurIPS, 2024\n\n[3] Simple ReFlow: Improved Techniques for Fast Flow Models, ICLR, 2025", "other_strengths_and_weaknesses": "- **originality and significance** : while the usage of adversarial learning to improve the perceptual quality of neural ODEs have been already introduced in works such as [1] or [2], the idea of optimizing forward process parameters via simulation-based training is novel and interesting. Still, I find the experimental setup rather limited, and as discussed in **Methods And Evaluation Criteria**, additional results on higher resolution or more complex datasets will significantly strengthen this submission.\n\n[1] Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion, ICLR, 2024\n\n[2] Constant Acceleration Flow, NeurIPS, 2024", "other_comments_or_suggestions": "- Right column, lines 43-48 : $x\\_{est,\\theta,\\phi}$ is undefined.\n- Right column, Eq. (2) : $\\hat{x}\\_{0,\\theta}$ and $\\hat{x}\\_{1,\\theta}$ are undefined.\n- Table 4 : please clarify the difference between $z \\sim \\rho_T$ and $x_T \\sim \\rho_T$.", "questions_for_authors": "I am willing to raise the score to **Weak Accept** if the authors can address the following concerns.\n\n- [Q1] Can the authors provide any clarifications regarding the comment in **Claims And Evidence**?\n- [Q2] Can the authors provide additional experiments on ImageNet $64 \\times 64$ as a demonstration of scalability?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Standard diffusion and flow-based models interpolate between noise and data, and the interpolating coefficients are usually shared across dimensions. This paper proposes a simple method for having dimension-dependent interpolation coefficients.

The paper received three reviews, yet one was flagged as being AI-generated. Upon inspecting the review, I believe the flag to be very plausible since the review explicitly mentioned the authors conducting experiments only on CIFAR-10 despite this not being the case. I thus ignored this review in my decision making process, and instead read the paper and engaged with the authors.

Reviewers found the method novel and sensible, but criticized the lack of experiments on ImageNet. The authors provided these experiments in their rebuttal in a way that I believe satisfactorily addressed the main criticisms raised. Overall I believe this paper should be accepted: it presents a simple yet well-motivated and novel method which achieves good empirical results.
Decision: Accept (poster)