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

TL;DR

This paper introduces a fast parallel backpropagation algorithm for inverting convolutions, significantly improving the efficiency of Normalizing Flows and related image processing tasks by reducing computation time from cubic to square root order.

Contribution

The authors develop a novel parallel backpropagation method for convolution inversion with $O(\sqrt{n})$ complexity and demonstrate its application in Normalizing Flows for faster sampling and training.

Findings

Achieved $O(\sqrt{n})$ backpropagation time for convolution inversion.

Implemented Inverse-Flow models with improved sampling speed.

Maintained comparable bits per dimension to previous models.

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.