LMD: Faster Image Reconstruction with Latent Masking Diffusion
Zhiyuan Ma zhihuan yu Jianjun Li Bowen Zhou
As a class of fruitful approaches, diffusion probabilistic models (DPMs) have
shown excellent advantages in high-resolution image reconstruction. On the
other hand, masked autoencoders (MAEs), as popular self-supervised vision
learners, have demonstrated simpler and more effective image reconstruction and
transfer capabilities on downstream tasks. However, they all require extremely
high training costs, either due to inherent high temporal-dependence (i.e.,
excessively long diffusion steps) or due to artificially low spatial-dependence
(i.e., human-formulated high mask ratio, such as 0.75). To the end, this paper
presents LMD, a faster image reconstruction framework with latent masking
diffusion. First, we propose to project and reconstruct images in latent space
through a pre-trained variational autoencoder, which is theoretically more
efficient than in the pixel-based space. Then, we combine the advantages of
MAEs and DPMs to design a progressive masking diffusion model, which gradually
increases the masking proportion by three different schedulers and reconstructs
the latent features from simple to difficult, without sequentially performing
denoising diffusion as in DPMs or using fixed high masking ratio as in MAEs, so
as to alleviate the high training time-consumption predicament. Our approach
allows for learning high-capacity models and accelerate their training (by 3x
or more) and barely reduces the original accuracy. Inference speed in
downstream tasks also significantly outperforms the previous approaches.