Geometry-Aware Time Reparameterization for Flow-Map Distillation

1 INTRODUCTION

• Generative ODE sampling typically requires many sequential network evaluations, making cost reduction without sacrificing sample quality a central challenge.Flow matching still requires numerical integration, and too few steps can introduce substantial discretization error.Source: §1 · p. 1
• The proposed shared clock allocates more student time to trajectory regions with large normal acceleration while preserving the teacher’s geometric paths and terminal distribution.The clock equalizes a population normal-acceleration statistic under suitable assumptions and is robustly estimated from teacher trajectories.Source: §1 · p. 2
• The clock conditions a single student trained by Lagrangian flow-map distillation, without teacher retraining, architecture changes, or extra inference-time network evaluations.It is estimated once before distillation.Source: §1 · p. 2
• At matched inference budgets, geometry-aware time reparameterization improves sample quality, including one-step generation where intermediate sampling times cannot be adjusted.The one-step gains highlight the role of time reparameterization during distillation itself.Source: §1 · p. 2

2 RELATED WORK

• Flow-map methods learn endpoint predictions or finite-time transitions, while Shortcut Models and MeanFlow use compositional consistency or average velocities; this work builds on Lagrangian flow-map distillation.Related approaches include progressive distillation (Salimans & Ho, 2022), consistency models, and flow map matching (Boffi et al., 2025b).Source: §2 · p. 2
• Trajectory-simplification methods alter the velocity field or source–target coupling, whereas this method preserves a pretrained teacher’s paths and changes their traversal speed.The cited approaches include rectified flow and reflow (Liu et al., 2023), minibatch optimal transport (Tong et al., 2024), and Lee et al. (2023).Source: §2 · p. 2
• Sampling schedules allocate numerical integration times using discrepancy, trajectory, entropy, or acceleration criteria; this method applies its normal-acceleration clock during student training.Sharpen Your Flow (Gupta et al., 2026) uses full material acceleration for sampling schedules, unlike this method’s normal component and training-time use.Source: §2 · pp. 2–3
• Adaptive training methods allocate effort across temporal stages or partition time among specialized networks, whereas this approach trains one student with a fixed teacher-derived clock.Curvature-Adaptive Consistency Flow Matching uses consistency-error feedback; Complexity-Balanced Diffusion Splitting (Issachar et al., 2026) uses acceleration or spatial regularity estimates.Source: §2 · p. 3
• Other time-coordinate methods shape velocity regularity, prescribe conditional profiles, or use pilot-run losses; this work derives a shared clock from a frozen teacher’s normal acceleration.The student is trained directly against the reparameterized teacher velocity in the Lagrangian distillation objective.Source: §2 · p. 3

3.1 FLOW MATCHING

• Flow matching regresses the time-dependent velocity field using interpolant velocities between coupled source and target samples; experiments use linear interpolation.The resulting teacher field defines ODE trajectories that generally differ from individual interpolation paths, so generation still requires numerical integration.Source: §3.1 · pp. 3–4

3.2 FLOW-MAP DISTILLATION

• Flow-map distillation learns two-time maps approximating the teacher ODE’s transitions, replacing repeated teacher-velocity evaluations with finite-time student updates.The Lagrangian objective minimizes a residual under the unique-solution assumption; stop-gradient treatment avoids backpropagation through the teacher’s spatial Jacobian and was important for training stability.Source: §3.2 · p. 4
• Sampling composes learned maps over a chosen time partition, with each update requiring one student evaluation and one-step generation applying the map directly to a source sample.The time-pair sampling distribution is specified in the experimental setup.Source: §3.2 · p. 4

4.1 NORMAL ACCELERATION AS A PROXY FOR DISTILLATION DIFFICULTY

• Normal acceleration is used as a proxy for temporal distillation difficulty because directional changes can alter both the magnitude and direction of a trajectory’s average velocity; this link remains an empirical hypothesis.The normal component captures directional change, unlike tangential acceleration, which changes speed.Source: §4.1 · p. 5

4.2 GEOMETRY-AWARE TIME REPARAMETERIZATION

• The shared time transformation slows traversal through high-hardness regions and accelerates it where hardness is low, aiming to equalize the profile while preserving paths and the terminal distribution.The equalization construction assumes C2 teacher trajectories with nonzero velocity and a positive, continuously differentiable hardness profile.Source: §4.2 · pp. 5–6
• The transformation’s strength is adjustable: δ = 0 gives the identity clock, 0 < δ < 1 softens it, and δ > 1 concentrates it more on high-hardness regions.The equalizing theorem corresponds to δ = 1 and ξ = 0.Source: §4.2 · p. 6
• The clock is estimated before distillation from teacher trajectories using finite-difference acceleration, normal projection, and trimming of the largest 5% of values at each time.A numerical clock density is integrated and fit with a monotonicity-preserving PCHIP interpolant; trimming and approximation mean the practical clock need not exactly equalize the population profile.Source: §4.2 · p. 6
• Lagrangian flow-map distillation conditions the student on reparameterized times while retaining physical-time displacement scaling, separating temporal allocation from velocity scaling.The teacher and clock remain fixed during training.Source: §4.2 · pp. 6–7

5 EXPERIMENTS

• The experiments assess whether training in the clock’s time coordinate improves flow-map distillation after examining temporal geometry and teacher sampling.The evaluation covers synthetic and image datasets.Source: §5 · p. 7

5.1 TEMPORAL GEOMETRY AND TEACHER SAMPLING

• Both teachers have U-shaped normal-acceleration profiles, with larger values near the endpoints; the estimated CelebA64 clock makes the profile more uniform across time.The original teacher’s velocity norm is approximately constant over most of the interval, while the warped velocity exhibits greater variation.Source: §5.1 · p. 7
• With 25 Heun steps, clock-adapted CIFAR-10 sampling lowers FID from 8.8 on the uniform grid to 3.5, with the advantage diminishing at larger evaluation budgets.The teacher-sampling comparison reports FID against velocity-network evaluations.Source: §5.1 · p. 8, Figure 3 · p. 8

5.2.1 TOY EXPERIMENTS

• Across 1, 2, 3, and 4 network evaluations, ϕ improves all three reported checkerboard metrics over Id; samples show sharper boundaries and fewer points outside the target support.The metrics are histogram-based forward KL, Gaussian-kernel MMD, and empirical Wasserstein-2 distance.Source: Figure 4 · p. 8, §5.2.1 · p. 9

5.2.2 IMAGE EXPERIMENTS

• At matched budgets, ϕ improves target FID on both image datasets; at 400K updates, one-step FID falls from 8.57 to 6.99 on CIFAR-10 and from 6.19 to 4.64 on CelebA64.The one-step improvement is consistent across CelebA64 training checkpoints, with a similar trend on CIFAR-10.Source: §5.2.2 · p. 9, Figure 5 · p. 9, Figure 6 · p. 9

6 CONCLUSION

• The proposed geometry-aware time reparameterization 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.Source: §6 · p. 10
• Experiments show improved few-step generation, including one-step gains that highlight the role of time coordinates in learning the map itself.Source: §6 · p. 10
• Establishing a theoretical link between normal acceleration and distillation difficulty remains future work.Source: §6 · p. 10

AI USE STATEMENT

• AI tools supported implementation and debugging, literature searches, manuscript drafting and revision, and figure preparation.Source: AI USE STATEMENT · p. 10
• The authors also used these tools to discuss and refine mathematical arguments and experimental methodology.Source: AI USE STATEMENT · p. 10
• The authors reviewed and revised AI-assisted text, checked mathematical arguments and cited sources, and tested AI-assisted code against the stated methods.They take responsibility for the final manuscript, code, references, and reported findings.Source: AI USE STATEMENT · p. 10

A.1 PROOF OF THE MAIN THEOREM AND CONSTRUCTION OF ϕ

• The proof denotes teacher time by t and student time by τ, with both ranging over [0, 1].Source: §A.1 · p. 13
• For positive, continuously differentiable h_p, the construction yields a strictly increasing C2 map ψ from 0 to 1, whose inverse ϕ is also C2 with ϕ′ > 0.Source: §A.1 · p. 13
• Reparameterization preserves the unit tangent at corresponding times because ϕ′ > 0, and the term containing ϕ′′ is tangential.Source: §A.1 · p. 13
• The same deterministic clock applies to every trajectory, with expectations taken over the same initial distribution; substitution gives an expression independent of τ.Source: §A.1 · p. 13

• The target is uniform on the occupied squares of a two-dimensional checkerboard, with a standard normal source.Source: B TOY EXPERIMENTS: IMPLEMENTATION… · p. 14
• Evaluation uses histogram-based forward KL, Gaussian-kernel MMD, and empirical Wasserstein-2, with comparisons to the target and teacher samples for the latter two.Source: B TOY EXPERIMENTS: IMPLEMENTATION… · p. 14
• The forward-KL estimate uses a 64 × 64 histogram and retains penalties for generated mass outside the domain; its value depends on resolution, sample count, and regularization.Source: B TOY EXPERIMENTS: IMPLEMENTATION… · pp. 14–15
• MMD uses a Gaussian kernel and reports a biased estimate of the squared population discrepancy.Source: B TOY EXPERIMENTS: IMPLEMENTATION… · p. 15
• Empirical Wasserstein-2 is computed by solving a minimum-cost assignment between equally sized point clouds using squared-distance costs.Source: B TOY EXPERIMENTS: IMPLEMENTATION… · p. 15

C.1 SHARED DISTILLATION SETUP

• For image experiments, the clock remains fixed while student inputs use transformed time coordinates paired with the corresponding interpolant state.Source: §C.1 · p. 15
• Sampling uniformly in physical time yielded further improvements over identity time and was used for all reported image results.Source: §C.1 · p. 15
• A learned weighting head rescales time-pair contributions and penalizes arbitrarily large weights; it is used only during training, while evaluation matches checkpoints and network-evaluation counts.Source: §C.1 · p. 16
• Each reported configuration uses one run, preliminary repeats had similar curves and negligible evaluation differences, and selected learning-rate settings matched across clocks.Source: §C.1 · p. 16

C.2 CIFAR-10

• CIFAR-10 uses its 50,000 native-resolution training images with horizontal flips and a guided-diffusion U-Net configured with two-time student conditioning.Source: §C.2 · p. 16
• The CIFAR-10 experiment uses a fixed released teacher and reports checkpoint results across sampling budgets; Table 1 covers both FID and KID comparisons.Source: §C.2 · p. 17, Table 1 · p. 17

C.3 CELEBA64

• CelebA64 uses 162,770 training images, a fixed validation subset, independent source-data coupling, and an EDM2 U-Net with approximately 124 million teacher parameters.Source: §C.3 · pp. 17–18
• The frozen CelebA64 teacher scores approximately FID 2.36 and KID 1.46 · 10−3; student training uses four A100 GPUs, and Table 2 spans checkpoints and sampling budgets.Source: Table 2 · p. 18, §C.3 · p. 19