Parallel Backpropagation for Inverse of a Convolution with Application to Normalizing Flows

TL;DR

This work tackles the computational bottleneck of backpropagating through the inverse of a convolution, a key operation in Normalizing Flows and related imaging tasks. It derives a fast parallel backpropagation algorithm with running time $O(mk^2)$ on square images and provides a CUDA GPU implementation. By pairing this with Inverse-Flow, a multi-scale flow that uses the inverse of convolution in the forward pass and conventional convolution for sampling, the authors achieve significantly faster sampling while maintaining competitive bits-per-dimension. Experiments on MNIST and CIFAR-10 demonstrate substantial speedups in sampling without sacrificing density estimation quality, highlighting the practical impact for scalable, fast generative modeling.

Abstract

The inverse of an invertible convolution is an important operation that comes up in Normalizing Flows, Image Deblurring, etc. The naive algorithm for backpropagation of this operation using Gaussian elimination has running time $O(n^3)$ where $n$ is the number of pixels in the image. We give a fast parallel backpropagation algorithm with running time $O(\sqrt{n})$ for a square image and provide a GPU implementation of the same. Inverse of Convolutions are usually used in Normalizing Flows in the sampling pass, making them slow. We propose to use the Inverse of Convolutions in the forward (image to latent vector) pass of the Normalizing flow. Since the sampling pass is the inverse of the forward pass, it will use convolutions only, resulting in efficient sampling times. We use our parallel backpropagation algorithm to optimize the inverse of the convolution layer, resulting in fast training times. We implement this approach in various Normalizing Flow backbones, resulting in our Inverse-Flow models. We benchmark Inverse-Flow on standard datasets and show significantly improved sampling times with similar bits per dimension compared to previous models.

Figures (5)

• Figure 1: a). Images from MNIST dataset. b). Reconstructed images using an Inverse-Flow model based on the inv-conv layer for a forward pass.
• Figure 2: Invertible convolution with zero padding (top, left) on input $x$ and masking of kernel $W_{k,k} = 1$
• Figure 3: Multi-scale architecture of Inverse-Flow model and Inv Flow step.
• Figure 4: The convolution of a top-left padded image ($x_{m,m}$) with filter ($W$) involves a linear transformation of the input ($x$) via a lower triangular convolution matrix ($M_{m^2,m^2}$), where the diagonal elements correspond to the filter weights ($W_{k,k}$). Each row of $M$ computes a pixel, and these rows can be inverted in parallel using the steps outlined in the inversion algorithm by Kallappa et al. (2023).
• Figure 5: Overview of a small image classification model with two inverse of convolution ( $3 \times 3$ inv-con) layers with $97.6\%$ accuracy on MNIST dataset.

Theorems & Definitions (4)

• Theorem 1
• Theorem 2
• proof
• proof