Flash-DMD: 효율적 증류와 결합 강화 학습을 통한 고품질 Few-Step 이미지 생성 연구

초록

디퓨전 모델은 생성 모델의 주요 계층으로 부상했지만, 반복적인 샘플링 과정은 계산 비용이 높은 문제를 지니고 있습니다. 타임스텝 디스틸레이션은 생성 속도를 높이기 위한 유망한 기술이지만, 종종 대규모 학습을 필요로 하고 이미지 품질 저하를 초래합니다. 더욱이 미적 매력이나 사용자 선호도와 같은 특정 목적을 위해 강화학습(RL)을 사용하여 이러한 디스틸레이션 모델을 미세 조정하는 것은 불안정하고 보장 해킹에 쉽게 빠지는 것으로 알려져 있습니다. 본 연구에서는 디스틸레이션과 결합된 RL 기반 정교화를 통해 빠른 수렴을 가능하게 하는 새로운 프레임워크인 Flash-DMD를 소개합니다. 구체적으로, 우리는 먼저 향상된 현실감과 함께 학습 비용을 크게 절감하는 효율적인 타임스텝 인식 디스틸레이션 전략을 제안하며, DMD2 대비 단 2.1%의 학습 비용만으로 더 우수한 성능을 달성합니다. 둘째, 타임스텝 디스틸레이션 학습이 지속되는 동시에 RL 목적 함수를 통해 모델을 미세 조정하는 결합 학습 방식을 도입합니다. 지속되는 디스틸레이션으로부터의 안정적이고 명확한 손실이 강력한 정규화자로 작용하여 RL 학습 과정을 효과적으로 안정화하고 정책 붕괴를 방지한다는 것을 입증합니다. 스코어 기반 및 플로우 매칭 모델에 대한 광범위한 실험을 통해, 제안된 Flash-DMD가 유의미하게 빠르게 수렴할 뿐만 아니라 Few-Step 샘플링 영역에서 최첨단 생성 품질을 달성하여 시각적 품질, 인간 선호도, 텍스트-이미지 정렬 메트릭에서 기존 방법들을 능가함을 보여줍니다. 우리의 연구는 효율적이고 높은 충실도를 가지며 안정적인 생성 모델을 학습시키는 효과적인 패러다임을 제시합니다. 코드는 곧 공개될 예정입니다.

English

Diffusion Models have emerged as a leading class of generative models, yet their iterative sampling process remains computationally expensive. Timestep distillation is a promising technique to accelerate generation, but it often requires extensive training and leads to image quality degradation. Furthermore, fine-tuning these distilled models for specific objectives, such as aesthetic appeal or user preference, using Reinforcement Learning (RL) is notoriously unstable and easily falls into reward hacking. In this work, we introduce Flash-DMD, a novel framework that enables fast convergence with distillation and joint RL-based refinement. Specifically, we first propose an efficient timestep-aware distillation strategy that significantly reduces training cost with enhanced realism, outperforming DMD2 with only 2.1% its training cost. Second, we introduce a joint training scheme where the model is fine-tuned with an RL objective while the timestep distillation training continues simultaneously. We demonstrate that the stable, well-defined loss from the ongoing distillation acts as a powerful regularizer, effectively stabilizing the RL training process and preventing policy collapse. Extensive experiments on score-based and flow matching models show that our proposed Flash-DMD not only converges significantly faster but also achieves state-of-the-art generation quality in the few-step sampling regime, outperforming existing methods in visual quality, human preference, and text-image alignment metrics. Our work presents an effective paradigm for training efficient, high-fidelity, and stable generative models. Codes are coming soon.

Flash-DMD: 효율적 증류와 결합 강화 학습을 통한 고품질 Few-Step 이미지 생성 연구

Flash-DMD: Towards High-Fidelity Few-Step Image Generation with Efficient Distillation and Joint Reinforcement Learning

초록

Support