Accelerating Inference of Discrete Autoregressive Normalizing Flows by Selective Jacobi Decoding

TL;DR

This paper introduces a selective Jacobi decoding method to accelerate inference in discrete autoregressive normalizing flows, achieving up to 4.7x speedup while maintaining quality.

Contribution

It proposes a novel acceleration technique that reduces inference time by leveraging low dependency redundancy and parallel iterative optimization.

Findings

Achieves up to 4.7 times faster inference on multiple datasets.

Theoretical analysis shows superlinear convergence and iteration bounds.

Empirical results confirm effectiveness across various models and datasets.

Abstract

Discrete normalizing flows are promising generative models with advantages such as analytical log-likelihood computation and end-to-end training. However, the architectural constraints to ensure invertibility and tractable Jacobian computation limit their expressive power and practical usability. Recent advancements utilize autoregressive modeling, significantly enhancing expressive power and generation quality. Nevertheless, such sequential modeling inherently restricts parallel computation during inference, leading to slow generation that impedes practical deployment. In this paper, we first identify that strict sequential dependency in inference is unnecessary to generate high-quality samples. We observe that sub-variables in sequential modeling can also be approximated without strictly conditioning on all preceding sub-variables. Moreover, the models tend to exhibit low dependency redundancy in the initial layer and higher redundancy in subsequent layers. Leveraging these observations, we propose to selectively use Jacobi decoding strategy that accelerates its autoregressive inference through parallel iterative optimization. Theoretical analyses demonstrate the method's superlinear convergence rate and guarantee that the number of iterations required is no greater than the original sequential approach. Empirical evaluations across multiple datasets validate the generality and effectiveness of our acceleration technique, achieving up to 4.7 times faster inference on modern normalizing flow models while preserving generation quality.