PaperlayerView exact revision on arXiv

Source-linked AI summary

Bidirectional Normalizing Flow: From Data to Noise and Back

Yiyang Lu, Qiao Sun, Xianbang Wang, Zhicheng Jiang, Hanhong Zhao, Kaiming He

arXiv:2512.10953v1cs.LGcs.CV

TL;DR

Standard Normalizing Flows depend on exact inverses, while autoregressive variants face slow causal sampling. BiFlow learns a separate reverse model to approximate the inverse; on ImageNet it reports improved generation quality and sampling up to two orders of magnitude faster than its explicit-inverse counterpart.

  • Problem

    Explicitly invertible NF architectures constrain model design, and autoregressive variants retain slow, difficult-to-parallelize causal decoding.

    Source: §1 · pp. 1–2

  • Method

    BiFlow learns a separate reverse model to approximate the forward model’s inverse, allowing a flexible feedforward, non-causal Transformer for generation.

    Source: §1 · p. 2

  • Results

    FID 2.39 on ImageNet 256×256 with DiT-B is reported alongside sampling two orders of magnitude faster than improved TARFlow, with state-of-the-art NF-based and competitive 1-NFE results.

    Source: §1 · p. 2

  • Takeaways & Limitations

    BiFlow demonstrates that learned bidirectional flows can deliver competitive NF generation quality with much faster inference than an explicit-inverse counterpart.

    Source: §6 · p. 8

    Output-only supervision may be insufficient because mapping pure noise to data in one step is highly under-constrained for learning a reliable inverse from a single reconstruction loss.

    Source: §4.1 · p. 4

Abstract

from arXiv · show

Normalizing Flows (NFs) have been established as a principled framework for generative modeling. Standard NFs consist of a forward process and a reverse process: the forward process maps data to noise, while the reverse process generates samples by inverting it. Typical NF forward transformations are constrained by explicit invertibility, ensuring that the reverse process can serve as their exact analytic inverse. Recent developments in TARFlow and its variants have revitalized NF methods by combining Transformers and autoregressive flows, but have also exposed causal decoding as a major bottleneck. In this work, we introduce Bidirectional Normalizing Flow ($\textbf{BiFlow}$), a framework that removes the need for an exact analytic inverse. BiFlow learns a reverse model that approximates the underlying noise-to-data inverse mapping, enabling more flexible loss functions and architectures. Experiments on ImageNet demonstrate that BiFlow, compared to its causal decoding counterpart, improves generation quality while accelerating sampling by up to two orders of magnitude. BiFlow yields state-of-the-art results among NF-based methods and competitive performance among single-evaluation ("1-NFE") methods. Following recent encouraging progress on NFs, we hope our work will draw further attention to this classical paradigm.

1. Introduction

Standard NFs rely on exact invertibility, while autoregressive TARFlow retains slow causal sampling. BiFlow learns a separate reverse model, and on ImageNet it reports stronger NF-based results with much faster generation.

1. Introduction

  • Standard NFs require explicitly invertible forward architectures with tractable, differentiable Jacobians, limiting use of general-purpose models such as U-Nets and Vision Transformers.TARFlow integrates Transformers with autoregressive flows but retains long causal chains that slow sampling and restrict reverse-model attention.Source: §1 · pp. 1–2
  • BiFlow’s F and G are learned separately, letting a flexible feedforward, non-causal reverse model approximate the inverse for 1-NFE generation.The forward model can remain a tractable NF such as improved TARFlow.Source: §1 · p. 2
  • BiFlow’s learned reverse model can outperform the explicit inverse, while hidden-state alignment, score-step removal, and perceptual loss contribute to the approach’s advantages.The paper reports that hidden-state alignment exceeds distillation of noise-to-data trajectories and that perceptual loss is difficult to use with an explicit inverse.Source: §1 · p. 2
  • FID 2.39 on ImageNet 256×256 with a DiT-B model accompanies sampling two orders of magnitude faster than improved TARFlow, setting an NF-based state of the art and a competitive 1-NFE result.Figure 2 compares generation quality with wall-clock time on 8 v4 TPU cores; its caption notes that VAE decoding is omitted.Source: §1 · p. 2, Figure 2 · p. 2
  • The authors argue that learning NF trajectories rather than pre-scheduling them can be advantageous without imposing inference-time limitations, and point to potential synergy with Flow Matching.They hope this work draws further attention to the NF paradigm.Source: §1 · p. 2

2. Related Work

Related work traces how autoregressive Transformers revived Normalizing Flows while retaining sequential inference, and contrasts their learned mappings with continuous-time approaches.

2. Related Work

  • TARFlow and STARFlow combine Transformers with autoregressive flows, helping improve NF generation quality and scalability.Earlier flows including Planar, NICE, Real NVP, Glow, IAF, and MAF developed expressivity under invertibility and tractable-likelihood constraints.Source: §2 · p. 2
  • Autoregressive flows impose causal ordering and sequential dependencies, constraining architecture and slowing inference; TARFlow’s 8-block, 256-token example requires 8×256 sequential evaluations.Each token depends on its predecessors within each block.Source: §2 · p. 2, Figure 3 · p. 3
  • Continuous Normalizing Flows model transformations as continuous-time ODE dynamics, while Flow Matching reformulates explicit maximum-likelihood training as an equivalent implicit objective.Diffusion models are described as a special case of Flow Matching with stochastic dynamics.Source: §2 · p. 2
  • The implicit formulations of Flow Matching and diffusion sacrifice the learnable bidirectional mapping that characterizes Normalizing Flows.Source: §2 · pp. 2–3

3. Background: Normalizing Flows

The background presents Normalizing Flows as bijections between data and a Gaussian prior, trained with a change-of-variables likelihood. It then describes TARFlow’s autoregressive Transformer design and its sequential sampling cost.

3. Background: Normalizing Flows

  • An NF maps data x to Gaussian-prior noise z with a bijective forward process, then uses the change-of-variables formula to assign data likelihood and trains by maximizing log likelihood.Generation applies the exact inverse to prior samples.Source: §3 · p. 3
  • Composing simpler bijections builds the forward transformation, and the likelihood calculation uses per-transformation log-determinants that must be computable and differentiable.The requirement motivates affine coupling and autoregressive-flow designs with tractable Jacobians.Source: §3 · p. 3
  • TARFlow combines Transformers with autoregressive flows, transforming tokens conditioned on predecessors and reversing sequence order in alternating blocks to capture bidirectional context.Source: §3 · p. 3
  • Autoregressive parameterization allows parallel likelihood evaluation and training, but sampling remains sequential and can require thousands of inverse transformations, causing substantial latency.The example gives 8×256 transformations.Source: §3 · p. 3

4. Bidirectional Normalizing Flow

BiFlow trains a separate reverse model to approximate a fixed forward model’s inverse, without requiring explicit invertibility. The section develops trajectory-alignment training, integrated denoising, flexible losses, norm control, and guidance.

4. Bidirectional Normalizing Flow

  • BiFlow trains a forward model to map data to noise, then a separate reverse model to approximate its inverse without requiring explicit invertibility.The reverse model can use arbitrary architectures and training objectives.Source: §4 · p. 3

4.1. Learning to Approximate the Inverse

  • Hidden distillation supervises the reverse trajectory against forward intermediate states, while naive distillation uses only final-output reconstruction loss, which may under-constrain one-step noise-to-data learning.Hidden distillation also constrains architecture by requiring repeated projections into input space and back to hidden space.Source: §4.1 · p. 4
  • Hidden alignment uses learnable projection heads to align reverse-model representations with forward states, supervising the full trajectory without repeated returns to input space.Unlike hidden distillation, it avoids requiring intermediate reverse states to match input-space representations, preserving architectural flexibility and avoiding potential semantic distortion.Source: §4.1 · p. 4, Figure 4 · p. 4

4.2. Eliminating Score-based Denoising

  • BiFlow integrates denoising into the reverse model with one additional block, producing clean samples in one pass without a separate score-based refinement step.This replaces TARFlow’s score calculation through a forward-backward pass, which nearly doubles inference cost.Source: §4.2 · pp. 4–5, Figure 5 · p. 5

4.3. Distance Metric

  • BiFlow supports varied or combined distance metrics, using MSE for intermediate alignment and perceptual loss at the final VAE-decoded image to enhance realism.The perceptual losses use VGG and ConvNeXt V2 feature spaces.Source: §4.3 · p. 5

4.4. Norm Control

  • Norm control clips forward transformation parameters and normalizes reverse intermediate states to stabilize norms and balance supervision across blocks.Applying either strategy alleviates imbalance in MSE loss across blocks and enhances performance.Source: §4.4 · p. 5

4.5. BiFlow with Guidance

  • BiFlow incorporates classifier-free guidance by extrapolating conditional and unconditional predictions at each hidden state, with guidance scales that can vary across blocks.Training with a range of guidance scales allows outputs at different strengths in a single forward pass, avoiding the doubled inference cost of directly applying CFG.Source: §4.5 · p. 5

5. Experiments

Experiments evaluate class-conditional ImageNet generation at 256×256, using an improved TARFlow as BiFlow’s forward-model baseline. BiFlow achieves a new state of the art among Normalizing Flows and performs on par with leading 1-NFE models.

5.1. Ablation: Learning to Approximate the Inverse

  • Naive MSE distillation beats the exact inverse, while hidden alignment performs best by removing repeated projections and retaining full trajectory-level supervision.Source: §5.1 · p. 6

5.2. Other Ablations

  • Training-time CFG halves inference cost versus standard CFG while improving FID; joint learned denoising improves quality over TARFlow’s score-based approach and avoids its 15.8×-flops extra pass.Source: §5.2 · p. 6
  • Norm control improves performance by balancing blockwise MSE supervision, while perceptual distance improves generation quality; scaling gains diminish after adding ConvNeXt features.Source: §5.2 · pp. 6–7

5.3. BiFlow vs. improved TARFlow

  • 1-NFE generation surpasses the improved TARFlow-XL/2 exact inverse and achieves up to a 42× speedup over 256×2 sequential decoding steps on similar-size TPU models.Source: §5.3 · p. 7, Table 3 · p. 7
  • Eliminating score-based denoising and integrating CFG into training reduce flops by roughly 4×; bidirectional attention enables parallel sequence computation, leaving VAE decoding as a dominant overhead.Source: §5.3 · p. 7

5.4. Comparison with Prior Works

  • 133 million parameters, FID 2.39, and 1-NFE establish BiFlow-B/2 as a new Normalizing Flow state of the art; it performs on par with leading 1-NFE models.Source: §5.4 · p. 8, Table 4 · p. 8

6. Conclusion

BiFlow challenges the requirement that a Normalizing Flow’s reverse process be the exact analytic inverse, proposing a learned, bidirectional alternative. The authors report competitive NF generation quality and inference up to two orders of magnitude faster than the explicit inverse.

6. Conclusion

  • The paper argues that a Normalizing Flow’s reverse process need not be the exact analytic inverse of its forward process.BiFlow replaces that constraint with a learnable reverse model.Source: §6 · p. 8
  • BiFlow achieves competitive generation quality among Normalizing Flows while delivering up to two orders of magnitude faster inference than its explicit inverse counterpart.The conclusion frames this as a step toward expanding the scope of Normalizing Flows.Source: §6 · p. 8
  • BiFlow supports direct, fully parallelized 1-NFE generation rather than autoregressive sampling.Selected BiFlow-B/2 samples are shown for ImageNet 256×256 using guidance scale 2.0.Source: §6 · p. 8, Figure 6 · p. 8

A. Implementation Details

The implementation details specify ImageNet evaluation protocols, inference-cost measurements, and the training and sampling procedures used for BiFlow. Experiments use JAX on TPU hardware, with additional inference comparisons on GPU and CPU.

A. Implementation Details

  • Experiments use JAX on TPU v4, v5p, and v6e cores, with adaptive-weighting exponent p = 2 for the MSE-only ablation and p = 1 otherwise.Training configurations and hyperparameters are provided in Table 5.Source: A. Implementation Details · p. 9
  • FID is computed between 50,000 generated images and training images without augmentation, using 50 samples per class across all 1,000 ImageNet classes.The evaluation uses an Inception-V3 model converted to a JAX-compatible implementation.Source: A. Implementation Details · p. 9
  • Inference-cost comparisons report average per-image runtime across GPU, TPU, and CPU, including CFG and VAE decoding overhead when applicable.The reported NF-model wall-clock comparison includes those overheads and uses KV-cache for autoregressive models.Source: A. Implementation Details · p. 9, Table 6 · p. 9
  • Training pseudocode adds noise to input batches, obtains a forward trajectory and prior, then trains the reverse model with alignment and reconstruction losses.The forward model is frozen, while the reverse model has B + 1 blocks and projection heads are used in the training procedure.Source: A. Implementation Details · p. 10

B. Method Details

BiFlow training aligns the forward and reverse trajectories and reconstructs the clean input. Its training-time guidance scheme is designed to retain 1-NFE sampling and inference-time guidance-scale flexibility.

B.1. Pseudocode

  • The training procedure obtains forward and reverse trajectories, then combines hidden-state alignment with reconstruction loss between the clean input and reconstructed output.The forward trajectory includes its prior, while the reverse trajectory includes intermediate hidden states and the reconstructed input.Source: §B.1 · p. 10

B.2. BiFlow with Guidance

  • Training-time CFG achieves similar or better performance than inference-time CFG while requiring half the inference compute.It also retains post-hoc tuning flexibility; the comparison uses BiFlow-B/2 at 160 epochs with an MSE-only baseline.Source: Figure 7 · p. 10, §B.2 · p. 10
  • Training-time CFG trains a guided reverse model to produce guided samples in a single forward pass while retaining the ability to sweep guidance scales at inference.The model is conditioned on the guidance scale, which is sampled during training.Source: §B.2 · p. 10

B.3. Distance Metric

The distance-metric details describe adaptive weighting, perceptual losses, and normalized trajectories used in BiFlow training. The section also reports reconstruction-fidelity comparisons among inverse-learning strategies.

B.3. Distance Metric

  • Adaptive weighting is applied to all training loss terms, with an adaptive weight p and stop-gradient operator in its definition.The paper reports that adaptively weighted MSE consistently outperforms naive MSE.Source: §B.3 · p. 10
  • The VGG perceptual loss follows LPIPS and is computed after decoding predicted latents into image space; ConvNeXt features provide a complementary perceptual distance.The ConvNeXt distance excludes the final classification head, and the final performance is particularly sensitive to its tunable weight.Source: §B.3 · pp. 10–11
  • Among the three inverse-approximation strategies, hidden alignment achieves the most accurate reconstruction by MSE and LPIPS.The reported setting uses no CFG and adaptive-weighted MSE loss only, with all strategies sharing the same forward model.Source: Table 7 · p. 11
  • Normalized trajectories divide intermediate points by their dataset-averaged squared norm to maintain scale consistency across blocks.They are not used in experiments except for one ablation; the authors note their relevance when using a pretrained NF without clipping.Source: §B.3 · p. 11

C.1. Learning to Approximate the Inverse

The reverse model is evaluated as an approximate inverse, with hidden-state alignment achieving the lowest regression loss across both reconstruction metrics. Guidance can be applied differently to the dedicated denoising block and the inverse blocks.

C.2. BiFlow with Guidance

  • A separate, larger guidance scale for the dedicated denoising block improves sample quality.The denoising block is jointly trained with the other reverse-model blocks and serves a different purpose from the rest of the NF reverse process.Source: Table 8 · p. 11, §C.2 · p. 11

C.1. Learning to Approximate the Inverse

  • Hidden-state alignment achieves the lowest regression loss across both reconstruction metrics, indicating a more accurate mapping between x and x′.Reconstruction fidelity is evaluated with MSE and LPIPS (VGG-based).Source: §C.1 · p. 11

C.2. BiFlow with Guidance

  • Training-time CFG uses one shared guidance scale across blocks, while MSE-only ablations decouple the inverse-block and denoising-block scales.In the decoupled setting, the denoising-block scale is typically much larger.Source: §C.2 · p. 11

C.3. Distance Metric

The ablations compare perceptual losses and adaptive MSE weighting, examine norm control for TARFlow, and test alternative TARFlow CFG designs. They find benefits from perceptual-feature combinations, clipping, and online guidance.

C.3. Distance Metric

  • Combining VGG and ConvNeXt features yields better FID scores than either individually, and ConvNeXt performs best among the tested individual extractors.The feature-extractor and combination comparisons are reported in the perceptual-loss ablation.Source: §C.3 · p. 11, Table 9 · p. 12
  • Adaptive-weighted MSE consistently outperforms naive MSE.The ablation uses BiFlow-B/2 at 160 epochs with training-time CFG.Source: §C.3 · p. 11, Table 10 · p. 12

C.4. Improved TARFlow Norm Control

  • Clipping TARFlow parameters µ and α stabilizes intermediate-state norms and improves final FID performance.Without clipping, norms diverge across blocks and continue growing; with clipping, they remain controlled, benefiting downstream reverse-model training.Source: §C.4 · p. 12, Figure 8 · p. 12

C.5. Improved TARFlow CFG

  • Online CFG significantly outperforms offline CFG at similar runtime; a linear schedule offers a slight advantage, while parameter-space and pixel-space guidance perform similarly.The ablation concludes that TARFlow’s original CFG formulation is close to optimal.Source: Table 11 · p. 12, §C.5 · p. 13

D. Training-free Image Editing

BiFlow’s bidirectional mapping supports training-free image editing by modifying noise and mapping it back to images. The paper demonstrates inpainting and class editing, each using a single forward and reverse pass.

D. Training-free Image Editing

  • BiFlow supports training-free image editing through its explicit bidirectional mapping between data and noise.The applications shown are inpainting and class editing.Source: D. Training-free Image Editing · p. 13
  • For inpainting, BiFlow maps a masked image to noise, resamples the masked portion of the prior, and maps the modified noise back to an image.The reverse mapping fills the masked region with content coherent with the context.Source: D. Training-free Image Editing · p. 13
  • Changing the class condition for a fixed noise variable modifies class-specific appearance while largely preserving global spatial structure.The paper demonstrates this as training-free class editing without retraining.Source: D. Training-free Image Editing · p. 13, Figure 10 · p. 13
  • Both editing applications require only one data-to-noise forward pass and one noise-to-data reverse pass.Source: D. Training-free Image Editing · p. 13

E. Visualizations

Related papers

Loading 2512.10953v1…