About MaskUnet

MaskUnet is an official PyTorch implementation of the CVPR 2025 paper Not All Parameters Matter: Masking Diffusion Models for Enhancing Generation Ability. It introduces a novel method that improves diffusion model generation quality by masking specific U-Net parameters, challenging the assumption that all parameters are necessary throughout the denoising process. The approach leverages the finding that certain parameters, including those with large magnitudes, can be selectively disabled to enhance performance. MaskUnet dynamically identifies timestep- and sample-dependent effective parameters to optimize the generation process. The framework supports two operational modes: a training-based optimization setting and a training-free setting that requires no additional parameter updates. Designed to be lightweight, it achieves significant performance gains on benchmarks like COCO and demonstrates strong generalization to downstream tasks while introducing negligible computational overhead. The tool is applicabl

g

Published by

gudaochangsheng

Visit View Profile

README.md

View on GitHub

🚀 MaskUnet: Not All Parameters Matter for Enhancing Diffusion Generation

If you like our project, please give us a star ⭐ on GitHub for the latest update.

CVPR 2025

🔥 Highlights

A surprising finding: masking a subset of U-Net parameters can improve diffusion generation quality.
We propose MaskUnet, a simple yet effective method that leverages timestep- and sample-dependent effective parameters.
Supports both:
- Training-based optimization
- Training-free optimization
Achieves strong performance on COCO and downstream tasks with negligible additional parameters.

🖼️ Motivation

Analysis of parameter distributions and denoising effects across different timesteps for Stable Diffusion 1.5 with and without random masking. The first column shows the parameter distribution of SD 1.5; the second to fifth columns show the distributions of removed parameters under random masking. The last two columns compare generated samples from SD 1.5 and the random mask.

📘 Overview

Diffusion models gradually generate images from coarse structure to fine details across timesteps. However, standard U-Nets use the same parameters throughout the whole denoising process, forcing them to handle both structural and textural information simultaneously.

This motivates us to study time-wise effective parameters in diffusion models.

We find that properly masking certain U-Net parameters—including large-magnitude ones—can actually help denoising and improve generation quality on the fly. Based on this observation, we propose MaskUnet, a lightweight method that dynamically exploits timestep- and sample-dependent effective parameters.

To optimize MaskUnet, we provide two practical settings:

Training-based MaskUnet
Training-free MaskUnet

MaskUnet achieves strong zero-shot generation performance on COCO and also generalizes well to downstream tasks.

🧠 Method

Pipeline of MaskUnet. G-Sig denotes the Gumbel-Sigmoid activation function. GAP denotes global average pooling.

✨ Qualitative Results

Qualitative comparison with existing methods.

📈 Quantitative Results

quantitative results

⚡ Quick Start

Installation

conda env create -f environment.yaml
conda activate maskunet

Dataset

fantasyfish/laion-art
- Hugging Face
- HF Mirror

🏋️ Training-based MaskUnet

Train

bash ./training/train_hyperunet.sh

Inference

bash ./training/infer_sd1-5_hardmask.sh

🎬 Training-free MaskUnet

bash ./training-free/infer_sd1-5_x0_optim_mask_fnal_para.sh

📌 Notes

The training-based version learns a masking strategy through optimization.
The training-free version directly performs mask optimization at inference time.
Both settings are built upon Stable Diffusion 1.5.

📚 Citation

If you find this project useful, please consider giving it a star ⭐ and citing our paper.

@inproceedings{wang2025not,
  title={Not All Parameters Matter: Masking Diffusion Models for Enhancing Generation Ability},
  author={Wang, Lei and Li, Senmao and Yang, Fei and Wang, Jianye and Zhang, Ziheng and Liu, Yuhan and Wang, Yaxing and Yang, Jian},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={12880--12890},
  year={2025}
}

🙏 Acknowledgement

This project is based on Diffusers.
Thanks to the Diffusers team for their awesome work.

📮 Contact

If you have any questions, please feel free to reach out to:[email protected]

MaskUnet