Normalizing Flow Losses¶

← Mixture Density Networks | Error-in-Variables →

Normalizing flows transform a simple base distribution (Gaussian) into a complex target distribution through learnable invertible transformations — enabling arbitrary conditional density estimation.

Mathematical Background & Formulation for Supervised Regression¶

Conditional Normalizing Flows model complex, non-Gaussian, multi-dimensional target distributions \(p(\mathbf{y} \mid \mathbf{x})\) for regression tasks.

1. Conditional Density Estimation in Regression¶

Given observed input features \(\mathbf{x} \in \mathbb{R}^D\) (e.g. tabular covariates or high-dimensional embeddings) and continuous target variables \(\mathbf{y} \in \mathbb{R}^M\):

1. Backbone Context Extractor: A neural network (e.g. 1D CNN, MLP, or ResNet) \(c_\theta(\mathbf{x})\) maps the input \(\mathbf{x}\) to a conditioning context vector \(\mathbf{c} \in \mathbb{R}^{d_{\text{ctx}}}\).
2. Invertible Flow Head: An invertible transformation \(\mathbf{f}_\phi(\mathbf{y}; \mathbf{c})\) conditioned on context \(\mathbf{c}\) maps the target \(\mathbf{y} \in \mathbb{R}^M\) into a simple base Gaussian variable \(\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_M)\).

By change of variables, the exact conditional likelihood of target \(\mathbf{y}\) given input \(\mathbf{x}\) is:

Training minimizes the exact Negative Log-Likelihood (NLL) via standard gradient descent:

2. Standard Regression vs. Simulation-Based Inference (SBI)¶

In standard regression, conditional normalizing flows act as a drop-in replacement for GaussianNLLLoss or MixtureDensityLoss: - Beyond Gaussianity: Unlike GaussianNLLLoss (which restricts \(p(\mathbf{y} \mid \mathbf{x})\) to symmetric ellipses), flows capture non-linear parameter degeneracies (such as banana-shaped \(T_{\text{eff}}\) vs. \(\log g\) contours in astrophysics), heavy tails, and skewness. - Beyond Fixed Components: Unlike MixtureDensityLoss (which requires choosing a fixed number \(K\) of Gaussian mixture modes and can suffer from mode collapse during training), flows transform a continuous base space smoothly, providing stable likelihood gradients. - Fast Sampling at Inference: Given a new input spectrum \(\mathbf{x}_{\text{test}}\), evaluating context \(\mathbf{c} = c_\theta(\mathbf{x}_{\text{test}})\) enables drawing \(N = 5,000\) exact joint posterior samples \(\mathbf{y}^{(s)} \sim p(\mathbf{y} \mid \mathbf{x}_{\text{test}})\) in milliseconds to produce full 1D and 2D corner plots.

See the NormalizingFlowLoss API for the training contract and create_flow_loss helper.

Choosing a flow family¶

For noisy labels, wrap the flow with NoiseAwareDensityLoss instead of fitting the observed label directly. For imbalanced targets, remember that target-density reweighting changes the population density (see imbalanced regression).

Usage¶

Step 1: Create a zuko flow¶

For bpf, n_transforms and hidden_features control the autoregressive architecture, while degree (at least 1) controls its polynomial coefficients. The factory uses a local bounded Bernstein transform with \(f(x)=x\) for \(|x|\geq 5\), continuous endpoint values and unit endpoint derivatives. Its inverse solver tolerance controls numerical inversion, not a tail shift. The corrected transform retains the parameter and buffer names of the upstream BPF, so existing weights load when the actual layer count and widths match. Older factory versions ignored the requested BPF layer count and widths; use the checkpoint's actual architecture rather than that ignored request when loading. Previously recorded predictions and benchmarks keep their original version and must not be relabeled as results of this corrected map.

Use the create_flow_model helper or build your own zuko flow:

from torchregress.losses import create_flow_model

flow = create_flow_model(
    n_features=2,        # target dimensionality
    context_dim=64,      # model's output dim (conditioning)
    flow_type="nsf",     # Neural Spline Flow
    n_transforms=5,      # number of invertible blocks
)

Step 2: Wrap in NormalizingFlowLoss¶

from torchregress.losses import NormalizingFlowLoss

loss_fn = NormalizingFlowLoss(flow=flow)

Or use the shortcut:

from torchregress.losses import create_flow_loss

loss_fn = create_flow_loss(
    n_features=2, context_dim=64, flow_type="nsf", n_transforms=5,
)

Contrastive Likelihood-Ratio Flow Variant (InfoNCE / NRE)¶

When the downstream task is hypothesis ranking or likelihood-ratio estimation rather than generic density estimation, use ContrastiveFlowLoss. It is a multiclass InfoNCE objective over the flow's conditional log-densities: the observed target is scored under the generating context and one or more alternate contexts, and the flow is trained so that the true hypothesis scores highest:

from torchregress.losses import ContrastiveFlowLoss

# Likelihood-ratio mode: T=1, margin=0.
loss_fn = ContrastiveFlowLoss(flow=flow)
loss = loss_fn(
    positive_context,
    target,
    negative_context=negative_context,  # [batch, n_negatives, context_dim]
)

This is a contrastive likelihood-ratio / neural-ratio-estimation (NRE) objective, not a density-calibration objective. Two regimes matter:

For a dedicated classifier-based NRE loss (cheaper: no flow Jacobians, and the standard ratio-estimation formulation), see NeuralRatioLoss and the NRE guide. The Elsharkawy & Kahn (2025) Contrastive Normalizing Flow is related work, but its objective (contrastive density shaping between signal and background, followed by a classifier and Neyman intervals) is different from the InfoNCE loss implemented here.

In practice, expect ContrastiveFlowLoss to help only when:

• the evaluation task is parameter ranking, scanning, or retrieval over hypotheses
• you can generate informative alternate contexts during training
• the parameter space is low-dimensional enough that discrimination between nearby hypotheses matters
• full-density calibration and sampling quality are secondary to getting the ordering right

If you mainly want a well-calibrated conditional density model, start with NormalizingFlowLoss. ContrastiveFlowLoss is a task-specific objective, not a stronger default flow loss.

Step 3: Train¶

class ContextModel(nn.Module):
    """Backbone outputs a context vector that conditions the flow."""
    def __init__(self, in_dim, context_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, 128), nn.ReLU(),
            nn.Linear(128, context_dim),
        )
    def forward(self, x):
        return self.net(x)

model = ContextModel(in_dim=10, context_dim=64)
optimizer = torch.optim.Adam(
    list(model.parameters()) + list(loss_fn.parameters()), lr=1e-4,
)

for x, y in train_loader:
    context = model(x)                  # [batch, context_dim]
    loss = loss_fn(context, y)          # NLL
    optimizer.zero_grad(); loss.backward(); optimizer.step()

Training recipe for NLL and calibration¶

Three things decide whether an NSF head is calibrated, in this order of importance (measured with the flow settings above: 3 transforms, [64, 64] hidden, Adam at 1e-3, batch 64, 5 seeds, 1-D targets; NLL is test negative log-likelihood in native units, IS the 90% interval score from 512 samples):

1. Stop on validation NLL. A fixed 60-epoch budget overfits small and medium tabular sets (the flow sharpens on the training rows, intervals under-cover). Hold out about 15% of the training rows, evaluate loss_fn on them each epoch, stop after 10 epochs without improvement and restore the best weights.
2. Standardise the target and derive the spline bound from it. The spline is the identity outside [-B, B], so a standardised skewed or heavy-tailed target with |y| > 5 carries no learning signal there. Pass bound=recommended_tail_bound(y_standardised) (data range with 25% head-room, clipped to [5, 25]; it is 5 for light tails, so nothing changes there).
3. Use bins=16 with a widened bound so the bulk of the data keeps its resolution.

import copy
import torch
from torchregress.losses import NormalizingFlowLoss, create_flow_model, recommended_tail_bound

y_std = (y - y.mean()) / y.std()
flow = create_flow_model(
    n_features=1, context_dim=64, flow_type="nsf", n_transforms=3,
    hidden_features=[64, 64], bins=16, bound=recommended_tail_bound(y_std),
)
loss_fn = NormalizingFlowLoss(flow=flow)
n_val = int(0.15 * len(x))
x_val, y_val, x_tr, y_tr = x[:n_val], y_std[:n_val], x[n_val:], y_std[n_val:]
params = list(model.parameters()) + list(loss_fn.parameters())
optimizer = torch.optim.Adam(params, lr=1e-3)

best, best_state, bad = float("inf"), None, 0
for epoch in range(200):
    model.train()
    for i in torch.randperm(len(x_tr)).split(64):
        optimizer.zero_grad()
        loss_fn(model(x_tr[i]), y_tr[i]).backward()
        optimizer.step()
    model.eval()
    with torch.no_grad():
        val = float(loss_fn(model(x_val), y_val))
    if val < best - 1e-4:
        best, bad = val, 0
        best_state = copy.deepcopy((model.state_dict(), loss_fn.state_dict()))
    else:
        bad += 1
        if bad >= 10:
            break
model.load_state_dict(best_state[0]); loss_fn.load_state_dict(best_state[1])

On the house_sales row of the harness pzflow_parity suite the torchregress NSF trains for a fixed 60 epochs on all training rows with no validation split, which is the fixed 60 epochs recipe above; its 0.756 coverage at a nominal 0.90 (pzflow: 0.915) is the over-sharpening signature that early stopping removes on the datasets above.

The same holds for joint calibration of multi-target flows. On the harness multivariate_intervals suite (five datasets with 2 or 3 targets: energy_efficiency and four synthetic sets, one seed), the share of test targets inside the 90% joint highest-density region (sample_hpd_coverage) moved as follows when the fixed 60-epoch budget was replaced by the recipe above:

Test NLL improved in 14 of the 15 comparisons. The ensemble of early-stopped flows had the smallest joint calibration error on every dataset: its mean absolute coverage error over the levels 0.5, 0.68, 0.9 and 0.95 was 0.008 to 0.038, and its largest single deviation was 0.084 (0.14 for single early-stopped NSF or TarFlow).

Inference¶

Sampling¶

with torch.no_grad():
    context = model(x_test)
    samples = loss_fn.sample(context, n_samples=1000)
    # samples shape: [batch, 1000, n_features]

For a challenge-style parameter scan built on top of ContrastiveFlowLoss, see Contrastive Flow Parameter Estimation. For shared-budget comparisons against Gaussian-summary and plain-flow baselines, see Comparison Against Baselines (--mode compare).

Stable affine layers (MAF, RealNVP, cINN)¶

An affine layer maps coordinate \(i\) as \(z_i = s_i\,x_i + t_i\) with \((\log s_i, t_i) = h(x_{<i}, \mathbf{c})\) for MAF (\(x_{<i}\) are the preceding coordinates in the layer's order) or \(h(x_A, \mathbf{c})\) for a coupling layer (\(x_A\) is the unchanged half). Sampling runs the inverse \(x_i = (z_i - t_i)/s_i\). The conditioner \(h\) is a ReLU network, which extrapolates linearly outside the training range. One far-tail base draw gives a large \(x_j\); \(h\) answers with a large shift and a small scale, so \(x_{j+1}\) is larger still and the error grows geometrically over coordinates and layers. The density stays correct (the test NLL looks fine), but a few in a million samples are astronomically large, finite numbers, and sample-based scores (CRPS, energy score, intervals) are ruined. torchregress observed samples of \(10^{12}\) in standardized units from a 6-target MAF trained on one million rows, and the same tail from zuko's MAF used through sbi.

create_flow_model therefore builds the affine flows (zuko layout, masks and state-dict keys) with two safeguards:

so that each layer stretches or shrinks a coordinate by at most \(e^{3}\) (zuko's default \(L = 6.9\) allows a factor \(10^3\)), and the conditioner sees softly clipped inputs

\(\operatorname{clip}_B\) is the identity to second order near zero and saturates at \(\pm B\), so the conditioner's outputs stay in the range it learned and the inverse cannot cascade. The same conditioner is used in both directions: the layers remain exact bijections with an exact triangular Jacobian, and log_prob and sample describe the same density. \(B\) is in the target's units, so standardize the target (as for NSF) or pass a wider bound=, e.g. recommended_tail_bound(y_standardised) for heavy tails.

import torch
from torchregress.losses import NormalizingFlowLoss, create_flow_model

torch.manual_seed(0)
y = torch.distributions.StudentT(3.0).sample((4096, 3))  # heavy tails
x = torch.randn(4096, 8)
y_std = (y - y.mean(0)) / y.std(0)
encoder = torch.nn.Sequential(torch.nn.Linear(8, 32), torch.nn.ReLU(), torch.nn.Linear(32, 32))
loss_fn = NormalizingFlowLoss(create_flow_model(3, 32, flow_type="maf", n_transforms=3))
params = list(encoder.parameters()) + list(loss_fn.parameters())
optimizer = torch.optim.Adam(params, lr=1e-3)
for _ in range(200):
    optimizer.zero_grad()
    loss = loss_fn(encoder(x), y_std)
    loss.backward()
    optimizer.step()
with torch.no_grad():
    samples = loss_fn.sample(encoder(x[:100]), n_samples=1000)  # [100, 1000, 3]
print(float(samples.abs().max()))

NSF and BPF need no such safeguard: their monotone maps are the identity outside \([-B, B]\) and map \([-B, B]\) onto itself, so a layer's inverse satisfies \(|x| \le \max(|z|, B)\).

Transformer autoregressive flows (TarFlow)¶

flow_type="tarflow" builds a tabular, conditional version of TarFlow (Zhai et al. 2025, ref. 14) and its deep-shallow scaling, STARFlow (Gu et al. 2025, ref. 15). Block \(k\) is an exact autoregressive bijection in a coordinate order \(\pi_k\):

where the token sequence is a context token followed by one token per preceding target, and \(f\) is a monotone rational-quadratic spline (as in NSF, ref. 16) or, with univariate="affine", TarFlow's affine map. Consecutive blocks reverse the order. The Jacobian is triangular, so

The difference from NSF and MAF is the conditioner: a masked MLP (MADE) gives every target the same fixed-width view of its predecessors, while attention lets each target weigh any of them. Density evaluation is one Transformer pass per block; sampling needs one pass per target coordinate per block.

The image-specific parts of TarFlow (patch tokens, noise augmentation with a score-based denoising step, guidance) and STARFlow's pretrained autoencoder latent space are left out, because tabular targets are continuous, low-dimensional and modelled directly.

import torch

from torchregress.losses import NormalizingFlowLoss, create_flow_model, recommended_tail_bound

torch.manual_seed(0)
x = torch.randn(2000, 2)
z = torch.randn(2000, 8).cumsum(dim=1) / torch.arange(1, 9).sqrt()  # correlated
y = z + 0.3 * torch.tanh(x[:, :1]) * z**2  # skewed, context-dependent, 8 targets
y = (y - y.mean(0)) / y.std(0)

flow = create_flow_model(
    n_features=8,
    context_dim=2,
    flow_type="tarflow",
    n_transforms=3,
    hidden_features=[64],
    n_layers=(4, 1, 1),  # STARFlow-style: one deep block, then shallow ones
    bound=recommended_tail_bound(y),
)
loss_fn = NormalizingFlowLoss(flow)
optimizer = torch.optim.Adam(flow.parameters(), lr=1e-3)
for epoch in range(3):
    for batch in torch.randperm(2000).split(256):
        optimizer.zero_grad()
        loss_fn(x[batch], y[batch]).backward()
        optimizer.step()

samples = loss_fn.sample(x[:5], n_samples=1000)  # [5, 1000, 8], for corner plots

TarFlow against NSF and MAF¶

Each flow was trained with Adam (\(10^{-3}\)), batch size 256 and early stopping on validation NLL (patience 15). The data were 8000 training, 2000 validation and 4000 test pairs, with standardized targets and the spline bound from recommended_tail_bound. Results are means over 3 seeds. The tasks have known conditional densities. KL gap is \(\mathbb{E}[\log p - \log \hat p]\) on the test set in nats (0 is perfect). Calibration is the mean absolute error of the sample-based joint HPD coverage (sample_hpd_coverage) at 0.5, 0.68, 0.9 and 0.95. Sampling is the wall time for 256 draws at each of 1000 contexts (8 CPU threads shared by 4 jobs). NSF and MAF use 5 blocks of MLP widths [128, 128]. TarFlow uses 3 blocks of width 64, with 2 layers per block, or 4, 1 and 1 layers in the deep-shallow variant.

† One job per 8-core session (2 threads per job), so the sampling times are not directly comparable with the other rows.

The affine variant (univariate="affine", five blocks) was only run on the first four tasks. It beat the spline variant on the banana (0.053 nats) but was worse on the other three (0.547, 0.117 and 0.263 nats), and worst of all flows on the mixture: one affine map per block cannot split a mode.

Flow Architectures¶

Limitations¶

1. Computational cost: Each training step evaluates \(K\) invertible transforms and their Jacobian determinants. For NSF with 5+ transforms and large context dimensions, training can be 5–20× slower than an equivalent MDN.
2. Mask semantics: Flow objectives model a joint density over all target dimensions. Only sample-level masking is supported — if one target dimension is missing, drop the whole sample.
3. Architecture choice matters: RealNVP is fast but less expressive; NSF is most expressive but slower. Match the flow type to your target: use NSF for complex multimodality, MAF for conditional density estimation, RealNVP for a quick baseline.
4. ContrastiveFlowLoss is task-specific: It optimises a contrastive likelihood-ratio/ranking objective, not density calibration. temperature != 1 or margin != 0 leave the calibrated-likelihood-ratio regime; negative contexts must follow a meaningful reference distribution for the ratios to be interpretable. If you need well-calibrated conditional densities for sampling or downstream UQ, start with NormalizingFlowLoss. For classifier-based NRE see NeuralRatioLoss.
5. Dependency: Requires zuko. Not installed by default; use pip install torchregress[flows].

When To Use Which Flow Objective¶

Next steps¶

• MDN losses — lighter-weight mixture models for simpler multimodality
• SLS regression — volume-optimal prediction regions (flow-based frontier)
• Contrastive flow estimation — parameter ranking example
• Flow comparison — Gaussian summary, flow NLL, contrastive flow, NRE and NPE on one budget

Ensembles of flows and joint calibration¶

A single flow captures aleatoric spread, not the uncertainty in its own fit. Train \(M\) flows independently (different seeds or bootstraps) and combine them with DensityEnsemble, the equal-weight mixture

whose entropy splits into a mean member entropy (aleatoric) and the members' mutual information with \(y\) (epistemic, \(\geq 0\)):

Unlike a variance decomposition, this holds for multimodal and multi-target densities.

Per-dimension PIT and interval coverage cannot detect a wrong joint shape (for example, a wrong correlation between targets). sample_hpd_level works on joint highest-density regions instead. For each target it counts how many predictive draws have a higher density than the observed \(y\):

For a calibrated model \(\ell\) is uniform on \([0, 1]\) in any dimension, and sample_hpd_coverage reports the coverage of the nominal 50/68/90/95% regions. Values pile up near 1 when the model is over-confident and near 0 when it is under-confident.

import torch
from torch import nn
from torchregress.ensemble import DensityEnsemble
from torchregress.losses import create_flow_loss, recommended_tail_bound
from torchregress.metrics import model_hpd_level

torch.manual_seed(0)
x = torch.randn(2048, 4)
y = torch.stack([x[:, 0] + 0.3 * torch.randn(2048), x[:, 0] * x[:, 1] + 0.3 * torch.randn(2048)], 1)

bound = recommended_tail_bound(y)  # keep the spline range wide enough for x0 * x1
heads = []
for seed in range(3):
    torch.manual_seed(seed)
    head = create_flow_loss(n_features=2, context_dim=4, flow_type="nsf", bound=bound)
    optimiser = torch.optim.Adam(head.parameters(), lr=3e-3)
    for _ in range(200):
        loss = head(x, y)            # conditional NLL with the features as context
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
    heads.append(head)

ensemble = DensityEnsemble(heads)
x_test = torch.randn(256, 4)
y_test = torch.stack(
    [x_test[:, 0] + 0.3 * torch.randn(256), x_test[:, 0] * x_test[:, 1] + 0.3 * torch.randn(256)], 1
)
with torch.no_grad():
    nll = -ensemble.log_prob(x_test, y_test).mean()
parts = ensemble.entropy_decomposition(x_test, n_samples=64)       # total / aleatoric / epistemic
levels = model_hpd_level(ensemble, x_test, y_test, n_samples=128)  # ~ Uniform(0, 1) if calibrated
coverage_90 = (levels <= 0.9).float().mean()  # well below 0.9 means over-confident

Short training budgets typically leave flow ensembles over-confident (90% joint coverage well below 0.9 in this example). Validation-NLL early stopping and more members both help; conformalising the ensemble density with ConformalizedDensity restores marginal coverage.

Fitted weights with error bars. Equal weights are a choice, not an estimate. ensemble.fit_weights(x_val, y_val) fits \(\pi\) by maximum likelihood on held-out data (concave, solved by EM) and stores a sandwich covariance; ensemble.log_prob_std(x, y) then returns the standard error of \(\log p(y \mid x)\) due to estimating \(\pi\). This applies the wifi ensemble idea (Benevedes & Thaler 2025, see NRE) to densities. It quantifies weight uncertainty only, not errors inside the frozen members, and assumes the weights are away from zero.

# Continuing the example above: fresh data that the heads were not trained on.
x_val = torch.randn(1024, 4)
y_val = torch.stack(
    [x_val[:, 0] + 0.3 * torch.randn(1024), x_val[:, 0] * x_val[:, 1] + 0.3 * torch.randn(1024)], 1
)
ensemble.fit_weights(x_val, y_val)
log_p = ensemble.log_prob(x_test, y_test)
log_p_std = ensemble.log_prob_std(x_test, y_test)  # [N]

References¶