PaperlayerView exact revision on arXiv

Source-linked AI summary

Flowing Backwards: Improving Normalizing Flows via Reverse Representation Alignment

Yang Chen, Xiaowei Xu, Shuai Wang, Chenhui Zhu, Ruxue Wen, Xubin Li, Tiezheng Ge, Limin Wang

arXiv:2511.22345v2cs.CV

TL;DR

Standard normalizing flows can learn semantically weak features under log-likelihood optimization, limiting their generative quality. The paper aligns intermediate features on the reverse generative pass and introduces a training-free test-time classifier; it reports simultaneous gains in generation and classification over a strong TARFlow baseline.

  • Problem

    Standard NFs optimized for log-likelihood often fail to learn semantically meaningful features, while their potential synergy between representation learning and generation remains underexploited.

    Source: §1 · pp. 1–2

  • Method

    R-REPA aligns intermediate features during the NF’s reverse generative pass, and a training-free test-time optimization method probes classification using the model’s own loss landscape.

    Source: §1 · p. 2

  • Results

    R-REPA delivers simultaneous gains in generative fidelity and classification over the strong TARFlow baseline, establishing a new state of the art for NFs on ImageNet.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The test-time classifier provides a more intrinsic probe of the semantic knowledge learned by an NF.

    Source: §6 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Standard NFs can fail to learn semantic features when optimized solely for log-likelihood, leaving their potential synergy between representation learning and generation underused. This work aligns features on the reverse generative path and introduces a training-free classifier; experiments report improvements in generative and discriminative performance.

1 Introduction

  • Standard NFs optimized solely for forward-pass log-likelihood often learn semantically weak features, limiting their generative quality.Prior representation-alignment work for diffusion models used pretrained visual features as guidance and reported gains in training efficiency and generation quality [Yu et al. 2024].Source: §1 · pp. 1–2
  • The paper adapts representation alignment to NFs by asking how their invertible structure can link representation learning and generation.Its reverse-path approach extends rather than directly applies REPA to this model class.Source: §1 · p. 2
  • A training-free, test-time optimization algorithm classifies with the NF’s own loss landscape rather than an auxiliary classifier.The authors introduce it to probe intrinsic discriminative ability and measure embedded semantic knowledge.Source: §1 · p. 2
  • Reverse representation alignment enforces semantic consistency on intermediate features during the true generative pass, using the invertibility of NFs.The method is presented as a strategy built from representation alignment and exploitation of the reverse pass.Source: §1 · p. 2
  • The approach improves both generative and discriminative performance, establishing a new state of the art for combining these capabilities in NFs.The experiments assess generative quality with FID and classification with the proposed test-time method.Source: §1 · p. 2

2 Related Work

Related work describes renewed interest in normalizing flows through modern architectures and representation alignment as a way to guide generative models with pretrained visual features. Prior alignment work has extended beyond diffusion models to other components and backbones.

Normalizing Flows

  • TARFlow introduced a Transformer-based NF with state-of-the-art likelihoods and sample quality reported as comparable to diffusion models [Zhai et al. 2024].JetFormer and FARMER use an NF as a jointly trained component for image-text generation, while STARFlow scales NFs in model capacity and task complexity [Tschannen, Pinto, and Kolesnikov 2024; Zheng et al. 2025; Gu et al. 2025].Source: §2 · p. 2

Representation Alignment for Generation

  • Representation alignment uses pretrained vision foundation model features to guide generative-model training, with REPA reporting faster convergence and improved sample quality [Yu et al. 2024].Subsequent work applied alignment to the latent diffusion pipeline, VAE latent spaces, and U-Net backbones [Lee, Park, and Kim 2024; Zheng et al. 2024; Tian et al. 2024b].Source: §2 · p. 2

3 Preliminary: Normalizing Flow and TARFlow

Normalizing flows learn invertible transformations from data to a tractable latent distribution, enabling exact likelihood training and generation by inversion. TARFlow builds on autoregressive flows by stacking Transformer-based blocks with alternating orderings.

Normalizing Flows

  • A normalizing flow maps data to a tractable base distribution with an invertible transformation and trains by maximizing data log-likelihood.The inverse transformation generates data from samples drawn from the base distribution.Source: §3 · p. 2

TARFlow

  • Figure 2 presents selected ImageNet 256 × 256 samples from L-TARFlow + R-REPA, using classifier-free guidance of 2.0.Source: Figure 2 · p. 3

Autoregressive Flows

  • An autoregressive flow computes each dimension’s parameters from preceding input dimensions and uses elementwise affine transformations.The forward and inverse equations specify encoding and sampling for each dimension.Source: §3 · p. 3
  • The autoregressive structure makes the flow Jacobian triangular, reducing its log-determinant to a sum.This simplifies the log-determinant term in the likelihood objective.Source: §3 · p. 3

TARFlow

  • TARFlow stacks Transformer AutoRegressive Blocks and alternates their autoregressive orderings to capture dependencies across dimensions.Its block parameters are modeled with causal Transformer layers.Source: §3 · p. 3
  • TARFlow also uses noise-augmented training and score-based denoising to improve modeling capability.Source: §3 · p. 3

4 Methodology

The methodology combines a training-free, test-time classifier with semantic feature alignment for TARFlow, then adapts the flow to latent space for high-resolution synthesis. Alignment uses pretrained visual features and explores gradient paths through the forward and generative computations.

4.1 Training-Free Classification with NFs

  • The classifier predicts from a single gradient of conditional log-likelihood with respect to soft class logits, avoiding a separately trained classifier.It selects the class with the largest gradient component after one forward and backward pass.Source: §4.1 · pp. 3–4, §4 · p. 3

4.2 Representation Alignment

  • The three alignment strategies differ in gradient reach: forward updates preceding blocks, detach updates the current block, and reverse updates subsequent blocks via the generative graph.Reverse alignment detaches the forward-computed latent and computes alignment loss in the inverse flow.Source: Figure 3 · p. 4, §4.2 · p. 4
  • A frozen pretrained vision encoder supplies patch features that a learnable MLP projects TARFlow features toward using patch-wise similarity.The alignment objective can be applied at any chosen set of layers.Source: §4.2 · p. 4
  • Training combines the NF loss with averaged alignment losses across selected layers, weighted by an alignment hyperparameter.The alignment-loss gradient follows one of the three backpropagation strategies.Source: §4.2 · p. 5

4.3 Moving to Latent Space

  • For high-resolution synthesis, TARFlow models VAE-compressed latents, samples a noisy latent through the inverse flow, refines it with score-based denoising, and decodes it into an image.The approach shifts pixel-level modeling to compact latent codes.Source: §4.3 · p. 5

5 Experiments

Experiments evaluate R-REPA’s classification and generation performance on ImageNet at 64×64 and 256×256, alongside ablations of alignment design and training or sampling choices. The results show improved performance over TARFlow baselines, with two-step generation at both resolutions.

5 Experiments

  • Experiments use ImageNet-1K at 64×64 and 256×256, evaluate generation with FID, sFID, and IS, and assess classification accuracy on the validation set using test-time optimization.Generation evaluation follows TARFlow by sampling 4096 images and choosing each model’s CFG scale by grid search.Source: §5 · p. 5

5.1 Ablation Studies

  • Reverse alignment outperforms the Forward and Detach strategies on FID, while Forward degrades sFID from 33.79 to 37.97.Reverse alignment updates only the generative path, refining synthesis without disrupting the density model.Source: §5.1 · pp. 5–6
  • Aligning Blocks 7 and 8 gives the best FID, whereas aligning Blocks 1 and 2 gives the best semantic accuracy but harms FID; deeper alignment improves fidelity from 11.93 to 11.71.Within the best-performing setup, moving alignment from the second to sixth layer improves FID; the fourth layer gives the best IS.Source: §5.1 · p. 6
  • The latent-noise standard deviation peaks at 0.20, and the final “Reverse 7&8” configuration achieves the lowest FID with an optimal CFG scale near 3.1.Figure 5 reports 3.3× faster training convergence on ImageNet 64×64.Source: Figure 5 · p. 6, §5.1 · p. 7

5.2 Main Results

  • 11.25 FID and 57.02% accuracy on ImageNet 64×64 improve on TARFlow’s 11.76 FID and 39.97% accuracy; at 256×256, FID improves from 13.05 to 12.79 and accuracy from 40.22% to 56.24%.The reported 64×64 model at 400k iterations also exceeds the fully trained baseline on FID and accuracy.Source: §5.2 · p. 7
  • 3.69 FID and 4.34 sFID on class-conditional ImageNet 64×64 improve on TARFlow’s 4.21 FID and 5.34 sFID, surpassing BigGAN’s FID while remaining behind iDDPM’s FID.The result uses two sampling steps.Source: §5.2 · p. 7, Table 3 · p. 7
  • 4.18 FID and 4.96 sFID on ImageNet 256×256 are reported for R-REPA with a 1×1 patch-size adjustment; direct R-REPA alone achieves 4.95 FID.The high-resolution results retain two-step sampling.Source: §5.2 · pp. 7–8, Table 4 · p. 7

6 Conclusion

The conclusion presents R-REPA as a way to improve NF semantic awareness by aligning features on the generative pass, with reported gains in generation, classification, and training efficiency. The appendix outlines supplementary algorithm, setup, ablation, and visualization material, including TARFlow’s autoregressive block computation.

6 Conclusion

  • R-REPA enforces semantic consistency on the generative z-to-x pass, achieving simultaneous FID and classification-accuracy gains over TARFlow on ImageNet and boosting training efficiency by over 3.3×.Source: §6 · p. 8

Appendix

  • The appendix covers algorithm details, experimental procedures, computational resources, additional ablations, and qualitative visualizations.Its algorithm review includes TARFlow block computation, accelerated implementation, and reverse- versus forward-pass gradient flow.Source: Appendix · p. 8

A.1 Review of the TARFlow Block Computation

  • A TARFlow block maps an input tensor to an output tensor of the same shape, with forward encoding parallelizable and inverse sampling sequential across tokens.For token d, inverse sampling depends on preceding tokens x<d.Source: §A.1 · p. 8

A.2 Accelerated R-REPA Implementation

The appendix makes reverse-pass alignment practical by replacing sequential autoregressive computation with a cached-context pseudo-reverse pass. It reports substantial speed and memory gains and explains how reverse alignment directs gradients differently from forward alignment.

A.2 Accelerated R-REPA Implementation

  • Caching detached forward-pass block inputs lets a pseudo-reverse pass use precomputed autoregressive context and compute inverse transformations for all tokens in parallel.The resulting intermediate features supply the R-REPA alignment loss while gradients flow through the aligned block and subsequent blocks, not preceding blocks.Source: §A.2 · p. 9

A.3 Computational Efficiency

  • Approximately 50× speedup and nearly 50% lower peak memory make accelerated reverse-pass alignment feasible for deep models.These gains are measured against a naive sequential autoregressive implementation.Source: §A.3 · p. 9
  • Forward and accelerated Reverse strategies incur minor throughput overhead relative to Forward, which adds a loss to the existing graph.The appendix describes the overhead as a small price for the generative-quality improvements reported in the main paper; Table 5 compares throughput under the same devices.Source: §A.3 · p. 9, Table 5 · p. 9

A.4 Detailed Gradient Flow Analysis for Forward-REPA and R-REPA

  • Forward alignment updates the aligned block and preceding blocks, whereas reverse alignment updates the aligned block and subsequent blocks in the original forward ordering.The reverse gradient pathway follows the generative chain from abstract latents toward concrete data.Source: §A.4 · pp. 9–10
  • The authors posit that injecting semantic guidance along the generative direction is key to R-REPA’s enhanced performance.R-REPA’s alignment loss guides how later-block representations are transformed into semantically aligned features by earlier blocks.Source: §A.4 · p. 10

B.1 Training Setup

The supplementary setup describes the training, sampling, and classification procedures, along with computational resources and additional experimental results. It specifies a frozen visual encoder and a training-free classifier whose predictions are robust to several implementation choices.

B.1 Training Setup

  • Training uses AdamW, EMA evaluation weights, resize-and-crop plus horizontal-flip augmentation, and frozen DINOv2-B targets with a learnable three-layer MLP projector.The encoder and projector settings are reported as consistent across ImageNet 64×64 and 256×256.Source: §B.1 · p. 10

B.2 Sampling Setup

  • Sampling follows TARFlow with a fixed random seed and CPU-generated Gaussian noise transferred to the GPU for the inverse transformation.The fixed seed is used to make evaluation sampling deterministic and reproducible.Source: §B.2 · p. 10

B.3 Classification Setup

  • The training-free classifier’s final prediction is robust to tested optimization settings because one gradient identifies the steepest-ascent direction in conditional log-likelihood.The tested variations include multiple optimization steps, early-stopping criteria, and logit learning rates.Source: §B.3 · p. 10

B.4 Computational Resources

  • Experiments use eight NVIDIA GPUs; full training takes approximately 6–7 days at 64×64 and 8–9 days at 256×256.The final models train for one million iterations, compared with 400,000 iterations for ablations.Source: §B.4 · p. 10

C Additional Experimental Results

  • Supplementary experiments include further ablations, and Table 8 compares R-REPA encoders against a baseline FID of 12.91.The table caption identifies 12.91 as the FID without R-REPA and says arrows indicate whether higher or lower values are better.Source: C Additional Experimental Results · p. 12, Table 8 · p. 12

C.1 Further Ablation on Alignment Layer for the Detach Strategy

Within the tested Detach configuration, changing the alignment layer has limited impact on generative quality. The best Detach configuration is outperformed by the Reverse strategy on both reported metrics.

C.1 Further Ablation on Alignment Layer for the Detach Strategy

  • The Detach ablation’s FID scores range from 12.14 to 12.37 across alignment layers, with only slight classification gains at deeper layers.The authors conclude that performance in this configuration is not highly sensitive to the precise layer choice.Source: §C.1 · p. 12
  • Reverse achieves FID 11.71 and Acc 57.35%, outperforming the best Detach configuration at FID 12.14 and Acc 54.98%.The cited Reverse result uses layer 6; the comparison is presented as evidence for leveraging the generative pathway.Source: §C.1 · p. 12

C.2 Ablation on Encoder Architectures

R-REPA performs robustly across the tested encoder architectures, and the results show a positive correlation between discriminative accuracy and generation quality.

C.2 Ablation on Encoder Architectures

  • The ablation applies R-REPA to layer 6 of blocks 7 and 8 while replacing default DINOv2-B with CLIP-B, MAE-B, and DINOv2-L.Quantitative results are presented in Table 8.Source: §C.2 · p. 12
  • Every tested encoder configuration significantly outperforms the baseline FID of 12.91, demonstrating R-REPA's robustness across encoder architectures.The tested alternatives are CLIP-B, MAE-B, and DINOv2-L, compared with the default DINOv2-B encoder.Source: §C.2 · p. 12
  • Encoders with better accuracy, such as DINOv2, also achieve the best FID scores, revealing a positive correlation between discrimination and generation quality.The authors describe this observation as empirical support for their motivation to improve generation by enhancing discriminative features.Source: §C.2 · p. 12

C.3 Ablation on Alignment Loss Weight λalign

On ImageNet 256×256, the ablation favors λalign=0.1 over 0.5, and the authors use 0.1 in their main experiments.

C.3 Ablation on Alignment Loss Weight λalign

  • λalign balances the standard Normalizing Flow objective LNF with representation alignment loss Lalign.The authors suggest that an overly strong alignment term can interfere with density estimation and slightly degrade final sample fidelity.Source: §C.3 · p. 12
  • FID is 13.26 at λalign=0.1 versus 13.74 at λalign=0.5 on ImageNet 256×256, supporting the smaller weight.The authors selected λalign=0.1 for all main experiments.Source: §C.3 · p. 12
  • In the alignment-weight ablation, Table 9 reports that λalign = 0.1 yields better FID (13.26) than λalign = 0.5 (13.74), so 0.1 was selected for all main experiments.Source: §C.3 · p. 12

D Additional Visualizations

The additional visualizations show samples from R-REPA models at 256×256 and 64×64 resolutions. The ethical statement discusses dual-use risks, possible inherited biases, computational impact, and research transparency.

D Additional Visualizations

Ethical Statement

  • The authors identify high-fidelity image generation as dual-use, including potential misuse to create misinformation or deceptive synthetic media.They advocate responsible development and deployment, alongside detection and mitigation strategies.Source: Ethical Statement · p. 13
  • ImageNet-1K training data and web-pretrained DINOv2 features may carry societal biases that the model could inherit or amplify.The authors urge fairness assessment before real-world deployment.Source: Ethical Statement · p. 13
  • The work uses publicly available ImageNet-1K and pretrained models under their respective licenses for academic research, without personally identifiable or sensitive private data.Source: Ethical Statement · p. 13
  • Training required significant computational resources, while the authors report over 3.3× faster convergence as reducing computational cost and associated energy consumption.Source: Ethical Statement · p. 13
  • The authors aim to advance scientific understanding and release findings to foster innovation and transparent discussion of generative technologies' capabilities and societal implications.Source: Ethical Statement · p. 13
Loading 2511.22345v2…