Grid & basis densities: Wasserstein-1 and Rank-N-Contrast¶

← Normalizing Flows | Evidential Regression →

Tools for heads that predict a whole 1D distribution over a fixed support (binned classification, spline bases) rather than a parametric family, plus a representation loss that orders the feature space by the continuous target.

#	Reference
1	Zha, Y.-J. et al. (2023). Rank-N-Contrast: Learning Continuous Representation Degrees for Improved Forecasting. NeurIPS 2023.
2	Ramsay, J. O. (1988). Monotone Regression Splines in Action. Statistical Science 3(4), 425–441. (M-splines)
3	Villani, C. (2009). Optimal Transport: Old and New. Springer. (Wasserstein distances)
4	de Boor, C. (2001). A Practical Guide to Splines. Springer. (B-spline basis, Cox–de Boor recursion)

Related methods for other 1D distribution heads: → See Normalizing Flows, SLS (structured least-squares quantiles), and MDN.

`BSplineDensityBasis` (`torchregress.utils`)¶

A density on \([lo, hi]\) is written as

\[ p(x) = \sum_{m=1}^{M} c_m\, M_m(x), \qquad c \in \Delta^{M-1}, \qquad M_m(x) = \frac{k\, B_{m,k}(x)}{t_{m+k} - t_m}, \]

where \(B_{m,k}\) are order-\(k\) (degree \(k-1\)) B-splines on a clamped knot vector and \(M_m\) are their unit-integral (M-spline) rescalings. Standard B-splines satisfy \(\sum_m B_m = 1\), not \(\int B_m = 1\); the rescaling is what makes a softmax over coefficients produce a normalised, non-negative density with no extra pass.

Method	Returns	Notes
`from_uniform(lo, hi, n_intervals, degree=3)`	basis	Equally spaced breakpoints
`from_quantiles(samples, n_interior, lo, hi)`	basis	Interior knots at empirical quantiles (balanced gradient flow)
`evaluate(x)`	`(..., M)`	Cox–de Boor recursion; zero outside `[lo, hi]`
`density(coeffs, grid)`	`(..., G)`	\(\sum_m c_m M_m(z_g)\)
`bin_integrals(edges)`	`(K, M)`	\(T_{km} = \int_{e_k}^{e_{k+1}} M_m\); exact via piecewise Gauss–Legendre
`cdf(x)` / `cumulative_moment(x, r)`	`(len(x), M)`	\(\int_{lo}^{x} u^r M_m(u)\,du\)
`absolute_deviation(x)`	`(len(x), M)`	\(\mathbb{E}_{M_m}\lvert U - x\rvert\) — exact \(W_1\) to a point mass

Bin masses are a fixed linear map of the coefficients, masses = coeffs @ T.T, so tomographic/binned summaries of a predicted density need no sampling.

import torch
from torchregress.utils import BSplineDensityBasis

basis = BSplineDensityBasis.from_quantiles(train_targets, n_interior=12, lo=0.0, hi=3.0)
coeffs = torch.softmax(head(features), dim=-1)          # [N, basis.n_basis]
nll = -(coeffs * basis.evaluate(y)).sum(-1).log()       # exact density at the labels
w1 = (coeffs * basis.absolute_deviation(y)).sum(-1)     # exact W1(p, delta_y) = E|U - y|

`DiscreteWasserstein1Loss`¶

For mass vectors \(p, q\) located at bin centres \(c_1<\dots<c_K\),

\[ W_1(p, q) = \sum_{k=1}^{K-1} \lvert F_k - G_k \rvert\,(c_{k+1} - c_k), \]

with \(F, G\) the cumulative sums. Unlike KL or cross-entropy, the penalty grows with how far misplaced mass is from the target, which complements a local divergence when training binned or basis heads. Targets may be [N, K] masses or [N] scalar values (quantised to the containing bin; for continuous parameterisations prefer BSplineDensityBasis.absolute_deviation).

from torchregress.losses import DiscreteWasserstein1Loss

loss = DiscreteWasserstein1Loss(bin_edges, from_logits=True)
value = loss(logits, target_masses)          # or loss(logits, scalar_targets)

`RankNContrastLoss`¶

Rank-N-Contrast (Zha et al., NeurIPS 2023) trains features \(h_i\) so that similarity is ordered by label distance. For anchor \(i\) and any \(j \ne i\),

\[ \mathcal{L}_i = \frac{1}{N-1}\sum_{j\ne i} -\log \frac{\exp(s_{ij}/\tau)}{\sum_{k\in S_{ij}} \exp(s_{ik}/\tau)}, \qquad S_{ij} = \{k \ne i : \lvert y_i - y_k\rvert \ge \lvert y_i - y_j\rvert\}, \]

with \(s_{ij} = -\lVert h_i - h_j\rVert_2\) (default) or cosine similarity. Ties in label distance are included in \(S_{ij}\). The implementation sorts by label distance per anchor and uses a cumulative log-sum-exp, so it is \(O(N^2 \log N)\) rather than \(O(N^3)\).

from torchregress.losses import RankNContrastLoss

rnc = RankNContrastLoss(temperature=2.0)
value = rnc(features, labels)                # features [N, D], labels [N] or [N, L]

mask removes samples from both the anchor set and every denominator; weights are per anchor.

Scope

These are building blocks. The library does not claim any particular downstream accuracy from them; validate on your own held-out data.

When to use which¶

Building block	Operates on	Best when	Not for
`BSplineDensityBasis`	Density output space (continuous variable)	You need a normalised, differentiable 1D density with exact evaluation, bin masses, CDF and moments from a softmax head	High-dimensional or multi-output densities (it is a 1D basis)
`DiscreteWasserstein1Loss`	Output space, binned masses	Penalising where misplaced mass lands, not just whether bins disagree (KL/CrossEntropy are blind to distance)	Continuous parameterisations — use `BSplineDensityBasis.absolute_deviation` instead
`RankNContrastLoss`	Representation space	Ordering features by target distance to improve downstream uncertainty/forecasting heads — pair it with a distribution head, not instead of one	Direct probabilistic prediction: RNC shapes features; it does not itself define a predictive distribution

→ Combine them: train features with RankNContrastLoss, then a BSplineDensityBasis head with an NLL plus a small DiscreteWasserstein1Loss term.

Generic cumulative-CDF objective¶

CumulativeCDFLoss and cumulative_cdf_loss retain the generic binned objective independently of an astronomy-specific recipe. For ordered centers \(c_j\), bin probabilities \(p_j=\operatorname{softmax}(a)_j\), and cumulative masses \(C_j=\sum_{k\le j}p_k\), the objective is

\[ L(y,a)=\bar h\sum_j\left(C_j-\mathbf 1[c_j\ge y]\right)^2, \qquad \bar h=(c_K-c_1)/(K-1). \]

The mean-width weighting is a fixed-grid training objective; it is not the exact piecewise-linear CRPS on nonuniform bins. Use the distribution metrics for that score.

import torch
from torchregress.losses import CumulativeCDFLoss

loss = CumulativeCDFLoss(torch.linspace(-2., 2., 9))
logits = torch.zeros((3, 9), requires_grad=True)
value = loss(logits, torch.tensor([-.2, .3, 1.1]))
value.backward()

For continuous scoring, BinnedDensityLoss.log_prob divides each bin mass by its width and returns negative infinity outside the finite support. MixtureDensityLoss.log_prob instead evaluates the normalized continuous mixture density. Binned classification cross-entropy alone has different units.