Marginal Contrastive Discrimination¶

MarginalContrastiveDensity estimates a conditional density by splitting it into an unconditional density and a classification problem (Meziani, Ndiaye & Riu 2022, ref. 1):

• \(p(\mathbf{y})\) is any unconditional flow over the targets, e.g. create_flow_model(n_features=D, context_dim=0).
• \(r\) is learned by a classifier that tells joint pairs \((\mathbf{x}_i, \mathbf{y}_i)\) from shuffled pairs \((\mathbf{x}_j, \mathbf{y}_i)\). It is trained with contrastive neural ratio estimation (NeuralRatioLoss(criterion="cnre"), ref. 2), whose logit estimates the absolute \(\log r\), not \(\log r\) up to a constant per \(\mathbf{x}\).

The training objective loss(x, y) is the marginal negative log-likelihood plus the CNRE loss, where each \(\mathbf{y}_i\) is paired with \(\mathbf{x}\) from n_negatives other rows of the batch:

Normalization and sampling¶

The exact ratio satisfies \(Z(\mathbf{x}) = \mathbb{E}_{\mathbf{y} \sim p(\mathbf{y})}[r(\mathbf{x}, \mathbf{y})] = 1\); a learned one does not. log_prob therefore subtracts a Monte-Carlo estimate by default,

with \(M = 4096\) (n_samples). Its error is about \(\operatorname{sd}(r(\mathbf{x}, \cdot)) / \sqrt{M}\), larger for contexts whose conditional sits in the tail of \(p(\mathbf{y})\). The normalizer is computed without gradients and once per distinct context, so log_prob (and forward, its negative) is a score, not a training objective: train with loss.

In eval mode the marginal draws are cached (one bank per size) and shared by log_prob, sample and effective_sample_size. Scores from separate calls, such as a target and its predictive draws in model_hpd_level, then use the same normalizer, and the effective sample size describes the bank that sample resampled. train(), load_state_dict and refresh_bank() discard the cache.

sample(x, n) uses sampling-importance-resampling: a bank of n_proposals marginal draws is resampled with weights \(r(\mathbf{x}, \tilde{\mathbf{y}}_m)\). When \(p(\mathbf{y} \mid \mathbf{x})\) is much narrower than \(p(\mathbf{y})\), a few draws carry the weight and are repeated; effective_sample_size(x) reports the Kish effective size of the bank, between 1 and n_proposals.

Example¶

import torch

from torchregress.losses import build_ratio_classifier, create_flow_model
from torchregress.metrics import model_hpd_level
from torchregress.models import MarginalContrastiveDensity

torch.manual_seed(0)
x = torch.randn(2000, 3)
y = torch.stack([x[:, 0] + 0.3 * torch.randn(2000), x[:, 0] * x[:, 1] + 0.3 * torch.randn(2000)], 1)
y = (y - y.mean(0)) / y.std(0)

model = MarginalContrastiveDensity(
    create_flow_model(n_features=2, context_dim=0, n_transforms=3),
    build_ratio_classifier(context_dim=3, observable_dim=2),
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(5):
    for batch in torch.randperm(2000).split(128):
        optimizer.zero_grad()
        model.loss(x[batch], y[batch]).backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    log_density = model.log_prob(x[:5], y[:5])  # normalized by default
    samples = model.sample(x[:5], n_samples=500)  # [5, 500, 2], for corner plots
    ess = model.effective_sample_size(x[:5])  # check before trusting the samples
levels = model_hpd_level(model, x[:200], y[:200])  # joint HPD calibration

To score several targets per context (e.g. predictive draws for HPD levels), compute log_normalizer(x) once and call log_prob(..., normalize=False).

When to use it¶

Harness comparison¶

On the torchregress-harness multivariate_intervals suite (five datasets with 2 or 3 targets, one seed), all rows trained with the same early-stopping recipe: a 15% validation carve, patience 10, recommended_tail_bound and 16 bins. MCD uses an NSF marginal and a CNRE classifier on the raw features. HPD90 is the share of test targets inside the 90% joint highest-density region; cal. err. is the mean absolute coverage error over the levels 0.5, 0.68, 0.9 and 0.95; ESS is the mean effective share of the 8192-draw bank.

MCD's NLL was worse than a single NSF on four of the five datasets. Its joint calibration was competitive, best of the three on the bimodal set, but it over-covered on energy_efficiency, where only 16% of the bank carried the weight. The ensemble of early-stopped flows had the best NLL on every dataset.

MCD's appeal is that the conditional part is plain classification, so the ratio model can be anything that produces a logit. Its weak point is structural: sampling and normalization both average over \(p(\mathbf{y})\), which wastes draws when the features are informative.

→ See also normalizing flow losses, neural ratio estimation and joint HPD calibration metrics.

References¶