Restoring Initial Noise Sensitivity in Text-to-Image Distillation through Geometric Alignment

Generative distillation significantly accelerates text-to-image (T2I) generation by compressing multi-step trajectories into few-step student models while preserving perceptual quality. However, existing methods primarily optimize efficiency and output fidelity, often neglecting critical properties of the original trajectory. In this work, we identify a key missing property: sensitivity to initial noise, whose degradation impairs downstream control methods relying on noise-based optimization and manipulation. We trace this issue to standard distillation objectives that enforce pointwise output alignment, inadvertently flattening the input-output landscape and suppressing the teacher’s local geometric structure. To address this, we propose Geometry-Aware Distillation (GAD), a sensitivity-preserving framework that aligns the local functional behavior of teacher and student models. Specifically, GAD matches Jacobian-vector products with respect to input noise, enabling the student to reproduce the teacher’s differential response to perturbations. Extensive experiments across multiple T2I paradigms and noise-driven control tasks demonstrate that GAD significantly restores sensitivity and improves diversity while maintaining high visual fidelity. Code is available at https://github.com/Hannah1102/GAD.

Modern AI image generators can create impressive pictures from text prompts, but the fastest versions of these systems are often produced by “compressing” large models into smaller and more efficient ones. During this compression process, the smaller models tend to lose sensitivity to their random starting noise, causing them to generate images that look repetitive and are harder to control. In this work, we show that the problem is not just a loss of randomness, but a loss of the underlying geometric behavior learned by the original model. Existing training methods focus mainly on matching final outputs, while ignoring how the model should react to small input changes. We introduce a new training method called Geometry-Aware Distillation (GAD), which teaches compact models to preserve these local response patterns during training. This allows the accelerated model to better follow the behavior of the original system while remaining efficient. Experiments across several state-of-the-art image generation frameworks show that our method improves image diversity, controllability, and overall generation quality without adding extra cost during inference. Our findings suggest that preserving local geometric structure is important for building faster generative AI systems that remain reliable and expressive.