Research map / Probabilistic and scientific modeling / Variational and generative sampling
Robustly Train Normalizing Flows via KL Divergence Regularization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a KL-divergence regularizer on the Jacobians of normalizing flows that is equivalent to adding identity-covariance-like samples, making NF training robust to outliers and small or high-dimensional training sets, evaluated on OoD detection with Glow-style coupling flows.
Problem
Normalizing flow training is sensitive to outliers and to small sample sizes or high dimensionality, where covariance estimates become low-rank or biased.
Approach
Regularize the Jacobian of each coupling block with a KL divergence to a positive definite matrix (identity in practice), which is shown to be equivalent to adding a set of in-distribution samples with covariance G to the training set; an efficient upper-bound form uses weight norms and the scale outputs so it is cheap to compute.
Key results
- On CIFAR10 with IB (gamma=infinity), adding KLJ regularization raised Rot-RGB OoD detection score from 67.7 to 88.2 (about 30% gain in the abstract's framing).
- With IB (gamma=1), KLJ improved Rot-RGB OoD score from 78.0-type baseline values (68.7 to 75.5 on Rot, 71.4 to 76.0 on QuickDraw), while ID accuracy rose from 78.0 to 85.3.
- KLJ regularization reduced in-distribution accuracy for cNLL in some settings and cNLL+KLJ lowered Noisy CIFAR10 score (47.1 to 43.2).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add KLJ regularization to IB (γ = 1) | Glow normalizing flow · Rot-RGB CIFAR10, QuickDraw, TinyImageNet, and Noisy CIFAR10 | OoD detection score: 78.0 -> 85.3; 68.7 -> 75.5; 71.4 -> 76.0; 77.4 -> 84.9 | Helped |
| Add KLJ regularization to IB (γ →∞) | Glow normalizing flow · Rot-RGB CIFAR10, QuickDraw, TinyImageNet, and Noisy CIFAR10 | OoD detection score: 67.7 -> 88.2; 50.5 -> 87.0; 66.1 -> 74.3; 48.7 -> 69.1 | Helped |
| Add KLJ regularization to cNLL | Glow normalizing flow · CIFAR10 OoD detection evaluation | OoD detection score: 69.7 -> 76.2 | Helped |
| Add KLJ regularization to IB (γ = 1) | Glow normalizing flow · CIFAR10 | accuracy: 89.3 -> 85.6 | Hurt |
| Add KLJ regularization to IB (γ →∞) | Glow normalizing flow · CIFAR10 | accuracy: 92.0 -> 88.6 | Hurt |
| Add KLJ regularization to IB (γ = 1) | Glow normalizing flow · Rot-RGB CIFAR10, QuickDraw, TinyImageNet, and Noisy CIFAR10 | OoD/ID entropy ratio: 1.45 -> 1.69; 0.76 -> 1.11; 2.30 -> 2.45; 0.15 -> 0.31 | Mixed |
| Add KLJ regularization to IB (γ →∞) | Glow normalizing flow · Rot-RGB CIFAR10, QuickDraw, TinyImageNet, and Noisy CIFAR10 | OoD/ID entropy ratio: -0.07 -> 0.03; -0.42 -> -0.13; 0.32 -> 0.07; -0.23 -> -0.19 | Mixed |
| Vary the KLJ coefficient α (ablation) | IB-trained Glow normalizing flow · Rot-RGB CIFAR10, TinyImageNet, QuickDraw, and Noisy CIFAR10 | OoD detection score | Helped |
| Add KLJ regularization under Gaussian-noise outliers in training data (ablation) | IB (γ = 1), IB (γ →∞), and cNLL · OoD detection with training outlier ratios r = 0.05, r = 0.1, and r = 0.15 | OoD detection score | Mixed |
| Add KLJ regularization when training data are reduced (ablation) | IB (γ = 1) and IB (γ = ∞) · Rot-RGB CIFAR10 and Noisy CIFAR10, reduced training-set sizes | OoD detection score | Helped |
Surprising findings
- The regularization reduced in-distribution accuracy for the generative classifier cNLL objective, a trade-off the authors attribute to fitting less label-correlated noise.
- L2 regularization makes the estimated density collapse in the linear 1D case, whereas KL regularization does not.
Limitations
- The equivalence analysis is for linear flows and extended to deep coupling flows via heuristic upper bounds and a tuned coefficient tau.
- Evaluation limited to OoD detection on CIFAR10-scale images with a single Glow-based architecture.
- Regularization reduced ID accuracy and hurt some OoD metrics (e.g., cNLL on Noisy CIFAR10).
Experimental setup
Datasets: CIFAR10 CIFAR100 QuickDraw Tiny-ImageNet Rot-RGB CIFAR10 Noisy CIFAR10
Benchmarks: CIFAR10 OoD detection (Rot-RGB, QuickDraw, Tiny-ImageNet, Noisy CIFAR10)
Metrics: OoD detection ROC-AUC score OoD/ID entropy ratio ID accuracy
Baselines: cNLL IB (information bottleneck) L2 regularization
Builds on: Glow information bottleneck objective for invertible networks ITML
Closest papers by meaning
- Understanding Likelihood of Normalizing Flow and Image Complexity through the Lens of Out-of-Distribution Detection (AAAI 2024 · similarity 0.85)
- Decoupled Kullback-Leibler Divergence Loss (NeurIPS 2024 · similarity 0.85)
- Fast and unified path gradient estimators for normalizing flows (ICLR 2024 · similarity 0.85)
- Optimal Ridge Regularization for Out-of-Distribution Prediction (ICML 2024 · similarity 0.84)
- Normalizing Flows are Capable Generative Models (ICML 2025 · similarity 0.84)
- An Error Analysis of Flow Matching for Deep Generative Modeling (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.