MONET: Scalable Monosemantic Experts in Transformers

• MONET is a Transformer architecture that integrates sparse dictionary learning into MoE pretraining, addressing polysemanticity by specializing experts towards single concepts.
• It employs horizontal and vertical decomposition techniques to scale expert counts up to 262,144 per layer while keeping parameter growth sublinear.
• Monet enables precise knowledge manipulation such as domain-specific unlearning and toxicity reduction, maintaining robust overall performance.

Mixture of Monosemantic Experts for Transformers (Monet) is a Transformer architecture that integrates sparse dictionary learning directly into end-to-end Mixture-of-Experts (MoE) pretraining. Monet aims to address the polysemanticity problem in LLMs—the phenomenon where single neurons or units respond to multiple, unrelated concepts—by scaling the number of interpretable, functionally disentangled experts per layer to extreme counts (up to 262,144 per layer) while keeping total model parameters sublinear in the number of experts. This enables direct mechanistic interpretability at scale and enables targeted knowledge manipulation such as domain or toxicity unlearning with minimal collateral degradation in model performance (Park et al., 2024).

1. Monet Architecture: End-to-End Monosemantic MoE Pretraining

Monet is based on a standard decoder-only Transformer backbone (e.g., LLaMA-style) where each feed-forward (MLP) sublayer is replaced by a sparse mixture of N experts. Each expert is itself a small MLP, and per-token gating selects a small subset (top-$k$) of experts to process the hidden state at each layer: $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$ where $g_i(x)$ is the router softmax score, $\mathcal{K}(x)$ is the top-$k$ selected expert indices, and $E_i(x)$ denotes the $i$-th expert's MLP block. Unlike Sparse Autoencoders (SAEs), Monet places sparse dictionary learning inside the main model, so each expert is optimized jointly with the language modeling loss across pretraining. The overall Monet learning objective is: $$\mathcal{L} = \mathcal{L}_\mathrm{LM} + \lambda\,\mathcal{L}_\mathrm{unif} + \lambda\,\mathcal{L}_\mathrm{amb},$$ which combines standard LM loss, load-balancing, and monosemanticity (specialization) regularizers.

2. Expert Decomposition and Sublinear Parameter Scaling

A central innovation in Monet is the ability to scale the number of experts, $N$, per layer up to 262,144 with only $O(\sqrt{N})$ growth in parameters per layer. This is achieved via two factorizations:

• MoHDE (Horizontal Decomposition): Store only $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$0 "bottom" $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$1 and $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$2 "top" $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$3 matrices. Any expert is on-the-fly composed as $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$4 for $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$5, routing via product key retrieval.
• MoVDE (Vertical Decomposition): Each expert's hidden dimension $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$6 is split vertically, and composition from $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$7 half-matrices per block is used to build up to $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$8 experts.

Each token is routed to a pair $$\mathrm{MoE}(x) = \sum_{i\in\mathcal{K}(x)} g_i(x) E_i(x), \quad E_i(x) = V_i\,\sigma(U_ix),$$9, selecting one "bottom" and one "top" matrix, yielding a full expert with $g_i(x)$0 storage rather than $g_i(x)$1. Routing and compute complexity are similarly sublinear.

Expert Count $g_i(x)$2	SMoE Params	Monet Params
$g_i(x)$3	$g_i(x)$4	$g_i(x)$5

This factorization enables practical training and inference with extreme expert counts.

3. Training Objectives and Monosemanticity Promotion

Monet includes auxiliary objectives designed for both diversity and exclusivity of expert usage. These are critical for each expert to become near-monosemantic (i.e., respond to a single concept):

• Uniformity loss $g_i(x)$6 encourages even routing probability usage across both the set of bottom and top experts:

$g_i(x)$7

• Ambiguity loss $g_i(x)$8 encourages specialization (sharper distribution):

$g_i(x)$9

Domain specificity is measured by whether an expert $\mathcal{K}(x)$0 has mean routing score on a domain more than double that in other domains.

4. Empirical Evaluation: Specialization and Knowledge Editing

Monet models with 262,144 experts per layer and up to 4.1B parameters are shown to match or closely approach baseline LLaMA language modeling performance on standardized metrics (e.g., 0-shot/5-shot MMLU, ARC, HellaSwag), while promoting substantially higher monosemanticity compared to baseline MoE or post-hoc SAE approaches.

Domain, Language, and Toxicity Knowledge Masking

Monet supports targeted removal of experts to ablate entire knowledge domains with minimal impact on out-of-domain generalization:

• Domain masking: On MMLU, removing all biology-specific experts in Monet-1.4B drops biology accuracy by –4.66 points, versus widespread performance loss in LLaMA or OLMoE models.
• Language masking: On MULTI​PL-E code generation, removing the top Python-specialized experts (∼57k) reduces Python pass@100 from 31.64% to 1.06%, leaving other languages minimally affected (–0.6 to –1.8 points).
• Toxic expert removal: Experts whose activation is strongly correlated with Perspective-API toxicity can be purged, reducing toxicity-scores while maintaining helpfulness.

Monet’s routing structure thus enables direct, interpretable interventions for alignment and content moderation.

Dictionary Analysis

Qualitatively, Monet yields experts interpretable as single concepts (e.g., chemical formulas, specific regions, or domain-specific terms). For example, one expert routes on “Bayesian” contexts while its counterpart routes on “Bay Area” geography, illustrating disambiguation.

5. Innovations, Implications, and Limitations

Innovations:

• Extreme expert scaling: Factorized expert parameterization supports $\mathcal{K}(x)$1262k monosemantic units per layer without prohibitive compute or memory cost.
• End-to-end monosemanticity: Monosemantic experts are learned during pretraining, not via post-hoc reconstruction, preserving LLM performance.
• Efficient routing: Product-key routing and block-wise factorization enable high-throughput, hardware-parallel deployment.

Implications for interpretability and alignment:

• Fine-grained, interpretable expert routing underpins direct and transparent knowledge editing.
• Enables surgical removal (“unlearning”) or grafting (“lifelong addition”) of domain-specific knowledge.
• Strong candidate architecture for regulatory-compliant or mission-critical LLMs where transparency is mandatory.

Limitations:

• Current monosemanticity metrics (e.g., domain specificity, 2× skewness) are recognized as preliminary; robust quantitative benchmarks are needed.
• Selection of which experts to modify/remove relies on empirical routing heuristics rather than proved criteria.
• Most reported studies focus on unlearning; dynamic, continual learning (e.g., new expert addition post-training) remains under active research.

6. Broader Context and Future Directions

Monet represents a significant advance in the mechanistic interpretability of LLMs, directly addressing polysemanticity and enabling explicit, localized control of model knowledge and behavior. The approach bridges the gap between post-hoc sparse dictionary methods and scalable, performant Transformer pretraining. Extension to multimodal contexts (e.g., VisionMonet) is suggested as a next logical step. Development of improved, semantically grounded routing and expert selection criteria, and extension to continual or online knowledge addition, remain central open challenges (Park et al., 2024).