Multi-Head Attention as a Source of Catastrophic Forgetting in MoE Transformers

• URL: http://arxiv.org/abs/2602.12587v1
• Date: Fri, 13 Feb 2026 03:59:09 GMT
• Title: Multi-Head Attention as a Source of Catastrophic Forgetting in MoE Transformers
• Abstract summary: Mixture-of-Experts (MoE) architectures are often considered a natural fit for continual learning.<n>We propose MH-MoE, which performs head-wise routing over sub-representations to increase routing granularity and reduce composition collisions.
• Score: 28.800065959523696
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Mixture-of-Experts (MoE) architectures are often considered a natural fit for continual learning because sparse routing should localize updates and reduce interference, yet MoE Transformers still forget substantially even with sparse, well-balanced expert utilization. We attribute this gap to a pre-routing bottleneck: multi-head attention concatenates head-specific signals into a single post-attention router input, forcing routing to act on co-occurring feature compositions rather than separable head channels. We show that this router input simultaneously encodes multiple separately decodable semantic and structural factors with uneven head support, and that different feature compositions induce weakly aligned parameter-gradient directions; as a result, routing maps many distinct compositions to the same route. We quantify this collision effect via a route-wise effective composition number $N_{eff}$ and find that higher $N_{eff}$ is associated with larger old-task loss increases after continual training. Motivated by these findings, we propose MH-MoE, which performs head-wise routing over sub-representations to increase routing granularity and reduce composition collisions. On TRACE with Qwen3-0.6B/8B, MH-MoE effectively mitigates forgetting, reducing BWT on Qwen3-0.6B from 11.2% (LoRAMoE) to 4.5%.

Related papers

• CRISP: Cliff-awaRe Input-adaptive Sparse Prefilling with Structural-Mass-Motivated Routing [90.63350360476294]
  We present CRISP (Cliff-awaRe Input-adaptive Sparse Prefilling), which identifies and addresses two structural challenges in this dynamic routing paradigm.<n>We show that the routing decision can be read directly off the structure of the proxy attention map.<n>We formalize the post-softmax mass cliff and demonstrate theoretically that strictly cumulative coverage thresholds accumulate O(n) background noise at long contexts.
  arXiv  Detail & Related papers  (2026-09-01T22:49:46Z)
• 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)
• ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts [13.754183593794131]
  ThAME is a 3D heterogeneous multi-chiplet architecture for MoE inference.<n>We show that ThAME outperforms state-of-the-art counterparts by up to 15.7x in terms of speedup.
  arXiv  Detail & Related papers  (2026-07-19T04:45:01Z)
• Sticky Routing: Training MoE Models for Memory-Efficient Inference [0.0]
  StickyMoE penalises abrupt expert switches between adjacent tokens.<n>Unlike post-hoc methods, StickyMoE allows expert representations and routing decisions to co-adapt from the first training step.
  arXiv  Detail & Related papers  (2026-06-12T23:47:19Z)
• 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)
• Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts [23.123497435269112]
  We study how routing decisions in SMoEs are formed mechanistically.<n>In a $1$B SMoE trained from scratch, higher router scores predict stronger expert neuron activations.<n>Our results explain how routers form assignment geometry that supports an effective division of labor.
  arXiv  Detail & Related papers  (2026-05-12T17:55:02Z)
• 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)
• TopoCurate:Modeling Interaction Topology for Tool-Use Agent Training [53.93696896939915]
  Training tool-use agents typically rely on Supervised Fine-Tuning (SFT) on successful trajectories and Reinforcement Learning (RL) on pass-rate-selected tasks.<n>We propose TopoCurate, an interaction-aware framework that projects multi-trial rollouts from the same task into a unified semantic quotient topology.<n>TopoCurate achieves consistent gains of 4.2% (SFT) and 6.9% (RL) over state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-03-02T10:38:54Z)
• Generalizing GNNs with Tokenized Mixture of Experts [75.8310720413187]
  We show that improving stability requires reducing reliance on shift-sensitive features, leaving an irreducible worst-case generalization floor.<n>We propose STEM-GNN, a pretrain-then-finetune framework with a mixture-of-experts encoder for diverse computation paths.<n>Across nine node, link, and graph benchmarks, STEM-GNN achieves a stronger three-way balance, improving robustness to degree/homophily shifts and to feature/edge corruptions while remaining competitive on clean graphs.
  arXiv  Detail & Related papers  (2026-02-09T22:48:30Z)
• Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts [32.65737144630759]
  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.
  arXiv  Detail & Related papers  (2026-01-05T14:16:11Z)
• Route-DETR: Pairwise Query Routing in Transformers for Object Detection [11.46025964297103]
  Detection Transformer (DETR) offers an end-to-end solution for object detection.<n>DETR suffers from inefficient query competition where multiple queries converge to similar positions.<n>We present Route-DETR, which addresses these issues through adaptive pairwise routing in decoder self-attention layers.
  arXiv  Detail & Related papers  (2025-12-15T20:26:58Z)
• Flash Multi-Head Feed-Forward Network [51.82159978122374]
  Multi-Head FFN (MH-FFN) is motivated by the structural similarity between single-head attention and FFN.<n>MH-FFN faces two challenges: memory consumption scaling with the head count, and an imbalanced ratio between the growing intermediate size and the fixed head dimension.<n>We propose Flash Multi-Head FFN (FlashMHF), with two key innovations: an I/O-aware fused kernel computing outputs online in akin to FlashAttention, and a design using dynamically weighted parallel sub-networks.
  arXiv  Detail & Related papers  (2025-12-07T20:50:20Z)
• Dr.LLM: Dynamic Layer Routing in LLMs [55.11953638340419]
  Dr.LLM is a retrofittable framework that equips pretrained models with lightweight per-layer routers deciding to skip, execute, or repeat a block.<n>On ARC (logic) and DART (math), Dr.LLM improves accuracy by up to +3.4%p while saving 5 layers per example on average.
  arXiv  Detail & Related papers  (2025-10-14T17:51:26Z)
• Optimizing MoE Routers: Design, Implementation, and Evaluation in Transformer Models [0.0]
  Mixture of Experts (MoE) architectures increase large language model scalability, yet their performance depends on the router module that moves tokens to specialized experts.<n>This work provides a comparative analysis of MoE router designs and offers insights into optimizing their performance for efficient and effective large-scale model deployment.
  arXiv  Detail & Related papers  (2025-06-19T15:55:43Z)
• Mixture of Routers [16.169900017745327]
  We propose an efficient fine-tuning method called Mixture of Routers (MoR)<n>MoR uses multiple sub-routers for joint selection and uses a learnable main router to determine the weights of the sub-routers.<n>Results show that MoR outperforms baseline models on most tasks, achieving an average performance improvement of 1%.
  arXiv  Detail & Related papers  (2025-03-30T08:39:09Z)
• Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural Networks [74.68583356645276]
  In deep learning, mixture-of-experts (MoE) activates one or few experts (sub-networks) on a per-sample or per-token basis. We show for the first time that pMoE provably reduces the required number of training samples to achieve desirable generalization.
  arXiv  Detail & Related papers  (2023-06-07T00:16:10Z)
• SMILE: Scaling Mixture-of-Experts with Efficient Bi-level Routing [47.11171833082974]
  We introduce SMILE, which exploits heterogeneous network bandwidth and splits a single-step routing into bi-level routing. Our experimental results show that the proposed method obtains a 2.5x speedup over Switch Transformer in terms of pretraining throughput on the Colossal Clean Crawled Corpus without losing any convergence speed.
  arXiv  Detail & Related papers  (2022-12-10T03:44:16Z)
• StableMoE: Stable Routing Strategy for Mixture of Experts [109.0602120199226]
  Mixture-of-Experts (MoE) technique can scale up the model size of Transformers with an affordable computational overhead. We propose StableMoE with two training stages to address the routing fluctuation problem. Results show that StableMoE outperforms existing MoE methods in terms of both convergence speed and performance.
  arXiv  Detail & Related papers  (2022-04-18T16:48:19Z)

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.