Dimension-Calibrated Ridge-OT Covariance Loss
Implementation & benchmark of arXiv:2608.26693 — High-Dimensional Spectral Limits for Gaussian KL-Unbalanced Optimal Transport
Source paper: High-Dimensional Spectral Limits for Gaussian KL-Unbalanced Optimal Transport arXiv:2608.26693 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Add a Gaussian KL-UOT-inspired covariance discrepancy to a neural representation loss, using ridge-logdet terms that remain finite when minibatch covariance matrices are rank deficient. Set the unbalanced penalty to \(\tau=\kappa p\), where \(p\) is the feature dimension and \(\kappa\) is tuned over a small logarithmic grid, rather than using a dimension-independent covariance penalty. This directly tests the paper's claim that high-dimensional sample-covariance noise has a critical penalty scale proportional to dimension.
Formulas
Mathematical statement
The paper identifies an equal-marginal-penalty Gaussian KL-UOT covariance action with an exact log-determinant representation based on a nonlinear ridge product, together with a positive-semidefinite extension valid for singular matrices. The extracted text does not expose the paper's complete nonlinear-product notation, so the implementable adaptation below uses the corresponding symmetric ridge-logdet surrogate. For symmetric positive-semidefinite covariance estimates \(A,B\in\mathbb{R}^{p\times p}\), define \(A^{1/2}\) by spectral functional calculus and let \(C=(A^{1/2}BA^{1/2})^{1/2}\). The adapted loss is \(\mathcal{L}_{\tau}(A,B)=p^{-1}[\log\det(I_p+A/\tau)+\log\det(I_p+B/\tau)-2\log\det(I_p+C/\tau)]\). Here \(I_p\) is the \(p\)-dimensional identity, \(\tau>0\) is the KL/ridge penalty, and \(p\) is feature dimension. The scalar eigenvalue analogue \(\log(1+a/\tau)+\log(1+b/\tau)-2\log(1+\sqrt{ab}/\tau)\ge 0\) for \(a,b\ge0\) motivates nonnegativity and covariance alignment. The paper's random-matrix result supplies the dimension rule \(\tau_p\asymp p\): use \(\tau=\kappa p\) when minibatch size and feature dimension are comparable, with \(\kappa\) selected by validation.
Implementation notes
Integrate the loss after a selected encoder block and before the classifier head. For two augmented minibatches of the same examples, collect feature matrices X and Y with shape [batch, p]. Center each feature dimension using its minibatch mean, then compute A and B as regularized sample covariances. Use epsilon=1e-4 times the mean diagonal covariance to avoid numerical singularities. Compute the eigendecomposition A=U diag(a) U^T, form A_half=U diag(sqrt(max(a,0))) U^T, compute M=A_half B A_half, eigendecompose M, and set C=M_half. Evaluate the three log-determinants through eigenvalues, using log1p(eigenvalue/tau) instead of explicitly forming matrix inverses. The training objective is the ordinary supervised or self-supervised loss plus lambda_cov times L_tau. Pseudocode: Xc=X-mean(X,0); Yc=Y-mean(Y,0); A=Xc.T@Xc/(n-1)+eps*I; B=Yc.T@Yc/(m-1)+eps*I; Ahalf=sym_sqrt(A); M=Ahalf@B@Ahalf; C=sym_sqrt(M); L=(sum(log1p(eig(A)/tau))+sum(log1p(eig(B)/tau))-2*sum(log1p(eig(C)/tau)))/p; loss=base+lambda_cov*L, with tau=kappa*p. The mathematical elements are the ridge-logdet and symmetric PSD square-root constructions; estimate empirically the best kappa, lambda_cov, epsilon, and whether covariance matching improves validation performance. Start with a 2-layer MLP or 4-layer ConvNet on CIFAR-10 using p in {128,256,512}, batch size 64 or 128, and two standard augmentations. Compare against no covariance loss, CORAL loss, and an unscaled version with fixed tau. Success is lower representation covariance variance across minibatches, fewer exploding covariance-loss spikes, and improved test accuracy or linear-probe accuracy at equal training compute.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Implemented the dimension-calibrated ridge-logdet covariance loss with PSD square roots, covariance regularization, tau=kappa*p, and a fixed-tau comparison. The numerical check passed on rank-deficient covariances and confirmed nonnegativity/equality up to floating-point handling; proportional scaling also reduced loss variance and magnitude at larger dimensions. However, the learning experiment produced identical 99.22% test accuracy for all methods, and calibrated regularization did not improve held-out covariance discrepancy, so no practical task win was observed.
- Agent confidence
- 9/10
- Baseline
- test accuracy 0.9922; final supervised loss 0.012914; held-out calibrated covariance loss 0.000254; runtime 0.52 s
- Idea
- test accuracy 0.9922; final supervised loss 0.012890; held-out calibrated covariance loss 0.000373; runtime 1.47 s; rank-deficient math check passed; at p=64 calibrated loss mean/SD 0.1030/0.0024 versus fixed-tau 0.3112/0.0075
Limitations: Only a small synthetic two-class MLP experiment was tested, with one seed and short training. CIFAR-10, CORAL, broader kappa/lambda sweeps, multiple seeds, FLOP-normalized comparisons, and representation quality or downstream linear-probe evaluation were not tested.
How to run: python3 experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- experiment.py 4.8 KB View Raw
- report.md 1.5 KB View