Mixture of Wide Experts: Sparse Routing for Interpretable Native Coordinates

AI summary

Overview

Research area: Mechanistic interpretability and language-model architecture design — specifically sparse mixture-of-experts feed-forward layers and the interpretability of a model's native activation coordinates.

Technical level: Advanced. The paper assumes familiarity with transformer feed-forward blocks, top-k mixture-of-experts routing, residual streams, sparse autoencoders, and activation/path patching.

Scope: The paper proposes Mixture of Wide Experts (MoWE), a sparse FFN layer whose selected experts read the full residual stream but write to fixed, disjoint output coordinates, and tests whether this design yields readable native coordinates without sacrificing language-modeling quality.

What This Paper Is About

Language models tend to spread a single concept across many internal dimensions, and individual dimensions often respond to several unrelated concepts, which makes the model's internal features hard to interpret. Existing fixes such as sparse autoencoders recover clearer features only after training, by learning a new dictionary over the model's activations. This paper asks whether an architectural change during training — routing tokens to experts that each own a private slice of the output vector — can make the model's own coordinates more interpretable in the first place, and whether those coordinates can then be used directly to trace and steer predictions.

Key Contributions

1. A read-all, write-disjoint FFN layer (MoWE). A router selects k of G experts per token; every selected expert reads the complete d-dimensional input, but each expert's output is confined to its own fixed, non-overlapping slice of the FFN update. In the final language models, d = G = 768 and slice width s = 1, so each expert writes exactly one coordinate, and two different experts' contributions have strictly disjoint support (inner product zero).

2. A controlled test with known hidden causes. The authors build a synthetic task with 32 latent causes (grouped into eight domains) that independently generate a 256-bit observation and a 128-bit prediction target, and measure whether causal information concentrates in identifiable coordinates, four-coordinate expert slices, and slice ablations, using purity and effective-cause-count metrics.

3. A native-coordinate readability metric ("top-context purity" / specificity). A language-model judge describes a coordinate from passages where it is strongly active and tests the description on held-out passages, with a correction that subtracts the score of a description borrowed from another coordinate in the same model, layer, and representation type. The paper also compares native coordinates against a fixed random orthogonal rotation and against TopK sparse-autoencoder features evaluated at the same site.

4. Two case studies using native coordinates directly. One traces four MoWE coordinates supporting the prediction "degrees" after "A right angle is 90" and transfers seven coordinate values to change the prediction for "A rectangle has four" from "sides" to "degrees"; the other uses 15 native coordinates as a hazard detector plus a fixed seven-coordinate steering packet to shift next-token predictions toward refusal.

Main Findings

• Predictive quality is roughly preserved, but not best. At approximately matched active compute (19 layers, 10 billion tokens, two runs per architecture), MoWE reaches a mean evaluation cross-entropy of 2.417 versus 2.452 for the FLOP-equivalent dense model. Standard MoE reaches the lowest loss at 2.22, but uses 3.3% more active MACs and 32.6% more parameters than MoWE.

• Parameter cost is higher for MoWE. MoWE has 274.2M parameters versus 103.5M for dense and 363.7M for standard MoE, while active MACs per token are 133.4M, 133.2M, and 137.8M respectively.

• Native coordinates are more specific under disjoint writing. On the coordinate-specificity measure, MoWE scores 0.523, compared with 0.375 for dense and 0.268 for standard MoE.

• MoWE's native coordinates rival a post-hoc SAE on the dense model. Site-matched TopK SAEs (12,288 features, exactly 38 active features per example) score 0.489 on the dense model and 0.556 on MoWE. Native MoWE is therefore comparable to the dense SAE, and its own SAE improves further.

• The controlled task separates routing from localization. The authors distinguish routing mutual information (do examples with the same cause follow similar routes?) from output localization (does the cause land in identifiable coordinates or slices?). The paper reports that writing to disjoint coordinates preserves predictive performance while localizing the known causes more clearly than dense and standard MoE models; the specific BCE, accuracy, purity, and effective-cause-count values are not reported in the provided text.

• Single-prediction tracing works. The frozen atlas labels all 19 × 768 = 14,592 coordinates, of which GPT-5.6 Luna produces usable labels for 14,562 (30 are inactive or have insufficient evidence). Of the 722 active layer-coordinate sites on the angle prompt, the authors identify a four-coordinate path supporting "degrees," compare each link against same-layer alternatives at the 95th percentile, and compare the joint ablation against 1,000 random four-coordinate controls.

• Coordinates can be reused to change another prompt. Adding the seven largest positive angle-coordinate values to "A rectangle has four" changes the predicted next token from "sides" to "degrees"; the control uses 500 random coordinate sets matched in layer, sign, value, and scale.

• A sparse native-coordinate policy can refuse hazardous requests. On a test set of 100 hazardous and 100 benign requests built from matched templates, the complete MoWE policy produced refusal-family next tokens for 99 hazardous requests and nine benign controls, versus two hazardous and no benign controls for the corresponding dense policy. The detector reads 15 native coordinates, and the steering packet is a fixed seven-value activation pattern stored from one refusal example.

• Forward-throughput benchmarking was set up but its results are not reported in the provided text. The protocol uses one NVIDIA RTX A6000, 16 sequences of 1,024 tokens, PyTorch 2.4.1, CUDA 12.4, BF16 weights, six alternating rounds, and 80 timed replays per model per round, with CUDA Graph replay and grouped kernels for the selected experts.

Methodology in Plain English

The authors start from a simple observation about standard mixture-of-experts layers: routing decides which experts run, but every selected expert still writes a full-width vector into the same output coordinates, so experts can overwrite each other's contributions. MoWE keeps the routing but changes the writing rule. Each expert is a complete SwiGLU network with its own up and down projections (in the final models: 768 inputs, 8 hidden activations, and 1 output value), and the down projection is restricted to that expert's own coordinate. Because 768 experts each own one coordinate, any two active experts at the same layer occupy different parts of the FFN update.

To make comparisons fair, the authors match active computation rather than parameter count: they widen the dense FFN (m = 455) until it executes about as many multiply–accumulates per token as MoWE. A standard MoE with the same G, k, H, data, token budget, and schedule serves as a routing-matched control that isolates the effect of disjoint writing.

For interpretability, they collect, for each coordinate, the single most activating token from each of 1,024 unseen sequences. A judge model (GPT-5.6 Luna) writes a description from the top 20 contexts and scores it against the next 20 on a 0–10 scale. Because a vague description might fit anything, they subtract the score of a description shuffled in from another coordinate, producing a normalized specificity between 0 and 1. They also rotate the activation space randomly — which preserves information but destroys the learned axes — and train TopK SAEs on the exact same FFN-update vectors, so only the representation differs between comparisons.

The case studies follow label-then-intervene discipline. Coordinates are labeled and frozen before any example is chosen, so example selection cannot bias the descriptions. Causal relevance is tested by zeroing a coordinate's FFN-update value and rerunning the model, with same-layer alternatives as controls. Finally, the steering experiments are built entirely offline: a detector and a seven-value packet are fitted and frozen on validation requests, and only then applied to held-out test templates.

Why This Matters

Impact on research. The dominant interpretability pipeline treats the trained model as fixed and fits a post-hoc dictionary (an SAE) on top of it. MoWE tests a different hypothesis: that the choice of output basis during training changes how readable the model is in its own coordinates. If native coordinates are already close to SAE quality, that shortens the path from "inspect the model" to "act on the model," because no encoder/decoder translation is needed to steer. It also gives a clean ablation: standard MoE and MoWE share routing hyperparameters, so the difference between them isolates the effect of where experts write rather than which experts run.

Real-world applications:

• Safety monitoring and refusal steering. A small set of native coordinates plus a fixed activation packet can flag and redirect hazardous requests, as shown by the 15-coordinate detector and seven-value packet.
• Cheap behavioral auditing. Readable coordinates give auditors a handle for checking what a deployed model is sensitive to, without training a separate feature dictionary per checkpoint.
• Targeted model editing and debugging. Knowing which coordinates carry a concept allows precise edits (e.g., correcting a factual association) rather than blunt prompt engineering.
• Transparent model comparison. Coordinate-level specificity offers a way to compare architectures on interpretability as well as loss, which matters when choosing between dense, MoE, and MoWE backbones.

Industry relevance. Sparse MoE is already a standard way to scale parameter count without scaling per-token compute, so an architectural variant that keeps sparse routing while improving the readability of the FFN update is directly relevant to teams that deploy MoE models and need monitoring, red-teaming, or controllable behavior. The released OLM-native architecture, CUDA and Triton forward kernels, and the 19-layer MoWE, dense, and standard-MoE checkpoints lower the barrier to reproducing and extending the results.

Future Directions

• Close the loss gap to standard MoE. Standard MoE reaches 2.22 versus MoWE's 2.417 and dense's 2.452, so the interpretability benefit currently comes with a quality cost relative to the best-performing baseline. Whether more experts, wider per-expert hidden layers, or a different slice layout can recover it is untested here.
• Test whether the locality survives composition. The paper states explicitly that the disjointness guarantee holds only within one FFN layer; attention and later layers can recombine coordinates into distributed representations. Measuring how far down the network the localization persists is an open question.
• Move beyond single-token and toy-scale evidence. The case studies trace one prediction and one refusal behavior in 19-layer, 768-dimensional models. Whether the atlas-and-packet approach works at frontier scale, on multi-step reasoning, and without a fresh label pass per checkpoint is not established.
• Compare against and combine with stronger post-hoc methods. Native MoWE (0.523) sits between the dense SAE (0.489) and the MoWE SAE (0.556), which suggests the two approaches are complementary rather than substitutive. Studying SAEs trained on MoWE updates, and SAEs that exploit the disjoint structure, is a natural next step.

Target Audience

Interpretability researchers, MoE architecture researchers, and safety engineers who work with open-weight language models. Readers will get the most from the paper with working knowledge of transformer FFNs, top-k routing, and sparse autoencoders; the writing is clear enough that a strong graduate student in machine learning can follow the experimental logic without prior familiarity with the specific interpretability literature.

Authors’ abstract

Language models often spread a concept across many internal dimensions, making their features difficult to interpret. Sparse autoencoders (SAEs) recover clearer features after training. We ask whether models can instead learn interpretable output coordinates directly. We introduce Mixture of Wide Experts (MoWE), a sparse feed-forward layer in which selected experts read the full representation but write to fixed, disjoint output coordinates. On a controlled task with known hidden causes, this design preserves predictive performance while localizing causes more clearly than dense and standard mixture-of-experts (MoE) models. At approximately matched compute, we train two 19-layer runs of each architecture for 10 billion tokens. MoWE achieves evaluation cross-entropy 2.417 versus 2.452 for dense, while standard MoE reaches 2.22. We measure native-coordinate readability by asking a language-model judge to describe each coordinate from passages where it is strongly active and test that description on new passages. Correcting for descriptions that also fit other coordinates gives a specificity score of 0.523 for MoWE, 0.375 for dense, and 0.268 for standard MoE. TopK SAEs trained on the same FFN updates score 0.489 for dense and 0.556 for MoWE, making native MoWE comparable to the dense SAE. In the first case study, we identify a four-coordinate causal path supporting the model's prediction of degrees after the prompt A right angle is 90. We then copy seven angle-related coordinate values into the model as it processes A rectangle has four, changing its next-token prediction from sides to degrees. In the second, one fixed coordinate set detects hazardous requests and another steers the next token toward refusal. These examples show that MoWE combines sparse conditional computation with coordinates that can be identified and reused directly.