Sparsely gated tiny linear experts

Abstract: Sparsity allows scaling model parameters without proportionally increasing computational cost. While mixture of experts (MoE) models are made increasingly sparse, individual experts typically remain large and dense. Here, we demonstrate that further increasing sparsity by shrinking each expert to consist of a single neuron and selecting a tiny fraction of many available neurons can improve compute efficiency and interpretability. Counterintuitively, the key to achieving both is removing the nonlinearity typically applied to the experts, resulting in a network of sparsely gated linear neurons (sgatlin). In an isoflop comparison, we find that replacing all transformer feedforward layers with sgatlin improves perplexity in LLMs across different compute budgets. At the same time, the sparsity and linearity of the resulting feedforward circuits present new opportunities for model interpretability. In a small-scale case study, we demonstrate that feedforward circuits in sgatlin can be interpreted without having to train additional replacement models. We find that they form semantically structured clusters and are causally implicated in factual recall. Our findings paint a possible path towards compute-efficient and interpretable transformer feedforward layers.

• The paper introduces a sparsely gated linear neuron (sgatlin) architecture that replaces dense or nonlinear MoE layers with single linear neurons activated via top-k gating.
• It achieves high parameter efficiency by using only ≤0.1% of feedforward parameters per token, matching or exceeding MoE baselines in language modeling perplexity.
• The approach enhances interpretability by mapping sparse gating outputs to distinct semantic circuits, enabling causal analysis of factual recall.

Sparsely Gated Tiny Linear Experts: Summary and Analysis

Introduction and Motivation

This work introduces the "sparsely gated linear neuron" (sgatlin) feedforward layer as a compute-optimal and interpretable replacement for standard dense or nonlinear mixture-of-experts (MoE) feedforward stages in transformers. While traditional transformer feedforward layers scale model capacity but with a linear increase in computational cost, and MoEs leverage sparsity but still activate massive numbers of dense parameters per example, sgatlin pushes this paradigm further—reducing each expert to a single linear neuron and entirely omitting expert nonlinearity. This design enables a higher ratio of parameters per FLOP and simplifies interpretability, addressing two persistent barriers in transformer scaling: computational efficiency and model transparency.

Sgatlin Architecture

Sgatlin layers consist of a large array of linear neurons (“experts”) and a gating network that, for each input, computes a sparse top-$k$ selection of these neurons to activate (typically $k=8$). The gating function is realized via product-top-$k$ over query and key projections of the input, similar to product-key mechanisms used in prior memory-augmented models. Notably, sgatlin entirely removes nonlinearities: neither the gating output nor the neuron activations are passed through elementwise nonlinearity functions, resulting in a linear, low-rank affine circuit determined by the gating configuration.

By keeping the number of active neurons per input fixed at $k$ and the gating bottleneck small, the increase in computational complexity with respect to the total number of neurons is sublinear: $O(\sqrt{d_{\text{ffw}}})$, where $d_{\text{ffw}}$ is the number of available neurons. This enables scaling parameter count without commensurate FLOP cost, directly contrasting with dense MLPs or coarse-grained MoEs.

Empirical Evaluation

Compute-Optimal Language Modeling

Sgatlin layers were benchmarked as drop-in replacements for transformer feedforward blocks in decoder-only LLMs trained autoregressively on the 627B-token SlimPajama corpus. Isoflop comparisons (controlling total training FLOPs across model configurations) included dense GeLU MLP, SwiGLU, two established MoE variants (GPT-OSS and PEER), and sgatlin, with architectures ranging up to 4B parameters and varying compute budgets. Across all budgets, MoE variants consistently outperformed dense feedforward architectures in test perplexity, in line with established scaling laws.

Sgatlin matched or exceeded the perplexity of all MoE baselines at comparable or higher compute sparsity (defined as parameters per FLOP activated per token). Critically, only $\leq$0.1% of feedforward parameters are used per token, a marked improvement in activation efficiency. In ablation, introducing nonlinearity (ReLU, GeLU, Swish) to sgatlin decreased performance, evidencing that nonlinear processing is effectively produced by the top-$k$ gating structure itself, not elementwise activation.

Interpretability Studies

The use of linear experts and sparse gating creates an architecture where model computations at each position can be exactly specified by the sparse vector of gating weights, defining a unique low-rank circuit. The paper leverages this property in a suite of interpretability analyses on a small sgatlin-based model trained on TinyStories:

• Semantic Structure of Circuits: Gating weights across layers and inputs form distinct clusters in metric space (e.g., via UMAP projection), frequently corresponding to semantics—pronouns, names, and syntactic roles. This suggests that neuron subpopulations are reused as composable, interpretable circuits mapping to abstraction in the data.
• Neighbourhood Analysis: Nearest-neighbour search among the circuits reveals that similar gating configurations are systematically reused for semantically related or functionally analogous inputs, further supporting the hypothesis of emergent monosemantic circuits.
• Causal Intervention: Directly modifying gating weights (i.e., patching the feedforward circuit at specific tokens and layers with those from counterfactual prompts) demonstrates that these circuits are causally implicated in factual recall. The normalized indirect effect is most prominent when patching at noun positions and in earlier (second) layers, indicating that facts are retrieved upstream of the output token.

Theoretical and Practical Implications

Scaling and Efficiency

Sgatlin enables models to scale parameter count and thus representational potential without a proportional increase in memory and compute requirements. Only the selected neurons incur computational cost per example; thus, the overall wall-clock speed hinges on hardware support for highly sparse activation patterns. While the experiments use isoflop benchmarks, such sparsity may not yield similar wall-time gains on current dense-matrix-optimized accelerators, an acknowledged limitation.

Interpretability

Unlike prior interpretability approaches that rely on training surrogate (replacement) models (e.g., sparse autoencoders, transcoders), sgatlin provides intrinsic access to activations’ causal structure without any post hoc extraction. The explicit, sparse, and input-selective nature of its gating means that the function computed at each token is directly analyzable in terms of a small subset of parameters, allowing circuit analysis at scale. The observed formation of semantically coherent and causally effective subpopulations suggests a path toward models with enhanced monosemanticity and reduced “polysemantic” feature entanglement.

Comparison to Related Architectures

While PEER and other ultra-fine-grained MoEs have shrunk experts to a single neuron, they retain expert nonlinearities and use less efficient routing mechanisms. Sgatlin offers computational and memory complexity improvements by further simplifying expert computations to linearity and sharing query projections.

The conceptual decomposition of transformer computation into “circuit formation” (gating/attention selection) and “circuit application” (linear transformation) echoes the decomposition of attention mechanisms. When both stages are linearized, the end-to-end forward pass reduces to log-linear (affine) computations, facilitating theoretical analysis and potentially bridging empirical advances with learning theory in ways that dense, nonlinear models could not.

Limitations and Future Directions

Sgatlin only addresses the sparsity and interpretability of the feedforward component; attention mechanisms remain dense and opaque. Real-world training speedups will depend on future hardware/software tuned for sparse activation. All interpretability experiments were performed at small scale; next steps include extension to larger models and integration with deconstructions of attention and other modules.

The architecture opens several avenues for systematizing and scaling up mechanistic interpretability, such as tracking circuit provenance, information flow, and circuit reuse across transformer layers—the sparse gating basis provides a natural coordinate system for these analyses.

Conclusion

Sgatlin demonstrates that sparsely gated, purely linear experts can match and sometimes exceed traditional and MoE-based dense feedforward networks in autoregressive language modeling tasks while enabling parameter scaling with sublinear compute cost. More importantly, this architecture offers a transparent mapping from input to sparse, interpretable circuits, empirically structured by semantics and causally responsible for factual recall. Sgatlin thus provides a compelling foundation for scaling both the performance and interpretability of transformer architectures, motivating further theoretical and empirical investigation as the field pushes the limits of model scale and understanding.

---

Reference: “Sparsely gated tiny linear experts” (2606.07414)