Heat-Wasserstein curvature preconditioner

Integrate this into the optimizer for a graph neural network or an embedding model, not into the raw high-dimensional parameter metric. At each training refresh, take a minibatch of feature vectors $z_i=f_\theta(x_i)\in\mathbb R^d$ with $d\leq64$, construct a $k$-nearest-neighbor graph, and define edge weights $w_{ij}=\exp(-\|z_i-z_j\|^2/\sigma^2)$. Let $L=D-W$ be its symmetric normalized graph Laplacian and approximate heat diffusion by $H_t=\exp(-tL)$, computed with 5--10 Lanczos steps or a truncated eigendecomposition. Treat row $H_t[i,:]$ as the discrete analogue of $\mu_{t,x_i}$. For nearby points $i,j$, estimate the induced metric tensor $\widetilde G_t$ by finite differences of feature coordinates and a Sinkhorn Wasserstein-2 distance between the corresponding heat rows; alternatively use the local covariance of transported displacements as a cheaper proxy. Use $G$ from the $t=0$ local feature metric and estimate the second-order coefficient with the central finite difference $C=(\widetilde G_{2t}-2\widetilde G_t+G)/(2t^2)$. The first-order curvature proxy is $R=(G-\widetilde G_t)/(2t)$; do not attempt to separately identify $\mathcal Q_g$ unless curvature tensors are available. During optimization, average local matrices over the minibatch, symmetrize, clip eigenvalues below $10^{-4}$, and apply $(\widetilde G_t+\lambda I)^{-1}$ to the gradient in feature coordinates or to the final projection-layer gradient. Pseudocode: build graph; compute $H_t,H_{2t}$; estimate $G,\widetilde G_t,\widetilde G_{2t}$; set $C$ by the displayed formula; set $\widetilde G_t^{(2)}=G-2tR+t^2C$; precondition the gradient and update parameters. The mathematics supplies the expansion and required scaling; graph bandwidth $\sigma$, diffusion time $t$, Sinkhorn iterations, and damping $\lambda$ are empirical. First test on MNIST or CIFAR-10 with a 2-layer GNN or MLP whose penultimate embedding has dimension 16, comparing AdamW, ordinary local-metric preconditioning, and this second-order metric at equal wall-clock budget. Success means fewer steps to the same accuracy, reduced sensitivity to learning rate, or better kNN neighborhood preservation under noisy labels; reject it if Sinkhorn and graph construction erase any optimization gain.