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

2024-10-18

arXiv.org

AI Summary

To address the high computational cost of backpropagation through invertible convolutions in normalizing flows—where conventional Gaussian elimination incurs O(n³) complexity—this paper proposes a GPU-accelerated parallel backpropagation algorithm, reducing inversion complexity to O(√n). It introduces the novel design of moving invertible convolution into the forward mapping, enabling sampling to rely solely on standard (non-invertible) convolutions; this decouples sampling efficiency from training. Based on this insight, we propose Invert-Flow, a new architectural paradigm that refactors RealNVP and Glow. The method integrates parallel linear algebra, CUDA-optimized kernels, and principled reversible network design. Experiments demonstrate significantly faster sampling, competitive bits-per-dimension (BPD) performance against state-of-the-art methods, accelerated training convergence, and we publicly release an efficient implementation of the invertible convolution layer.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Search and Optimization: Non-convex Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work

Abstract

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 Convolutions are usually used in Normalizing Flows in the sampling pass, making them slow. We propose to use Inverse 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 for optimizing the inverse convolution layer resulting in fast training times also. 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.

Problem

• Efficient inversion of convolutions in Normalizing Flows.
• Parallel backpropagation reduces training time significantly.
• Improved sampling times in Inverse-Flow models.

Innovation

• Parallel backpropagation algorithm
• GPU implementation for efficiency
• Inverse-Flow models optimization