Hermite-fiber MoE router initialization
Implementation & benchmark of arXiv:2608.12701 — Sharp proper estimation of fixed-component Gaussian location mixtures in polynomial time
Source paper: Sharp proper estimation of fixed-component Gaussian location mixtures in polynomial time arXiv:2608.12701 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace random or k-means initialization of a k-expert router with a moment-based range finder on a calibration batch of hidden states. Estimate a low-dimensional second-moment subspace, enlarge it using one-free-index third-Hermite contractions, and fit the router's expert centroids and weights only in this resulting subspace. The router can then operate on projected hidden states while retaining an optional small residual adapter.
Formulas
Mathematical statement
The paper models observations as X=U+Z, where U is drawn from a mixing distribution Gamma=sum_{j=1}^s w_j delta_{mu_j}, s<=k, w_j>=0, sum_j w_j=1, ||mu_j||<=R, and Z is standard Gaussian N(0,I_d). Its guarantee is E H(P_{hat Gamma^E},P_Gamma)<=C_{k,R} min{1,sqrt(d/n)}, where H^2(P,Q)=integral (sqrt(p)-sqrt(q))^2 is squared Hellinger distance and P_Gamma=Gamma*N(0,I_d). The adaptation uses the Gaussian-Hermite identity E[H_2(X)]=sum_j w_j mu_j mu_j^T, with H_2(x)=xx^T-I_d. Let M_2 be its empirical estimate and let V=[v_1,...,v_r] contain its leading eigenvectors. The one-free-index third-Hermite contractions are b_{ab}=E[H_3(X)[v_a,v_b,:]], where H_3(x)=x^{otimes 3}-x otimes I-I otimes x-I^{(3)} otimes x componentwise, equivalently H_3(x)[v_a,v_b,:]=(v_a^T x)(v_b^T x)x-(v_a^T x)v_b-(v_b^T x)v_a-(v_a^T v_b)x. For a mixture, b_{ab}=sum_j w_j(v_a^T mu_j)(v_b^T mu_j)mu_j, so every contraction lies in the component-mean span and can expose directions missed by V. Define W=span(V,{b_{ab}}), project h to z=W^T h, and fit k component means nu_j and weights pi_j in this space. The router logits are ell_j(h)=log pi_j-||W^T h-nu_j||^2/(2 sigma^2).
Implementation notes
1. Integration point: apply this during initialization of a top-1 or top-2 MoE router, or before fitting a Gaussian-mixture density head, using a calibration buffer of hidden states h in R^d from a pretrained encoder. Standardize the buffer by subtracting its mean and whitening with a regularized covariance estimate so that the approximate unit-covariance assumption is reasonable. The projected router can replace the full d-dimensional distance computation; retain a residual linear correction only if accuracy drops. 2. Pseudocode: collect n hidden states x_i; compute M2=(1/n) sum_i (x_i x_i^T-I); obtain the top r=min(k,d) eigenvectors v_a by randomized Lanczos; for every a,b compute b_hat_ab=(1/n) sum_i[((v_a^T x_i)(v_b^T x_i))x_i-(v_a^T x_i)v_b-(v_b^T x_i)v_a-(v_a^T v_b)x_i]; form W by QR orthogonalization of all v_a and b_hat_ab; set z_i=W^T x_i; fit k weighted centroids (nu_j,pi_j) with EM or k-means-plus-weighted refinement; initialize router logits as log pi_j-||W^T h-nu_j||^2/(2 sigma^2). 3. Computed from the mathematics: the H2/H3 estimators, eigenspace, contraction range, and projected logits. Estimated empirically: r, sigma, whitening regularization, and whether the residual adapter is needed. 4. First experiment: use a 4- or 8-expert Transformer MoE on a small language-model or CIFAR-100 encoder benchmark, comparing random router initialization, k-means initialization, and this method at identical training FLOPs. Measure early loss reduction, router load balance, expert utilization, final validation accuracy or perplexity, and projection dimension dim(W). A useful success signal is 20-30% lower loss after the first 1-5% of training with no worse final loss, or a 2x reduction in router projection FLOPs at matched validation quality. Use a held-out calibration split to detect overfitting of moment estimates.
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.