Gaussian Wasserstein bound surrogate¶

← Faithful Gaussian | Robust Losses →

This loss supervises a predicted Gaussian $ \mathcal{N}(\hat\mu, \hat\Sigma) $ against a target mean $ \mu $ and target covariance $ \Sigma $ (or a pseudo-covariance) using a standard Frobenius surrogate built from principal matrix square roots. It is useful for joint mean–covariance learning; it is not a drop-in replacement for GaussianNLLLoss on raw targets.

When to use which¶

Mathematical definition¶

Let $ \hat S $ and $ S $ denote symmetric principal square roots of $ \hat\Sigma $ and $ \Sigma $ (implemented via eigen-decomposition with a small eigenvalue floor). The per-sample objective is

In diagonal mode, the covariance term is $ \sum_i \bigl(\sqrt{\hat v_i} - \sqrt{v_i}\bigr)^2 $ for positive variances $ \hat v_i, v_i $.

Parameterisations¶

The helper symmetric_spd_matrix_sqrt (exported from torchregress.losses) applies the same root used internally for full matrices; see GaussianWassersteinBoundLoss.

API¶

from torchregress.losses import GaussianWassersteinBoundLoss, gaussian_wasserstein_bound_loss

loss_fn = GaussianWassersteinBoundLoss(
    covariance_parameterization="covariance",
    mean_weight=1.0,
    covariance_weight=1.0,
    jitter=1e-6,
    reduction="mean",
)
loss = loss_fn(pred_mean, target_mean, pred_cov, target_cov, mask=mask, weights=w)

create_loss_from_config accepts "type": "gaussian_wasserstein_bound".

Runnable example¶

See Gaussian Wasserstein bound demo and examples/gaussian_wasserstein_bound_demo.py --mode loss.

For a two-phase pseudo-covariance pretrain then Gaussian NLL fine-tune, see Wasserstein-bound hybrid pretraining (--mode hybrid).

Covariance pseudo-labels (experimental)¶

When you do not have oracle covariance targets, the research plan suggests neighbour-weighted local covariance of y in input space as a heuristic supervision signal. v1 lives under algorithms (not losses):

• torchregress.algorithms.NeighborhoodCovariancePseudoLabeler with fit_predict(x, y) -> [n, d, d]
• torchregress.algorithms.mahalanobis_covariance_pseudo_labels functional wrapper

Metrics supported: pooled Mahalanobis in x (default) or Euclidean distances; weights are a softmax over neighbour distances. Outputs are symmetrised and eigenvalue-clamped for numerical SPD-ish matrices. Treat as experimental and validate on your modality before relying on gradients.

References¶

Limitations¶

1. Surrogate, not exact Wasserstein: This objective is a Frobenius-norm upper-bound surrogate related to Gaussian 2-Wasserstein ideas. It is not the exact Wasserstein-2 distance in all non-commutative cases. Treat it as a training signal, not an interpretable distance metric.
2. Requires covariance targets: Unlike NLL which learns from raw targets \((x, y)\), this loss requires a target covariance \(\Sigma\) per sample (or a shared \(\Sigma\)). Covariance targets are rarely available — they must come from problem structure, pseudo-labeling, or external estimation.
3. Pseudo-label quality: The NeighborhoodCovariancePseudoLabeler in torchregress.algorithms provides heuristic covariance targets from neighbour-weighted statistics. These are experimental — validate on your modality before relying on gradients.
4. Eigenvalue floor: The matrix square root requires eigenvalue decomposition with a floor (jitter) to guarantee positive-semidefiniteness. If jitter is too small, gradients become unstable; if too large, covariance information is washed out. Repeated eigenvalues (e.g. an isotropic \(\sigma^2 I\) initialisation) are safe: symmetric_spd_matrix_sqrt back-propagates by solving the Sylvester equation \(X\,dX + dX\,X = d\Sigma\) in the eigenbasis (gradient entries \(1/(s_i + s_j)\)) instead of differentiating through eigh, whose backward is NaN at degenerate eigenvalues.

Recommendations¶

• Use only with covariance targets: This loss is for supervised or pseudo-supervised covariance learning. For standard likelihood training on \((x, y)\) pairs, use GaussianNLLLoss.
• Start with diagonal parameterisation: covariance_parameterization="diagonal" is simpler, faster, and numerically more stable than full covariance or Cholesky modes.
• Two-phase strategy: Consider a pseudo-covariance pretrain with GaussianWassersteinBoundLoss followed by GaussianNLLLoss fine-tuning. See the hybrid pretraining example (examples/gaussian_wasserstein_bound_demo.py --mode hybrid).
• Validate pseudo-labels: When using NeighborhoodCovariancePseudoLabeler, check that the produced covariance matrices are plausible (positive eigenvalues, reasonable scale) before training.

Next steps¶

• Gaussian losses — NLL, CRPS, and multivariate likelihoods for standard Gaussian training
• Beta-NLL — stabilized heteroscedastic likelihood training without covariance targets
• TIC-TAC — Taylor-Induced Covariance for improved covariance estimation
• Wasserstein demo — runnable example with pseudo-covariance pretraining