Mixture of Experts Made Intrinsically Interpretable

Neurons in large language models often exhibit \emph{polysemanticity}, simultaneously encoding multiple unrelated concepts and obscuring interpretability. Instead of relying on post-hoc methods, we present \textbf{MoE-X}, a mixture-of-experts (MoE) language model designed to be \emph{intrinsically} interpretable. Our approach is motivated by the observation that, in language models, wider networks with sparse activations are more likely to capture interpretable factors. however, directly training such large sparse networks is computationally prohibitive. MoE architectures offer a scalable alternative by activating only a subset of experts for any given input, inherently aligning with interpretability objectives. In MoE-X, we establish this connection by rewriting the MoE layer as an equivalent sparse, large MLP. This approach enables efficient scaling of the hidden size while maintaining sparsity. To further enhance interpretability, we enforce sparse activation within each expert and redesign the routing mechanism to prioritize experts with the highest activation sparsity. These designs ensure that only the most salient features are routed and processed by the experts. We evaluate MoE-X on chess and natural language tasks, showing that it achieves performance comparable to dense models while significantly improving interpretability. MoE-X achieves a perplexity better than GPT-2, with interpretability surpassing even sparse autoencoder (SAE)-based approaches.

AI That "Thinks" Clearer: New Model Aims to Make Language AI Easier to UnderstandImagine trying to follow a conversation where each word has a dozen different meanings, all jumbled together. That's a bit like how current AI language models, the brains behind chatbots and translation tools, can sometimes work internally. Their "neurons" – the tiny processing units in these AI brains – often handle multiple unrelated concepts at once. This makes it incredibly difficult for us to understand exactly how they arrive at an answer, a problem scientists call "polysemanticity."This research paper introduces a new AI language model called MoE-X that's designed to be understandable from the get-go. Instead of trying to decode a jumbled mess after the fact, MoE-X is built to think in a more organized way.The key idea is inspired by how a very wide, but sparsely used, network might be easier to interpret – think of a massive library where only a few specific books are pulled out for any given topic. However, building such enormous networks directly is too computationally expensive.MoE-X cleverly sidesteps this by using a "mixture-of-experts" (MoE) approach. Imagine a team of specialist consultants. When a question comes in, the system smartly routes it to only the most relevant one or two experts, rather than every expert trying to chime in on every topic. This keeps things focused.MoE-X takes this a step further:- It's structured so that this "team of experts" system is equivalent to a very large, but sparsely used, internal network, making it efficient to train.- It encourages each "expert" to activate only for very specific features or concepts.- The system that decides which "expert" sees which piece of information is redesigned to pick experts that are very selective about what they respond to.The result? MoE-X performs as well as other popular language models (even outperforming the well-known GPT-2 in some measures) on tasks like understanding chess moves and general language. Crucially, it's significantly easier to look inside MoE-X and see which "expert" is handling which concept, making its decision-making process much clearer than traditional models, and even better than other methods specifically designed to improve this kind of transparency.In essence, MoE-X is a step towards AI that not only provides answers but can also show its working in a more straightforward and understandable way.