{"id":"5fccd68b-ab80-4c7f-bd12-2cb33b6ca2aa","arxiv_id":"2505.00792","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"Similarity-weighted and attention-weighted routing scores reduce expert-switching instability in sparse MoE and yield small but consistent gains on Wikitext-103, ImageNet-1K, and sentiment/intent fine-tuning.","lead":"This paper proposes two cheap router modifications that make expert selection in sparse Mixture of Experts models less erratic, by blending each token's routing score with scores from similar tokens or from the tokens it attends to. Tested on language modeling, image classification, and fine-tuning, the changes reduce late-training expert switching and give small accuracy or perplexity gains.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Proposition 1's entropy-reduction guarantee is proven only for J_i-restricted sums in the tau->0/sigma->0 limit, not for the tau=sigma=1 all-token implementation.","rationale":"The reader's weakest assumption identifies exactly the gap between the theoretical guarantee in Prop. 1 and the implemented algorithms: the proof uses the restricted set J_i and the limits tau->0/sigma->0, while Defs. 2 and 5 sum over all tokens at tau=sigma=1. My stress-test agrees and adds that the Attention-Aware implementation also approximates the posterior by head selection h* and the point estimate U=MHA(X), creating an additional mismatch. This concern is load-bearing because the paper's abstract and Sec. 4 claim a theoretical proof of entropy reduction, which is a key part of the central claim. However, the empirical evidence (consistent perplexity gains, lower fluctuation, robustness improvements across several tasks) is plausible and could support a more modest claim. The appropriate resolution is to qualify the theoretical statement or extend the proof; the conditional verdict already captures this need. Therefore no change to the reader's verdict is required.","tokens_in":23834,"tokens_out":6343,"duration_ms":60422,"concrete_test":"On the trained Wikitext-103 models from Table 1, take a fixed validation batch and compute, for every token, H(pi) for the actual implemented pi = sum_j S[i,j] r(uj) (Def. 2, tau=1) and similarly for Attention-Aware (Def. 5, sigma=1), and compare against H(ri). Report the mean and max difference H(pi)-H(ri) and the fraction of tokens exceeding 0. If a substantial fraction (say >5%) have H(pi)>H(ri), Prop. 1's guarantee fails in the implemented regime; if the difference is systematically negative, the empirical claim holds despite the proof gap, and the paper's conclusion could stand with a corrected qualifier.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's central theoretical claim is that Similarity/Attention-Aware routing reduces expert-selection entropy (Prop. 1). The proof in Appendix B.2 defines pi = sum_{j in J_i} s(i,j) r_j with J_i = {j : H(e_j) <= H(e_i)}, and establishes H(pi) <= H(ri) only when H(s_i) -> 0, i.e., tau->0 (Def. 1) or sigma->0 (Def. 4). However, the implemented Def. 2 and Def. 5 sum over all j = 1..N at tau = sigma = 1 (Appendix D.1.2). At finite temperature, H(s_i) is positive (up to log N), and tokens outside J_i can have H(e_j) > H(e_i), so the chain H(pi) <= H(s_i) + sum_j s(i,j) H(e_j) <= H(ri) is not guaranteed. Moreover, the Attention-Aware implementation further approximates the posterior by selecting a single lowest-entropy head h* and the point estimate U = MHA(X), so even the sigma->0 limit of the full posterior is not what runs in practice. Thus the theoretical justification does not cover the actual algorithm; only the empirical entropy curves (Fig. 3) support the claim, and those are reported without error bars or multiple seeds. This is load-bearing because the abstract and Sec. 4 advertise a proof of entropy reduction as the basis for stable routing.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes two router modifications for sparse Mixture of Experts: Similarity-Aware (S)MoE, which replaces the per-token router score r(u_i) with a similarity-weighted average sum_j S[i,j] r(u_j) (Def. 2, Eq. 8), and Attention-Aware (S)MoE, which uses a posterior attention matrix A^p_{h*}[i,j] in place of S (Def. 5, Eq. 16). The authors introduce probabilistic graphical models (G1-G4) to motivate these designs, and Proposition 1 claims that the entropy of the resulting expert-selection distribution satisfies H(p_i) <= H(r_i) in the limits tau->0 (similarity) or sigma->0 (attention). Experiments on Wikitext-103, ImageNet, and three downstream tasks report improved perplexity/accuracy, lower routing fluctuation in the final epochs, and lower average routing entropy, with roughly 5% extra compute.","tokens_in":24184,"tokens_out":7593,"duration_ms":78885,"significance":"If the entropy-reduction guarantee held for the implemented algorithm, the contribution would be significant: it offers a nearly drop-in router change that stabilizes late-training SMoE routing, with thorough experiments across language, vision, and fine-tuning settings and with plug-and-play integrations into X-MoE and SMoE-dropout. The paper also includes useful ablations on temperature, number of experts, TopK values, and compute/memory overhead. A strength is the breadth of empirical validation and the clear formulation of the averaging mechanism. However, the advertised theoretical guarantee does not cover the actual algorithm: the proof is restricted to a subset J_i of tokens and to one-hot limits, while the implementation sums over all tokens at tau=sigma=1 and applies TopK and additional approximations. The PGM framework is a modeling choice rather than an independent derivation, but because the paper does not claim to fit the PGM, I do not view that as a correctness error. The significance is therefore conditional on reconciling the theorem with the implementation or reframing the theoretical claim.","major_comments":[{"comment":"The entropy bound is proved for p_i = sum_{j in J_i} s(i,j) r_j, where J_i = {j : H(e_j|u_j) <= H(e_i|u_i)}, and the inequality H(p_i) <= H(r_i) is obtained only by taking tau->0 or sigma->0 so that H(s_i) -> 0. The implemented models, however, sum over all j = 1..N and are run at tau = sigma = 1 as stated in Appendix D.1.2. At finite temperature H(s_i) is positive, and tokens outside J_i may have H(e_j) > H(e_i), so the chain H(p_i) <= H(s_i) + sum_j s(i,j) H(e_j) <= H(r_i) is not guaranteed. Since the abstract and Section 4 advertise a proof of entropy reduction as the basis for stable routing, the theoretical claim does not currently cover the algorithm as implemented.","section":"Section 4 (Prop. 1, Eq. 17) vs. Defs. 2, 5 and Appendix D.1.2"},{"comment":"The Attention-Aware implementation further deviates from the generative model used in Proposition 1. Equation 14 selects a single head h* by lowest average attention entropy, and Equation 15 replaces the expectation over U in Eq. 13 with the point estimate U = MHA(X). Neither step is part of the sigma->0 argument in Appendix B.2, which concerns the full posterior A^p_h under the generative model. Consequently, the claimed sigma->0 limit does not apply to the actual Def. 5 computation, and the proposition cannot be cited as a guarantee for Eq. 16.","section":"Section 3.2, Eqs. 14-16"},{"comment":"The implemented methods apply TopK to the averaged scores, whereas Proposition 1 bounds the entropy of the un-thresholded mixture p_i. TopK is a nonlinear operation, and no argument is given that the entropy or the top-k assignment process inherits the bound; additionally, Fig. 3 measures fluctuation as changes in TopK expert assignments, while the proposition concerns score entropy. The manuscript therefore does not rigorously connect the proven statement to the reported fluctuation reduction. A separate top-k analysis, or an explicit statement that the fluctuation results are empirical only, is needed.","section":"Section 2.2, Eq. 8 and Section 4, Prop. 1"}],"minor_comments":[{"comment":"The proof writes H(e_i) > H(e_j) for j in J_i, but J_i is defined by H(e_j|u_j) <= H(e_i|u_i); the strict inequality should be non-strict so that equal-entropy tokens are not excluded.","section":"Appendix B.2"},{"comment":"The definition of h* should state how ties in the lowest average attention entropy are broken, since this affects reproducibility of the Attention-Aware results.","section":"Section 3.2, Eq. 14"},{"comment":"No standard errors or number of seeds are reported; this is especially important for Fig. 3, where the central claims of reduced entropy and reduced fluctuation are made.","section":"Section 5, Fig. 3 and Tables 1-3"},{"comment":"The text says 'selecting Top-1 or Top-2 experts (E = 2)', but E denotes the total number of experts (16) elsewhere; the notation should be corrected to K, the number of selected experts.","section":"Appendix D.1.2"},{"comment":"Several typos remain, including 'Techincal Proofs', 'poterior head selection matrx', 'Comuputation and memory', and 'refered'; a proofreading pass is needed.","section":"Supplement headings"},{"comment":"The authors state that their implementation builds on publicly available codebases but do not release their own code; providing code would materially strengthen reproducibility.","section":"General"}],"recommendation":"major_revision","confidential_remarks":"The main theorem as stated does not cover the implemented algorithm, so the paper's advertised theoretical guarantee currently overstates what is proved. I would ask the editor to weigh whether the authors can supply a finite-temperature, all-token statement or whether the manuscript should be reframed as an empirical study with a conjectured entropy mechanism; either path is defensible, but the current combination of a strong proof claim and an unverified empirical curve is not."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Briefly: this paper proposes a simple and practical modification to SMoE routing — replace each token's router score with a similarity-weighted or attention-weighted average of other tokens' scores. It shows consistent, if modest, gains in language modeling, image classification, and fine-tuning, at about 5% extra compute. The specific formulas (Eqs. 8 and 16) are new, and the empirical work is decent: they compare against X-MoE and SMoE-dropout and show their method works as a plug-in.\n\nThe main soft spot is the theory. Proposition 1 claims an entropy-reduction guarantee, but the proof is for a restricted set J_i of tokens and only in the limits tau→0 or sigma→0, where the weights become one-hot and there is no mixing. The implemented Definitions 2 and 5 sum over all tokens at tau=sigma=1. At finite temperature the bound can fail; the paper just assumes it doesn't. The Attention-Aware version also approximates the posterior by a single low-entropy head and a point estimate of U, so even the sigma→0 limit of the full posterior isn't what runs. This is load-bearing because the abstract and Section 4 advertise a proof. The empirical entropy curves in Fig. 3 support the direction, but they are reported without error bars or multiple seeds.\n\nA second concern is the ImageNet comparison. The V-MoE baseline has SMoE layers only in the final two blocks, while the proposed variants alternate SMoE layers more broadly. That's not an apples-to-apples architecture comparison. And the PGM 'derivation' is somewhat circular: the generative model is constructed so its expectation reproduces the routing formula, so it frames the method rather than independently deriving it. That's fine as exposition, but it shouldn't be sold as a derivation.\n\nNone of this kills the paper. The method is plausible, cheap, and seems to help. But the contribution should be reframed as an empirical smoothing technique with a heuristic entropy argument, or the proof needs to be extended to finite temperature and the actual all-token sum. If either happens, this is a solid contribution.\n\nWho it's for: people working on MoE training stability. It deserves a serious referee, but the right outcome is major revision, not acceptance as-is.","headline":"A cheap, effective router-smoothing method for SMoE whose advertised entropy proof only covers a degenerate limit, not the actual algorithm.","tokens_in":24685,"tokens_out":2350,"would_cite":true,"duration_ms":22870,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Routing in sparse mixture-of-experts models becomes more stable when each token's expert score is replaced by a similarity- or attention-weighted average of other tokens' scores; the paper proves this lowers selection entropy and reports…","keywords":["sparse mixture of experts","routing fluctuation","token routing","expert selection entropy","probabilistic graphical model","similarity-aware routing","attention-aware routing","mixture of experts robustness"],"falsifier":"Run the deployed Similarity-Aware and Attention-Aware routers (Definitions 2 and 5, with $\\tau = \\sigma = 1$) on Wikitext-103 at the final epoch, compute $H(p_i)$ and $H(r_i)$ per token, and count how many tokens — especially low-entropy tokens outside $J_i$ — have strictly higher entropy after mixing; a non-negligible fraction would refute the theoretical justification at the operating point. A second check: shuffle the similarity matrix so each token is weighted by other tokens' scores at random; if routing fluctuation falls as much under shuffled weights as under true similarities, then plain averaging, not the graph structure, is doing the work.","tokens_in":23623,"feed_emoji":"🗳️","tokens_out":10768,"duration_ms":95134,"temperature":0.7,"pith_summary":"Sparse mixture-of-experts (SMoE) models assign each token to experts independently, through the router score $r(u_i)$ and nothing else; the paper argues this per-token independence is why routing decisions keep changing at the end of training, when up to a third of tokens still switch experts between final epochs. It proposes two nearly drop-in fixes, Similarity-Aware and Attention-Aware routing, in which the score used for token $u_i$ becomes a weighted mixture $\\sum_j S[i,j] r(u_j)$ of its neighbors' router scores, with weights given by token similarity or by the attention matrix. The paper proves that this mixing lowers the entropy of expert selection (in the limits $\\tau \\to 0$ or $\\sigma \\to 0$), making indecisive tokens more confident, and reports reduced routing fluctuation, lower perplexity, and higher accuracy across language, vision, and fine-tuning benchmarks at roughly 5% extra compute. If correct, stable routing is available as a small router-level modification to existing SMoE layers rather than a new architecture.","feed_headline":"Similarity-weighted expert votes steady sparse MoE routing","feed_subtitle":"Weighted token votes lower routing entropy and cut late-training expert switches, at ~5% extra compute.","key_machinery":"The load-bearing mechanism is the replacement of the standalone router score with a weighted average over tokens: $p_i = \\sum_j S[i,j] r(u_j)$, computed either from token similarities (Definition 2, Eq. 8) or from the attention matrix restricted to the head with the lowest average attention entropy (Definition 5, Eq. 16). This coupling is what breaks the conditional independence of expert selections, and it is also what makes the entropy argument work: writing the mixed selection as a mixture of the neighbors' categorical distributions and applying the decomposition $H(e_i^s, t_i) = H(t_i) + \\sum_j s(i,j) H(e_j) \\geq H(e_i^s)$ gives the Proposition 1 bound, with $H(t_i)$ vanishing as the weighting becomes one-hot. The PGM graphs G1–G4 carry the conceptual load: they make the independence assumption explicit and show both proposed routers arise as conditional expectations under a generative story.","core_discovery":"Read through a probabilistic graphical model, the (S)MoE generative process draws each token's expert from $\\mathrm{Cat}(r(u_i))$ and treats the draws as conditionally independent given the input; the paper identifies this independence, $e_i \\perp e_j \\mid U$, as the source of routing fluctuation, since nothing forces similar tokens toward the same expert. To break it, the paper places a graph over the tokens and replaces the per-token router score with a mixture $p_i = \\sum_j S[i,j] r(u_j)$, where $S$ is either the temperature-softmax similarity matrix $S[i,j] = \\mathrm{Softmax}(u_i^\\top W_s u_j / \\tau)$ (Similarity-Aware) or the posterior attention matrix of the lowest-entropy head (Attention-Aware). Proposition 1 bounds $H(p_i) \\leq \\sum_j s(i,j) H(r_j) + H(s_i)$ and shows $H(p_i) \\leq H(r_i)$ as $\\tau \\to 0$ or $\\sigma \\to 0$, giving the theoretical claim that the entropy of expert selection decreases. Empirically, the variants cut the share of tokens switching experts between the last training epochs, lower average routing entropy across layers, improve load balance, and outperform softmax-gated SMoE baselines on Wikitext-103, ImageNet and its corrupted and adversarial variants, and fine-tuning tasks.","pith_inferences":["A consequence the paper leaves implicit: if lower selection entropy is the active ingredient, then any cheap pairwise coupling — for example, averaging router scores within hash-bucketed blocks of tokens — should yield similar stability at far less than the $O(N^2)$ cost of a full similarity matrix, which would matter for very long sequences.","A testable extension suggested by the appendix's temperature study: scheduling $\\tau$ (and $\\sigma$) from large to small during training might preserve the stability gains while sharpening the mixture toward the entropy-reduced limit where Proposition 1 provably holds.","The paper notes as its own limitation that it does not model how tokens are generated; a generative model that includes the token-construction step would place the independence problem in the prior rather than patching it at the router, so the present work can be read as evidence for that more principled fix.","For deployment, the honest reading is empirical: the proven guarantee holds at one-hot limits, so choosing $\\tau = \\sigma = 1$ rests on the observed correlation between smoother routing, lower entropy, and fewer switches rather than on Proposition 1 itself."],"forward_implications":["Because the change is confined to how router scores are combined, the method layers onto existing SMoE routers; the paper demonstrates this by improving X-MoE, SMoE-dropout, and GLaM-style baselines with the same weighted-mixing trick.","If the entropy reduction transfers to training dynamics, late-epoch routing fluctuation — the non-robustness signal the paper measures as tokens switching experts between epochs 59 and 60 — should shrink, and the reported plots show exactly that for both variants.","More confident expert selection should make predictions less sensitive to small input perturbations, consistent with the reported gains on adversarially attacked Wikitext-103 and on ImageNet-C, ImageNet-A, and ImageNet-R.","Attention-Aware routing reuses a quantity the transformer already computes (the attention matrix), so its overhead is small — reported forward-pass ratios of about 1.07x computation and 1.06x memory over the baseline — making it a cheap upgrade in MoE transformers."],"supporting_citations":[{"why":"Defines the sparse mixture-of-experts layer and sparse gating that the paper's routers modify.","marker":"(Shazeer et al., 2017)"},{"why":"Supplies the Switch Transformer, the MoE-Transformer block structure whose PGM is extended and the main language-modeling baseline.","marker":"(Fedus et al., 2022)"},{"why":"The X-MoE baseline targeting routing consistency; the paper compares against it and improves it with weighted mixing.","marker":"(Chi et al., 2022)"},{"why":"The V-MoE vision baseline used for ImageNet and robustness benchmarks.","marker":"(Riquelme et al., 2021)"},{"why":"StableMoE, the prior routing-stability method that frames the routing-fluctuation problem the paper addresses.","marker":"(Dai et al., 2022)"},{"why":"GLaM, a second language-modeling baseline whose Similarity/Attention-Aware variants are evaluated.","marker":"(Du et al., 2022)"},{"why":"SMoE-dropout, a stability-oriented SMoE baseline included in perplexity comparisons.","marker":"(Chen et al., 2023)"}],"fun_headline_variants":["Graph of tokens steadies sparse MoE routing","Token similarity kills expert-switching in sparse MoE","Similarity-aware routing stabilizes sparse mixture of experts","Attention-aware experts reduce late-training route flips"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"Proposition 1's entropy reduction $H(p_i) \\leq H(r_i)$ is proved only for tokens whose every neighbor has no larger selection entropy than they do (the set $J_i$), and only in the limits $\\tau \\to 0$ or $\\sigma \\to 0$, where the weighting matrix becomes one-hot and no mixing occurs; the deployed Definitions 2 and 5 average over all tokens at $\\tau = \\sigma = 1$, and the paper assumes without proof that the guarantee survives both relaxations.","fun_headline_variants_meta":{"raw":{"variants":["Graph of tokens steadies sparse MoE routing","Token similarity kills expert-switching in sparse MoE","Similarity-aware routing stabilizes sparse mixture of experts","Attention-aware experts reduce late-training route flips"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000605,"raw_usage":{"total_tokens":2917,"prompt_tokens":1139,"completion_tokens":1778,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":755,"completion_tokens_details":{"reasoning_tokens":1717}},"tokens_in":755,"tokens_out":1778,"duration_ms":13260,"temperature":1.0,"reasoning_tokens":1717,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T04:35:59.130260+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the deployed Similarity-Aware and Attention-Aware routers (Definitions 2 and 5, with $\\tau = \\sigma = 1$) on Wikitext-103 at the final epoch, compute $H(p_i)$ and $H(r_i)$ per token, and count how many tokens — especially low-entropy tokens outside $J_i$ — have strictly higher entropy after mixing; a non-negligible fraction would refute the theoretical justification at the operating point. A second check: shuffle the similarity matrix so each token is weighted by other tokens' scores at random; if routing fluctuation falls as much under shuffled weights as under true similarities, then plain averaging, not the graph structure, is doing the work.","supporting_citations":[{"cited_title":"On the representation collapse of sparse mixture of experts","cited_arxiv_id":null,"evidence_quote":"The X-MoE baseline targeting routing consistency; the paper compares against it and improves it with weighted mixing."},{"cited_title":"K., Liu, S., and Wang, Z","cited_arxiv_id":null,"evidence_quote":"SMoE-dropout, a stability-oriented SMoE baseline included in perplexity comparisons."}],"review_version":1}