Quantile-space Huber-Wasserstein loss
Implementation & benchmark of arXiv:2608.13131 — Huber-Wasserstein barycenters for robust distribution-valued data
Source paper: Huber-Wasserstein barycenters for robust distribution-valued data arXiv:2608.13131 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
For models that predict probability distributions, replace the usual Wasserstein-2 loss or a Huber penalty on the final Wasserstein distance with a Huber penalty on quantile-by-quantile prediction errors. This suppresses gradients from localized outliers while retaining quadratic gradients on the majority of the distribution, which is useful for uncertainty prediction, histogram prediction, and distributional distillation.
Formulas
Mathematical statement
Let Q_theta(x,u) be the predicted quantile function for input x and quantile level u, and let Q_y(u) be the target quantile function. The displacement-wise robust loss is L_delta(theta;x,y)=integral_0^1 h_delta(Q_theta(x,u)-Q_y(u))du, approximated by a uniform sum over m quantile levels. Its derivative with respect to the predicted quantile is psi_delta(e_j)=clip(e_j,-delta,delta), where e_j=Q_theta(x,u_j)-Q_y(u_j). Therefore each quantile has an independently bounded influence, unlike distance-wise Huberization, which computes H_delta(W_2(Q_theta,Q_y)) and multiplies the gradient for every quantile by one shared scalar. To ensure that predictions are valid quantile functions, emit an offset b and unconstrained increments d_k and define Q_theta(x,u_j)=b(x)+sum_{k=1}^j softplus(d_k(x)). This guarantees monotonicity without a sorting operation.
Implementation notes
Use this as the output loss for a distributional predictor, uncertainty head, histogram decoder, or teacher-student distillation model. The network should emit b and m unconstrained increment logits d_1,...,d_m; construct monotonically increasing predicted quantiles with Q_j=b+sum_{k<=j}softplus(d_k), avoiding invalid CDF inversions. Convert every training target distribution to the same m quantile levels once during preprocessing. Compute e_j=Q_pred,j-Q_target,j and apply the piecewise Huber loss independently to every j. Backpropagation then supplies the clipped gradient clamp(e_j,-delta,delta)/m, rather than a single scalar robustification of the whole Wasserstein distance. Start with delta=1.345 times a running robust scale s, where s=1.4826*median(abs(e_j-median_k(e_k))) over the current minibatch; detach s from the gradient and update it using an exponential moving average. Compare fixed delta values {0.25,0.5,1,2} and adaptive scaling. The mathematical operation is exact numerical quadrature of the displacement-wise Huber-Wasserstein objective; target quantile estimation and scale tracking are empirical approximations. First run a synthetic contaminated-mixture regression experiment and a small heteroscedastic regression network. Compare MSE on quantiles, ordinary W2 loss, distance-wise Huber-W2, and this loss under 10-40% target outliers. Record clean Wasserstein error, contaminated-test Wasserstein error, calibration or coverage, gradient-norm distributions, and convergence at equal optimization steps. Success is lower contaminated-test error and better calibration while retaining near-MSE clean performance when contamination is absent.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.