Not All Parameters Matter: Masking Diffusion Models for Enhancing Generation Ability

TL;DR

This work shows that diffusion models can be substantially improved without updating pretrained U-Net weights by learning where to activate parameters at each timestep and for each input. The proposed MaskUNet framework introduces a learnable mask (via a lightweight mask generator) that yields timestep- and sample-dependent parameter selection, with a training-based path using diffusion loss and a training-free path guided by reward models. Empirically, MaskUNet delivers notable gains in zero-shot COCO generation (e.g., reduced FID by about 1.1 points) and strengthens performance on downstream tasks such as image customization, relation inversion, and text-to-video generation, while also improving semantic-binding benchmarks. The approach preserves generalization, requires minimal additional parameters, and offers practical benefits for applications demanding controllable, high-fidelity image synthesis without full fine-tuning of the U-Net.

Abstract

The diffusion models, in early stages focus on constructing basic image structures, while the refined details, including local features and textures, are generated in later stages. Thus the same network layers are forced to learn both structural and textural information simultaneously, significantly differing from the traditional deep learning architectures (e.g., ResNet or GANs) which captures or generates the image semantic information at different layers. This difference inspires us to explore the time-wise diffusion models. We initially investigate the key contributions of the U-Net parameters to the denoising process and identify that properly zeroing out certain parameters (including large parameters) contributes to denoising, substantially improving the generation quality on the fly. Capitalizing on this discovery, we propose a simple yet effective method-termed ``MaskUNet''- that enhances generation quality with negligible parameter numbers. Our method fully leverages timestep- and sample-dependent effective U-Net parameters. To optimize MaskUNet, we offer two fine-tuning strategies: a training-based approach and a training-free approach, including tailored networks and optimization functions. In zero-shot inference on the COCO dataset, MaskUNet achieves the best FID score and further demonstrates its effectiveness in downstream task evaluations. Project page: https://gudaochangsheng.github.io/MaskUnet-Page/

Figures (9)

• Figure 1: The motivation of our method.
• Figure 2: The pipeline of the MaskUnet. G-Sig represents the Gumbel-Sigmoid activate function. GAP is global average pooling.
• Figure 3: Quality results compared to other methods.
• Figure 4: Quality results compared to other methods.
• Figure 5: Quality results by Textual Inversion galimage with or without mask.