Collision-free energy-kernel prototype flow
Implementation & benchmark of arXiv:2608.01182 — Wasserstein gradient flows of Maximum Mean Discrepancy with energy kernels
Source paper: Wasserstein gradient flows of Maximum Mean Discrepancy with energy kernels arXiv:2608.01182 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace ordinary codebook or prototype updates with the paper's diagonal-free particle dynamics in feature space. Each prototype is attracted toward the current minibatch feature distribution, while pairwise negative-distance interactions spread prototypes apart and suppress codebook collapse. The exponent q in (0,2) controls the interaction profile, with q less than 1 producing a stronger near-collision force but requiring numerical regularization.
Formulas
Mathematical statement
The paper uses the energy kernel K(z)=-|z|^q, where z is a vector in R^d and 0<q<2. For z not equal to zero, the gradient is grad K(z)=-q|z|^(q-2)z. The continuum MMD flow is partial_t rho_t=-div(rho_t v_t), with v_t=-grad K*(rho_t-mu), where rho_t is the particle distribution and mu is the target distribution. For N particles, the diagonal-free dynamics are dot{x}_i=-(1/N) sum_{j not equal to i} grad K(x_i-x_j)+grad(K*mu)(x_i). Substitution gives dot{x}_i=(q/N) sum_{j not equal to i} |x_i-x_j|^(q-2)(x_i-x_j)-q E_{y sampled from mu}[|x_i-y|^(q-2)(x_i-y)]. The first term spreads particles and the second matches their distribution to the target. Self-interaction must be excluded because the gradient is singular at zero for q<1. In a minibatch implementation, mu is approximated by the empirical distribution of encoder features.
Implementation notes
Integrate the flow into a prototype bank, VQ-style codebook, or latent-particle module rather than directly into network weights. Let an encoder f_theta produce a minibatch Y={y_b=f_theta(s_b)} in R^d, and maintain N trainable prototype vectors X={x_i}. At each training step, compute the vector field in the final formula. The first sum is over all j not equal to i, explicitly excluding self-interaction; the second sum treats the current feature minibatch as the empirical target mu. Update prototypes with Euler steps, either once per network step or with K=2 to 10 inner prototype steps. Use epsilon between 10^-3 and 10^-2 times the median pairwise feature distance to prevent exploding gradients for q<1, and optionally clip each prototype update to a fixed norm. Stop gradients through y_b for a codebook-only experiment; in a joint experiment, permit encoder gradients through the assignment or reconstruction loss but not through the prototype-flow update. Normalize features before applying the flow. The paper-derived components are the negative-distance kernel, the diagonal-free interaction, and the repulsive-attractive vector field. Estimate the target expectation with the minibatch and do not assume a theoretical global convergence rate. First test on a CIFAR-10 or MNIST encoder with N=64 prototypes. Compare vanilla EMA or k-means codebook updates, standard MMD-gradient updates, and this flow at matched update FLOPs. Measure dead-prototype fraction, minimum pairwise distance, feature reconstruction error, downstream linear-probe accuracy, and MMD to held-out features. Success means lower collapse and better feature coverage at equal prototype count without unstable norms or slower encoder loss descent.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.