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

Mixture-of-Experts (MoE) architectures are appealing 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_{\mathrm{eff}}$ and find that higher $N_{\mathrm{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 across multiple backbones, MH-MoE consistently improves the retention--accuracy trade-off over LoRA-MoE variants.

Large language models are often updated over time so that they can learn new tasks or new knowledge, but this can cause them to forget what they previously learned. A popular way to reduce this problem is to divide the model into multiple "experts" so that different inputs can use different parts of the model. However, we find that this design still suffers from forgetting because the model often mixes several different types of information together before deciding which expert to use. This can send very different learning signals to the same expert, making old knowledge easier to overwrite. We propose MH-MoE, a new expert-routing method that makes decisions separately over different parts of the model’s internal representation. This gives the model a finer way to decide where each piece of information should go, reducing harmful interference between different kinds of knowledge. Experiments on a continual-learning benchmark show that MH-MoE better preserves earlier abilities while maintaining strong performance on later tasks. This work helps explain why expert-based language models can still forget, and offers a practical design for making them more reliable when they are updated over time.