Neural Ratio Estimation (NRE)¶

← Normalizing Flows | Loss catalogue →

Neural Ratio Estimation trains a classifier to estimate the likelihood-to-evidence ratio \(r(x,\theta) = p(x \mid \theta)/p(x)\) directly, without ever modelling a normalized density. A classifier that separates true parameter–observation pairs from mismatched pairs estimates \(\log r(x,\theta)\) as its logit at its optimum, so ranking candidate parameters by that logit is maximum-likelihood estimation. The training prior cancels in the ratio, which makes NRE attractive for simulator-based parameter inference, likelihood-ratio tests, and amortized regression heads on problems where a normalized likelihood is unavailable or expensive.

The implementation is in torchregress.losses.nre: NeuralRatioLoss wraps any nn.Module scorer, build_ratio_classifier constructs the usual MLP, and neural_ratio_logits / neural_ratio_loss are the functional forms.

When to use which objective¶

Mathematical background¶

The joint-to-marginal ratio¶

Let \(\theta\) be a parameter/hypothesis (the context in the API) and \(x\) an observation. NRE targets

Consider a classifier \(s(\theta, x)\) trained with balanced labels to distinguish samples from the joint \(p(\theta, x)\) (label 1) from samples from the marginal \(p(\theta)\,p(x)\) (label 0). Its Bayes-optimal posterior is

so the logit of the optimal classifier is exactly the log likelihood ratio,

This is the likelihood-ratio trick (Cranmer, Pavez & Louppe 2015): classification replaces density estimation, and the classifier can be any differentiable scorer, e.g. build_ratio_classifier.

Losses¶

For one observation with a matched positive context \(\theta^+\) and \(K\) negative candidate contexts \(\{\theta^-_k\}_{k=1}^K\), write \(s_+ = s(\theta^+, x)\) and \(s^-_k = s(\theta^-_k, x)\). The balanced binary NRE objective is

where the \(1/K\) factor balances the single positive against the \(K\) negatives. The multiclass (categorical) NRE objective treats the positive and the \(K\) negatives as the \(K+1\) classes of a softmax classifier,

with \(s_0 = s_+\). This is the standard categorical cross-entropy, i.e. the NRE-B / InfoNCE-style form: the softmax is normalized over the sampled candidate set, and the objective is a contrastive bound rather than an exact finite-\(K\) ratio estimator. As \(K \to \infty\) its optimum satisfies \(\exp(s(\theta, x)) \propto r(x,\theta)\), i.e. the logit approaches \(\log r\) up to an additive constant; at finite \(K\) NRE-B has an intrinsic bias term (Miller, Weniger & Forré 2022). Use criterion="cnre" (below) or criterion="bce" (NRE-A) when an accurate absolute log-ratio is needed. Use the multiclass form only for ranking with many candidates.

Contrastive NRE (CNRE) (Miller, Weniger & Forré 2022) removes that bias by adding a class \(y = 0\) in which none of the \(K\) candidates generated \(x\). With prior odds \(\gamma\) between "candidate \(k\) generated \(x\)" (\(y \geq 1\), uniform over \(k\)) and \(y = 0\), the class posterior is

and the loss mixes a dependent set (the true \(\theta_+\) and \(K-1\) negatives) with an independent set (all \(K\) negatives):

The \(y = 0\) class fixes the normalisation, so the optimum is \(s = \log r(x, \theta)\) exactly for every \(K\). \(K = 1\) recovers binary NRE with prior odds \(\gamma\). The negatives must be drawn from the reference prior independently of \(x\), so negative_target is rejected for criterion="cnre".

In both objectives the classifier sees a positive pair and negative pairs. The reduction over the batch is handled by BaseLoss ('none' | 'mean' | 'sum', mask, weights).

Why the training prior cancels for the MLE¶

Insert \(p(\theta, x) = p(\theta)\,p(x \mid \theta)\) into the ratio:

The parameter prior \(p(\theta)\) cancels identically. Consequently

which is the maximum-likelihood estimate; \(\log p(x)\) does not depend on \(\theta\). The requirement is not that the prior be uniform: it is that the negative contexts are drawn from the same reference prior \(p(\theta)\) (or reweighted to it), so that the marginal class really is \(p(\theta)\,p(x)\). Under that condition the training prior is irrelevant for ranking \(\theta\).

If the training negatives instead come from a proposal \(q(\theta)\), the optimal logit becomes

so the proposal can be corrected by subtracting \(\log p(\theta) - \log q(\theta)\) when the two densities are known. A MAP estimate adds the target prior back explicitly: \(\log p(\theta \mid x) = \log r(x,\theta) + \log p(\theta) + \text{const}\).

This argument is exact for the balanced binary objective. The finite-\(K\) multiclass objective superimposes the contrastive bias discussed above on the same prior-free target.

API¶

import torch

from torchregress.losses.nre import (
    NeuralRatioLoss,
    build_ratio_classifier,
    neural_ratio_logits,
    neural_ratio_loss,
)

# Class form: the loss owns (and optimizes) the classifier parameters.
classifier = build_ratio_classifier(context_dim=3, observable_dim=8, hidden_dims=(64, 64))
loss_fn = NeuralRatioLoss(classifier, criterion="bce", reduction="mean")
loss = loss_fn(
    positive_context,            # [N, 3] candidate parameters (e.g. labels or model estimates)
    observations,                # [N, 8] data
    negative_context=negative,   # [N, K, 3], [1, K, 3] or [K, 3] (shared bank)
    negative_target=None,        # optional [N, K, 8] mismatched observations
    mask=sample_mask,            # optional [N] boolean
    weights=sample_weights,      # optional [N] non-negative
)

# Functional form (same math, classifier passed explicitly).
loss = neural_ratio_loss(classifier, positive_context, observations, negative)

# Log-ratio estimation and plain scoring.
ratios = loss_fn.log_likelihood_ratio(positive_context, observations, negative)  # [N, K]
logits = neural_ratio_logits(classifier, candidate_context, observations)        # [N]

criterion="bce" is the balanced binary loss; criterion="multiclass" is the categorical / InfoNCE-style loss; criterion="cnre" is the bias-free contrastive loss (tune gamma, default 1.0). The classifier is a plain nn.Module mapping concat([context, observable]) to one logit; build_ratio_classifier tags it with context_dim / observable_dim, and NeuralRatioLoss infers the dimensions lazily for custom modules.

Negative-sampling conventions¶

negative_context is required; its shape determines how negatives are shared across the batch:

For a 2-D [N, C] input whose leading dimension equals the batch size, set shared_negative_context=True (shared bank) or False (per-sample, one negative per sample). The same rules apply to negative_target via shared_negative_target.

negative_target selects the NRE variant:

• None (paired negatives) — the negative hypotheses are scored against the same observation \(x\): \(s^-_k = s(\theta^-_k, x)\). The classifier learns the ratio, and log_likelihood_ratio then returns \(\log p(x \mid \theta^+) - \log p(x \mid \theta^-_k)\), a likelihood ratio between hypotheses.
• Given (mismatched pairs) — each negative is a full pair \(s^-_k = s(\theta^-_k, x^-_k)\) with \(x^-_k\) sampled independently of \(\theta^-_k\) (e.g. shuffled observations). This is the canonical joint-versus-marginal NRE training setup.

Runnable example: Gaussian MLE by grid argmax¶

The simulator is \(c \sim \mathcal{N}(0, 1)\) and \(x \mid c \sim \mathcal{N}(c, \sigma^2)\) with \(\sigma = 0.5\). We train the classifier on paired negatives drawn from the reference prior, then invert a new observation by scoring a dense candidate grid and taking the argmax (the MLE).

import torch

from torchregress.losses.nre import (
    NeuralRatioLoss,
    build_ratio_classifier,
    neural_ratio_logits,
)

torch.manual_seed(0)
sigma = 0.5


def simulate(n: int) -> tuple[torch.Tensor, torch.Tensor]:
    c = torch.randn(n, 1)                       # latent parameter (NRE context)
    x = c + sigma * torch.randn(n, 1)           # observation
    return c, x


# 1. Train the ratio classifier. Negative contexts are alternative hypotheses
#    c' for the same observation, drawn from the same prior as c.
classifier = build_ratio_classifier(context_dim=1, observable_dim=1, hidden_dims=(64, 64))
loss_fn = NeuralRatioLoss(classifier, criterion="bce")
optimiser = torch.optim.Adam(classifier.parameters(), lr=3e-3)

for _ in range(600):
    c, x = simulate(256)
    negative_c = torch.randn(256, 8, 1)         # reference-prior alternatives
    loss = loss_fn(c, x, negative_context=negative_c)
    optimiser.zero_grad()
    loss.backward()
    optimiser.step()
print(f"final training loss: {float(loss.detach()):.3f}")

# 2. Maximum-likelihood inversion on a grid: score every candidate for new data.
grid = torch.linspace(-3.0, 3.0, 121).reshape(-1, 1)
torch.manual_seed(123)
c_true, x_obs = simulate(5)
n_obs, n_candidates = x_obs.shape[0], grid.shape[0]
candidate_context = grid.unsqueeze(0).expand(n_obs, n_candidates, 1).reshape(-1, 1)
observations = x_obs.unsqueeze(1).expand(n_obs, n_candidates, 1).reshape(-1, 1)
with torch.no_grad():
    scores = neural_ratio_logits(classifier, candidate_context, observations)
scores = scores.reshape(n_obs, n_candidates)
c_hat = grid[scores.argmax(dim=1)].squeeze(-1)

# The flat-prior MLE of a Gaussian location model is x_obs (nearest grid point).
print("true c:     ", c_true.squeeze(-1))
print("grid MLE c: ", c_hat)
print(f"mean |c_hat - x_obs| = {float((c_hat - x_obs.squeeze(-1)).abs().mean()):.3f}")
print(f"mean |c_hat - c_true| = {float((c_hat - c_true.squeeze(-1)).abs().mean()):.3f}")

# 3. Pairwise log likelihood ratio between two hypotheses for one observation.
hypothesis_a = torch.tensor([[0.0]])
hypothesis_b = torch.tensor([[1.0]])
log_ratio = loss_fn.log_likelihood_ratio(
    hypothesis_a, x_obs[:1], hypothesis_b.unsqueeze(0)
)
print(f"estimated log-ratio at c=0 vs c=1: {float(log_ratio.detach()):.3f}")

On a typical run the grid-MLE estimate tracks the analytic MLE \(x_{\text{obs}}\) within about one grid step (a mean absolute deviation of about \(0.03\) on a grid with spacing \(0.05\)) and the generating \(c\) within the observation noise (\(\approx 0.2\) for \(\sigma = 0.5\)).

Frequentist uncertainties: wifi ensembles¶

A trained ratio estimator has no error bar. wifi ensembles (Benevedes & Thaler, PRD 112, 056024, 2025) add one. Train \(M\) ratio estimators \(f_i(x) \approx \log r(x)\) independently (bootstrap resamples), freeze them, and model

fitting \(w\) on held-out samples \(x \sim p\) (numerator) and \(x \sim q\) (denominator) by the symmetrised maximum-likelihood-classifier loss

which is convex in \(w\) and minimised pointwise by \(s = \log p/q\). Because \(\hat w\) is an M-estimator, its covariance is the sandwich \(\Sigma_w = H^{-1} V H^{-1}\) (\(H\) the Hessian of \(L\), \(V\) the covariance of the per-sample scores), and

For a mixture \(x \sim \kappa p + (1-\kappa) q\), the pseudo-MLE of \(\kappa\) maximises \(\sum_j \log\bigl(1 + \kappa(\hat r_j - 1)\bigr)\) and its variance adds the ratio uncertainty (Gong & Samaniego 1981):

WifiRatioEnsemble works on stacked member outputs, so the members can be NeuralRatioLoss classifiers, plain binary classifiers, or flow log-density differences.

import torch
import torch.nn.functional as F
from torch import nn

from torchregress.inference import WifiRatioEnsemble

torch.manual_seed(0)
# Numerator p = N([1, 0], I), denominator q = N([0, 0], I) in 2-D.
def draw_p(n): return torch.randn(n, 2) + torch.tensor([1.0, 0.0])
def draw_q(n): return torch.randn(n, 2)

# 1) Train M classifiers on bootstrap resamples; logit ~ log p(x)/q(x).
x_train = torch.cat([draw_p(4000), draw_q(4000)])
label = torch.cat([torch.ones(4000), torch.zeros(4000)])
members = []
for _ in range(8):
    index = torch.randint(len(x_train), (len(x_train),))
    net = nn.Sequential(nn.Linear(2, 32), nn.SiLU(), nn.Linear(32, 32), nn.SiLU(), nn.Linear(32, 1))
    optimiser = torch.optim.Adam(net.parameters(), lr=3e-3)
    for _ in range(300):
        loss = F.binary_cross_entropy_with_logits(net(x_train[index]).squeeze(-1), label[index])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
    members.append(net)

# 2) Freeze the members and fit the wifi weights on fresh held-out samples.
features = lambda x: WifiRatioEnsemble.member_features(members, x)  # [N, M]
wifi = WifiRatioEnsemble().fit(features(draw_p(5000)), features(draw_q(5000)))

# 3) Log ratio with a frequentist standard error.
x_test = torch.tensor([[0.0, 0.0], [1.0, 0.0], [2.0, 1.0]])
log_r, log_r_std = wifi.log_ratio(features(x_test)), wifi.log_ratio_std(features(x_test))
true_log_r = x_test[:, 0] - 0.5

# 4) Mixture fraction kappa in x ~ kappa p + (1 - kappa) q, with the ratio
#    uncertainty propagated (Gong-Samaniego).
from_p = torch.rand(10_000) < 0.2
x_mix = torch.where(from_p.unsqueeze(1), draw_p(10_000), draw_q(10_000))
estimate = wifi.mixture_fraction(features(x_mix))
print(f"kappa = {estimate.fraction:.3f} +/- {estimate.std:.3f} "
      f"(stat {estimate.std_statistical:.3f}, ratio {estimate.std_ratio:.3f})")

The equations are a reconstruction from the paper's text, validated by coverage tests on problems with a known ratio rather than against the authors' code.

Comparison with related objectives¶

In short: NRE is the lightest of these when you have a simulator and care about ratios and rankings, but it is not a replacement for a calibrated density or posterior surrogate. For finite-sample coverage of prediction intervals, pair any inference with conformal methods and calibration diagnostics (see calibration metrics).

Caveats and failure modes¶

References¶

Next steps¶

• Normalizing flow losses — the flow-based alternative, including ContrastiveFlowLoss
• Loss catalogue — where NRE sits among regression objectives
• Gaussian losses — plain NLL when a normalized Gaussian likelihood is acceptable
• Likelihood-ratio posterior utilities — turn log-ratios into posterior weights, effective sample sizes and MLE/MAP estimates
• SBC and TARP — validate posterior calibration before reporting intervals
• Calibration metrics — diagnose ratio/logit calibration downstream
• Conformal prediction — distribution-free prediction intervals