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

TL;DR

This paper identifies a pre-routing bottleneck in MoE Transformers caused by multi-head attention, which leads to catastrophic forgetting, and proposes MH-MoE with head-wise routing to mitigate this issue.

Contribution

The paper reveals how multi-head attention causes routing collisions in MoE Transformers and introduces MH-MoE, a head-wise routing method that reduces forgetting in continual learning.

Findings

MH-MoE reduces BWT from 11.2% to 4.5% on Qwen3-0.6B.

Higher effective composition number $N_{eff}$ correlates with increased forgetting.

Routing collisions are a key factor in catastrophic forgetting in MoE Transformers.

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%.