Unverified 2026

Early-EM KL Router

Implementation & benchmark of arXiv:2609.01688 — On the discretization of the object space in inverse problems with application to cryo-electron microscopy

Usefulness6/10

Difficulty4/10

Novelty4/10

Regularization Stability Generalization Accuracy Moe routing Optimizer Regularization Optimization Probability Statistics Information theory

Source paper: On the discretization of the object space in inverse problems with application to cryo-electron microscopy arXiv:2609.01688 ⓘ · analyzed Sep 3, 2026

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

Idea description

Replace direct optimization of noisy simplex-valued mixture or MoE gate weights with a fixed number of EM responsibility updates. The finite iteration count acts as an implicit reverse-KL regularizer toward the uniform gate distribution, preserving low-frequency experts without selecting an arbitrary entropy coefficient.

Formulas

$$F(\alpha):=\frac{1}{n}\sum_{i=1}^{n}\log\left(\sum_{m=1}^{M}p_{\sigma}(y_i\mid x_m)\alpha_m\right),\qquad \alpha\in\Delta_{M-1}:=\{\alpha_m\geq 0,\ \sum_m\alpha_m=1\}.$$

$$r_{im}^{(t)}=\frac{\alpha_m^{(t)}p_{\sigma}(y_i\mid x_m)}{\sum_{j=1}^{M}\alpha_j^{(t)}p_{\sigma}(y_i\mid x_j)},\qquad \alpha_m^{(t+1)}=\frac{1}{n}\sum_{i=1}^{n}r_{im}^{(t)}.$$

$$\widetilde{\alpha}_{\sigma,T}=\arg\max_{\alpha\in\Delta_{M-1}}\left[F(\alpha)-\frac{1}{T+1}\operatorname{KL}(u\|\alpha)\right],\qquad \operatorname{KL}(u\|\alpha)=\sum_{m=1}^{M}u_m\log\frac{u_m}{\alpha_m},\quad u_m=\frac{1}{M}.$$

$$a_{\sigma}:=\frac{R_x}{\sigma}+\frac{A_xR_x}{\sigma^{2}},\qquad \frac{\|\alpha^{(T_{\sigma}+1)}-\widetilde{\alpha}_{\sigma,T_{\sigma}}\|_2}{\|\alpha^{(T_{\sigma}+1)}-u\|_2}=O_p\bigl((T_{\sigma}+1)a_{\sigma}\bigr)=o_p(1)\ \text{when }(T_{\sigma}+1)a_{\sigma}\to0.$$

Mathematical statement

The paper models observations with candidate states x_m and likelihoods p_sigma(y_i|x_m), using simplex weights alpha in Delta_{M-1}. Its empirical log-likelihood is F(alpha) = (1/n) sum_i log(sum_m p_sigma(y_i|x_m) alpha_m). Let u=(1/M,...,1/M) be the uniform simplex point and KL(u||alpha)=sum_m u_m log(u_m/alpha_m). The EM responsibilities are r_im^(t) = alpha_m^(t) p_sigma(y_i|x_m) / sum_j alpha_j^(t) p_sigma(y_i|x_j), and the M-step is alpha_m^(t+1)=(1/n) sum_i r_im^(t). Proposition 6 states that for Y_i^(sigma)=mu_i+sigma Z_i, with Z_i iid N(0,I_d), A_x=max_{i,j}||x_j-mu_i||, a_sigma=R_x/sigma + A_x R_x/sigma^2, and (T_sigma+1)a_sigma -> 0, the relative discrepancy between alpha^(T_sigma+1) and the maximizer of F(alpha)-KL(u||alpha)/(T_sigma+1) is O_p((T_sigma+1)a_sigma). The transferable property is that early stopping creates an explicit reverse-KL barrier with approximate strength 1/(T+1), rather than being merely an informal optimization heuristic.

Implementation notes

(1) Integrate this into a mixture-of-experts router or a latent mixture head whose per-example component evidence is available. For a batch of tokens x_i and experts m, let s_im be the router logit or frozen-expert log-likelihood. Define l_im=exp(s_im/tau), where tau is the routing temperature, and maintain a positive simplex vector alpha, either globally per layer or separately per task or domain. Do not optimize alpha with Adam in the first experiment; update it with the EM recursion.

(2) Pseudocode: initialize alpha_m=1/M. For t=0,...,T-1, compute log_r_im=log(alpha_m)+s_im/tau, normalize over experts with logsumexp, and set r_im=softmax_m(log_r_i). Then set alpha_m=(1/B) sum_i r_im. For streaming training, use alpha <- beta alpha +(1-beta) batch_mean(r), followed by normalization. Use alpha as the gate prior, or multiply token-level router probabilities by alpha and renormalize. T is the regularization knob: small T gives a strong approximate KL(u||alpha) barrier, while large T permits specialization. If the router is trainable, alternate K gradient steps on router logits with one alpha update, keeping alpha detached during the router gradient step.

(3) Computed directly from the paper's mathematics are the responsibilities, simplex-preserving M-steps, and implicit regularization coefficient 1/(T+1). The quantity a_sigma can be computed only when a Gaussian-noise interpretation is appropriate; estimate effective noise and evidence scales empirically otherwise. The theorem should be treated as motivation for the schedule, not as an unverified guarantee for neural logits.

(4) First cheap experiment: train a 4- or 8-expert sparsely gated MLP on CIFAR-100 or a small language-model classification task. Compare standard softmax routing, standard explicit uniform entropy or load-balancing loss, and Early-EM KL Router with T in {1,2,4,8,16}. Keep FLOPs, expert capacity, optimizer, and total router updates fixed. Measure validation accuracy, held-out mixture likelihood, routing entropy, dead-expert fraction, gradient variance, and loss at equal optimizer steps. Success means fewer dead experts and equal-or-better validation loss than standard routing, with less sensitivity to an entropy coefficient; in a noisy or frozen-expert setup, the strongest signal is improved held-out likelihood at the same computation.

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.