Paper deep dive
Geometry-Aware Time Reparameterization for Flow-Map Distillation
Félix Dedek, Makoto Yamada
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 96%
Last extracted: 10/5/2026, 6:03:13 AM
Summary
This paper proposes a geometry-aware time reparameterization method for flow-map distillation to improve one- and few-step generative model performance. The authors hypothesize that trajectory segments with high normal acceleration are harder to distill. They derive a shared time clock that equalizes the population normal-acceleration statistic, effectively slowing down the traversal of complex trajectory regions. This clock is estimated once from teacher trajectories and used to condition the student model during Lagrangian flow-map distillation, without requiring teacher retraining or additional inference costs. Experiments on synthetic data, CIFAR-10, and CelebA-64 demonstrate improved sample quality over identity-time distillation, particularly in one-step generation.
Entities (11)
Relation Signals (7)
Félix Dedek → affiliatedwith → Okinawa Institute of Science and Technology Graduate University
confidence 99% · Affiliation: Machine Learning and Data Science Unit Affiliation: Okinawa Institute of Science and Technology Graduate University
Makoto Yamada → affiliatedwith → Okinawa Institute of Science and Technology Graduate University
confidence 99% · Affiliation: Machine Learning and Data Science Unit Affiliation: Okinawa Institute of Science and Technology Graduate University
Geometry-Aware Time Reparameterization → isappliedin → Lagrangian Flow-Map Distillation
confidence 96% · We incorporate this clock into Lagrangian flow-map distillation, using the transformed time coordinate to condition the student.
Geometry-Aware Time Reparameterization → uses → Normal Acceleration
confidence 95% · Motivated by the hypothesis that trajectory segments with large normal acceleration are harder to distill, we propose a geometry-aware time reparameterization
Geometry-Aware Time Reparameterization → improves → Flow-Map Distillation
confidence 94% · Experiments on synthetic data, CIFAR-10, and CelebA-64 show improved sample quality over identity-time distillation
Flow-Map Distillation → enables → One-step generation
confidence 93% · Flow-map distillation enables one- and few-step generation by learning finite-time transitions of a pretrained generative ODE.
Normal Acceleration → isproxyfor → Distillation Difficulty
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Flow-map distillation enables one- and few-step generation by learning finite-time transitions of a pretrained generative ODE. We investigate whether changing the teacher's time parameterization can make these transitions easier to learn. Motivated by the hypothesis that trajectory segments with large normal acceleration are harder to distill, we propose a geometry-aware time reparameterization that allocates more student time to these regions while preserving the teacher's geometric paths and terminal distribution. We derive a shared clock that equalizes a population normal-acceleration statistic under suitable assumptions, and construct a practical approximation from robust, regularized estimates across teacher trajectories. We incorporate this clock into Lagrangian flow-map distillation, using the transformed time coordinate to condition the student. The clock is estimated once before distillation and requires neither teacher retraining nor additional student parameters or inference-time network evaluations. Experiments on synthetic data, CIFAR-10, and CelebA-64 show improved sample quality over identity-time distillation at matched inference budgets, including improvements in one- and two-step image generation. The gains in one-step generation, where no intermediate sampling times can be adjusted, highlight the benefits of time reparameterization during distillation.
Tags
Links
- Source: https://arxiv.org/abs/2610.02427v1
- Canonical: https://arxiv.org/abs/2610.02427v1
Trouble viewing inline? Open PDF directly →
Full Text
66,329 characters extracted from source content.
Expand or collapse full text
Geometry-Aware Time Reparameterization for Flow-Map Distillation Félix Dedek & Makoto Yamada Affiliation: Machine Learning and Data Science Unit Affiliation: Okinawa Institute of Science and Technology Graduate University Affiliation: Okinawa, Japan Abstract Flow-map distillation enables one- and few-step generation by learning finite-time transitions of a pretrained generative ODE. We investigate whether changing the teacher’s time parameterization can make these transitions easier to learn. Motivated by the hypothesis that trajectory segments with large normal acceleration are harder to distill, we propose a geometry-aware time reparameterization that allocates more student time to these regions while preserving the teacher’s geometric paths and terminal distribution. We derive a shared clock that equalizes a population normal-acceleration statistic under suitable assumptions, and construct a practical approximation from robust, regularized estimates across teacher trajectories. We incorporate this clock into Lagrangian flow-map distillation, using the transformed time coordinate to condition the student. The clock is estimated once before distillation and requires neither teacher retraining nor additional student parameters or inference-time network evaluations. Experiments on synthetic data, CIFAR-10, and CelebA-64 show improved sample quality over identity-time distillation at matched inference budgets, including improvements in one- and two-step image generation. The gains in one-step generation, where no intermediate sampling times can be adjusted, highlight the benefits of time reparameterization during distillation. 1 Introduction Diffusion and score-based generative models have emerged as a leading paradigm for learning complex data distributions (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021b), enabling high-fidelity image synthesis and flexible conditioning at scale (Dhariwal and Nichol, 2021; Rombach et al., 2022). These models learn to reverse a gradual noising process, with generation described by reverse-time stochastic dynamics or an associated probability-flow ordinary differential equation (ODE). Their success, however, comes with substantial inference cost: producing a sample typically requires many sequential neural-network evaluations. Reducing this cost while preserving sample quality remains a central challenge, motivating extensive work on deterministic sampling and specialized numerical solvers (Song et al., 2021a; Lu et al., 2022). Flow matching (FM) (Lipman et al., 2023) provides a flexible framework for learning continuous-time transport through a simulation-free regression objective. Together with closely related rectified-flow and stochastic-interpolant formulations (Liu et al., 2023; Albergo and Vanden-Eijnden, 2023), it learns a time-dependent velocity field that transports a simple source distribution to the data distribution, using analytically accessible conditional paths to construct training targets. This viewpoint accommodates diffusion probability paths as well as simpler alternatives, such as linear interpolation between noise and data, offering greater flexibility in the design of the generative dynamics. Importantly, however, straight conditional interpolation paths do not generally imply straight trajectories of the resulting marginal velocity field (Frans et al., 2025). Sampling therefore still requires numerical ODE integration, and using too few steps can introduce substantial discretization error. One- and few-step generative models address this limitation by learning finite-time transport directly, either through distillation from a pretrained model (Salimans and Ho, 2022) or through direct training from data. Consistency models (Song et al., 2023) learn to map points along a probability-flow trajectory to a common endpoint, while consistency trajectory models (Kim et al., 2024) and flow map matching (FMM) (Boffi et al., 2025a) learn two-time maps connecting intermediate states along the dynamics. Related formulations include Shortcut Models (Frans et al., 2025), which learn step-size-conditioned updates by enforcing consistency between one large step and two smaller steps, and MeanFlow (Geng et al., 2025), which learns average velocities through a differential identity relating them to instantaneous velocities. By learning finite-time updates, these approaches replace many evaluations of an instantaneous velocity field with one or a few evaluations of a learned map. Their effectiveness consequently depends on how accurately this map can be learned. We investigate whether changing a pretrained teacher’s time parameterization can make its flow map easier to learn. Motivated by the hypothesis that segments with large normal acceleration are harder to distill, we construct a shared clock that allocates more student time to these regions while preserving the teacher’s geometric paths and terminal distribution. The clock is estimated once from robust, regularized normal-acceleration statistics. We then train the student against the reparameterized velocity b~τ(x)=ϕ′(τ)bϕ(τ)(x) b_τ(x)=φ (τ)b_φ(τ)(x) using Lagrangian flow-map distillation. This requires neither teacher retraining nor changes to the student architecture or additional inference-time network evaluations. Our main contributions are: • A normal-acceleration-based clock. We derive a shared time transformation that equalizes a population normal-acceleration statistic under suitable assumptions, and estimate it robustly from teacher trajectories. • Geometry-aware flow-map distillation. We incorporate this time coordinate into Lagrangian flow-map distillation, training a single student directly against the reparameterized teacher dynamics. The time map is estimated once before distillation and requires neither teacher retraining nor changes to the student architecture, with no additional inference-time network evaluations. • Improved few-step generation. Through synthetic and image-generation experiments, we demonstrate improved sample quality over identity-time distillation at matched inference budgets. In particular, we observe meaningful improvements to one-step generation, where no intermediate sampling times can be adjusted, highlighting the role of time reparameterization in the distillation problem itself. 2 Related work Flow-map learning. Progressive distillation (Salimans and Ho, 2022) compresses successive teacher updates, while consistency models and consistency trajectory models learn endpoint predictions or transitions between intermediate noise levels (Song et al., 2023; Kim et al., 2024). Flow map matching (Boffi et al., 2025a) provides a mathematical framework for these constructions, with objectives for both teacher distillation and direct training. Subsequent work develops self-distillation formulations (Boffi et al., 2025b) and practical continuous-time distillation objectives (Sabour et al., 2025). Shortcut Models and MeanFlow offer complementary formulations based on compositional consistency and average velocities (Frans et al., 2025; Geng et al., 2025). Recent extensions include parallel prediction of interval velocities (Shaul et al., 2026) and distribution-based distillation using average-velocity information (Zhao et al., 2026). Our work builds on Lagrangian flow-map distillation and investigates the time coordinate in which the teacher’s finite-time transport is learned. Trajectory geometry and coupling. Several approaches improve generative dynamics by simplifying their trajectories. Rectified flow and reflow encourage straighter transport (Liu et al., 2023), while minibatch optimal transport modifies source–target couplings to obtain simpler dynamics (Tong et al., 2024). Lee et al. (2023) learn couplings that reduce trajectory curvature and also investigate the benefits of straighter teachers for distillation. Isokinetic Flow Matching introduces a lookahead velocity-consistency regularizer motivated by reducing material acceleration (Khan, 2026). These methods modify the learned velocity field or its training coupling. Our transformation instead retains the geometric trajectories of a pretrained teacher and adjusts their traversal speed before flow-map distillation. Sampling schedules and trajectory diagnostics. Sampling efficiency also depends on the placement of numerical integration steps. Align Your Steps optimizes schedules through a stochastic-calculus-based discrepancy criterion (Sabour et al., 2024), while Geometry-Inspired Time Scheduling uses trajectory information and dynamic programming to select integration times (Chen et al., 2024). Entropic Time Schedulers allocate sampling points according to conditional-entropy progress (Stančević et al., 2025). Particularly close to our construction, Sharpen Your Flow estimates the full material-acceleration magnitude of a pretrained flow and derives a square-root sampling density implemented through inverse-CDF scheduling (Gupta et al., 2026). Our method uses the normal component of acceleration and applies the resulting clock to the teacher dynamics during student training. This distinction is particularly relevant to one-step flow-map generation, where no intermediate inference times can be adjusted. Adaptive allocation of training effort. Temporal difficulty has also motivated adaptive training procedures. Curvature-Adaptive Consistency Flow Matching uses consistency-error feedback and reinforcement learning to allocate training across temporal stages (Tian et al., 2026). Its boundary-focused allocation provides qualitative evidence of nonuniform distillation difficulty, although its diagnostic differs from normal acceleration. Complexity-Balanced Diffusion Splitting uses acceleration or spatial regularity estimates to partition time among specialized networks (Issachar et al., 2026). Our approach trains a single student using a fixed time transformation estimated from the pretrained teacher. Time reparameterization and learnability. Time parameterization can affect learning as well as numerical integration. Tsimpos et al. (2025) optimize the schedule of a fixed transport interpolation to control the spatial Lipschitz constant of its velocity field. Chen et al. (2025) develop Lipschitz-guided interpolation schedules and transfer formulas for applying them to pretrained models. Other approaches prescribe conditional velocity profiles (Bondar, 2026) or construct an inverse-CDF interpolation schedule from regression losses measured during a pilot training run (Tania and Khan, 2026). Time transformations also appear in Flow Map Language Models, where a decoding-error-based clock guides temporal sampling and flow-map learning (Lee et al., 2026). These works establish the relevance of temporal coordinates to generative learning. We investigate a shared clock derived from the normal acceleration of a frozen teacher, and train the student directly against the reparameterized teacher velocity in the Lagrangian distillation objective. 3 Background In this section, we first formulate the flow matching problem and then introduce flow-map distillation methods. 3.1 Flow Matching Let (ˇ0,ˇ1)∼π∈Π(ρ0,ρ1)( x_0, x_1) π∈ ( _0, _1) denote a pair of samples from a coupling of the source and target distributions. Following the stochastic-interpolant formulation (Albergo and Vanden-Eijnden, 2023), we define intermediate states through a time-differentiable interpolant: It(ˇ0,ˇ1):=αtˇ0+βtˇ1,I_t( x_0, x_1):= _t x_0+ _t x_1, with boundary conditions I0(ˇ0,ˇ1)=ˇ0,I1(ˇ0,ˇ1)=ˇ1.I_0( x_0, x_1)= x_0, I_1( x_0, x_1)= x_1. For brevity, we write It:=It(ˇ0,ˇ1)I_t:=I_t( x_0, x_1), where (ˇ0,ˇ1)∼π( x_0, x_1) π. Flow matching (Lipman et al., 2023) learns a time-dependent velocity field by regressing the interpolant velocity: ℒFM()=t,(ˇ0,ˇ1)∼π[∥bt(It)−I˙t∥2],L_FM( θ)=E_t,( x_0, x_1) π [ b θ_t(I_t)- I_t ^2 ], where t∼([0,1])t ([0,1]) independently of the endpoint pair and I˙t=dtIt I_t= ddtI_t. This loss is minimized by bt⋆()=[I˙t∣It=].b_t ( x)=E [ I_t I_t= x ]. Under suitable regularity assumptions, the ODE driven by bt⋆b_t and initialized from ρ0 _0 reproduces the interpolant marginals. In all our experiments, we use linear interpolation, It(ˇ0,ˇ1)=(1−t)ˇ0+tˇ1,I˙t(ˇ0,ˇ1)=ˇ1−ˇ0.I_t( x_0, x_1)=(1-t) x_0+t x_1, I_t( x_0, x_1)= x_1- x_0. We denote the trained teacher field by btb_t and reserve tx_t for its ODE trajectories: ˙t=bt(t),0∼ρ0. x_t=b_t(x_t), x_0 _0. These trajectories generally differ from individual interpolation paths. Generation requires numerical ODE integration, typically involving many network evaluations and motivating few-step samplers and distillation methods. 3.2 Flow-Map Distillation Consistency models (Song et al., 2023) learn to map states along a probability-flow trajectory to a common data endpoint. Flow-map distillation (Boffi et al., 2025a) learns two-time maps connecting intermediate states of a pretrained teacher ODE. Let Xs,tX_s,t denote its solution map from time s to time t: Xs,t(s)=t.X_s,t(x_s)=x_t. (1) Approximating this map allows a student to replace many evaluations of the teacher velocity btb_t with a single finite-time update. Following Boffi et al. (2025a), we parameterize the student as X^s,t()=+(t−s)v(,s,t), X θ_s,t( x)= x+(t-s)v_ θ( x,s,t), (2) where v_ θ approximates the average velocity over the interval [s,t][s,t]. This parameterization enforces the identity condition X^s,s()= X θ_s,s( x)= x. Assuming the teacher ODE has unique solutions, the flow-map relation equation 1 is equivalently characterized by the Lagrangian equation ∂tXs,t()=bt(Xs,t()),Xs,s()=. _tX_s,t( x)=b_t (X_s,t( x) ), X_s,s( x)= x. Since our parameterization enforces the identity condition, Lagrangian flow-map distillation (LMD) (Boffi et al., 2025a) trains the student by minimizing the residual of this equation. We use the following stop-gradient formulation: ℒLMD()=s,t,(ˇ0,ˇ1)∼π[‖∂tX^s,t(Is)−sg(bt(X^s,t(Is)))‖2].L_LMD( θ)=E_s,t,( x_0, x_1) π [ \| _t X θ_s,t(I_s)-sg\! (b_t\! ( X θ_s,t(I_s) ) ) \|^2 ]. (3) The time-pair sampling distribution is specified in the experimental setup. Following the stop-gradient treatment of Lagrangian distillation discussed by Boffi et al. (2025b), we apply sg(⋅)sg(·) to the teacher velocity evaluated at the student output. This avoids backpropagation through the teacher’s spatial Jacobian, while the teacher parameters remain fixed throughout training. We also found this treatment important for training stability in our experiments; further details are provided in the appendix. Sampling from the flow map Given a partition 0=t0<⋯<tK=10=t_0<·s<t_K=1, we generate samples by composing the learned maps: ^tk+1=X^tk,tk+1(^tk),k=0,…,K−1,^0∼ρ0. x_t_k+1= X θ_t_k,t_k+1( x_t_k), k=0,…,K-1, x_0 _0. The terminal state ^1 x_1 is the generated sample. Each update requires one student-network evaluation, allowing the sampling budget K to be chosen after training. In particular, one-step generation applies X^0,1 X θ_0,1 directly to a source sample. 4 Proposed Method We first motivate normal acceleration as a proxy for flow-map learning difficulty. We then construct a time transformation that redistributes this quantity across time and incorporate the resulting clock into Lagrangian distillation. Figure 1: Time warping represented on the 2d trajectory t=(t,sin(2πt)) x_t=(t, (2π t)). The orange dots represent the reparametrized trajectory ϕ(τK) x_φ( _K) with uniform time steps τk=k/(N+1) _k=k/(N+1). The colorscale represents the density q=ψ′q=ψ of ϕ(t)φ( t) with t∼[0,1] t [0,1]. 4.1 Normal Acceleration As A Proxy For Distillation Difficulty Our working hypothesis is that trajectory segments with large normal acceleration are harder to represent through finite-time flow maps. To motivate this, consider the exact average velocity along a teacher trajectory: t−st−s=1t−s∫stbu(u)u,s<t. x_t-x_st-s= 1t-s _s^tb_u(x_u)\,du, s<t. If the velocity is constant on [s,t][s,t], this average equals bs(s)b_s(x_s). More generally, if its direction is fixed, the average remains aligned with that direction, and only its magnitude varies with the interval. When the direction changes, both the magnitude and direction of the average can depend on the two endpoints. This motivates using normal acceleration, the component associated with directional changes, as a diagnostic of temporal learning difficulty. The resulting connection to distillation error is an empirical hypothesis, rather than a consequence of the flow-map identity. Normal-acceleration diagnostic. For the teacher field btb_t, define the material acceleration at():=∂tbt()+Dbt()bt().a_t( x):= _tb_t( x)+D_ xb_t( x)\,b_t( x). Along a teacher trajectory, ¨t=at(t) x_t=a_t(x_t). Where bt()≠0b_t( x)≠ 0, let b^t()=bt()/∥bt()∥ b_t( x)=b_t( x)/ b_t( x) and decompose at∥()=(b^t()⊤at())b^t(),at⟂()=at()−at∥().a_t ( x)= ( b_t( x) a_t( x) ) b_t( x), a_t ( x)=a_t( x)-a_t ( x). The tangential component changes speed, while the normal component changes velocity direction. We define nt():=∥at⟂()∥.n_t( x):= a_t ( x) . Geometrically, ntn_t measures how strongly a trajectory bends at its current traversal speed: it equals the geometric curvature multiplied by the squared speed. Thus, slowing down reduces normal acceleration without altering the path taken by a trajectory. To summarize this diagnostic across teacher trajectories, we define the hardness profile hp(t):=([nt(t)p])1/p,p>0,h_p(t):= (E [n_t(x_t)^p ] )^1/p, p>0, where the expectation is over teacher trajectories initialized from 0∼ρ0x_0 _0. We use this profile as a proxy for time-dependent distillation difficulty. 4.2 Geometry-Aware Time Reparameterization The hardness profile identifies where additional temporal resolution may be useful. We therefore seek a shared change of time that slows traversal through regions of high hardness and accelerates it where hardness is low. Such a transformation preserves the teacher’s geometric paths and terminal distribution, while changing its normal acceleration. Our goal is to choose the new time coordinate so that the resulting hardness profile is constant. The following theorem gives an explicit construction. Theorem 4.1 (Equalizing the hardness profile). Assume the teacher trajectories are C2C^2 with nonzero velocity, and hph_p is positive and continuously differentiable. Define ψ(t):=∫0thp(u)u∫01hp(u)u,ϕ:=ψ−1.ψ(t):= _0^t h_p(u)\,du _0^1 h_p(u)\,du, φ:=ψ^-1. Then the reparameterized trajectories τ:=ϕ(τ)y_τ:=x_φ(τ) have constant hardness: h~p(τ):=([n~τ(τ)p])1/p=(∫01hp(u)u)2,τ∈[0,1], h_p(τ):= (E[ n_τ(y_τ)^p] )^1/p= ( _0^1 h_p(u)\,du )^2, τ∈[0,1], where n~τ n_τ denotes the normal-acceleration magnitude of the reparameterized dynamics. Proof. See Appendix A.1 ∎ Remark 4.2 (Tempering the time transformation). The strength of the transformation can be controlled by choosing ψ′(t)∝(hp(t)+ξ)δ/2ψ (t) (h_p(t)+ξ)^δ/2, where δ≥0δ≥ 0 and ξ≥0ξ≥ 0 is an optional numerical floor. The choice δ=0δ=0 recovers the identity clock, 0<δ<10<δ<1 softens the transformation, and δ>1δ>1 increases its concentration on high-hardness regions. The theorem corresponds to δ=1δ=1 and ξ=0ξ=0. Practical construction. We estimate the time map once before distillation using M teacher trajectories evaluated on a fine grid tk=k/Nt_k=k/N, k=0,…,Nk=0,…,N, with N≥100N≥ 100. We estimate material acceleration by finite differences of the teacher velocities along each trajectory, then project onto the normal direction to obtain nk,mn_k,m. At each time, we discard the largest 5%5\% of these values before computing the hardness profile. This limits the influence of rare sharp turns, which can produce particularly large normal accelerations (phenomenon which we mainly observed in 2d experiments). We integrate the resulting clock density numerically and fit a monotonicity-preserving PCHIP interpolant to the inverse pairs (τk,tk)( _k,t_k) (Fritsch and Butland, 1984). The resulting C1C^1 map has an analytically evaluable derivative. A positive floor ξ ensures a strictly increasing cumulative clock; in our experiments, the estimated hardness already remained bounded away from zero. Trimming, tempering, and numerical approximation mean that the practical clock need not exactly equalize the population profile. Algorithm 1 summarizes the procedure. Algorithm 1 Offline construction of the shared time map 1: Teacher btb_t, source ρ0 _0, trajectory count M, grid size N, aggregation p>0p>0, tempering δ≥0δ≥ 0, floor ξ≥0ξ≥ 0 2: Generate M teacher trajectories-velocity pairs (k,m,k,m)(x_k,m,b_k,m) from 0,m∼ρ0x_0,m _0 on tk=k/Nt_k=k/N, k=0,…,Nk=0,…,N, m=1,…,Mm=1,…,M 3: k,m←FDt[⋅,m]ka_k,m _t[b_·,m]_k 4: ^k,m←k,m/∥k,m∥ b_k,m←b_k,m/ b_k,m ; nk,m←∥k,m−(^k,m⊤k,m)^k,m∥n_k,m← a_k,m-( b_k,m a_k,m) b_k,m 5: for k=0,…,Nk=0,…,N do 6: ℐk←I_k← indices of the smallest ⌈0.95M⌉ 0.95M values of nk,mm=1M\n_k,m\_m=1^M 7: h^p,k←(1|ℐk|∑m∈ℐknk,mp)1/p h_p,k← ( 1|I_k| _m _kn_k,m^p )^1/p; wk←(h^p,k+ξ)δ/2w_k←( h_p,k+ξ)^δ/2 8: end for 9: S0←0S_0← 0; Sk←Sk−1+tk−tk−12(wk−1+wk)S_k← S_k-1+ t_k-t_k-12(w_k-1+w_k), k=1,…,Nk=1,…,N 10: τk←Sk/SN _k← S_k/S_N, k=0,…,Nk=0,…,N 11: ϕ←PCHIP((τk,tk)k=0N)φ (\( _k,t_k)\_k=0^N ) 12: return ϕφ and its derivative ϕ′φ Flow-map distillation in warped time Given the clock ϕ=ψ−1φ=ψ^-1 (Theorem 4.1), we train a student to approximate the reparameterized flow map Yα,β:=Xϕ(α),ϕ(β),0≤α≤β≤1.Y_α,β:=X_φ(α),φ(β), 0≤α≤β≤ 1. The trajectories τ=ϕ(τ)y_τ=x_φ(τ) satisfy dτdτ=b~τ(τ),b~τ():=ϕ′(τ)bϕ(τ)(). dy_τdτ= b_τ(y_τ), b_τ( x):=φ (τ)b_φ(τ)( x). For image experiments, we condition the student on the reparameterized times while retaining the physical-time interval as its displacement scale: Y^α,β()=+(ϕ(β)−ϕ(α))v(,α,β). Y θ_α,β( x)= x+ (φ(β)-φ(α) )v_ θ( x,α,β). (4) This choice separates temporal allocation from velocity scaling. Empirically, the original teacher’s velocity norm is approximately constant over most of the time interval, whereas reparameterization introduces substantial variation through the factor ϕ′φ (see Figure 2). The parameterization also enforces Y^α,α()= Y θ_α,α( x)= x. For the toy experiments, we instead use displacement scale β−αβ-α; the warped-time distillation objective below applies to either parameterization. Both scalings coincide under the identity clock. We apply the Lagrangian distillation objective to the reparameterized teacher: ℒLMDϕ()=α,β,(ˇ0,ˇ1)∼π[‖∂βY^α,β(Iϕ(α))−sg(b~β(Y^α,β(Iϕ(α))))‖2].L_LMD^φ( θ)=E_α,β,( x_0, x_1) π [ \| _β Y θ_α,β(I_φ(α))-sg\! ( b_β\! ( Y θ_α,β(I_φ(α)) ) ) \|^2 ]. (5) For image experiments, we sample two independent uniform physical times and sort them to obtain s≤ts≤ t. The student receives α=ψ(s)α=ψ(s) and β=ψ(t)β=ψ(t), with input state Is=Iϕ(α)I_s=I_φ(α). Thus, both clocks use the same physical-time pair distribution, while ϕφ changes the student’s temporal coordinates. Both the teacher and the clock remain fixed during training. We additionally apply learned loss weighting, as detailed in Appendix C. We sample analogously by composing Y^τk,τk+1 Y θ_ _k, _k+1 over a student-time partition 0=τ0<⋯<τK=10= _0<·s< _K=1, starting from ^0∼ρ0 y_0 _0. The overhead that comes from computing the scaling factor (ϕ(τk+1)−ϕ(τk))(φ( _k+1)-φ( _k)) is negligible compared with neural-network evaluation, and these differences can be precomputed once for a fixed partition. 5 Experiments We first examine the temporal distribution of normal acceleration and the effect of our clock on teacher sampling. We then evaluate whether training in the resulting time coordinate improves flow-map distillation on synthetic and image datasets. 5.1 Temporal Geometry and Teacher Sampling Figure 2: CelebA64 flow matching model (teacher) measurements plotted against their respective coordinates. Left: estimated hardness before and after reparametrization. Middle: estimated map ϕφ. Right: original and reparameterized velocity norms. Normal-acceleration profiles. We evaluate normal-acceleration profiles along trajectories of the CIFAR-10 and CelebA64 teachers. Both exhibit a similar U-shaped temporal profile, with larger normal acceleration near the endpoints and smaller values in the middle. Figure 2 shows CelebA64 as a representative example; the CIFAR-10 profiles exhibit the same qualitative behavior, with different magnitudes. This concentration motivates allocating more student time to the boundary regions. It is also qualitatively consistent with the U-shaped distillation difficulty reported by Tian et al. (2026), who learn an adaptive time-sampling curriculum using reinforcement learning. This agreement motivates our diagnostic, although it does not establish a direct correspondence between normal acceleration and distillation error. Figure 2 illustrates the resulting transformation for CelebA64. The estimated clock makes the normal-acceleration profile more uniform across time. Meanwhile, the original teacher’s velocity norm is approximately constant over most of the interval, whereas the warped velocity exhibits greater variation. This observation motivates the physical-time displacement scaling used for image distillation. Teacher sampling. Figure 3: CIFAR-10 Flow Matching model (teacher) evaluated FID using different solvers and number of steps. Before evaluating distillation, we test whether the estimated clock improves numerical sampling from the frozen CIFAR-10 teacher. For fixed-step solvers with K integration steps, we integrate the original field btb_t using either uniform nodes tk=k/Kt_k=k/K or clock-adapted nodes tk=ϕ(k/K)t_k=φ(k/K), for k=0,…,Kk=0,…,K. For the adaptive Dopri5 solver, we instead compare integration of btb_t with integration of the reparameterized field b~τ=ϕ′(τ)bϕ(τ) b_τ=φ (τ)b_φ(τ). We report FID against the number of velocity-network evaluations (NFE), which dominate sampling cost. Figure 3 shows that clock-adapted grids improve sample quality at low evaluation budgets. With 2525 Heun steps, FID decreases from 8.88.8 on the uniform grid to 3.53.5 on the clock-adapted grid, approaching the teacher’s performance at larger budgets. Naturally, the advantage diminishes as the evaluation budget increases and as we get closer to the real ODE’s solutions. These results support the usefulness of the estimated clock for allocating a limited fixed-step integration budget. The following experiments test whether training in the corresponding time coordinate also improves flow-map distillation. 5.2 Flow-Map Distillation We compare distillation under the identity clock, Id(t)=tId(t)=t, with distillation under our estimated clock ϕφ. For each dataset, we use a fixed flow-matching teacher and train two students with matched architectures, training budgets, and optimization settings. The IdId student learns from b, while the ϕφ student learns from its reparameterized counterpart b~ b. Before distillation, we estimate ϕφ from 256256 teacher trajectories generated with an adaptive ODE solver. The clock remains fixed throughout training. We use student-time displacement scaling for toy data and physical-time displacement scaling for images, as described in the method section. All comparisons use matched numbers of student-network evaluations. Implementation details and additional results are provided in the appendix. 5.2.1 Toy Experiments Figure 4: Quantitative results of the flow matching teacher and its distilled flow maps trained on the 2D checkerboard distribution. We evaluate on a two-dimensional checkerboard distribution, whose sharp boundaries and disconnected support provide a test of sample fidelity. The teacher and students use four-layer MLPs of width 512512, trained for 50,00050,000 updates with a batch size of 81928192. The source distribution is ρ0=(,) _0=N(0, I). Figure 7 shows that ϕφ produces sharper boundaries and fewer samples outside the target support than IdId. We quantify sample quality using a histogram-based forward KL statistic, Gaussian-kernel MMD with bandwidth σ=0.05σ=0.05, and empirical Wasserstein-2 distance. The KL computation exploits the known target density; its discretization and normalization are specified in Appendix B. MMD and Wasserstein-2 are evaluated against both the target distribution and teacher-generated samples, measuring target fidelity and distillation fidelity, respectively. Across the evaluated budgets of 11, 22, 33, and 44 network evaluations, ϕφ improves all three reported metrics relative to IdId (Figure 4). 5.2.2 Image Experiments Figure 5: Left: CIFAR-10 Flow Map FID after 300k training steps. Right: CelebA64 Flow Map FID after 350k training steps. We evaluate unconditional generation on CIFAR-10 (Krizhevsky et al., 2009) and CelebA64 (Liu et al., 2015). Both experiments use RAdam (Liu et al., 2020) and learned loss weighting following EDM2 (Karras et al., 2024) and its adaptation to flow maps (Boffi et al., 2025b). We report FID here and provide full FID/KID tables, including additional checkpoints, in Appendix C. Figure 6: CelebA64 1 step FID vs training budget. CIFAR-10. We use approximately 3535M-parameter U-Nets and a pretrained TorchCFM teacher (Tong et al., 2023; Tong et al., 2024). Students are trained for 400,000400,000 updates with a global batch size of 196196. CelebA64. We use approximately 124124M-parameter EDM2 U-Nets. The teacher is trained on the official CelebA training partition for 350,000350,000 updates, achieving FID 2.362.36 with 100100 Heun integration steps. Students are trained for 400,000400,000 updates with a global batch size of 256256. Results. Figure 5 compares IdId and ϕφ at matched inference budgets. On both datasets, ϕφ improves FID against the target distribution at every displayed budget. This does not imply uniformly better agreement with the teacher distribution: at the displayed checkpoints, IdId achieves lower teacher-relative FID at eight steps on CIFAR-10 and at two, four, and eight steps on CelebA64. The one-step comparison is particularly informative: both clocks have identical endpoints and displacement scales, so improvements reflect the learned map rather than the placement of intermediate sampling times. At 400400K training updates, one-step target FID decreases from 8.578.57 to 6.996.99 on CIFAR-10 and from 6.196.19 to 4.644.64 on CelebA64 (Appendix C).Moreover Figure 6 shows that the 1 step improvement is consistent across training steps for the CelebA64 model (CIFAR-10 shows a very similar trend). 6 Conclusion We proposed a geometry-aware time reparameterization that redistributes normal acceleration across time to improve flow-map distillation. The clock is estimated once, preserves teacher trajectories, and requires neither teacher retraining nor additional inference-time network evaluations. Experiments show improved few-step generation, in particular one-step gains that highlight the role of time coordinates in learning the map itself. Establishing a theoretical link between normal acceleration and distillation difficulty remains a direction for future work. AI Use Statement We used ChatGPT and AI coding assistants to support implementation and debugging, literature searches, manuscript drafting and revision, and figure preparation. We also used these tools to discuss and refine mathematical arguments and experimental methodology. The authors take responsibility for the final manuscript, code, references, and reported findings. The authors reviewed and revised AI-assisted text, checked mathematical arguments and cited sources, and tested AI-assisted code against the stated methods. References Albergo and Vanden-Eijnden (2023) M. S. Albergo and E. Vanden-Eijnden Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations, External Links: Link Cited by: §1, §3.1. Bińkowski et al. (2021) M. Bińkowski, D. J. Sutherland, M. Arbel, and A. Gretton Demystifying mmd gans. External Links: 1801.01401, Link Cited by: §C.1. Boffi et al. (2025a) N. M. Boffi, M. S. Albergo, and E. Vanden-Eijnden Flow map matching with stochastic interpolants: a mathematical framework for consistency models. Transactions on Machine Learning Research. External Links: Link Cited by: §C.2, §1, §2, §3.2, §3.2, §3.2. Boffi et al. (2025b) N. M. Boffi, M. S. Albergo, and E. Vanden-Eijnden How to build a consistency model: learning flow maps via self-distillation. External Links: 2505.18825, Link Cited by: §C.1, §2, §3.2, §5.2.2. Bondar (2026) V. Bondar Velocity scheduled flow matching. External Links: 2607.11442, Link Cited by: §2. Chen et al. (2024) D. Chen, Z. Zhou, C. Wang, C. Shen, and S. Lyu On the trajectory regularity of ode-based diffusion sampling. External Links: 2405.11326, Link Cited by: §2. Chen et al. (2025) Y. Chen, E. Vanden-Eijnden, and J. Xu Lipschitz-guided design of interpolation schedules in generative models. External Links: 2509.01629, Link Cited by: §2. Dhariwal and Nichol (2021) P. Dhariwal and A. Nichol Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, Vol. 34, p. 8780–8794. External Links: Link Cited by: §C.2, §1. Frans et al. (2025) K. Frans, D. Hafner, S. Levine, and P. Abbeel One step diffusion via shortcut models. In International Conference on Learning Representations, External Links: Link Cited by: §1, §1, §2. Fritsch and Butland (1984) F. N. Fritsch and J. Butland A method for constructing local monotone piecewise cubic interpolants. SIAM Journal on Scientific and Statistical Computing 5 (2), p. 300–304. External Links: Document, Link, https://doi.org/10.1137/0905021 Cited by: §4.2. Geng et al. (2025) Z. Geng, M. Deng, X. Bai, J. Z. Kolter, and K. He Mean flows for one-step generative modeling. In Advances in Neural Information Processing Systems, Vol. 38. External Links: Link Cited by: §1, §2. Gupta et al. (2026) A. Gupta, S. H. Lim, A. Yu, and N. B. Erichson Sharpen your flow: sharpness-aware sampling for flow matching. External Links: 2605.11547, Link Cited by: §2. Heusel et al. (2018) M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter GANs trained by a two time-scale update rule converge to a local nash equilibrium. External Links: 1706.08500, Link Cited by: §C.1. Ho et al. (2020) J. Ho, A. Jain, and P. Abbeel Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, Vol. 33, p. 6840–6851. External Links: Link Cited by: §1. Issachar et al. (2026) N. Issachar, D. Lischinski, and R. Fattal Complexity-balanced diffusion splitting. External Links: 2606.06477, Link Cited by: §2. Karras et al. (2024) T. Karras, M. Aittala, J. Lehtinen, J. Hellsten, T. Aila, and S. Laine Analyzing and improving the training dynamics of diffusion models. External Links: 2312.02696, Link Cited by: §C.1, §C.3, §5.2.2. Khan (2026) T. Khan Isokinetic flow matching for pathwise straightening of generative flows. External Links: 2604.04491, Link Cited by: §2. Kim et al. (2024) D. Kim, C. Lai, W. Liao, N. Murata, Y. Takida, T. Uesaka, Y. He, Y. Mitsufuji, and S. Ermon Consistency trajectory models: learning probability flow ODE trajectory of diffusion. In International Conference on Learning Representations, External Links: Link Cited by: §1, §2. Krizhevsky et al. (2009) A. Krizhevsky G. Hinton et al. Learning multiple layers of features from tiny images. Technical report Technical report, University of Toronto. Cited by: §C.2, §5.2.2. Lee et al. (2026) C. Lee, J. Yoo, M. Agarwal, S. Shah, J. Huang, A. Raghunathan, S. Hong, N. M. Boffi, and J. Kim Flow map language models: one-step language modeling via continuous denoising. External Links: 2602.16813, Link Cited by: §2. Lee et al. (2023) S. Lee, B. Kim, and J. C. Ye Minimizing trajectory curvature of ODE-based generative models. In Proceedings of the 40th International Conference on Machine Learning, A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (Eds.), Proceedings of Machine Learning Research, Vol. 202, p. 18957–18973. External Links: Link Cited by: §2. Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In International Conference on Learning Representations, External Links: Link Cited by: §1, §3.1. Liu et al. (2020) L. Liu, H. Jiang, P. He, W. Chen, X. Liu, J. Gao, and J. Han On the variance of the adaptive learning rate and beyond. In Proceedings of the Eighth International Conference on Learning Representations (ICLR 2020), Cited by: §C.2, §5.2.2. Liu et al. (2023) X. Liu, C. Gong, and Q. Liu Flow straight and fast: learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, External Links: Link Cited by: §1, §2. Liu et al. (2015) Z. Liu, P. Luo, X. Wang, and X. Tang Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), Cited by: §C.3, §5.2.2. Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter Decoupled weight decay regularization. External Links: 1711.05101, Link Cited by: Appendix B. Lu et al. (2022) C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu DPM-Solver: a fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, Vol. 35, p. 5775–5787. External Links: Link Cited by: §1. Rombach et al. (2022) R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 10684–10695. External Links: Link Cited by: §1. Sabour et al. (2024) A. Sabour, S. Fidler, and K. Kreis Align your steps: optimizing sampling schedules in diffusion models. External Links: 2404.14507, Link Cited by: §2. Sabour et al. (2025) A. Sabour, S. Fidler, and K. Kreis Align your flow: scaling continuous-time flow map distillation. In Advances in Neural Information Processing Systems, Vol. 38. External Links: Link Cited by: §2. Salimans and Ho (2022) T. Salimans and J. Ho Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, External Links: Link Cited by: §1, §2. Shaul et al. (2026) N. Shaul, C. Liu, A. Vahdat, and J. Berner Parallel decoding distillation for fast image and video generation. External Links: 2607.26004, Link Cited by: §2. Sohl-Dickstein et al. (2015) J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 37, p. 2256–2265. External Links: Link Cited by: §1. Song et al. (2021a) J. Song, C. Meng, and S. Ermon Denoising diffusion implicit models. In International Conference on Learning Representations, External Links: Link Cited by: §1. Song et al. (2023) Y. Song, P. Dhariwal, M. Chen, and I. Sutskever Consistency models. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, p. 32211–32252. External Links: Link Cited by: §1, §2, §3.2. Song et al. (2021b) Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, External Links: Link Cited by: §1. Stančević et al. (2025) D. Stančević, F. Handke, and L. Ambrogioni Entropic time schedulers for generative diffusion models. In Advances in Neural Information Processing Systems, Vol. 38. External Links: Link Cited by: §2. Tania and Khan (2026) A. A. Tania and M. R. Khan Difficulty-calibrated interpolation paths for conditional flow matching. External Links: 2608.21286, Link Cited by: §2. Tian et al. (2026) S. Tian, G. Chen, B. Li, J. Ma, and Z. Yu Curvature-adaptive consistency flow matching: autonomous trajectory optimization via reinforcement learning. External Links: 2606.22394, Link Cited by: §2, §5.1. Tong et al. (2024) A. Tong, K. FATRAS, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research. Note: Expert Certification External Links: ISSN 2835-8856, Link Cited by: §C.2, §C.2, §2, §5.2.2. Tong et al. (2023) A. Tong, N. Malkin, K. Fatras, L. Atanackovic, Y. Zhang, G. Huguet, G. Wolf, and Y. Bengio Simulation-free schrödinger bridges via score and flow matching. arXiv preprint 2307.03672. Cited by: §C.2, §C.2, §5.2.2. Tsimpos et al. (2025) P. Tsimpos, Z. Ren, J. Zech, and Y. Marzouk Optimal scheduling of dynamic transport. In Proceedings of the Thirty-Eighth Conference on Learning Theory, Proceedings of Machine Learning Research, Vol. 291, p. 5441–5505. External Links: Link Cited by: §2. Zhao et al. (2026) A. Zhao, S. Zhang, Z. Sun, Y. Zhou, Z. Li, L. Yang, T. Chen, and L. Sun Mean flow distillation: robust and stable distillation for flow matching models. External Links: 2606.11155, Link Cited by: §2. Appendix A Mathematical Details and Proofs A.1 Proof of the main Theorem and construction of ϕφ Proof of Theorem 4.1. Let t∈[0,1]t∈[0,1] denote teacher time and τ∈[0,1]τ∈[0,1] denote student time. Write Z:=∫01hp(u)u.Z:= _0^1 h_p(u)\,du. Since hph_p is positive and continuously differentiable, 0<Z<∞0<Z<∞ and ψ′(t)=hp(t)Z>0.ψ (t)= h_p(t)Z>0. Thus ψ is a strictly increasing C2C^2 map with ψ(0)=0ψ(0)=0 and ψ(1)=1ψ(1)=1. Its inverse ϕ=ψ−1φ=ψ^-1 is also C2C^2, with ϕ′>0φ >0, and satisfies t=ϕ(τ),τ=ψ(t).t=φ(τ), τ=ψ(t). Consider a teacher trajectory x and its reparameterized version τ=ϕ(τ).y_τ=x_φ(τ). By the chain rule, ˙τ=ϕ′(τ)˙ϕ(τ),¨τ=ϕ′(τ)˙ϕ(τ)+ϕ′(τ)2¨ϕ(τ). y_τ=φ (τ) x_φ(τ), y_τ=φ (τ) x_φ(τ)+φ (τ)^2 x_φ(τ). Since ϕ′(τ)>0φ (τ)>0, the two trajectories have the same unit tangent at corresponding times. The term containing ϕ′φ is tangential, so a~τ⟂(τ)=ϕ′(τ)2aϕ(τ)⟂(ϕ(τ)),n~τ(τ)=ϕ′(τ)2nϕ(τ)(ϕ(τ)). a_τ (y_τ)=φ (τ)^2a_φ(τ) (x_φ(τ)), n_τ(y_τ)=φ (τ)^2n_φ(τ)(x_φ(τ)). The same deterministic clock applies to every trajectory, hence h~p(τ)=ϕ′(τ)2hp(ϕ(τ)). h_p(τ)=φ (τ)^2h_p(φ(τ)). Here both expectations are taken over the same initial distribution of teacher trajectories. Finally, differentiating ψ(ϕ(τ))=τψ(φ(τ))=τ gives ψ′(ϕ(τ))ϕ′(τ)=1,ϕ′(τ)=Zhp(ϕ(τ)).ψ (φ(τ))φ (τ)=1, φ (τ)= Z h_p(φ(τ)). Substituting this expression into the preceding identity, we obtain h~p(τ)=Z2hp(ϕ(τ))hp(ϕ(τ))=(∫01hp(u)u)2. h_p(τ)= Z^2h_p(φ(τ))h_p(φ(τ))= ( _0^1 h_p(u)\,du )^2. which is independent of τ. ∎ Appendix B Toy Experiments: Implementation Details Figure 7: 8192 samples from the target distribution (checkerboard, top left, black), the flow matching teacher (bottom left, green), the IdId-TP student (baseline, top middle and right) and the ϕφ-TP student (ours, bottom middle and right). Dataset and training. The target is a uniform distribution on the occupied squares of a two-dimensional checkerboard. The source distribution is ρ0=(,) _0=N(0, I). The teacher and students use four-layer MLPs of width 512512 and SiLU activation function. Each model is trained for 50,000 optimizer updates with a batch size of 8192. The teacher is optimized using AdamW (Loshchilov and Hutter, 2019) with learning rate 10−310^-3, β1=0.9 _1=0.9, β2=0.999 _2=0.999, weight decay 10−410^-4, and ϵ=10−8ε=10^-8. Its learning rate follows a cosine schedule down to 10−610^-6, and its exponential moving average uses decay 0.9980.998. The students use AdamW with learning rate 10−210^-2, the same Adam coefficients, weight decay, and numerical stabilization, followed by an exponential learning-rate schedule with multiplicative decay 0.99990.9999 after each optimizer update. Their exponential moving average uses decay 0.9990.999. The maximum gradient norms are 1010 for the teacher and 11 for the students. EMA weights are used for validation and evaluation. We use a constant loss weight 12 12. The ϕφ student uses displacement scale β−αβ-α. Evaluation protocol. We evaluate generated samples using a histogram-based forward KL, Gaussian-kernel MMD, and empirical Wasserstein-2 distance. KL uses 64,00064,000 generated samples and the known target density. For MMD and Wasserstein-2, we use 81928192 samples per distribution and compare against both target samples and teacher-generated samples, measuring target fidelity and distillation fidelity, respectively. Histogram-based forward KL. We estimate the forward divergence DKL(p∥q)D_KL(p\|q) using the known checkerboard density p and a histogram estimate of the generated density q. A fixed domain containing the target support is partitioned into 64×6464× 64 equal-area bins. For bin center cijc_ij, bin area ΔA A, and generated count nijn_ij, we set q^(cij)=nijNΔA, q(c_ij)= n_ijN A, where N is the total number of generated samples, including those outside the histogram domain. We compute D^KLε=∑i,j:p(cij)>0p(cij)ΔAlog(p(cij)q^(cij)+ε),ε=10−10. D_KL^\, = _i,j:\,p(c_ij)>0p(c_ij) A \! ( p(c_ij) q(c_ij)+ ), =10^-10. Normalization by N retains the penalty associated with generated mass outside the domain. The density floor prevents infinite values from empty target-positive bins. This statistic approximates the forward KL at the chosen spatial resolution; its value depends on binning, sample count, and regularization. We use identical evaluation settings for all models. Maximum mean discrepancy. We use the Gaussian kernel k(,)=exp(−∥−∥22σ2),σ=0.05.k( x, y)= \! (- x- y ^22σ^2 ), σ=0.05. For distributions μ and ν, the squared population discrepancy is MMDk2(μ,ν)=[k(X,X′)]+[k(Y,Y′)]−2[k(X,Y)],MMD_k^2(μ,ν)=E[k(X,X )]+E[k(Y,Y )]-2E[k(X,Y)], where X,X′∼μX,X μ and Y,Y′∼νY,Y ν are independent. We report the biased estimate MMD^b=(1n2∑i,jk(xi,xj)+1m2∑i,jk(yi,yj)−2nm∑i,jk(xi,yj).)1/2. MMD_b= ( 1n^2 _i,jk(x_i,x_j)+ 1m^2 _i,jk(y_i,y_j)- 2nm _i,jk(x_i,y_j). )^1/2. Empirical Wasserstein distance. We compare empirical sample measures using the quadratic transport cost: W22(μ^,ν^)=min∫γ∈Π(μ^,ν^)∥−∥2γ(,).W_2^2( μ, ν)= _γ∈ ( μ, ν) x- y ^2\,dγ( x, y). For two equally sized point clouds with n samples, we solve the assignment problem with cost matrix Cij=∥xi−yj∥2C_ij= x_i-y_j ^2 using SciPy’s linear_sum_assignment. If σ is the resulting optimal permutation, we report W2(μ^,ν^)=(1n∑i=1n∥xi−yσ(i)∥2)1/2.W_2( μ, ν)= ( 1n _i=1^n x_i-y_σ(i) ^2 )^1/2. Appendix C Image Experiments: Implementation and Additional Results C.1 Shared distillation setup Training inputs and time sampling. Both image experiments use unconditional students trained with the stop-gradient Lagrangian distillation objective. The teacher and, where applicable, the estimated clock remain fixed throughout distillation. Input states are sampled directly from the linear interpolant: IαI_α for IdId and Iϕ(α)I_φ(α) for ϕφ. We draw u,v∼([0,1])u,v ([0,1]) independently and set s=min(u,v),t=max(u,v).s= (u,v), t= (u,v). For IdId, the student receives (s,t)(s,t); for ϕφ, it receives (α,β)=(ψ(s),ψ(t))(α,β)=(ψ(s),ψ(t)). Both use the input state IsI_s. The physical-time pair (s,t)(s,t) is uniform over the upper triangle; the transformed pair (α,β)(α,β) generally is not. We do not explicitly include diagonal samples. In preliminary experiments, directly sampling and sorting uniform student times also improved one- and two-step target FID over IdId. Sampling uniformly in physical time yielded further improvements and is used for all image results reported here. Learned loss weighting. We use the learned weighting introduced for diffusion training in EDM2 (Karras et al., 2024) and adapted to two-time flow-map objectives by Boffi et al. (2025b). Let ℓα,β _α,β denote the unweighted per-example squared distillation residual. A scalar head with parameters η predicts wα,β:=wη(α,β−α).w_α,β:=w_η(α,β-α). We jointly optimize the student and weighting head using ℒweighted(,η)=[e−wα,βℓα,β+wα,β].L_weighted( θ,η)=E [e^-w_α,β _α,β+w_α,β ]. (6) The exponential factor rescales the contribution of each time pair, while the additive term penalizes arbitrarily suppressing the residual through large positive weights. We found this weighting important for training stability. Each student learns its own weighting function, using the same weighting scheme for IdId and ϕφ. The weighting head is used only during training and adds no inference-time network evaluations. Hyperparameter selection. For each image dataset, we compared candidate learning rates and decay schedules using FID and KID computed from 50,000 generated samples at multiple training checkpoints and sampling budgets. Both metrics consistently selected the same learning-rate and decay configuration across these checkpoints and budgets. The selected configuration was also the same for IdId and ϕφ. Student evaluation. We evaluate EMA student weights, using decay 0.99990.9999 in both image experiments. Comparisons between IdId and ϕφ use matched training checkpoints and numbers of student network evaluations. The main text reports FID (Heusel et al., 2018); the additional tables report both FID and KID (Bińkowski et al., 2021) at multiple training stages. Both statistics are computed using 50,000 generated samples and legacy TensorFlow-compatible Inception-v3 features. For the KID we use a standard cubic polynomial kernel k(x,y)=(x⊤yd+1)3,k(x,y)= ( x yd+1 )^3, where d is the Inception feature dimension. We report the mean unbiased estimate of MMDk2MMD_k^2 over 100100 subsets of 10001000 samples per distribution. Each reported configuration uses a single training run. Preliminary repetitions of selected configurations yielded very similar training curves and negligible differences between evaluation results. Construction of ϕφ We used the parameters p=2p=2, δ=1δ=1, ξ=10−5ξ=10^-5, N=100N=100 as described in 1. In practice, we found that the parameter p does not have a major effect on the resulting ϕφ (for tested values in [1,4][1,4]). Moreover the numerical floor ξ has negligible effect in all our experiments since the estimate hph_p has very high values (in range 10210^2 and above). C.2 CIFAR-10 Training Model FID to target FID to teacher 1 2 4 8 20 1 2 4 8 20 150k Id 10.09 6.31 5.54 5.49 5.65 4.44 1.83 1.44 1.52 1.75 150k ϕφ (teacher time) 7.94 5.60 5.34 5.45 5.17 2.85 1.49 1.44 1.69 1.66 300k Id 8.83 5.80 5.09 4.99 5.08 3.56 1.64 1.32 1.35 1.48 300k ϕφ (teacher time) 7.16 5.06 4.88 4.89 4.85 2.45 1.32 1.27 1.43 1.48 400k Id 8.57 5.71 4.98 4.89 5.02 3.40 1.62 1.30 1.31 1.45 400k ϕφ (teacher time) 6.99 5.01 4.80 4.80 4.79 2.38 1.32 1.25 1.37 1.45 Training Model KID to target (×103× 10^3) KID to teacher (×103× 10^3) 1 2 4 8 20 1 2 4 8 20 150k Id 6.07 3.37 2.91 3.04 3.24 2.43 0.61 0.33 0.45 0.68 150k ϕφ (teacher time) 4.38 2.63 2.79 3.08 2.73 1.20 0.26 0.37 0.58 0.45 300k Id 5.43 3.16 2.71 2.75 2.90 1.91 0.52 0.28 0.31 0.48 300k ϕφ (teacher time) 3.90 2.37 2.53 2.72 2.73 0.97 0.19 0.23 0.40 0.45 400k Id 5.21 3.13 2.71 2.62 2.92 1.78 0.52 0.27 0.28 0.45 400k ϕφ (teacher time) 3.78 2.34 2.45 2.66 2.71 0.92 0.20 0.21 0.35 0.42 Table 1: CIFAR-10 distillation results across training checkpoints and sampling budgets. KID values are multiplied by 10310^3. For each checkpoint, reference distribution, and sampling budget, the lower value across the two clocks is highlighted. Data. We use the 50,00050,000 training images of CIFAR-10 (Krizhevsky et al., 2009) at their native 32×3232× 32 resolution, without class conditioning. The source distribution is ρ0=(,) _0=N(0, I). We add horizontal flips. The images are then mapped to [-1,1]. Architecture. The teacher and students use the guided-diffusion U-Net architecture of Dhariwal and Nichol (2021), as implemented in TorchCFM (Tong et al., 2023; Tong et al., 2024), with base width 128128 and channel multipliers (1,2,2,2)(1,2,2,2). As in Boffi et al. (2025a), the student adapts the time-conditioning pathway to accept two times. The corresponding resolution–channel pairs are (322,128),(162,256),(82,256),(42,256).(32^2,128), (16^2,256), (8^2,256), (4^2,256). The backbone uses two encoder residual blocks and three skip-connected decoder blocks per resolution, with self-attention at 16×1616× 16 and in the bottleneck. Attention uses four heads with 6464 channels per head. The models contain approximately 3535M parameters. The teacher embeds physical time t, whereas the student embeds α and β separately, using 128128-dimensional Fourier features with time-input scale 10001000. These features are projected into the residual-block conditioning stream. Pretrained teacher. We use a publicly released TorchCFM checkpoint (Tong et al., 2023; Tong et al., 2024). The supplied training configuration uses Adam with learning rate 2⋅10−42· 10^-4, default optimizer parameters, a batch size of 128128, 5,0005,000 learning-rate warmup updates, gradient-norm clipping at 1.01.0, EMA decay 0.99990.9999, and 400,000400,000 training updates. The exact checkpoint is TorchCFM release v1.0.4, cfm_cifar10_weights_step_400000.pt, trained using Independent Conditional Flow Matching (I-CFM) with independent source–target coupling and noise σ=0σ=0. The teacher achieves approximately 50k-FID 3.53.5 (see Fig 3) Student optimization. Both students use RAdam (Liu et al., 2020), a global batch size of 196196, and 400,000400,000 optimizer updates. The learning rate at optimizer update k is η(k)=2⋅10−4 0.992k/1000.η(k)=2· 10^-4\,0.992^k/1000. Both students use the learned weighting from Appendix C.1. Evaluation. The main-text figure uses the checkpoint at 300k training steps. Full results across inference budgets and training steps are reported in Table 1. C.3 CelebA64 Training Model FID to target FID to teacher 1 2 4 8 20 1 2 4 8 20 150k Id 8.41 5.73 4.70 4.35 4.15 3.93 1.73 1.48 1.76 1.98 150k ϕφ (teacher time) 6.00 4.18 3.84 4.05 4.03 2.69 1.46 1.97 2.60 2.68 200k Id 7.56 5.14 4.18 3.92 3.82 3.35 1.44 1.25 1.57 1.81 200k ϕφ (teacher time) 5.38 3.72 3.44 3.63 3.62 2.46 1.39 1.89 2.45 2.52 250k Id 6.98 4.78 3.93 3.74 3.66 2.93 1.26 1.14 1.51 1.75 250k ϕφ (teacher time) 5.20 3.67 3.37 3.56 3.61 2.13 1.15 1.61 2.24 2.42 300k Id 6.56 4.54 3.75 3.60 3.58 2.60 1.11 1.06 1.46 1.74 300k ϕφ (teacher time) 4.72 3.29 3.08 3.33 3.38 2.04 1.15 1.65 2.29 2.43 350k Id 6.35 4.47 3.74 3.59 3.56 2.43 1.07 1.05 1.46 1.75 350k ϕφ (teacher time) 4.68 3.30 3.07 3.29 3.33 1.95 1.09 1.53 2.16 2.30 400k Id 6.19 4.35 3.63 3.51 3.51 2.34 1.02 1.01 1.44 1.74 400k ϕφ (teacher time) 4.64 3.36 3.12 3.28 3.26 1.81 1.05 1.49 2.11 2.26 Training Model KID to target (×103× 10^3) KID to teacher (×103× 10^3) 1 2 4 8 20 1 2 4 8 20 150k Id 5.43 3.55 3.06 3.01 2.91 2.72 1.05 1.17 1.69 1.92 150k ϕφ (teacher time) 3.79 2.65 2.76 3.08 3.06 1.97 1.14 2.03 2.73 2.71 200k Id 4.84 3.13 2.68 2.71 2.72 2.32 0.86 1.01 1.54 1.79 200k ϕφ (teacher time) 3.43 2.38 2.49 2.77 2.74 1.91 1.17 1.99 2.61 2.57 250k Id 4.43 2.93 2.53 2.60 2.60 1.99 0.75 0.93 1.48 1.73 250k ϕφ (teacher time) 3.17 2.21 2.31 2.67 2.73 1.50 0.83 1.60 2.34 2.44 300k Id 4.15 2.78 2.43 2.53 2.59 1.74 0.64 0.86 1.44 1.74 300k ϕφ (teacher time) 2.92 2.03 2.16 2.52 2.57 1.50 0.88 1.63 2.37 2.44 350k Id 4.01 2.76 2.43 2.54 2.59 1.61 0.62 0.86 1.44 1.75 350k ϕφ (teacher time) 2.84 2.01 2.10 2.46 2.51 1.37 0.79 1.50 2.22 2.29 400k Id 3.89 2.66 2.33 2.45 2.54 1.56 0.58 0.82 1.42 1.74 400k ϕφ (teacher time) 2.80 2.04 2.10 2.41 2.43 1.31 0.80 1.49 2.18 2.26 Table 2: CelebA64 distillation results across training checkpoints and sampling budgets. KID values are multiplied by 10310^3. For each checkpoint, reference distribution, and sampling budget, the lower value across the two clocks is highlighted. Data and preprocessing. We use the official CelebA training partition (Liu et al., 2015), containing 162,770162,770 images. Validation uses 3,0723,072 fixed images sampled with seed 00 from the official validation partition. The official test partition is not used for training. Images are center-cropped or padded to 140×140140× 140, resized to 64×6464× 64 using area interpolation, and converted from uint8 RGB to [−1,1][-1,1]. We use no horizontal flipping or other online augmentation. The source distribution is ρ0=(,0.25) _0=N(0,0.25 I). Source and data samples are coupled independently. Teacher training uses uniformly sampled physical times and linear interpolation without additional noise. Architecture. We use the NVLabs EDM2 U-Net backbone (Karras et al., 2024) with base width 128128 and channel multipliers (1,2,3,4)(1,2,3,4), giving (642,128),(322,256),(162,384),(82,512).(64^2,128), (32^2,256), (16^2,384), (8^2,512). There are three encoder residual blocks and four skip-connected decoder blocks per resolution, with attention at 16×1616× 16 and 8×88× 8. Teacher dropout is zero; student dropout is 0.10.1. The teacher and student U-Nets contain 124,198,821124,198,821 and 124,264,357124,264,357 trainable parameters, respectively. The backbone uses magnitude-preserving convolutions, pixel normalization, magnitude-preserving residual and skip mixing, normalized attention, and learnable residual and embedding gains. An internal constant-one plane is appended to RGB, so the first convolution receives four channels. No spatial time-conditioning planes are used. The teacher uses a deterministic 128128-dimensional Fourier embedding of physical time, z(t)=2[cos(tω),sin(tω)],z(t)= 2\,[ (tω), (tω)], with logarithmically spaced frequencies and maximum period 10410^4. A magnitude-preserving projection maps this embedding to a 512512-dimensional conditioning stream. Students instead concatenate separate 128128-dimensional embeddings of α and β and project the result to 512512 dimensions. After magnitude-preserving SiLU, the student embedding modulates every residual block through the original conditioning pathway. Teacher optimization. The teacher minimizes the independent conditional flow-matching loss, averaged over batch, channels, and spatial dimensions, without learned loss weighting. We use RAdam with (β1,β2)=(0.9,0.999)( _1, _2)=(0.9,0.999) and ϵ=10−8ε=10^-8, a global batch size of 256256, and trained for 350,000350,000 updates. The learning-rate schedule is η(k)=10−2max(k/35000,1).η(k)= 10^-2 (k/35000,1). Training uses four A100 GPUs with 6464 images per GPU, BF16 computation, gradient-norm clipping at 1.01.0, and EMA decay 0.99990.9999. EDM2 weights are renormalized after every optimizer update. We freeze the step-350,000350,000 EMA state as an immutable, hashed snapshot for distillation. Its evaluation with 100100 Heun steps and 50,00050,000 generated samples gives approximately FID 2.362.36 and KID 1.46⋅10−31.46· 10^-3. Student optimization and weighting implementation. Both students use one RAdam optimizer with separate parameter groups for the U-Net and weighting head. The groups share the same learning rate, (β1,β2)=(0.9,0.999)( _1, _2)=(0.9,0.999), ϵ=10−8ε=10^-8, and zero weight decay. Students are trained for 400,000400,000 updates with learning-rate schedule η(k)=3⋅10−3 0.992k/1000.η(k)=3· 10^-3\,0.992^k/1000. The global batch size is 256256, with 6464 examples per GPU on four GPUs. We use BF16 computation, gradient-norm clipping at 1.01.0, and EMA decay 0.99990.9999. EDM2 weights are renormalized after every optimizer update and every EMA update. There is no early stopping. For Equation 6, ℓα,β _α,β is a sum over all 3⋅64⋅643· 64· 64 image coordinates. Weighted losses are then averaged over the batch. The weighting head embeds α and β−αβ-α separately with 128128-dimensional sinusoidal features, concatenates them, and applies a zero-initialized linear projection to one scalar. It contains 257257 parameters in addition to the U-Net.