Langevin Escape for Codebook Lloyd Traps

Apply this optimizer only to codebook or prototype parameters, not to all neural-network weights. At each outer step, obtain embeddings $z_b$, solve the approximate balanced Laguerre assignment using the dual-weight procedure above, compute the minibatch energy $E_k=B^{-1}\sum_b\|z_b-x_{a_b}\|^2/2$ with capacity correction, and form prototype gradients $g_{i,k}=p_i(x_i-b_i)$. Maintain $E_{\mathrm{best},k}=\min(E_{\mathrm{best},k-1},E_k)$ and calculate $q_k=(E_k-E_{\mathrm{best},k})/(|E_{\mathrm{best},k}|+\epsilon)$. Set $T_k=T_{\min}+(T_{\max}-T_{\min})\operatorname{sigmoid}(\alpha q_k)$. Draw independent $\xi_{i,k}\sim\mathcal N(0,I_d)$ and update each prototype by $x_{i,k+1}=x_{i,k}-h g_{i,k}+\sqrt{2hT_k}\xi_{i,k}$. If embeddings lie in a bounded domain, project prototypes into the same bounding box after the update. During joint encoder training, freeze encoder weights during exploration or perform one Langevin prototype step every 20 Adam steps. Always checkpoint the prototype set with the lowest validation or smoothed minibatch energy, because the paper explicitly distinguishes raw-iterate behavior from best-so-far convergence. The paper supplies the Euler structure and nonzero-temperature principle; $h$, $T_{\min}$, $T_{\max}$, and $\alpha$ are empirical. Test on a 32-code VQ-VAE on CIFAR-10 against Adam/Lloyd updates with equal prototype-update FLOPs. Compare best reconstruction error, final validation error, code usage entropy, dead-code frequency, and run-to-run variance. The desired signal is lower best-so-far quantization error and fewer collapsed-code runs, not convergence of the noisy iterates.