# Balanced Laguerre Codebook Layer

- ID: 3106
- Canonical URL: https://synthcore.org/idea/3106/balanced-laguerre-codebook-layer
- API JSON: https://synthcore.org/api/idea/3106.json
- API Markdown: https://synthcore.org/api/idea/3106.md
- Verification status: unverified
- Source: [arXiv:2609.03405](https://arxiv.org/abs/2609.03405)
- Category: architecture
- Solves: accuracy, stability, generalization
- ML areas: moe-routing, quantization, embedding, regularization
- Math tags: optimal-transport, geometry, convex-analysis
- Ratings: usefulness 6/10; difficulty 6/10; novelty 6/10

## 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.

## 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.

## Key formulas

- $$E(X)=\frac{1}{2}W_2^2\left(\nu,\sum_{i=1}^{N}p_i\delta_{x_i}\right),\qquad \nu(dy)=\rho(y)\mathbf{1}_{\Omega}(y)dy.$$
- $$L_i(X,w)=\left\{y\in\Omega:\|y-x_i\|^2-w_i\leq\|y-x_j\|^2-w_j\ \text{for all }j\right\},\qquad \int_{L_i(X,w)}\rho(y)dy=p_i.$$
- $$b_i(X)=\frac{1}{p_i}\int_{L_i(X,w)}y\rho(y)dy,\qquad \nabla_{x_i}E(X)=p_i\bigl(x_i-b_i(X)\bigr).$$
- $$x_i^{+}=x_i-\eta p_i(x_i-b_i)=(1-\eta p_i)x_i+\eta p_i b_i.$$

## 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.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.