Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts

• URL: http://arxiv.org/abs/2601.02144v1
• Date: Mon, 05 Jan 2026 14:16:11 GMT
• Title: Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts
• Abstract summary: Mixture-of-Experts (MoE) architectures scale large language models efficiently by employing a parametric "router" to dispatch tokens to a sparse subset of experts.<n>We introduce kNN-MoE, a retrieval-augmented routing framework that reuses optimal expert assignments from a memory of similar past cases.<n>Experiments show kNN-MoE outperforms zero-shot baselines and rivals computationally expensive supervised fine-tuning.
• Score: 32.65737144630759
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Mixture-of-Experts (MoE) architectures scale large language models efficiently by employing a parametric "router" to dispatch tokens to a sparse subset of experts. Typically, this router is trained once and then frozen, rendering routing decisions brittle under distribution shifts. We address this limitation by introducing kNN-MoE, a retrieval-augmented routing framework that reuses optimal expert assignments from a memory of similar past cases. This memory is constructed offline by directly optimizing token-wise routing logits to maximize the likelihood on a reference set. Crucially, we use the aggregate similarity of retrieved neighbors as a confidence-driven mixing coefficient, thus allowing the method to fall back to the frozen router when no relevant cases are found. Experiments show kNN-MoE outperforms zero-shot baselines and rivals computationally expensive supervised fine-tuning.

Related papers

• Share First, Route What Remains: A Unified Framework for Token-Adaptive MoE Computation [49.8459545291965]
  Mixture-of-experts (MoE) models have recently moved beyond routing a fixed number of complete experts.<n>We study this dependency by decomposing sparsely upcycled feed-forward experts into key-value channels.<n>We instantiate it in UniF-MoE, a unified framework for token-adaptive MoE computation.
  arXiv  Detail & Related papers  (2026-08-11T02:40:58Z)
• Agent-as-a-Router: Agentic Model Routing for Coding Tasks [54.00153674547733]
  Agent-as-a-router (AC) formalizes routing as a C-A-F loop (Context->Action->FeedbackContext)<n>AC achieves the lowest cumulative regret on in-distribution tasks and generalizes to out-of-distribution agentic-programming tasks.
  arXiv  Detail & Related papers  (2026-06-22T06:37:31Z)
• Redesign Mixture-of-Experts Routers with Manifold Power Iteration [60.55368504665436]
  We propose a router redesign with Manifold Power Iteration (MPI)<n>MPI drives router rows to converge toward the principal singular directions of associated experts.<n>We show that MPI drives router rows to converge toward the principal singular directions of associated experts.
  arXiv  Detail & Related papers  (2026-06-10T17:57:36Z)
• When Model Merging Breaks Routing: Training-Free Calibration for MoE [67.27839627141972]
  Hessian-Aware Router (HARC) is a training-free framework that leverages second-order curvature information to realign the merged router.<n> Experiments on mathematical reasoning and code generation tasks show that HARC effectively mitigates routing breakdown across diverse MoE merging baselines.
  arXiv  Detail & Related papers  (2026-06-02T09:33:33Z)
• Adaptive Inverted-Index Routing for Granular Mixtures-of-Experts [49.09151538536423]
  Mixture-of-experts (MoE) models enable scalable transformer architectures by activating only a subset of experts per token.<n>Recent evidence suggests that performance improves with increasingly granular experts, i.e., many small experts instead of a few large ones.<n>We introduce adaptive inverted-index routing for MoE (AIR-MoE), an inverted-index-inspired routing architecture based on vector quantization (VQ)
  arXiv  Detail & Related papers  (2026-05-06T14:15:10Z)
• Self-Routing: Parameter-Free Expert Routing from Hidden States [14.018297453424454]
  Mixture-of-Experts (MoE) layers increase model capacity by activating only a small subset of experts per token.<n>We propose Self-Routing, a parameter-free routing mechanism that uses a designated subspace of the token hidden state directly as expert logits.
  arXiv  Detail & Related papers  (2026-04-01T03:05:20Z)
• Path-Constrained Mixture-of-Experts [23.80197713200086]
  MoE routing selects each layer's experts independently, creating NL possible expert paths -- for N experts across L layers.<n>We propose pathmoe, which shares router parameters across consecutive layers.<n>Experiments on 0.9B and 16B parameter models demonstrate consistent improvements on perplexity and downstream tasks over independent routing.
  arXiv  Detail & Related papers  (2026-03-18T21:35:53Z)
• Routing Manifold Alignment Improves Generalization of Mixture-of-Experts LLMs [24.791817951102487]
  We show that aligning the manifold of routing weights with that of task embedding can effectively reduce the gap.<n>In experiments, we finetune routers in OLMoE, DeepSeekMoE, and Qwen3-MoE using RoMA.
  arXiv  Detail & Related papers  (2025-11-10T18:59:53Z)
• Robust Nearest Neighbour Retrieval Using Targeted Manifold Manipulation [0.0]
  Nearest-neighbour retrieval is central to classification and explainable-AI pipelines.<n>We propose Targeted Manifold Manipulation-Nearest Neighbour (TMM-NN), which reconceptualises retrieval by assessing how readily each sample can be nudged into a designated region of the feature manifold.<n>TMM-NN implements this through a lightweight, query-specific trigger patch.
  arXiv  Detail & Related papers  (2025-11-09T07:37:05Z)
• Rewiring Experts on the Fly:Continuous Rerouting for Better Online Adaptation in Mixture-of-Expert models [52.502867924372275]
  Mixture-of-Experts (MoE) models achieve efficient scaling through sparse expert activation, but often suffer from suboptimal routing decisions due to distribution shifts in deployment.<n>We propose textita data-free, online test-time framework that continuously adapts MoE routing decisions during text generation without external supervision or data.
  arXiv  Detail & Related papers  (2025-10-16T16:24:36Z)
• From Score Distributions to Balance: Plug-and-Play Mixture-of-Experts Routing [52.01745035243826]
  Mixture-of-Experts (MoE) models can scale parameter capacity by routing each token to a subset of experts.<n> conditional routing shifts the burden on inference memory, limiting the number of experts per device.<n>We present LASER, a plug-and-play, inference-time routing algorithm that balances load while preserving accuracy.
  arXiv  Detail & Related papers  (2025-09-29T16:29:17Z)
• RouteMark: A Fingerprint for Intellectual Property Attribution in Routing-based Model Merging [69.2230254959204]
  We propose RouteMark, a framework for IP protection in merged MoE models.<n>Our key insight is that task-specific experts exhibit stable and distinctive routing behaviors under probing inputs.<n>For attribution and tampering detection, we introduce a similarity-based matching algorithm.
  arXiv  Detail & Related papers  (2025-08-03T14:51:58Z)
• MoQAE: Mixed-Precision Quantization for Long-Context LLM Inference via Mixture of Quantization-Aware Experts [29.11217299899888]
  MoQAE is a mixed-precision quantization method via mixture of quantization-aware experts.<n>We show that MoQAE outperforms state-of-the-art KV cache quantization approaches in both efficiency and effectiveness.
  arXiv  Detail & Related papers  (2025-06-09T08:16:24Z)
• Tight Clusters Make Specialized Experts [1.7597562616011944]
  Sparse Mixture-of-Experts (MoE) architectures have emerged as a promising approach to decoupling model capacity from computational cost.<n>We present a novel router that learns the underlying clustering structure of the input distribution in order to send input tokens to appropriate experts.<n>Our AC router enables the MoE model to obtain three connected benefits: 1) faster convergence, 2) better robustness to data corruption, and 3) overall performance improvement.
  arXiv  Detail & Related papers  (2025-02-21T09:10:54Z)
• Noisy Correspondence Learning with Self-Reinforcing Errors Mitigation [63.180725016463974]
  Cross-modal retrieval relies on well-matched large-scale datasets that are laborious in practice. We introduce a novel noisy correspondence learning framework, namely textbfSelf-textbfReinforcing textbfErrors textbfMitigation (SREM)
  arXiv  Detail & Related papers  (2023-12-27T09:03:43Z)
• Tractable Bounding of Counterfactual Queries by Knowledge Compilation [51.47174989680976]
  We discuss the problem of bounding partially identifiable queries, such as counterfactuals, in Pearlian structural causal models. A recently proposed iterated EM scheme yields an inner approximation of those bounds by sampling the initialisation parameters. We show how a single symbolic knowledge compilation allows us to obtain the circuit structure with symbolic parameters to be replaced by their actual values.
  arXiv  Detail & Related papers  (2023-10-05T07:10:40Z)
• Sparse Backpropagation for MoE Training [118.31785160874024]
  We introduce SparseMixer, a scalable gradient estimator that bridges the gap between backpropagation and sparse expert routing. Grounded in a numerical ODE framework, SparseMixer harnesses the mid-point method, a second-order ODE solver, to deliver precise gradient approximations. Applying SparseMixer to Switch Transformer on both pre-training and machine translation tasks, SparseMixer showcases considerable performance gain.
  arXiv  Detail & Related papers  (2023-10-01T22:43:57Z)
• Soft Merging of Experts with Adaptive Routing [38.962451264172856]
  We introduce Soft Merging of Experts with Adaptive Routing (SMEAR) SMEAR avoids discrete routing by using a single "merged" expert constructed via a weighted average of all of the experts' parameters. We empirically validate that models using SMEAR outperform models that route based on metadata or learn sparse routing through gradient estimation.
  arXiv  Detail & Related papers  (2023-06-06T15:04:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.