Geometry-Aware Time Reparameterization for Flow-Map Distillation

2026-10-01

AI Summary

This study addresses the challenge in flow matching distillation where excessive normal acceleration along trajectories impedes the learning of finite-time transitions. To overcome this, we propose a geometry-aware temporal reparameterization method. By introducing a shared clock mechanism, our approach preserves the geometric paths of the teacher model while balancing the statistics of normal acceleration. Furthermore, we construct an efficient approximation scheme by integrating Lagrangian flow matching distillation with robust regularized estimation. Experimental results demonstrate that the proposed method significantly improves image generation quality for both single-step and few-step settings on benchmarks such as CIFAR-10, without incurring any additional inference overhead.

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.

Problem

• Flow-map distillation
• Generative ODE
• Time reparameterization
• Few-step generation
• Normal acceleration

Innovation

• Flow-map distillation
• Geometry-aware time reparameterization
• Normal acceleration
• Shared clock
• Lagrangian flow-map distillation