Flowing Backwards: Improving Normalizing Flows via Reverse Representation Alignment

AI summary

Overview

Research area: Generative modeling in computer vision, specifically Normalizing Flows (NFs), representation alignment, and class-conditional image synthesis.

Technical level: Advanced. The paper assumes familiarity with normalizing flows, autoregressive models, change-of-variables likelihood, computational graphs, and standard generative metrics (FID, sFID, IS).

Scope: The paper proposes R-REPA, a training strategy that aligns the internal features of a Normalizing Flow's generative (reverse) pass with features from a pretrained vision foundation model, and evaluates it on class-conditional ImageNet generation at 64×64 and 256×256 resolution.

What This Paper Is About

Normalizing Flows are generative models with an exactly invertible structure: a forward pass maps images to noise for likelihood training, and a reverse pass turns noise back into images. Although this duality should let representation learning and generation reinforce each other, standard NFs trained only on log-likelihood learn weak semantic features, which caps their image quality. This paper asks whether aligning the NF's features during the reverse (generative) pass — rather than the forward pass, as prior work like REPA does for diffusion models — can improve both image fidelity and the model's internal semantic knowledge.

Key Contributions

1. Reverse Representation Alignment (R-REPA). The authors design and systematically compare three gradient-backpropagation strategies for aligning NF features with a pretrained vision encoder — Forward (F-REPA), Detach (D-REPA), and Reverse (R-REPA) — and show that the reverse strategy, which updates parameters only along the generative computational graph, performs best.

2. A training-free, test-time classification algorithm for NFs. Rather than training a linear probe, the method optimizes a set of classification logits at inference time by taking a single gradient of the conditional log-likelihood with respect to those logits, then selecting the class with the largest gradient component. It requires only one forward and one backward pass.

3. Demonstration of joint improvement in generation and discrimination. The method improves FID and classification accuracy simultaneously, which the authors present as evidence of NF invertibility being exploited to its full potential.

4. New state-of-the-art results for Normalizing Flows on ImageNet 64×64 and 256×256, achieved while preserving two-step sampling and accelerating training convergence by 3.3×.

Main Findings

• Reverse alignment beats forward and detach alignment. In ablations at 400k iterations on ImageNet 64×64 (all blocks, 2 layers), Reverse achieved FID 12.21 / sFID 33.80 / IS 42.08 / accuracy 49.91%, versus Detach at 12.19 / 34.31 / 41.98 / 49.06% and Forward at 12.25 / 37.97 / 40.85 / 46.97%, against the TARFlow baseline of 12.91 / 33.79 / 36.62 / 37.43%. The authors attribute the Forward strategy's sFID degradation to tension between the maximum-likelihood objective and the alignment loss in early blocks.

• A trade-off exists between where alignment is applied and what improves. Aligning blocks 7 & 8 (the first to operate on the latent variable during generation) gave the best FID of 11.93, while aligning blocks 1 & 2 gave the best semantic accuracy of 61.63% but harmed FID.

• Deeper alignment layers improve fidelity. Within R-REPA on blocks 7 & 8, FID improved from 11.93 (2nd layer) to 11.84 (4th layer) to 11.71 (6th layer). The 4th layer gave the best IS at 46.06, while the 6th layer was chosen for best image quality.

• Faster convergence and better results at every checkpoint. On ImageNet 64×64, R-REPA at 400K iterations already reached FID 11.71 and accuracy 57.76%, surpassing the fully trained 1M-iteration TARFlow baseline (FID 11.76, accuracy 39.97%). At 1M iterations R-REPA reached FID 11.25 and accuracy 57.02%. On ImageNet 256×256 in latent space, FID improved from 13.05 to 12.79 and accuracy from 40.22% to 56.24% at 1M iterations.

• State-of-the-art NF generation at 64×64. With 50K samples, R-REPA reduced FID from TARFlow's 4.21 to 3.69 and sFID from 5.34 to 4.34, surpassing BigGAN (FID 4.06) and approaching iDDPM (FID 2.92), while using only two sampling steps.

• Strong high-resolution results. On ImageNet 256×256, combining R-REPA with a 1×1 patch size reached FID 4.18 and sFID 4.96, versus Latent-TARFlow's 5.15 and 6.78. Direct R-REPA application gave FID 4.95 and sFID 6.89.

• Optimal hyperparameters identified. The best classifier-free guidance scale was near 3.1, and the optimal latent noise standard deviation was σ = 0.20. Samples in Figure 2 were generated with classifier-free guidance equal to 2.0.

• The test-time classifier is a valid probe. Comparing the single-score test-time evaluation against layer-by-layer linear probing showed that the new metric consistently reflects the overall performance ceiling for both TARFlow and +R-REPA models, avoiding expensive per-layer evaluation.

Methodology in Plain English

The starting point is TARFlow, a high-performing Normalizing Flow built from stacked Transformer Autoregressive Blocks (TARBlocks). The authors keep TARFlow's architecture but add an auxiliary training signal.

A frozen pretrained vision encoder extracts semantic features from each training image. A small learnable MLP projects the NF's intermediate features into that same feature space, and a cosine-similarity loss encourages the two to match. The novelty is which part of the network this loss is allowed to update:

• Forward REPA backpropagates through the normal encoding path, updating all blocks before the alignment layer. This hurt sFID.
• Detach REPA stops the gradient at the block input, so only the current block and the projector update.
• Reverse REPA runs the forward pass to get the latent, detaches it, then builds a fresh computational graph by running the inverse flow. The alignment loss is computed on this generative graph, so gradient updates flow only into the parameters after the alignment layer.

Because a true reverse pass is sequential and slow on GPUs, the authors introduce an accelerated version: features from the forward pass are cached (with gradients stopped) and reused as conditioning during a "pseudo-reverse" pass, making the alignment loss trainable in parallel.

Separately, to scale to 256×256, the NF operates in the latent space of a pretrained VAE, with noise of standard deviation 0.20 added to latents during training and a score-based denoising step applied after sampling. For classification, the paper's training-free algorithm initializes logits to zero, forms a softmax-weighted class embedding, evaluates the conditional log-likelihood, takes one gradient with respect to the logits, and predicts the class with the largest gradient.

Implementation specifics: at 64×64 the model uses 8 TARBlocks, each with 8 causal attention layers, channel dimension 1024, and 4×4 patches. At 256×256 it uses a pretrained VAE-ft-EMA with 2×2 patches, plus Rotary Position Embeddings and SwiGLU, keeping the same block count, layer count, and channel dimension.

Why This Matters

Impact on research. The paper extends the representation-alignment paradigm — previously established for diffusion models by REPA — to Normalizing Flows, and shows that the invertible architecture permits an alignment strategy (reverse-pass alignment) that has no direct analogue in non-invertible models. It also offers a cheap alternative to linear probing for measuring what a generative model has learned.

Real-world applications:

• Efficient high-resolution image synthesis, since competitive FID at 256×256 is achieved with only two sampling steps rather than the many iterative steps used by diffusion models.
• Latent-space generation pipelines that already use a VAE, where this alignment can be added without changing the compression stage.
• Model auditing and diagnostics: the training-free classification probe can be used to check how semantically meaningful a generative model's internal features are without training an auxiliary classifier.
• Image classification in settings where a class-conditional generative model is already deployed, using the same network for both synthesis and prediction.

Industry relevance. The reported 3.3× training acceleration and two-step inference point directly at reduced compute cost for training and serving generative image models. The paper's author affiliations list the State Key Laboratory for Novel Software Technology at Nanjing University, Alibaba Group, and Shanghai AI Lab, and the work was supported in part by Alibaba Group through the Alibaba Innovative Research Program.

Future Directions

• How far does alignment generalize? The paper demonstrates results on ImageNet at 64×64 and 256×256; whether R-REPA transfers to other datasets, other NF architectures (the authors cite JetFormer, FARMER, and concurrent work STARFlow as related but untested), or other modalities such as video remains unverified in the reported content.

• Reconciling the fidelity–semantics trade-off. Aligning early blocks maximized accuracy (61.63%) while aligning late blocks maximized FID (11.93). A single alignment configuration that achieves both is an open problem the paper identifies but does not solve.

• Closing the gap to diffusion and autoregressive models. R-REPA's FID of 4.18 at 256×256 still trails models listed such as VAR (1.73), MAR (1.55), and SiT (2.06); how much of that gap is addressable through further representation alignment is not reported.

• Extending the test-time classification algorithm. The method is demonstrated on TARFlow variants; its behavior on non-invertible generative models, and whether multi-step or iterative test-time optimization improves over the single-step gradient, are not explored in the provided content.

Target Audience

This paper is most useful to generative modeling researchers and graduate students working on normalizing flows, density estimation, or representation alignment, and to practitioners seeking faster-training, few-step image synthesis. It will also interest engineers evaluating whether flow-based generative models are competitive with diffusion and autoregressive alternatives at high resolution, and researchers looking for lightweight diagnostic methods to assess the semantic quality of generative representations.

Authors’ abstract

Normalizing Flows (NFs) are a class of generative models distinguished by a mathematically invertible architecture, where the forward pass transforms data into a latent space for density estimation, and the reverse pass generates new samples from this space. This characteristic creates an intrinsic synergy between representation learning and data generation. However, the generative quality of standard NFs is limited by poor semantic representations from log-likelihood optimization. To remedy this, we propose a novel alignment strategy that creatively leverages the invertibility of NFs: instead of regularizing the forward pass, we align the intermediate features of the generative (reverse) pass with representations from a powerful vision foundation model, demonstrating superior effectiveness over naive alignment. We also introduce a novel training-free, test-time optimization algorithm for classification, which provides a more intrinsic evaluation of the NF's embedded semantic knowledge. Comprehensive experiments demonstrate that our approach accelerates the training of NFs by over 3.3$\times$, while simultaneously delivering significant improvements in both generative quality and classification accuracy. New state-of-the-art results for NFs are established on ImageNet 64$\times$64 and 256$\times$256. Our code is available at https://github.com/MCG-NJU/FlowBack.