DSA-Diff: Dynamic schedule alignment for training-Inference consistent modality translation in x-prediction diffusion model.

Abstract

For modality translation tasks, diffusion models based on x-prediction offer faster and more accurate image generation compared to traditional ϵ-prediction. However, they often suffer from training-inference inconsistency (TII), which arises from a mismatch between the Gaussian distribution assumed by the preset noise schedule and the true data distribution. To address this, we propose DSA-Diff, a novel framework that employs dual noise schedules to decouple the training and inference processes. Our approach decomposes the noise schedule along three dimensions: noise sequence, timestep, and correction matrix, and introduces a Bayesian-Greedy Alignment Scheduler (BGAS) to dynamically reconstruct the inference schedule. BGAS combines greedy initialization and Bayesian optimization to align the generated data manifold with the true one. Additionally, we introduce progressive target prediction and multi-scale perceptual alignment to enhance the robustness and detail fidelity of the x-prediction model. Experiments on four datasets show that DSA-Diff achieves high-fidelity image synthesis in only 4-10 adaptive inference steps, with minimal computational cost (68 GFLOPS). It improves the SSIM metric by up to 2.56% in TFW dataset using only one additional algorithmic module, effectively mitigating TII. Code and models are available at: https://github.com/ElephantOH/DSA-Diff.