Wasserstein Poincare-deficit regularizer

Integrate the method at the loss level for a classifier or encoder. Begin with two endpoint minibatches \(x_{0,b}\sim\mu_0\) and \(x_{1,b}\sim\mu_1\) with a known coupling. For a cheap controlled experiment, choose Gaussian endpoints \(\mu_i=\mathcal{N}(m_i,\Sigma_i)\). Their exact Brenier map is \(T(x)=m_1+A(x-m_0)\), with \(A=\Sigma_0^{-1/2}(\Sigma_0^{1/2}\Sigma_1\Sigma_0^{1/2})^{1/2}\Sigma_0^{-1/2}\). Compute \(x_{1,b}=T(x_{0,b})\), sample \(t\sim\operatorname{Uniform}(0,1)\), and set \(x_{t,b}=(1-t)x_{0,b}+tx_{1,b}\). For Gaussian endpoints use the exact curvature lower bound \(\kappa_i=1/\lambda_{\max}(\Sigma_i)\); for other distributions use a conservative analytic or estimated lower bound. Define the scalar output as the class margin \(f_\theta(x)=z_y(x)-\max_{k\neq y}z_k(x)\). Compute its minibatch variance and per-example input gradient with autodiff, then optimize loss = cross_entropy + lambda * relu(var(f)-Bt**2*mean(sum(grad_x(f)**2))). Stop gradients through sampled transport points and \(B_t\). The theorem supplies the coefficient; minibatch variance, Jacobian energy, and transport quality are empirical. First compare against cross-entropy and ordinary input-gradient regularization on a two-dimensional Gaussian classification problem using a small MLP. Measure clean accuracy, accuracy on interpolated points, perturbation robustness, and the ratio \(\widehat{\operatorname{Var}}/(B_t^2\widehat{\mathbb{E}}\|\nabla f\|^2)\). A successful signal is a reduced ratio and improved transport-path robustness at equal FLOPs without clean-accuracy degradation. Then test the same loss on a frozen CIFAR-10 feature space with Gaussian endpoint clouds.