🤖 machine learning

Accelerating Inference of Discrete Autoregressive Normalizing Flows by Selective Jacobi Decoding

This paper proposes a selective Jacobi decoding strategy that accelerates the inference of discrete autoregressive normalizing flows by leveraging observed dependency redundancy to enable parallel iterative optimization, achieving up to 4.7 times faster generation without compromising quality.

Original authors: Jiaru Zhang, Juanwu Lu, Xiaoyu Wu, Ziran Wang, Ruqi Zhang

Published 2026-05-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiaru Zhang, Juanwu Lu, Xiaoyu Wu, Ziran Wang, Ruqi Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Slow Writer" Problem

Imagine you have a very talented AI artist (a Discrete Autoregressive Normalizing Flow) that can create beautiful pictures from scratch. This artist is famous for two things:

Precision: It can calculate exactly how likely a specific image is to exist (a mathematical superpower).
Quality: The pictures it makes are sharp and realistic.

However, this artist has a major flaw: they are incredibly slow.

Why? Because this artist works like a strict, old-fashioned writer. To write a sentence, they must write the first word, then the second word, then the third, and so on. They cannot write the second word until the first is finished. They cannot write the third until the second is done.

In the world of AI, this is called sequential inference. If you want to generate an image with 1,000 "words" (pixels or patches), the AI has to do 1,000 steps one after another. It's like waiting for a snail to cross a highway. This slowness makes it hard to use the AI in real-time applications.

The Discovery: "Do We Really Need to Wait?"

The researchers (Zhang, Lu, et al.) asked a simple question: "Does the artist really need to wait for the previous word to be perfectly finished before guessing the next one?"

They found that the answer is no.

The Analogy: Imagine you are reading a mystery novel. If you miss the very first sentence, you might be confused. But if you miss the 50th sentence, you can probably guess what happens next based on the general vibe of the story. The story has "redundancy." You don't need 100% perfect information from the past to predict the future; a good guess is often enough.

The researchers discovered that in these AI models, the "dependency" on previous steps is often weaker than we thought, especially as the AI gets deeper into the generation process.

The Solution: The "Group Guessing" Strategy (Jacobi Decoding)

To speed things up, the researchers introduced a technique called Jacobi Decoding.

The Old Way (Sequential): One person writes a sentence. Then they pass the pen to the next person. Then the next. It takes a long time.
The New Way (Jacobi): Imagine a team of 10 people sitting in a circle. Instead of passing the pen, they all write their part of the sentence at the same time, based on what they think the others wrote in the previous round.
- Round 1: Everyone guesses their word based on a rough draft.
- Round 2: They look at everyone's guesses from Round 1, refine their own word, and write again.
- Round 3: They refine again.

Because they are all working at the same time (parallel processing), this is much faster. Usually, after just a few rounds of "guessing and refining," everyone is writing the exact same sentence the slow, single writer would have produced.

The Twist: "Selective" Decoding

The researchers realized that you can't just use this "Group Guessing" method for everything.

The First Step is Critical: The very first part of the image (the seed) is like the foundation of a house. If you guess the foundation wrong, the whole house collapses. The first step needs to be done carefully and sequentially.
The Rest is Flexible: Once the foundation is set, the walls and roof can be built using the "Group Guessing" method because the structure is already there to support them.

So, they created a Selective strategy:

Step 1: Do it the slow, careful, sequential way (to get the foundation right).
Steps 2 to End: Switch to the fast, parallel "Group Guessing" method.

The Results: Speed Without Sacrificing Quality

The paper tested this on several datasets (CIFAR-10, CIFAR-100, and AFHQ, which are collections of images).

Speed: The new method was up to 4.7 times faster than the old slow method. On some tasks, it was even faster.
Quality: The pictures looked almost identical to the ones made by the slow method. The "Group Guessing" was so good that the human eye couldn't tell the difference.
Math Proof: The authors didn't just guess; they proved mathematically that this method will always converge (finish correctly) and that it does so very quickly (superlinear convergence).

Summary Analogy

Think of the AI model as a relay race.

The Old Way: The runners pass the baton one by one. Runner 1 runs, passes to Runner 2, who runs, passes to Runner 3. It takes a long time.
The Paper's Way: The researchers realized that for most of the race, the runners don't need the baton to be perfectly in hand to start sprinting. They can all start running at the same time, adjusting their speed as they see where the others are.
The Catch: The first runner still needs to start perfectly (the "Selective" part). But once the race is underway, everyone runs in parallel, finishing the race in a fraction of the time.

In short: The paper found a way to make a very smart but slow AI artist work much faster by letting it "guess and refine" multiple parts of an image simultaneously, without ruining the final picture.

Technical Summary: Accelerating Inference of Discrete Autoregressive Normalizing Flows by Selective Jacobi Decoding

Problem Statement
Discrete normalizing flows are powerful generative models capable of analytical log-likelihood computation and end-to-end training. Recent advancements, such as TarFlow, have enhanced their expressive power by employing autoregressive modeling, which splits data into sequences and uses masked autoregressive transformations. However, this sequential architecture imposes a critical bottleneck during inference: the inverse transformation requires generating each sub-variable strictly conditioned on all previously generated sub-variables. This inherent sequential dependency prohibits parallel computation, resulting in slow generation speeds that hinder practical deployment. While inference acceleration techniques exist for Diffusion Models, VAEs, and GANs, they typically rely on specific architectural properties or training paradigms not directly transferable to autoregressive normalizing flows.

Methodology
The authors propose a method to accelerate inference by selectively applying Jacobi decoding, a parallel iterative optimization strategy, to discrete autoregressive normalizing flows. The approach is grounded in two key empirical and theoretical observations:

Sequential Redundancy: The strict dependency on all preceding sub-variables for generating the current sub-variable contains significant redundancy. The paper demonstrates that sub-variables can be reasonably approximated without conditioning on the nearest preceding elements, suggesting that the strict sequential constraint is not always necessary for high-fidelity generation.
Depthwise Heterogeneity: The degree of this redundancy varies across the layers of the flow. The initial layer (transforming Gaussian noise) exhibits strong dependencies and low redundancy, whereas subsequent layers, which refine informative outputs, exhibit higher redundancy and weaker dependencies on immediate predecessors.

Based on these observations, the authors introduce Selective Jacobi Decoding (SJD):

Parallel Iteration: Instead of the standard sequential (Gauss-Seidel-like) update, the method reformulates the inference task as solving a system of non-linear equations. It employs Jacobi iterations where all sub-variables in a layer are updated simultaneously using values from the previous iteration step.
Selective Application: To avoid the computational overhead of parallel updates in layers with strong dependencies (specifically the first layer), the method uses standard sequential decoding for the first layer and switches to parallel Jacobi iterations for the remaining layers where redundancy is higher.
Convergence Guarantees: The method utilizes a stopping threshold $\tau$ to determine when iterations cease. Theoretically, the authors prove that the iteration exhibits superlinear convergence under appropriate initialization and is guaranteed to converge to the exact solution within at most $L$ iterations (where $L$ is the sequence length), providing a strict upper bound on computational steps.

Key Contributions

Discovery of Redundancy: The paper identifies that strict sequential dependency in discrete autoregressive normalizing flows contains exploitable redundancy, which varies significantly across different layers (low in the first layer, high in subsequent layers).
Selective Acceleration Framework: It introduces a novel inference strategy that selectively applies Jacobi decoding to high-redundancy layers while retaining sequential decoding for the dependency-heavy first layer. This approach requires no additional training or architectural modifications to the original model.
Theoretical and Empirical Validation: The work provides theoretical proofs for superlinear convergence and finite convergence guarantees. Comprehensive experiments on CIFAR-10, CIFAR-100, and AFHQ validate the method's effectiveness, showing significant speedups with minimal impact on generation quality.

Experimental Results
The proposed method was evaluated on TarFlow models across three datasets:

Speedup: The approach achieved up to 4.7 times faster inference compared to standard sequential decoding. Specifically, on CIFAR-100, the speedup was 4.7x; on CIFAR-10, 3.6x; and on AFHQ, 4.5x.
Quality Preservation: Quantitative metrics (FID, CLIP-IQA, BRISQUE) indicate that generation quality is largely preserved. For instance, on CIFAR-100, the FID score changed from 8.22 (sequential) to 8.19 (SJD), and visual inspections confirmed high perceptual fidelity.
Comparison with Uniform Jacobi: A uniform application of Jacobi decoding (UJD) across all layers failed to accelerate inference on the larger AFHQ dataset (actually slowing it down to 0.8x), highlighting the necessity of the selective strategy to avoid overhead in high-dependency layers.
Additional Tasks: The method also demonstrated effectiveness on Masked Autoregressive Flows (MAFs) for Boltzmann distribution approximation (15.7x speedup) and binary MNIST generation (18.4x speedup).

Significance and Claims
The paper claims to be the first work to explore inference acceleration specifically for normalizing flow models. Its primary significance lies in overcoming the sequential bottleneck of autoregressive normalizing flows without compromising the model's invertibility or requiring retraining. By leveraging the observed structural redundancy and layer-wise heterogeneity, the method enables practical deployment of these high-quality generative models. The authors emphasize that their approach offers a "worst-case guarantee" on convergence steps and achieves a favorable trade-off between inference speed and sample quality, making discrete autoregressive normalizing flows more viable for real-world applications.