Binarized Diffusion Model for Image Super-Resolution

TL;DR

A novel binarized diffusion model, BI-DiffSR, for image SR is introduced and the consistent-pixel-downsample and consistent-pixel-upsample are proposed to maintain dimension consistent and facilitate the full-precision information transfer.

Abstract

Advanced diffusion models (DMs) perform impressively in image super-resolution (SR), but the high memory and computational costs hinder their deployment. Binarization, an ultra-compression algorithm, offers the potential for effectively accelerating DMs. Nonetheless, due to the model structure and the multi-step iterative attribute of DMs, existing binarization methods result in significant performance degradation. In this paper, we introduce a novel binarized diffusion model, BI-DiffSR, for image SR. First, for the model structure, we design a UNet architecture optimized for binarization. We propose the consistent-pixel-downsample (CP-Down) and consistent-pixel-upsample (CP-Up) to maintain dimension consistent and facilitate the full-precision information transfer. Meanwhile, we design the channel-shuffle-fusion (CS-Fusion) to enhance feature fusion in skip connection. Second, for the activation difference across timestep, we design the timestep-aware redistribution (TaR) and activation function (TaA). The TaR and TaA dynamically adjust the distribution of activations based on different timesteps, improving the flexibility and representation alability of the binarized module. Comprehensive experiments demonstrate that our BI-DiffSR outperforms existing binarization methods. Code is released at: https://github.com/zhengchen1999/BI-DiffSR.

Figures (8)

• Figure 1: Visual comparison ($\times$4) of binarization methods. Some methods (e.g., BNN hubara2016binarized) cannot work on diffusion models. Several methods (e.g., BBCU xia2022basic) suffer from blurring and artifacts. In contrast, our proposed BI-DiffSR outperforms other methods with accurate results.
• Figure 2: The overall structure of the noise estimation network. (a) UNet: The model consists of ResBlock, CP-Down, CP-Up, and CS-Fusion. It predicts noise $\boldsymbol{\epsilon}_t$ with the upscaled LR image $\mathbf{y}$, noise image $\mathbf{x}_t$, and timestep $t$. (b) ResBlock: Residual block, utilizes the binarized convolution (BI-Conv) block. The input and output dimensions of the block remain consistent, making it suitable for binarization. (c) TE: Time encoding, encoders timestep $t$ to produce the timestep embedding $\mathbf{t}_{em}$.
• Figure 3: (a) CP-Down: Consistent-pixel-downsample. (b) CP-Up: Consistent-pixel-upsample. (c) CS-Fusion: Channel-shuffle fusion. (d) In the skip connection, the value ranges of two features ($\mathbf{x}_{1}$, $\mathbf{x}_{2}$) may be significant differences, which impedes effective fusion. (e) The illustration of channel shuffle. the shuffled features ($\mathbf{x}_{1}^{sh}$, $\mathbf{x}_{2}^{sh}$) have closely matched value ranges.
• Figure 4: Visualization of the changes in activation distribution across 50 timesteps.
• Figure 5: (a) The basic binarized convolutional (BI-Conv) block. The learnable bias $\mathbf{b}$ and the activation function RPReLU adjust the activations. (b) In timestep-aware redistribution (TaR) and activation function (TaA), multiple pairs of $\mathbf{b}$ and RPReLU are applied to adapt to the multi-step in DM. At each step $t$, only one pair of $\mathbf{b}$ and RPReLU is used (the darker modules with solid lines).