Jump-aware Wasserstein particle dynamics

Integrate the method at the particle-state update rather than inside the ordinary parameter-gradient operation. Maintain K particles \(x_s^k\in\mathbb R^d\), representing an optimizer ensemble, VAE latent particles, or K MoE routing prototypes at checkpoint s, and define \(\mu_s=K^{-1}\sum_k\delta_{x_s^k}\). Apply the unmodified model update to obtain provisional particles \(\tilde x_{s+1}^k\). Compute a minibatch Sinkhorn approximation to \(W_1(\mu_s,\tilde\mu_{s+1})\) using cost matrix \(C_{ij}=\|x_s^i-\tilde x_{s+1}^j\|_2\), uniform marginals, and entropic temperature \(\varepsilon\). Let \(P_{ij}\) be the resulting transport plan and estimate the velocity of particle i by \(v_s^i=\sum_jP_{ij}(\tilde x_{s+1}^j-x_s^i)/\sum_jP_{ij}\). Add the first term of \(\mathcal L_{\mathrm{jump}}\) to the task loss. When resampling, top-k router replacement, or particle birth/death occurs, mark the step as \(s\in\mathcal J\), match old and new particles with Hungarian matching or the Sinkhorn plan, and add the second term, where \(m_{s,k}\) is matched mass and \(\pi_s\) is the assignment. For a continuity diagnostic, use Gaussian test functions \(\varphi_r(t,x)=\exp(-\|x-c_r\|^2/(2\sigma^2))\), approximate \(\partial_t\varphi\) by finite differences, and compare the empirical weak residual with the matched velocity term. Sinkhorn cost, jump cost, and residual are computed from the paper's structure; \(\lambda,\beta,\varepsilon\), test centers \(c_r\), and bandwidth \(\sigma\) are empirical choices. First test a 2D Gaussian-mixture VAE or four-expert toy MoE with K=64 particles, comparing Adam plus unregularized resampling against the proposed loss at equal steps, parameters, and accounting for Sinkhorn FLOPs. Pre-register that total W1 variation decreases by at least 20% at fixed task loss, the weak residual after a replacement event is no worse than baseline and recovers within two steps, and held-out likelihood or routing balance improves without increasing K. Remove only the W1/jump term for the ablation. Falsify the transfer if variation does not decrease, residuals do not recover, or task quality drops by more than 1% at equal compute across three seeds.