Improved diffusion model fine-tuning with principal timestep initialization
Principal-timestep Restricted Init via Sparse Matrix-decomposition in Flow-matching
Machine Learning
Summary
Fine-tuning diffusion models for image generation can be very slow and inefficient. The authors found that existing methods that limit updates to low-rank changes don’t work well because gradients from different training steps vary too much. They introduced a technique that picks important training steps and filters out irrelevant features to align these updates better. This makes fine-tuning faster and leads to better results on tasks like personalized image generation and image restoration.
What this means in practice
- •For computer vision engineers: Speed up and improve fine-tuning of diffusion models for subject-driven and controllable image generation tasks.
- •For software developers in image restoration: Enhance convergence and final quality when adapting diffusion models for image deblurring applications.
Authors
Jiayang Gu, Zheng Fang, Lichaun Xiang, Fanghui Liu, Xu Cai, Hongkai Wen
Abstract
Flow-matching diffusion models have recently emerged as a strong paradigm for high-fidelity visual generation. However, their prohibitively high fine-tuning cost limits scalability to downstream tasks. While Low-Rank Adaptation (LoRA) combined with spectral initialization has demonstrated accelerated convergence and improved performance in autoregressive language models by better aligning gradient directions, we find that it fails to deliver similar gains in diffusion fine-tuning, often yielding marginal or even negative improvements over vanilla LoRA.We attribute this discrepancy to a fundamental mismatch between LoRA's low-rank parameterization and the intrinsically high-rank gradients induced by the flow-matching objective. In particular, stochastic timestep sampling introduces directionally heterogeneous gradient signals across training steps, leading to misaligned updates under low-rank constraints.To address this issue, we propose Prism-LoRA,a Principal-timestep Restricted Init via Sparse Matrix-decomposition framework that improves gradient alignment during fine-tuning. Our method consists of two key components: (i) principal timestep selection, which restricts initialization gradients to a subset of dominant timesteps to suppress effective gradient rank, and (ii) principal channel filtering, which removes task-irrelevant channels, enabling the one-step spectral initialization gradient to better align with the long-horizon optimization trajectory. Extensive experiments demonstrate that our method consistently improves both convergence speed and final performance across multiple diffusion fine-tuning benchmarks, including subject-driven generation, controllable generation, and deblurring, achieving not only performance improvement but also earlier stages of convergence over baseline LoRA and other spectral-init methods.