Polarized Gaussian bottleneck

Integrate this at the output of a stochastic bottleneck, VAE latent block, or MoE router before the next linear layer. Let \(Z\in\mathbb{R}^{d\times B}\) contain one centered latent vector per batch item, \(\bar z\) be the batch mean, and \(C=(B-1)^{-1}(Z-\bar z\mathbf{1}^{\mathsf T})(Z-\bar z\mathbf{1}^{\mathsf T})^{\mathsf T}\). To make the Gaussian scale meaningful, first normalize each latent coordinate or apply a learned diagonal scale constrained to \([0,1]\), then compute the symmetric eigendecomposition \(C=U\operatorname{diag}(\lambda_1,\ldots,\lambda_d)U^{\mathsf T}\). Add \(\beta d^{-1}\sum_i\lambda_i(1-\lambda_i)\) to the task loss, optionally with a variance-budget term \(\alpha(\sum_i\lambda_i-r)^2\) targeting effective rank \(r\). Pseudocode: z = bottleneck(x); zc = z - mean(z, batch); C = zc.T @ zc / (B-1); lam = eigvalsh(C); loss = task_loss + beta*mean(lam*(1-lam)) + alpha*(lam.sum()-r)**2; loss.backward(). For large \(d\), estimate only the extreme eigenvalues with 5-10 Lanczos iterations, or use a Hutchinson estimator for the trace of \(C-C^2\). The paper supplies the Gaussian-convex spectral target and endpoint interpretation; covariance estimation, unit-variance normalization, and the rank budget are engineering choices. First test on MNIST or CIFAR-10 with a 64-to-16 autoencoder bottleneck, comparing no regularizer, ordinary \(\|C-I\|_F^2\), and the polarization loss at equal parameter count. Measure reconstruction or classification accuracy, eigenvalue histograms, the number of eigenvalues below 0.05 or above 0.95, and post-training pruning of low-variance coordinates. Success means a bimodal spectrum with no accuracy loss and at least 2x removable bottleneck coordinates, or improved accuracy at the same effective rank. A second test can apply the loss to router assignment embeddings and measure expert load balance and routing entropy.