Balanced Laguerre Codebook Layer
Implementation & benchmark of arXiv:2609.03405 — Semi-discrete quadratic Wasserstein energy and state-dependent Langevin exploration
Source paper: Semi-discrete quadratic Wasserstein energy and state-dependent Langevin exploration arXiv:2609.03405 ⓘ · analyzed Sep 4, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace ordinary nearest-neighbor or softmax codebook assignment with a capacity-constrained Laguerre assignment whose cells have prescribed masses. Optimize the codebook using the semi-discrete quadratic Wasserstein energy, whose gradient moves every site toward the barycenter of its balanced cell. This directly prevents prototype collapse and gives explicit control over expert or codeword utilization.
Formulas
Mathematical statement
The paper studies $E(X)=\frac{1}{2}W_2^2(\nu,\mu_X)$, where $\nu(dy)=\rho(y)\mathbf{1}_{\Omega}(y)dy$ is a target probability measure on compact convex $\Omega\subset\mathbb{R}^d$, $\rho$ is continuous and bounded by $0<\underline\rho\leq\rho\leq\overline\rho$, $X=(x_1,\ldots,x_N)$ are codebook sites, and $\mu_X=\sum_{i=1}^Np_i\delta_{x_i}$ with prescribed masses $p_i>0$ satisfying $\sum_i p_i=1$. For dual weights $w_i$, the Laguerre cell is $L_i(X,w)=\{y\in\Omega:\|y-x_i\|^2-w_i\leq\|y-x_j\|^2-w_j\ \forall j\}$, with weights chosen so that $\int_{L_i}\rho(y)dy=p_i$. Its barycenter is $b_i=\frac{1}{p_i}\int_{L_i}y\rho(y)dy$. On collision-free configurations, the gradient is $\nabla_{x_i}E(X)=p_i(x_i-b_i)$. Hence a descent step is a mass-weighted Lloyd update. Local Lipschitz continuity at collisions and global semiconcavity provide stability diagnostics, while collision-free global minimizers motivate repulsion or validity checks for degenerate assignments.
Implementation notes
Integrate this as a prototype or codebook module after an embedding tensor and before a quantizer, router, or decoder. Let a minibatch produce vectors $z_1,\ldots,z_B\in\mathbb{R}^d$ and let $x_1,\ldots,x_N$ be learnable prototypes. Choose target masses $p_i$; uniform masses use $p_i=1/N$, while desired expert capacities use nonuniform values. Maintain dual weights $w$. For several inner iterations, assign each sample using $a_b=\arg\min_i(\|z_b-x_i\|^2-w_i)$, estimate empirical cell masses $\widehat m_i=B^{-1}\sum_b\mathbf{1}[a_b=i]$, and update weights by $w_i\leftarrow w_i+\gamma(\widehat m_i-p_i)$, optionally with damping or a zero-mean projection. Reassign after the weight iterations, compute $\widehat b_i=\sum_{b:a_b=i}z_b/\#\{b:a_b=i\}$, and update prototypes with $x_i\leftarrow x_i-\eta p_i(x_i-\widehat b_i)$. Use straight-through assignments for the encoder path, and add $\lambda\sum_i p_i\|x_i-\widehat b_i\|^2/2$ to the training objective. The exact paper mathematics supplies the transport objective and barycentric gradient; empirical cell masses, finite-batch barycenters, and dual-weight iterations are approximations. First test a 32- or 64-code VQ-VAE on CIFAR-10 against standard VQ-VAE and k-means initialization at equal parameters and FLOPs. Record reconstruction loss, codebook perplexity, maximum utilization imbalance, dead-code count, and collision rate. Success means substantially lower utilization variance and no reconstruction degradation; a stronger result is improved validation reconstruction at fixed codebook size.
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.