# Convex-Order Distributional Distillation

- ID: 2836
- Canonical URL: https://synthcore.org/idea/2836/convex-order-distributional-distillation
- API JSON: https://synthcore.org/api/idea/2836.json
- API Markdown: https://synthcore.org/api/idea/2836.md
- Verification status: unverified
- Source: [arXiv:2608.29565](https://arxiv.org/abs/2608.29565)
- Category: regularization
- Solves: accuracy, stability, generalization
- ML areas: loss, distillation, regularization
- Math tags: optimal-transport, convex-analysis, probability
- Ratings: usefulness 5/10; difficulty 6/10; novelty 7/10

## Idea description

Represent each neural prediction as a finite probability distribution and project it, under an optimal-transport cost, onto the set of distributions dominated by a teacher or target distribution in convex order. This enforces a global spread and risk relationship across all convex observables rather than adding separate variance, tail, and calibration penalties. Use a periodically refreshed projection during training and test whether it improves uncertainty calibration and robustness at equal predictive accuracy.

## Mathematical statement

The paper defines the backward Wasserstein projection as $T_c(\mu,P_{\leqslant\nu})=\inf_{\eta\in P_{\leqslant\nu}}T_c(\mu,\eta)$, where $\mu$ is the unconstrained model distribution, $\nu$ is a reference distribution, $P_{\leqslant\nu}=\{\eta:\eta\leq_{cx}\nu\}$ is the cone of probability measures dominated by $\nu$ in convex order, and $T_c(\mu,\eta)$ is the optimal-transport cost. Convex order means $\eta\leq_{cx}\nu$ if and only if $\mathbb E_\eta[\phi(X)]\leq\mathbb E_\nu[\phi(Y)]$ for every convex function $\phi$ for which the expectations exist. For one-dimensional distributions with finite first moments, convex order additionally requires equal means and can be checked using call functions: $\mathbb E_\eta[(X-t)_+]\leq\mathbb E_\nu[(Y-t)_+]$ for every threshold $t$. For discrete predictions on a common sorted grid, checking these inequalities at all grid breakpoints yields a finite constraint system. The neural target is $\hat\eta_\theta=\arg\min_{\eta\leq_{cx}\nu}\operatorname{OT}_h(\mu_\theta,\eta)$, with $h(u)=u^2$ or $h(u)=|u|$.

## Key formulas

- $$T_{c}\mathopen{}\mathclose{{\left(\mu,P_{\leqslant\nu}}}\right)=\inf_{\eta\in P_{\leqslant\nu}}T_{c}\mathopen{}\mathclose{{\left(\mu,\eta}}\right),\qquad T_{c}\mathopen{}\mathclose{{\left(P_{\mu\leqslant},\nu}}\right)=\inf_{\xi\in P_{\mu\leqslant}}T_{c}\mathopen{}\mathclose{{\left(\xi,\nu}}\right).$$
- $$\eta\leq_{cx}\nu\quad\Longleftrightarrow\quad \mathbb E_\eta[\phi(X)]\leq\mathbb E_\nu[\phi(Y)]\ \text{for every convex }\phi,$$
- $$\mathbb E_\eta[X]=\mathbb E_\nu[Y],\qquad \mathbb E_\eta[(X-t)_+]\leq\mathbb E_\nu[(Y-t)_+]\quad\forall t\in\mathbb R,$$
- $$\hat\eta_\theta=\arg\min_{\eta\leq_{cx}\nu}\operatorname{OT}_{h}(\mu_\theta,\eta),\qquad h(u)=u^2\ \text{or}\ |u|.$$

## Implementation notes

Integrate this at the distributional prediction head, not in the shared feature extractor. Let the student emit fixed ordered support locations $x_1<\cdots<x_m$ and logits $a_i$, with probabilities $p_i=\operatorname{softmax}(a)_i$ and student distribution $\mu_\theta=\sum_i p_i\delta_{x_i}$. Obtain a teacher or target distribution $\nu=\sum_j r_j\delta_{y_j}$ on the same one-dimensional grid. Every $K$ training steps, solve for projected masses $q$ by minimizing discrete OT cost over a coupling $\pi$: minimize $\sum_{i,j}\pi_{ij}h(x_i-y_j)$ subject to nonnegative coupling, its source marginal being $p$, its target marginal being $q$, and the convex-order constraints $\sum_jq_jy_j=\sum_jr_jy_j$ and $\sum_jq_j(y_j-t_k)_+\leq\sum_jr_j(y_j-t_k)_+$ for every sorted grid breakpoint $t_k$. Use a linear program for exact projection or entropic OT plus a differentiable constrained solver. Feed $q$ into the task loss, or add $\lambda\operatorname{KL}(q\|p)$ as a distillation loss. For a first implementation, stop gradients through the projection and use a straight-through estimator; later compare implicit differentiation and a soft penalty. The mathematical quantities are the OT objective and convex-order constraints; projection frequency $K$, solver tolerance, entropic temperature, and loss weight $\lambda$ are empirical choices. The first cheap experiment should use a small MLP on UCI heteroscedastic regression with a 32-bin output grid, comparing NLL-only, variance matching, and convex-order distillation. Measure NLL, calibration error, CRPS, and tail-risk error under covariate shift. Success means equal or better NLL with lower calibration and tail-risk error, while recording projection time and solver failures.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.