{"id":"40a8054f-5e25-4cfb-a402-8cdd41278589","arxiv_id":"2601.02144","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"Retrieving stored, likelihood-optimized expert assignments for similar router inputs improves MoE routing accuracy over zero-shot and is competitive with supervised fine-tuning.","lead":"kNN-MoE adds a memory of past router decisions to mixture-of-experts language models, retrieving similar cases at inference and blending their expert assignments with the frozen router. The method lifts accuracy over zero-shot routing on several benchmarks and roughly matches fine-tuning without changing model weights.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Untested premise: stored 'optimal' assignments depend on the reference target token, but retrieval keys are router inputs only; a one-step TopK gradient may not even beat the frozen router.","rationale":"The reader's verdict is already CONDITIONAL with MODERATE confidence, citing the locality assumption and lack of code/data. My concern sharpens that assumption: the memory values depend on the reference target token, which is absent from the retrieval key. This makes the core premise less secure than a generic 'similar inputs have similar outputs' assumption. However, the paper reports consistent empirical gains across 15 model-benchmark comparisons, so the appropriate disposition is to keep the conditional verdict and require the proposed test before treating the mechanism as established. The one-step TopK gradient and missing reference-likelihood verification further support the conditional status, but do not by themselves invalidate the empirical claim. No change to the reader's verdict is needed.","tokens_in":12067,"tokens_out":10279,"duration_ms":124069,"concrete_test":"Take a held-out set D_val with ground-truth tokens. For each token t, run the same memory-construction optimization (S=1, η=2e-2) on x_t and y_t to get an oracle assignment a*_oracle(t). Then compare three routing policies on D_val: (1) original router a(x_t); (2) kNN-MoE's retrieved a_mem(x_t) from Dref; (3) oracle a*_oracle(t). Measure average NLL (or task accuracy). If policy (2) is not substantially closer to (3) than to (1), or if (3) is much better than (2), the retrieval locality assumption is the bottleneck. As a control, permute the stored values across keys; if shuffled memory retains the gains, the specific key-value association is not load-bearing.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central mechanism requires that the stored value v_t = π(r_t^(S)) is a better expert assignment for test tokens whose router input x is close to x_t. But v_t is optimized against the reference target token y_t: the gradient ∇_r L_t in Eq. (2) depends on y_t through the output cross-entropy. At inference, retrieval keys are only router inputs; the target token is unknown. Thus the method implicitly assumes the mapping x ↦ optimal assignment is nearly single-valued (y_t predictable from x_t). This is not derived or measured. Additionally, only S=1 step through a non-differentiable TopK-softmax is used, so the stored values may barely differ from the frozen router's own output; no reference-set likelihood comparison is reported to show v_t is actually better than a(x_t). If optimal assignments are target-token-dependent or the one-step values do not improve on the router, then kNN-MoE's gains cannot be attributed to 'routing by analogy' and may instead come from a generic bias introduced by mixing with the reference set.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes kNN-MoE, an inference-time routing refinement for Mixture-of-Experts language models. For each MoE layer, a memory is built offline from a labeled reference set: router inputs are stored as keys, and values are expert assignments obtained by taking one or more gradient steps on token-level negative log-likelihood with respect to the routing logits. At test time, for each router input the method retrieves the K nearest keys, aggregates their stored assignments with an RBF similarity kernel, and interpolates this memory-based proposal with the frozen router's output using the average similarity as a confidence coefficient. Experiments on OLMoE, GPT-OSS, and Qwen3 across GPQA, MMLU, SuperGPQA, USMLE, and MedMCQA compare against zero-shot, 5-shot, SFT, and router-only SFT. The paper reports that kNN-MoE improves over zero-shot in all 15 model-benchmark cells and is competitive with SFT while requiring much less preparation time.","tokens_in":12383,"tokens_out":7316,"duration_ms":77304,"significance":"The idea is timely and pragmatic: augmenting frozen MoE routers with a retrievable store of locally optimized assignments offers a parameter-free (in the sense of no gradient updates to the model) alternative to supervised fine-tuning. The paper is careful in several respects: it uses disjoint reference/test splits, honestly frames the comparison with SFT as competitive rather than uniformly superior, and includes ablations over reference size, neighbor count, gradient steps, similarity function, and perplexity buckets. The limitations section explicitly acknowledges the reliance on a labeled, distribution-matched reference set. If the empirical claims are confirmed with appropriate statistical controls and a clearly specified optimization procedure, this would be a useful contribution to inference-time adaptation of MoE models. The main weaknesses are statistical (single runs, test-set-tuned hyperparameters) and a lack of verification of the central mechanism—the 'optimality' of stored assignments and the router-input-to-assignment mapping.","major_comments":[{"comment":"All reported accuracies are single-run point estimates without confidence intervals or significance tests. Several key improvements over zero-shot are very small, e.g., GPT-OSS MMLU +0.08, Qwen3 MMLU +0.27, OLMoE SuperGPQA +0.25. Since the abstract and §5.2 claim consistent gains across all models and benchmarks, the reader cannot assess whether these differences are meaningful. Please provide multiple seeds or bootstrap confidence intervals, and state significance for the aggregate comparisons.","section":"§5.2, Table 2"},{"comment":"The default hyperparameters K=1 and S=1 are selected using accuracy on MedMCQA, which is one of the test benchmarks reported in Table 2 (see Tables 6 and 7). This is test-set tuning for exactly the dataset where gains are largest (e.g., OLMoE MedMCQA +1.44). This compromises the validity of the MedMCQA result and makes it difficult to know how K and S generalize. Use a separate validation split disjoint from both D_ref and D_test, or report results across all K/S values without selecting the best on the test set.","section":"§5.1, §6.3, §6.4"},{"comment":"The optimization in Eq. (2) minimizes L_t(r) where the dependence on r goes through π(r)=TopK(Softmax(r)). TopK is a non-differentiable selection operator, so the gradient ∇_r L_t does not exist in the usual sense; with S=1, the update r^(1)=r^(0)-η∇L would leave r unchanged if the exact TopK derivative (zero a.e.) is used. The paper does not specify whether a straight-through estimator, a softmax relaxation, or a gradient through the softmax probabilities is used. This is central: if the stored values v_t equal the original router assignments, the method reduces to the zero-shot baseline. Please specify the surrogate gradient and report reference-set likelihoods of v_t vs a(x_t) to show that the 'optimal' assignments are actually better than the frozen router's.","section":"§4.1, Eq. (2)"},{"comment":"Memory values v_t are optimized against reference target tokens y_t, but retrieval keys are router inputs only. The method therefore assumes that, for a given router input x, the beneficial expert assignment is approximately independent of the target token. This assumption is not directly tested. The PPL-bucket analysis in §6.1 shows that gains concentrate in high-perplexity inputs, which is consistent with the mechanism, but does not rule out the alternative that the gain comes from a generic bias introduced by mixing with a reference-derived assignment. Please include a diagnostic, e.g., compare kNN-MoE against a random-memory baseline or measure the within-neighborhood variance of stored optimal assignments for nearby router inputs.","section":"§4.1–4.2, §6.1"}],"minor_comments":[{"comment":"Typo: 'We equip each MoE' is rendered as 'We equipeachMoE'.","section":"§4.1"},{"comment":"The header 'Ref. Source' is ambiguous; clarify that the column specifies the origin of the reference split.","section":"Table 1"},{"comment":"Please specify how perplexity is computed: per-token, per-sequence, or over answer tokens only.","section":"§6.1"},{"comment":"The RBF kernel width γ is 'set heuristically based on the average nearest neighbor distance in the memory.' Give the exact formula and state whether it is computed per layer or globally.","section":"§5.1"},{"comment":"The claim that Qwen3 router inputs are more uniform in magnitude is speculative; if kept, support it with statistics of router input norms.","section":"§6.5"}],"recommendation":"major_revision","confidential_remarks":"The paper proposes a clear, useful approach and the experimental scope is respectable, but the current validation is not yet at the standard needed to support the headline claims. The lack of error bars, the use of a test benchmark for hyperparameter selection, and the underspecified gradient through TopK all need to be addressed. The idea is not fundamentally flawed, so I recommend major revision rather than rejection. Since no code is provided, asking for code or at least detailed pseudocode for the surrogate gradient would also help reproducibility."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague, here's my take on arXiv:2601.02144. The core idea is genuinely new as far as I know: build a memory of router inputs paired with expert assignments that maximize next-token likelihood on a reference set, then at test time retrieve nearest neighbors and interpolate with the frozen router using a similarity-based confidence. That's not in the cited literature—kNN-LM interpolates output distributions, C3PO updates logits online, and the router fine-tuning methods change parameters. So credit for the combination.\n\nWhat the paper does well: the method is well specified, the ablations are thoughtful (number of neighbors, gradient steps, similarity function, reference set size), and the perplexity analysis in §6.1 is genuinely informative: gains concentrate in high-perplexity inputs, which supports the intuition that retrieval helps when the router is uncertain. The comparison to SFT is honestly framed—they don't claim to beat it uniformly, just to be competitive while using far less compute. The limitations section is straight about the dependence on a labeled, distribution-matched reference set and the inference latency.\n\nThe soft spots are real but addressable. Every number is a single run with no error bars; several differences are under 0.5 points, so we can't yet tell which of the small gains are signal. Hyperparameters K and gamma are selected using benchmark performance, and with K=1 the method is essentially trusting one nearest neighbor, which is fragile. The stress-test concern about target-token dependence is worth taking seriously: the memory values are optimized against the reference token, but retrieval keys are only router inputs, so the method implicitly assumes the oracle assignment is mostly a function of the input. That's not proven, though the perplexity bucketing gives some supporting evidence. They also don't release code or data, so reproduction is on trust.\n\nOverall: this is a useful incremental contribution, not a breakthrough. A serious referee could help push them to add error bars or confidence intervals, test the target-token dependence explicitly, and consider a more robust aggregation than K=1. If those are addressed, the paper would be a solid addition to the routing-adaptation literature.\n\nI'd send it to peer review. It's clear, honest, and the experiments are relevant even if the effect sizes are small.","headline":"A clean, honest paper on retrieval-augmented MoE routing; the gains are modest and single-run, but the method is clearly described and the comparison to SFT is fair.","tokens_in":12813,"tokens_out":2536,"would_cite":false,"duration_ms":25717,"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":"A retrieval memory of past optimal expert assignments can improve Mixture-of-Experts routing at inference, matching supervised fine-tuning without parameter updates.","keywords":["Mixture-of-Experts","retrieval-augmented routing","kNN memory","expert assignment","router","test-time adaptation","confidence-based mixing","MoE routing"],"falsifier":"Build the memory from a reference set in one domain and evaluate on a test set from a different domain; if kNN-MoE accuracy drops below the zero-shot baseline, the core distribution-match assumption fails. More directly, for a sample of router inputs, compute the expert assignment that truly maximizes target-token likelihood and check whether the nearest neighbor in router-input space has significantly higher assignment overlap than a random neighbor; if overlap is not significantly higher, the retrieval signal is absent.","tokens_in":1369,"feed_emoji":"🧠","tokens_out":2062,"duration_ms":51610,"temperature":0.7,"pith_summary":"The paper introduces kNN-MoE, a method that augments the frozen router of a Mixture-of-Experts model with a retrieval-based memory. Offline, for each token in a labeled reference set, it optimizes the routing logits to maximize the likelihood of the ground-truth token, storing the router input and the resulting expert assignment as a key-value pair. At inference, it retrieves the nearest neighbors of the current router input and interpolates their expert assignments with the original router prediction, using the average neighbor similarity as a confidence coefficient. Across three MoE models and five benchmarks, kNN-MoE outperforms the zero-shot baseline and is competitive with supervised fine-tuning, with a fraction of the offline cost. The central belief is that similar router inputs tend to share beneficial expert assignments, so past optimal choices can be reused by analogy.","feed_headline":"Reuse past expert picks to improve MoE routing","feed_subtitle":"A similarity-based memory of optimal expert assignments beats zero-shot and rivals fine-tuning without weight updates.","key_machinery":"The key object is a key-value memory per MoE layer, where keys are router input hidden states and values are per-token 'optimal' expert assignments obtained by gradient-based likelihood maximization on reference tokens. The second key mechanism is confidence-aware adaptive mixing: the average RBF similarity of the retrieved neighbors (with K=1 found best) acts as a scalar mixing coefficient that linearly interpolates the parametric router output and the memory proposal, allowing automatic fallback to the frozen router when no similar past case exists.","core_discovery":"The central claim is that expert assignment can be treated as a retrievable entity: rather than learning a fixed routing policy, store the actual assignments that would have maximized next-token likelihood for past tokens, and reuse them for similar future tokens. The memory is built by taking one gradient step (or a few) on per-token routing logits to maximize the likelihood of the reference token, holding all other model parameters frozen. At test time, the aggregate similarity of retrieved neighbors provides a principled mixing weight between the memory proposal and the frozen router, so the method degrades gracefully to the original router when retrieval is uncertain. The paper reports c","pith_inferences":["Editorial extension: The success of K=1 over K>1 hints that the router-input space may be locally structured such that one exemplar is enough; if so, memory pruning or quantization could cut storage and latency with little loss, which the paper leaves as future work.","Editorial extension: The mixing coefficient is an uncalibrated average similarity; a testable extension is to calibrate it or learn a threshold on a small validation set, which might reduce the slight degradation observed on low-perplexity inputs.","Editorial extension: The 'optimal' assignments are derived from a single gradient step and are not guaranteed to be globally optimal; a direct test would check whether the stored assignment actually improves likelihood over the router on the reference token, and whether this improvement correlates with downstream gains.","Editorial extension: The method depends on labeled, distribution-matched reference data; the paper's proposed unlabeled extension (using pseudo-labels from a judge model) could be tested directly and would determine whether the mechanism works without ground-truth answers."],"forward_implications":["If correct, any frozen MoE model can be adapted to a new domain by constructing a memory from a few hundred labeled examples, without retraining or modifying weights.","The method is especially beneficial for high-perplexity (hard, out-of-distribution) inputs, where the router is unreliable; low-perplexity inputs see little change or slight degradation.","The offline memory construction is substantially faster than LoRA-based supervised fine-tuning (about one-third the time in the reported setting), with moderate per-token latency overhead at inference.","The finding that K=1 beats larger K indicates that in routing, a single most similar past case carries the signal; aggregating more neighbors dilutes it.","In data-scarce settings where supervised fine-tuning overfits and degrades below zero-shot, kNN-MoE still improves over zero-shot, suggesting a robustness advantage."],"fun_headline_variants":["Retrieve past expert picks to improve MoE routing","kNN-MoE: reuse optimized routing decisions from memory","Memory-guided expert assignment rivals fine-tuning","Confidence-weighted retrieval boosts MoE routing"],"cache_read_input_tokens":14208,"weakest_assumption_plain":"The method assumes that if two router inputs are close in Euclidean (RBF) distance, then the expert assignment that would have maximized likelihood for one is beneficial for the other, and that the labeled reference set is representative of the test distribution; if similar router inputs do not share optimal experts, the memory proposal adds noise that the mixing coefficient cannot fully suppress.","fun_headline_variants_meta":{"raw":{"variants":["Retrieve past expert picks to improve MoE routing","kNN-MoE: reuse optimized routing decisions from memory","Memory-guided expert assignment rivals fine-tuning","Confidence-weighted retrieval boosts MoE routing"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000709,"raw_usage":{"total_tokens":2985,"prompt_tokens":654,"completion_tokens":2331,"prompt_tokens_details":{"cached_tokens":256},"prompt_cache_hit_tokens":256,"prompt_cache_miss_tokens":398,"completion_tokens_details":{"reasoning_tokens":2280}},"tokens_in":398,"tokens_out":2331,"duration_ms":18476,"temperature":1.0,"reasoning_tokens":2280,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-03T12:35:48.202985+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Build the memory from a reference set in one domain and evaluate on a test set from a different domain; if kNN-MoE accuracy drops below the zero-shot baseline, the core distribution-match assumption fails. More directly, for a sample of router inputs, compute the expert assignment that truly maximizes target-token likelihood and check whether the nearest neighbor in router-input space has significantly higher assignment overlap than a random neighbor; if overlap is not significantly higher, the retrieval signal is absent.","supporting_citations":[],"review_version":1}