Skip to main content
eScholarship
Open Access Publications from the University of California

CausalMoE: Enhancing Expert Specialization and Interpretability through Cognitively-Inspired Causal Routing

Creative Commons 'BY' version 4.0 license
Abstract

Mixture-of-Experts (MoE) architectures enable scalable language modeling via sparse, input-dependent computation, but suffer from weak expert specialization and limited routing interpretability. From a cognitive perspective, this contrasts with theories of modular cognition that emphasize selective recruitment of functionally specialized systems. We propose CausalMoE, a causality-aware MoE framework that aligns expert activation with causal structure in language. CausalMoE introduces a specialized Causal Expert and a Causal Routing mechanism that identifies tokens along causally salient dependency structures inferred from attention patterns. The Causal Expert is specialized in processing sequences with strong causal and logical dependencies. To further encourage functional differentiation, we introduce a Causal Auxiliary Loss that reduces representational redundancy between experts. Integrated into canonical MoE architectures, including Switch Transformer and GLaM, CausalMoE achieves an average relative perplexity reduction of 18.42% on logic-intensive code datasets and a mean improvement of 6.83% on causal reasoning benchmarks, demonstrating structure-sensitive and interpretable computation allocation.