[Paper Review] Learning Likelihoods with Conditional Normalizing Flows

Christina Winkler, Daniel E. Worrall+2 more|arXiv (Cornell University)|Nov 29, 2019

Advanced Image Processing TechniquesComputer Science39 references55 citations

TL;DR

The paper demonstrates learning conditional distributions pY|X using Conditional Normalizing Flows (CNFs) and applies them to super-resolution and retinal vessel segmentation, achieving competitive likelihoods and standard metrics while modeling multimodal, correlated outputs.

ABSTRACT

Normalizing Flows (NFs) are able to model complicated distributions p(y) with strong inter-dimensional correlations and high multimodality by transforming a simple base density p(z) through an invertible neural network under the change of variables formula. Such behavior is desirable in multivariate structured prediction tasks, where handcrafted per-pixel loss-based methods inadequately capture strong correlations between output dimensions. We present a study of conditional normalizing flows (CNFs), a class of NFs where the base density to output space mapping is conditioned on an input x, to model conditional densities p(y|x). CNFs are efficient in sampling and inference, they can be trained with a likelihood-based objective, and CNFs, being generative flows, do not suffer from mode collapse or training instabilities. We provide an effective method to train continuous CNFs for binary problems and in particular, we apply these CNFs to super-resolution and vessel segmentation tasks demonstrating competitive performance on standard benchmark datasets in terms of likelihood and conventional metrics.

Motivation & Objective

• Motivate learning conditional distributions pY|X without strong handcrafted losses in high-dimensional outputs.
• Introduce CNFs that condition both the base distribution and the invertible mapping on X.
• Demonstrate the ability to model multimodal and correlated outputs in structured prediction tasks.
• Present a variational dequantization framework for binary variables to enable CNFs on discrete data.

Proposed method

• Use a conditional normalizing flow pY|X(y|x)=pZ|X(z|x)|det dfφ(y,x)/dy| with z=fφ(y;x) and x-conditioned prior pZ|X.
• Condition prior, split priors, and affine coupling layers on x via a shared representation h=g(x).
• Train by flowing y to z and maximizing the log-likelihood under the conditional prior, enabling likelihood-based optimization.
• Extend variational dequantization to binary variables, introducing half-infinite noise to align supports between v and y.

Experimental results

Research questions

• RQ1 Can CNFs efficiently model conditional distributions pY|X(y|x) with multimodal and correlated outputs in high-dimensional prediction tasks?
• RQ2 Does conditioning both the base density and the invertible transform on x improve likelihoods and task-specific metrics over factorized baselines?
• RQ3 How does variational dequantization affect learning and calibration for binary segmentation with CNFs?
• RQ4 Are CNFs competitive with state-of-the-art per-pixel loss methods in super-resolution and vessel segmentation when evaluated on standard benchmarks?

Key findings

• CNFs achieve higher conditional log-likelihoods than a comparable factorized baseline on ImageNet-based super-resolution (e.g., ImageNet32: CNF 3.01 vs baseline 4.00; ImageNet64: CNF 2.90 vs baseline 3.61).
• CNF samples display sharper high-frequency content and crisper details than factorized baselines, with perceptual tuning via sampling temperature.
• On standard SR benchmarks, CNFs are competitive or superior in PSNR and comparable in SSIM to pixel-wise methods, illustrating the advantage of modeling output correlations.
• For retinal vessel segmentation, a binary-dequantized CNF achieves competitive bits-per-dimension and F-scores compared to strong CNN-based methods like DRIU and HED, while offering well-calibrated probabilistic outputs.
• The approach demonstrates that learning likelihoods with CNFs can outperform handcrafted loss baselines in multi-dimensional, structured prediction tasks.