Sparsely gated tiny linear experts

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Sparsely gated tiny linear experts

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

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 language models 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.1 MLP

SwiGLU

MoE

PEER

Compute budget 1e17 FLOPS

6e17 FLOPS

1e18 FLOPS

6e18 FLOPS

45 40

Test perplexity

arXiv:2606.07414v1 [cs.LG] 5 Jun 2026

Simon Schug Department of Computer Science Princeton University [email protected]

35 30 25 20

100μ 2

5 0.001 2

Compute sparsity

100μ 2

5 0.001 2

Compute sparsity

100μ 2

5 0.001 2

Compute sparsity

100μ 2

5 0.001 2

Compute sparsity

Figure 1: Sparsity improves compute efficiency. Comparing feedforward layers in transformerbased language models pretrained on SlimPajama 627B in a compute-optimal setting reveals that increasing compute sparsity – the ratio of parameters per FLOP – improves test perplexity across compute budgets. Our sparsely gated linear neuron layer (sgatlin) further increases this ratio. 1 Code available at https://github.com/smonsays/sparsely-gated-linear

Preprint.

Transformer block

Sparsely gated linear neurons

Gating network Gating weights

Product-top-k

sgatlin

Scores A

Norm

Linear neurons

Dense

Attention Norm

Scores B

Query

Gating network

Dense

Figure 2: Sparsely gated linear neurons layer. Left We replace the feedforward layer within each transformer block with a sgatlin layer. Center In sgatlin, a gating network selects a sparse subset of k linear neurons whose outputs are linearly combined. Right The gating network can efficiently select among a large number of neurons by combining a bottleneck with an efficient top-k operation over the product set of two sets of scores.

Introduction

Scale is a key ingredient for creating capable transformer-based models and has played an important part in their success. At the same time, it demands large amounts of computational resources and yields complex and opaque models that are difficult to understand. Reconciling scale with compute efficiency and model interpretability is a major challenge. Feedforward layers in particular consume a considerable amount of the total compute of transformer models as their scale increases. Since feedforward layers are thought to be the main repository of knowledge, increasing their number of parameters is an important scaling axis. However, scaling the number of parameters in classical dense feedforward layers leads to a proportional increase in computational cost. The large number of parameters also presents an obstacle for model interpretability, a challenge compounded by nonlinearities that preclude analysis using standard linear algebra tools. Sparsity can lower scale-induced computational cost and complexity: By activating only a subset of parameters, it reduces the compute requirements of parameter scaling and helps localize the subcircuits causally implicated in a given computation [1–4]. However, while mixture of experts (MoE) models are popular and increasingly sparse in current large-scale architectures, they continue to activate billions of parameters per token [5–10]. Here, we take a step towards scalable feedforward layers that further realize the potential of compute sparsity by drastically increasing the number of parameters per FLOP. In our biggest models, less than 0.1% of feedforward layer parameters are activated per token. We build on prior work by He [11] that demonstrated how equipping MoEs with an efficient routing mechanism allows scaling parameters at increasing sparsity with fine-grained experts. By removing the element-wise nonlinearity applied to the experts, eliminating expert parameter sharing across channels, and further decreasing the computational cost of the router, we obtain a layer of sparsely gated linear neurons (sgatlin) that improves language modelling perplexity in a compute-matched setting and opens avenues for model interpretability that do not rely on training additional replacement models. We summarize our main contributions in the following: • In Section 2, we present sgatlin, a new feedforward layer which sparsely gates a large number of linear neurons and can efficiently scale the number of parameters at sublinear time complexity. • In Section 3, we evaluate the language modelling perplexity of transformers using various feedforward architectures on the SlimPajama 627B dataset [12]. In an isoflop comparison with models of up to 4B parameters, we find that sgatlin performs competitively across compute budgets. • In Section 4, we conduct an interpretability study, demonstrating that the feedforward circuits of sgatlin form a semantically structured metric space and are causally implicated in factual recall. 2

Sparsely gated linear neurons

We now introduce our sparsely gated linear neurons layer (sgatlin), which we will use to replace the feedforward layer in each transformer block (see Figure 2 for an overview). It consists of two components, a gating network that creates a sparse vector of gating weights, and a large pool of linear neurons – the experts – that are linearly combined according to the gating weights. The gating network computes scores for all available neurons and selects the k largest ones, g = product-top-k(W key W query z), √ 2 dffw ×dkey

(1) √

where W query ∈ Rdkey ×dmodel and W key ∈ R . The product-top-k(x) : R2 dffw → Rdffw operation [13] selects the k√largest entries of the product set of two sets of scalars, obtained by slicing the input vector, x ∈ R2 dffw , into two parts of equal size, and returns a sparse vector of gating weights with exactly k non-zero entries, g ∈ Rdffw with ∥g∥0 = k. The gating weights are then used to linearly combine the output of a sparse subset of k linear neurons, X sgatlin(z) = gi wiout wiin⊤ z, (2) i:gi ̸=0

where wiin , wiout ∈ Rdmodel . The full sgatlin layer sums over C parallel channels, each of which has distinct sets of parameters except for a shared query projection. To keep the exposition in the main text concise, we defer the full equations for multiple channels to Appendix A. The resulting layer has several notable properties: Given the gating weights, the effective circuit processing the layer input is linear. Apart from the top-k selection, we do not apply any nonlinearities – neither to the gating weights nor to the neurons (see Section 3.2 for an ablation study). Given the gating weights, the effective circuit processing the layer input is purely linear. Specifically, each neuron i consists of a rank-one out in⊤ matrix, Wineuron each token is effectively processed by a (at most) rank-k matrix, P:= wi wi out, and circuit W (z) = i:gi ̸=0 gi wi wiin⊤ . Scaling the number of available neurons only incurs a sublinear increase in computational cost. Since the effective circuit uses a fixed number of neurons (we use k = 8 in all experiments), increasing the number of available neurons only affects the time complexity of the gating network. By keeping the key dimension small and constant (we use dkey = 128 in all experiments),√the time complexity of the gating network scales sublinearly in the number of available neurons, O( dffw ). A comparison of the time and memory complexities of the different feedforward layer variants studied here is shown in Table 2 in the appendix.

Language modelling

In this section, we evaluate decoder-only transformer models in a standard autoregressive language modelling task where we replace the feedforward layers in all transformer blocks. We will begin by comparing different feedforward layer choices in a compute optimal setting, followed by an ablation study on various architectural components in sgatlin. 3.1 Isoflop model comparison We evaluate how the choice of transformer feedforward layer affects language modelling performance in a compute-matched (isoflop) comparison. We consider four different compute budgets, {1 × 1017 , 6 × 1017 , 1 × 1018 , 6 × 1018 } FLOPs, for each of which we train models of multiple sizes on the SlimPajama text dataset consisting of 627B tokens [12]. We compare sgatlin to both dense feedforward layers, a standard GeLU MLP [14] and a SwiGLU layer [15], as well as two types of mixture of experts layers, the coarse-grained MoE layer used in GPT-OSS [6] and the very fine-grained PEER MoE layer [11]. To systematically vary model size for a given compute budget, we define a scaling ladder that varies both transformer depth (number of layers) and width (number of feedforward neurons, number of experts), which is detailed in Appendix B.2. We can then determine the number of training tokens that satisfy a particular compute budget based on the model type and its size. Throughout this section, we consider a context window size of 2048 tokens and a batch size of 128 sequences, resulting in 262 144 tokens per batch. We use causally masked multi-head attention in each transformer block, where we set the size of each attention head to 64 and choose the number 3

MLP

SwiGLU

MoE

PEER

Compute budget

Test perplexity

1e17 FLOPS

6e17 FLOPS

1e18 FLOPS

6e18 FLOPS

50 40 30 20 10M 2

5100M2

1B 2

Parameters

10M 2

5100M2

1B 2

Parameters

10M 2

5100M2

1B 2

Parameters

10M 2

5100M2

1B 2

Parameters

Figure 3: Sparsely gated linear neurons improve language modelling performance. Isoflop comparison of transformer-based language models trained on SlimPajama across four different compute budgets. We compare the choice of transformer feedforward layer type by training models of varying size for a given compute budget and report the resulting perplexities on the test set.

of heads according to the model dimension, which is scaled as a multiple of 128. We optimize the cross-entropy loss using the AdamW optimizer [16] with a warmup-stable-decay learning rate scheduler [17–19] that warms up the learning rate from zero for 1000 steps, followed by a stable learning rate of 1 × 10−3 and a final square-root decay over the last 20% of training tokens back to zero. The full set of hyperparameters and the procedure for selecting them are provided in Section B.1 of the appendix. Figure 3 compares the test perplexities of the different feedforward layer types after training on a given compute budget and model size. Consistent with prior work [e.g., 5, 20–22], we find that MoE-based models are more compute optimal than dense architectures. MoE-based models have higher degrees of compute sparsity, allocating more parameters per FLOP, and correspondingly larger model sizes tend to be optimal for a given compute budget. We plot the best performing model size per compute budget and feedforward type in Figure 1, illustrating the relationship between compute sparsity and final test perplexity. sgatlin has a particularly high compute sparsity and consistently achieves a competitive test perplexity across compute budgets. 3.2

Ablation study

The absence of a nonlinear activation function in Table 1: Model ablations. Neither adding a sgatlin might appear counterintuitive given that it nonlinearity to the neurons nor replacing the is a standard component in transformer feedforward gating network with the router from PEER layers. Given that MoE layers typically employ an [11] improves language modelling perforelementwise nonlinear activation function to the ex- mance of sgatlin. perts, we run an ablation study to verify how adding such an activation function in sgatlin affects lanAblation Test perplexity guage modelling performance on our highest com+ ReLU activation 21.0880 pute budget of 6 × 1018 FLOPs using 1.2B-parameter + GeLU activation 20.8445 models. In Table 1, we report the test perplexity of + Swish activation 20.6424 adding either a ReLU, GeLU, or Swish activation function to sgatlin and find that each choice leads ↔ PEER router 18.2964 to a decrease in modelling performance. These find∅ sgatlin 18.1910 ings suggest that the top-k gating by itself suffices to enable nonlinear processing in sgatlin. We additionally ablate our gating network and swap it with the router used in PEER, confirming that despite our gating networks’ reduced time and memory complexity, it leads to comparable performance.

Interpretability

The structure of sgatlin has interesting implications for model interpretability that we will explore in the following section. Central to our analysis is the observation that sgatlin linearly combines 4

1. Collect circuit database across example sequences

Feedforward circuits Gating weights are unique identiﬁers for the feedforward circuits they conﬁgure.

Gating weights

Each circuit is a low-rank linear transformation

Feedfoward circuit

Circuit similarity

sgatlin

Attention

sgatlin

Attention

sgatlin

Attention

[Once]

[upon]

[a]

[.]

2. Analyse circuits by nearest neighbour search

Feedforward circuits live in a metric space induced by the similarity of their corresponding gating weights.

0.6

0.1

sgatlin Attention

Circuit database sgatlin

Nearest neighbor circuits

Attention

1 2

sgatlin Attention

Cosine distance Circuit A

Nearest neighbour search

sgatlin Attention

Circuit B

Circuit C

The [red] ball ...

4 5

... and Lilly took the red ball ...

... Tim looked at the red collar ...

... and Lilly took the green doll ...

... played with the yellow stick ...

... and Lilly took the big box ...

Figure 4: Feedforward circuit analysis. The effective feedforward circuits that process information at a given position in sgatlin are uniquely determined by their gating weights. As a result, the similarity between two circuits can be calculated from the distance between their gating weights. We use this property to analyse a given feedforward circuit by comparing it to similar circuits. For this purpose, we collect a circuit database over a set of reference sequences and retrieve the nearest neighbours from this database to interpret the role of a new circuit. a small number of linear neurons, effectively applying a low-rank linear circuit for each position in the sequence. Do these feedforward circuits share reusable and possibly interpretable structure? Instrumental to answering this question is the insight that each circuit is uniquely identified by the sparse gating weights that created it. We can therefore analyse the set of feedforward circuits in terms of the structure of the metric space of the corresponding gating weights. To keep this analysis tractable, we will consider a small, sgatlin-based language model trained on the TinyStories corpus [23], a collection of short stories for small children which use simple language. Despite the simplicity of this dataset, models trained on it have been shown to produce coherent and diverse stories with proper grammatical and semantic structure [23]. We use the same training pipeline as for our larger models above, albeit with a small, custom tokenizer with a vocabulary size of 8192, to train a language model with four layers and a context window size of 1024 tokens. The full training details can be found in Appendix B.1.2. 4.1 Neuron subpopulations form reusable feedforward circuits We begin by exploring whether the feedforward circuits activated across novel sequences are reused and form semantically meaningful structure. For this purpose, we collect the gating weights at each layer as well as the corresponding input tokens for a random subset of 128 unseen stories from the TinyStories validation set. Each gating weight is a dffw -dimensional vector with k non-zero entries. To visualize the collection of gating weights, we embed them into a two-dimensional space using UMAP [24] equipped with the cosine distance. The result of this procedure is shown in Figure 5 for the second feedforward layer, filtered by the 100 most frequent input tokens in the random subset of unseen stories. A corresponding plot for all four feedforward layers of the model is shown in Figure 9 in the appendix. Intriguingly, the resulting embeddings have discernible structure with distinct clusters of varying size. Some – but not all – of the structure is related to the input tokens that activated the corresponding circuits. To illustrate this point, we partially colour points in the two-dimensional embedding according to the input tokens. The colours are chosen to group semantically related groups of input tokens. This reveals, for instance, that the larger clusters in 5

UMAP embedding of gating weights Input tokens Other

Be are is was were

Names Ben Lily Max Timmy Tom

Pronouns he her his it she their they

Punctuation ! !" " , ," . ." ?" Ċ ĠĊ

Figure 5: Feedforward circuits form semantic clusters. We collect the gating weights of sgatlin of the second layer of a language model trained on TinyStories for the 100 most frequent tokens and embed them into 2D space using UMAP. The resulting embedding forms clusters that are partially explained by the semantics of the corresponding input tokens. For instance, names used in the stories cluster together and pronouns form clusters in their vicinity. An extended version of this figure showing the remaining layers can be found in Figure 9. the outer perimeter of the plot are made up of circuits that activate for different punctuation tokens. The different tenses of the verb "to be" are similarly grouped together. Particularly interesting is that names of the characters in the stories group together in the vicinity of pronouns. If we track these groups across the different layers in Figure 9, we find that after the second layer, pronouns and names are processed with similar circuits, suggesting that the corresponding circuits operate on the semantic role of these tokens. Overall, these findings provide some indication that the gating network systematically combines neuron subpopulations into reusable circuits. This motivates our next analysis. 4.2

Neighbouring feedforward circuits perform semantically related functions

Next, we investigate the extent to which a particular feedforward circuit can be understood by relating it to similar circuits. The overall procedure is outlined in Figure 4. Once again, we exploit the fact that the gating weights corresponding to each circuit form a metric space. This allows us to measure the similarity of two circuits by computing the cosine similarity between their corresponding gating weights. We collect the gating weights corresponding to each input token from unseen sequences in the TinyStories validation set and store them in a database alongside the corresponding input sequence and top output token predictions. For a novel sequence and its feedforward circuits, we can then search for the closest circuits in the database. Through this procedure, we aim to establish whether similar circuits have been observed before and whether the circuit is activated by semantically related input tokens or elicits related output predictions. 6

big

dog named Max

Max

had

red

collar that

1. Neighbor

big

dog

named

Max

Mia

had

red

dress

that

wore every of

every

castle

day

2. Neighbor

big

dog

named

Spot

Max

had

red

shirt

that

she

could

every

face

Circuit distance

3. Neighbor

big

dog

named

Max

Spot

had

red

doors

that

every

pie

0.6

4. Neighbor

big

dog

named

Jack

Amy

had

red

wheel

that

who

looks

every

5. Neighbor

old

bunny

named

Max

had

purple

shield

that

Every wardrobe

6. Neighbor

old

kitten

named

Amy

she

had

big

swing

that

named

every

7. Neighbor

eggs

cat

named

Mittens

Tim

had

big

label

that

who

had

every

party

8. Neighbor

long

kitten

named

Mia

had

big

zipper

that

wear

any

snacks

9. Neighbor

yard

mouse

was

Kitty

she

had

big

fur

that

Max

little

cheek

10. Neighbor

little

robot

You

Jack

Timmy

had

big

ball

that

they

Every

juice

0.4

castle 0.2

Figure 6: Feedforward circuit neighbours. Following the procedure shown in Figure 4 of collecting reference gating weights across input sequences, we can analyse the feedforward circuits for a new input sequence by comparing them to their nearest neighbours among the reference set. In this example, we analyse the feedforward circuits at the penultimate layer for the new input sentence excerpt "... big dog named Max. Max had a red collar that he wore every day". Each column below a respective input token shows the nearest neighbours’ input tokens. The colour encodes the distance of a particular neighbour’s gating weight to the gating weight of the corresponding input token. For example, the token "he" in our new sentence excerpt has a close neighbouring gating weight that also activated for the input token "he", followed by a neighbour that activated for the input token "she". We show the context input tokens as well as the top output predictions for all neighbours in Table 3. Figure 6 shows the input tokens of the ten nearest neighbour circuits of the penultimate layer, along with their distance for a reference sequence. The extended Table 3 in the appendix additionally shows input excerpts and output predictions. We deliberately choose the penultimate feedforward layer here, where we expect the model to perform more abstract computations that go beyond the immediate input or output features. Some of the reference tokens elicit highly stereotyped circuits that simply activate for exactly the same token across all neighbours. Intriguingly, however, for several of the tokens in the reference sequence, neighbour circuits activate in different but semantically similar contexts. For instance, among the neighbours of the feedforward circuit activated for the dog token are circuits that activate for other animals. Consistent with what we observed in the feedforward circuit embedding in Figure 5, the feedforward circuit activating for the name "Max" is very similar to circuits that activate for other names. For some reference tokens, like the "wore" token, the closest neighbours found in the database have a comparably high distance and accordingly bear no apparent semantic relation.

Normalized indirect eﬀect

4.3 Feedforward circuits are causally implicated in factual recall Patched position: 2 Patched position: 3 Finally, we examine how causal interventions 0.5 on the feedforward circuits impact model predictions. The fact that sgatlin explicitly separates 0.4 the circuit formation from its application gives 0.3 us the chance to delineate a circuit’s function from the intermediate data it produces. Gener0.2 ally speaking, the feedforward layers in transformers are believed to be the primary store of 0.1 a model’s knowledge about the world. We can use feedforward circuit interventions to assess 0 where such factual knowledge is represented in All layers 1 2 3 4 All layers 1 2 3 4 the feedforward circuits. Patched layer(s) Patched layer(s) To ensure the knowledge we probe for exists in TinyStories and we can reasonably expect the model to have learned it, we ex- Figure 7: Causal intervention on feedforward tract frequently co-occurring noun-adjective circuits. Normalized indirect effect of patchpairs that match the syntactic template "The ing layers at different depth and position in the [noun] was [adjective]" (e.g., "The turtle sequence on counterfactual knowledge pairs exwas slow") from the TinyStories training set. We tracted from the TinyStories dataset. select the largest set of noun-adjective pairs that ensures no two nouns have the same target adjective. For the causal intervention, we then create pairs of prompts "The [noun] was" to serve as clean 7

and counterfactual inputs (xclean , xpatch ) with divergent nouns and target adjectives. We can intervene on the forward pass of the clean input by patching the feedforward circuit with the gating weights from the counterfactual input at different token positions and layers in the model. We quantify the influence of a feedforward circuit, identified by the gating weights g, with respect to a pair of input sequences, xclean and xpatch , using the normalized indirect effect [25]. Given the corresponding clean and counterfactual target tokens, yclean and ypatch , we measure the logit difference for an input x as m(x) = logit(yclean | x) − logit(ypatch | x). The normalized indirect effect can then be computed as NIE(m; g; xclean , xpatch ) =

m(xclean | do(g = gpatch )) − m(xclean ) , m(xpatch ) − m(xclean )

where do(g = gpatch ) denotes the do-intervention on the underlying computational graph, corresponding to clamping the gating weights at specific layers and token positions to the counterfactual gating weights. We report the normalized indirect effect on the logit difference between the clean and counterfactual adjectives in Figure 7. The overall effects we observe are only of moderate magnitude, potentially reflecting that the simple factoids probed for are not only encoded in the feedforward layers but also in other parts of the network such as the input/output embeddings. Nevertheless, we find that patching the feedforward circuits at the position of the noun has a stronger impact on the normalized indirect effect than patching at the final token prior to the prediction, suggesting that associated factoids are retrieved right when the noun is encountered. When patching individual layers, the effect is strongest for the second layer of the network, although the outsized effect of patching all layers indicates that feedforward circuits across layers work in conjunction to retrieve a particular factoid.

Related work

5.1 Feedforward layers The feedforward network architectures we have considered here can be categorized as dense, where all parameters are activated in each forward pass, or mixture of experts (MoE), where only a subset of parameters are active per token. Interestingly, sgatlin can be conceptually related to both. The original transformer uses dense ReLU multilayer perceptrons [26], which implicitly gate the subset of non-negative neurons. While many transformers replace the ReLU activation with a GeLU due to better modelling performance [27], the activation sparsity implicitly induced by ReLU might be useful to improve inference efficiency [28]. We can think of sgatlin as a multilayer perceptron with a special kind of input-dependent nonlinearity that guarantees a fixed activation sparsity. A similar comparison can be drawn to the dense SwiGLU feedforward layer [15], in which a second linear input transformation with an ensuing nonlinearity (e.g. sigmoid) is used to softly gate the intermediate neurons via a component-wise product [29]. In a similar vein, sgatlin can be interpreted as hard gating via a component-wise binary mask of fixed sparsity. As outlined in Section 2, sgatlin can be equally motivated from the perspective of mixture of experts (MoE) layers that select a subset of experts, each of which is a full feedforward network. Classically, such MoE layers were coarse-grained with large experts, replicating dense layers [20]. More recently, there has been a trend towards fine-grained MoEs with many, smaller experts [2]. PEER takes an extreme stance in terms of granularity, making experts as small as possible, with each expert consisting of a single neuron [11]. sgatlin is closely related to PEER, albeit with several important architectural changes: In contrast to PEER, sgatlin uses distinct expert parameters for each channel and removes the activation function applied to the experts as well as the softmax nonlinearity applied to the gating scores. 5.2 Interpretability Understanding the internal mechanisms of large-scale transformer-based models is a pressing area of active research given their widespread deployment. The prevailing approach attempts to automatically extract sparse, interpretable circuits that faithfully characterize a particular behaviour [30–32]. Typically, this involves training replacement models that substitute components of the original model with more interpretable counterparts. Replacement models such as sparse autoencoders, transcoders, or crosscoders emulate the original model using high-dimensional and sparse representations [33–35]. Their shared goal is to decompose the likely polysemantic features of the original model into more easily interpretable, monosemantic features [36]. However, in practice, training such models is costly, surrogate features often only imperfectly approximate the original model [37–39], and sparsity 8

constraints can be hard to reconcile with hierarchical structure, leading to uninterpretable features [36, 40]. An alternative line of work therefore attempts to identify interpretable structure directly in the model under investigation without relying on a replacement model [41]. Arora et al. [4] demonstrate that gradient-based attribution can in several instances successfully extract sparse circuits directly from dense feedforward layers. Herbst et al. [42] show that as MoEs become more sparsely activated, expert neurons tend to become less polysemantic. These findings suggest that it may be possible to modify the model architecture itself in order to improve its interpretability. While the idea of modifying the architecture to improve interpretability is not new, and several attempts have been made to this extent [e.g., 43, 44], such modifications typically lead to a loss in modelling performance or scalability and thus struggle to be adopted.

Discussion

We have introduced sgatlin, a new feedforward architecture that can serve as a simple drop-in replacement for the feedforward layer in transformers to improve compute efficiency through high levels of sparsity. In an isoflop comparison, we found competitive language modelling performance across scales and validated key architectural decisions in an ablation study. The sparse and linear feedforward circuits configured by sgatlin offer interesting avenues for model interpretability, which we explored in a small-scale analysis that revealed that neuron subpopulations form reusable feedforward circuits that live in a semantically structured metric space and are causally implicated in factual recall. Beyond interpreting feedforward circuits, replacing the feedforward layers with sgatlin might enable new approaches for understanding the inner workings of transformers as a whole. Similar to how the attention mechanism can be decomposed into a key-query circuit that computes the attention pattern and an output-value circuit that applies a linear transformation to the inputs [45], sgatlin can be decomposed into the gating circuit that configures a linear feedforward circuit, which in turn processes the inputs. This means that conditioned on a particular input’s fixed gating weights, attention patterns and layer normalization statistics, the full transformer forward pass becomes effectively a single, large affine transformation followed by the final softmax (a log-linear model). Linear networks more generally and linearized transformer models in particular have served as important analytical devices and advanced deep learning theory, but ultimately have to concede with simplifying assumptions [46–48]. The decomposition of the transformer into circuit formation and circuit application that can be applied both to the attention mechanism and to sgatlin might open new directions for bridging the gap between theory and practice. Limitations Here, we focused on compute sparsity of the feedforward layers in transformer networks, leaving other dense components – like the attention mechanism – unchanged, despite their equally important role in terms of compute efficiency and interpretability. Our model comparison is conducted in an isoflop setting with matched FLOP budgets. While we believe this is a fair principle for the purpose of our study, it does not directly translate to practical wall clock times during training on current hardware accelerators that are highly optimized for compute dense workloads. Finally, our interpretability study is conducted on a comparably small dataset and model to provide a controlled setting that allowed us to develop novel techniques for analysing our feedforward layers. An important next step is to validate these techniques on larger models and integrate them with the analysis of other model components such as the attention mechanism. Broader impacts This paper suggests a model architecture that has the potential to improve compute efficiency and interpretability of transformers. While we foresee no immediate negative societal impact, we hope that it may improve our understanding of this widely deployed technology. Acknowledgments We are grateful to Xu Owen He for insightful conversations and feedback early on in the project. In addition, we would like to thank Yassir Akram, Marc Kaufmann, Nicolas Zucchet, and Seijin Kobayashi for valuable discussions. Simon Schug was supported by Postdoc.Mobility grant P500PT_225369 from the Swiss National Science Foundation. This work was supported with 10 000 GPU hours on the Alps supercomputing infrastructure by the Swiss National Supercomputing Centre under project ID a118 as part of the Swiss AI Initiative. 9

References [1] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, January 2017. URL http://arxiv.org/abs/1701.06538. arXiv:1701.06538 [cs]. [2] Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michał Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóźdź, Piotr Sankowski, Marek Cygan, and Sebastian Jaszczur. Scaling Laws for Fine-Grained Mixture of Experts, February 2024. URL http://arxiv.org/abs/2402.07871. arXiv:2402.07871 [cs]. [3] Andres Potapczynski, Shikai Qiu, Marc Finzi, Christopher Ferri, Zixi Chen, Micah Goldblum, Bayan Bruss, Christopher De Sa, and Andrew Gordon Wilson. Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices, October 2024. URL http://arxiv. org/abs/2410.02117. arXiv:2410.02117 [cs]. [4] Aryaman Arora, Zhengxuan Wu, Jacob Steinhardt, and Sarah Schwettmann. Language Model Circuits Are Sparse in the Neuron Basis, January 2026. URL http://arxiv.org/abs/2601. 22594. arXiv:2601.22594 [cs]. [5] Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. OLMoE: Open Mixture-of-Experts Language Models, September 2024. URL https://arxiv.org/abs/2409.02060v2. [6] OpenAI, Haiming Bao, Lama Ahmad, Jason Ai, Sam Altman, Andy Applebaum, Edwin Arbus, Rahul K. Arora, Yu Bai, Bowen Baker, Haiming Bao, Boaz Barak, Ally Bennett, Tyler Bertao, Nivedita Brett, Eugene Brevdo, Greg Brockman, Sebastien Bubeck, Che Chang, Kai Chen, Mark Chen, Enoch Cheung, Aidan Clark, Dan Cook, Marat Dukhan, Casey Dvorak, Kevin Fives, Vlad Fomenko, Timur Garipov, Kristian Georgiev, Mia Glaese, Tarun Gogineni, Adam Goucher, Lukas Gross, Katia Gil Guzman, John Hallman, Jackie Hehir, Johannes Heidecke, Alec Helyar, Haitang Hu, Romain Huet, Jacob Huh, Saachi Jain, Zach Johnson, Chris Koch, Irina Kofman, Dominik Kundel, Jason Kwon, Volodymyr Kyrylov, Elaine Ya Le, Guillaume Leclerc, James Park Lennon, Scott Lessans, Mario Lezcano-Casado, Yuanzhi Li, Zhuohan Li, Ji Lin, Jordan Liss, Lily, Liu, Jiancheng Liu, Kevin Lu, Chris Lu, Zoran Martinovic, Lindsay McCallum, Josh McGrath, Scott McKinney, Aidan McLaughlin, Song Mei, Steve Mostovoy, Tong Mu, Gideon Myles, Alexander Neitz, Alex Nichol, Jakub Pachocki, Alex Paino, Dana Palmie, Ashley Pantuliano, Giambattista Parascandolo, Jongsoo Park, Leher Pathak, Carolina Paz, Ludovic Peran, Dmitry Pimenov, Michelle Pokrass, Elizabeth Proehl, Huida Qiu, Gaby Raila, Filippo Raso, Hongyu Ren, Kimmy Richardson, David Robinson, Bob Rotsted, Hadi Salman, Suvansh Sanjeev, Max Schwarzer, D. Sculley, Harshit Sikchi, Kendal Simon, Karan Singhal, Yang Song, Dane Stuckey, Zhiqing Sun, Philippe Tillet, Sam Toizer, Foivos Tsimpourlas, Nikhil Vyas, Eric Wallace, Xin Wang, Miles Wang, Olivia Watkins, Kevin Weil, Amy Wendling, Kevin Whinnery, Cedric Whitney, Hannah Wong, Lin Yang, Yu Yang, Michihiro Yasunaga, Kristen Ying, Wojciech Zaremba, Wenting Zhan, Cyril Zhang, Brian Zhang, Eddie Zhang, and Shengjia Zhao. gpt-oss-120b & gpt-oss-20b Model Card, 2025. URL https://arxiv.org/abs/2508.10925. _eprint: 2508.10925. [7] Google. Welcome Gemma 4: Frontier multimodal intelligence on device, April 2026. URL https://huggingface.co/blog/gemma4. [8] Qwen Team. Qwen3-Next: Towards Ultimate Training & Inference Efficiency, September 2025. URL https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct. [9] GLM-5-Team, Aohan Zeng, Xin Lv, Zhenyu Hou, Zhengxiao Du, Qinkai Zheng, Bin Chen, Da Yin, Chendi Ge, Chenghua Huang, Chengxing Xie, Chenzheng Zhu, Congfeng Yin, Cunxiang Wang, Gengzheng Pan, Hao Zeng, Haoke Zhang, Haoran Wang, Huilong Chen, Jiajie Zhang, Jian Jiao, Jiaqi Guo, Jingsen Wang, Jingzhao Du, Jinzhu Wu, Kedong Wang, Lei Li, Lin Fan, Lucen Zhong, Mingdao Liu, Mingming Zhao, Pengfan Du, Qian Dong, Rui Lu, Shuang-Li, Shulin Cao, Song Liu, Ting Jiang, Xiaodong Chen, Xiaohan Zhang, Xuancheng 10

Huang, Xuezhen Dong, Yabo Xu, Yao Wei, Yifan An, Yilin Niu, Yitong Zhu, Yuanhao Wen, Yukuo Cen, Yushi Bai, Zhongpei Qiao, Zihan Wang, Zikang Wang, Zilin Zhu, Ziqiang Liu, Zixuan Li, Bojie Wang, Bosi Wen, Can Huang, Changpeng Cai, Chao Yu, Chen Li, Chengwei Hu, Chenhui Zhang, Dan Zhang, Daoyan Lin, Dayong Yang, Di Wang, Ding Ai, Erle Zhu, Fangzhou Yi, Feiyu Chen, Guohong Wen, Hailong Sun, Haisha Zhao, Haiyi Hu, Hanchen Zhang, Hanrui Liu, Hanyu Zhang, Hao Peng, Hao Tai, Haobo Zhang, He Liu, Hongwei Wang, Hongxi Yan, Hongyu Ge, Huan Liu, Huanpeng Chu, Jia’ni Zhao, Jiachen Wang, Jiajing Zhao, Jiamin Ren, Jiapeng Wang, Jiaxin Zhang, Jiayi Gui, Jiayue Zhao, Jijie Li, Jing An, Jing Li, Jingwei Yuan, Jinhua Du, Jinxin Liu, Junkai Zhi, Junwen Duan, Kaiyue Zhou, Kangjian Wei, Ke Wang, Keyun Luo, Laiqiang Zhang, Leigang Sha, Liang Xu, Lindong Wu, Lintao Ding, Lu Chen, Minghao Li, Nianyi Lin, Pan Ta, Qiang Zou, Rongjun Song, Ruiqi Yang, Shangqing Tu, Shangtong Yang, Shaoxiang Wu, Shengyan Zhang, Shijie Li, Shuang Li, Shuyi Fan, Wei Qin, Wei Tian, Weining Zhang, Wenbo Yu, Wenjie Liang, Xiang Kuang, Xiangmeng Cheng, Xiangyang Li, Xiaoquan Yan, Xiaowei Hu, Xiaoying Ling, Xing Fan, Xingye Xia, Xinyuan Zhang, Xinze Zhang, Xirui Pan, Xu Zou, Xunkai Zhang, Yadi Liu, Yandong Wu, Yanfu Li, Yidong Wang, Yifan Zhu, Yijun Tan, Yilin Zhou, Yiming Pan, Ying Zhang, Yinpei Su, Yipeng Geng, Yong Yan, Yonglin Tan, Yuean Bi, Yuhan Shen, Yuhao Yang, Yujiang Li, Yunan Liu, Yunqing Wang, Yuntao Li, Yurong Wu, Yutao Zhang, Yuxi Duan, Yuxuan Zhang, Zezhen Liu, Zhengtao Jiang, Zhenhe Yan, Zheyu Zhang, Zhixiang Wei, Zhuo Chen, Zhuoer Feng, Zijun Yao, Ziwei Chai, Ziyuan Wang, Zuzhou Zhang, Bin Xu, Minlie Huang, Hongning Wang, Juanzi Li, Yuxiao Dong, and Jie Tang. GLM-5: from Vibe Coding to Agentic Engineering, February 2026. URL http://arxiv.org/abs/2602.15763. arXiv:2602.15763 [cs]. [10] DeepSeek-AI. DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence, 2026. URL https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash. [11] Xu Owen He. Mixture of A Million Experts, July 2024. URL http://arxiv.org/abs/2407. 04153. arXiv:2407.04153. [12] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, June 2023. URL https://huggingface.co/datasets/cerebras/SlimPajama-627B. [13] Guillaume Lample, Alexandre Sablayrolles, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large Memory Layers with Product Keys, December 2019. URL http://arxiv. org/abs/1907.05242. arXiv:1907.05242. [14] Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), 2023. URL https://arxiv.org/abs/1606.08415. _eprint: 1606.08415. [15] Noam Shazeer. GLU Variants Improve Transformer, 2020. URL https://arxiv.org/abs/ 2002.05202. arXiv: 2002.05202. [16] Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization, January 2019. URL http://arxiv.org/abs/1711.05101. arXiv:1711.05101 [cs, math]. [17] Shengding Hu, Yuge Tu, Xu Han, Ganqu Cui, Chaoqun He, Weilin Zhao, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Xinrong Zhang, Zhen Leng Thai, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. August 2024. URL https://openreview.net/ forum?id=3X2L2TFr0f. [18] Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations. November 2024. URL https://openreview.net/forum?id=Y13gSfTjGr. [19] Kaiyue Wen, Zhiyuan Li, Jason S. Wang, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape View. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=m51BgoqvbP. 11

[20] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. ISSN 1533-7928. URL http://jmlr.org/papers/v23/21-0998. html. [21] Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giridharan Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O’Horo, Jeffrey Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, and Veselin Stoyanov. Efficient Large Scale Language Modeling with Mixtures of Experts. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 11699–11732, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.emnlp-main.804. URL https://aclanthology.org/2022.emnlp-main.804/. [22] Xi Victoria Lin, Akshat Shrivastava, Liang Luo, Srinivasan Iyer, Mike Lewis, Gargi Ghosh, Luke Zettlemoyer, and Armen Aghajanyan. MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts, July 2024. URL https://arxiv.org/abs/2407.21770v3. [23] Ronen Eldan and Yuanzhi Li. TinyStories: How Small Can Language Models Be and Still Speak Coherent English?, May 2023. URL http://arxiv.org/abs/2305.07759. arXiv:2305.07759 [cs]. [24] Tim Sainburg, Leland McInnes, and Timothy Q Gentner. Parametric UMAP Embeddings for Representation and Semisupervised Learning. Neural Computation, 33(11):2881–2907, 2021. [25] Judea Pearl. Direct and indirect effects. In Proceedings of the Seventeenth conference on Uncertainty in artificial intelligence, UAI’01, pages 411–420, San Francisco, CA, USA, August 2001. Morgan Kaufmann Publishers Inc. ISBN 978-1-55860-800-9. URL https://dl.acm. org/doi/10.5555/2074022.2074073. [26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, August 2023. URL http: //arxiv.org/abs/1706.03762. arXiv:1706.03762 [cs]. [27] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, and others. Improving language understanding by generative pre-training, 2018. [28] Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C. Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models, 2023. URL https://arxiv.org/abs/2310. 04564. [29] Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language Modeling with Gated Convolutional Networks. In Proceedings of the 34th International Conference on Machine Learning, pages 933–941. PMLR, July 2017. URL https://proceedings.mlr. press/v70/dauphin17a.html. [30] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small, November 2022. URL http://arxiv.org/abs/2211.00593. arXiv:2211.00593 [cs]. [31] Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models. October 2024. URL https://openreview.net/forum?id=I4e82CIDxv. [32] Emmanuel Ameisen, Jack Lindsey, Adam Pearce, Wes Gurnee, Nicholas L. Turner, Brian Chen, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Marcus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, Andrew Persic, Zhenyi Qi, T. Ben Thompson, Sam Zimmerman, Kelley Rivoire, Thomas Conerly, Chris Olah, and Joshua Batson. Circuit Tracing: Revealing Computational Graphs in Language Models, 2025. 12

[33] Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse Autoencoders Find Highly Interpretable Features in Language Models, October 2023. URL http://arxiv.org/abs/2309.08600. arXiv:2309.08600 [cs]. [34] Jacob Dunefsky, Philippe Chlenski, and Neel Nanda. Transcoders Find Interpretable LLM Feature Circuits, November 2024. URL http://arxiv.org/abs/2406.11944. arXiv:2406.11944 [cs]. [35] J. Lindsey, A. Templeton, J. Marcus, T. Conerly, J. Batson, and C. Olah. Sparse Crosscoders for Cross-Layer Features and Model Diffing, 2024. URL https://transformer-circuits. pub/2024/crosscoders/. [36] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nicholas L. Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E. Burke, Tristan Hume, Shan Carter, Tom Henighan, and Chris Olah. Towards Monosemanticity: Decomposing Language Models With Dictionary Learning, 2023. URL https://transformer-circuits.pub/2023/ monosemantic-features. [37] Alon Jacovi and Yoav Goldberg. Towards Faithfully Interpretable NLP Systems: How Should We Define and Evaluate Faithfulness? In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4198–4205, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-main.386. URL https://aclanthology.org/2020.acl-main. 386/. [38] Wes Gurnee. SAE reconstruction errors are (empirically) pathological — AI Alignment Forum, March 2024. URL https://www.alignmentforum.org/posts/rZPiuFxESMxCDHe4B/ sae-reconstruction-errors-are-empirically-pathological. [39] Joshua Engels, Logan Riggs Smith, and Max Tegmark. Decomposing The Dark Matter of Sparse Autoencoders. Transactions on Machine Learning Research, December 2024. ISSN 2835-8856. URL https://openreview.net/forum?id=sXq3Wb3vef. [40] David Chanin, James Wilken-Smith, Tomáš Dulka, Hardik Bhatnagar, Satvik Golechha, and Joseph Isaac Bloom. A is for Absorption: Studying Feature Splitting and Absorption in Sparse Autoencoders. October 2025. URL https://openreview.net/forum?id=R73ybUciQF. [41] Yutong Gao, Qinglin Meng, Yuan Zhou, and Liangming Pan. Towards Intrinsic Interpretability of Large Language Models:A Survey of Design Principles and Architectures, April 2026. URL http://arxiv.org/abs/2604.16042. arXiv:2604.16042 [cs]. [42] Jeremy Herbst, Jae Hee Lee, and Stefan Wermter. The Expert Strikes Back: Interpreting Mixture-of-Experts Language Models at Expert Level, April 2026. URL http://arxiv.org/ abs/2604.02178. arXiv:2604.02178 [cs]. [43] Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer El Showk, Nicholas Joseph, Nova DasSarma, Ben Mann, Danny Hernandez, Amanda Askell, Kamal Ndousse, Andy Jones, Dawn Drain, Anna Chen, Yuntao Bai, Deep Ganguli, Liane Lovitt, Zac Hatfield-Dodds, Jackson Kernion, Tom Conerly, Shauna Kravec, Stanislav Fort, Saurav Kadavath, Josh Jacobson, Eli Tran-Johnson, Jared Kaplan, Jack Clark, Tom Brown, Sam McCandlish, Dario Amodei, and Christopher Olah‡ AFFILIATION Anthropic PUBLISHED June 27. Softmax Linear Units, 2022. URL https://transformer-circuits.pub/2022/ solu/index.html. [44] Lee Sharkey. A technical note on bilinear layers for interpretability, May 2023. URL http: //arxiv.org/abs/2305.03452. arXiv:2305.03452 [cs]. [45] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, 13

and Chris Olah. A Mathematical Framework for Transformer Circuits, 2021. URL https: //transformer-circuits.pub/2021/framework/index.html. [46] Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2014. URL https://arxiv.org/abs/ 1312.6120. _eprint: 1312.6120. [47] Andrew M. Saxe, Shagun Sodhani, and Sam Lewallen. The Neural Race Reduction: Dynamics of Abstraction in Gated Networks, July 2022. URL http://arxiv.org/abs/2207.10430. arXiv:2207.10430 [cs]. [48] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy Models of Superposition, 2022. URL https://transformer-circuits.pub/2022/toy_ model/index.html. [49] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax. [50] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github.com/google/flax. [51] Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http://github.com/deepmind. [52] Marvin Ritter, Ihor Indyk, Aayush Singh, Andrew Audibert, Anoosha Seelam, Camelia Hanes, Eric Lau, Jacek Olesiak, Jiyang Kang, and Xihui Wu. Grain - Feeding JAX Models, 2023. URL http://github.com/google/grain. [53] Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, Joe Davison, Mario Šaško, Gunjan Chhablani, Bhavitvya Malik, Simon Brandeis, Teven Le Scao, Victor Sanh, Canwen Xu, Nicolas Patry, Angelina McMillan-Major, Philipp Schmid, Sylvain Gugger, Clément Delangue, Théo Matussière, Lysandre Debut, Stas Bekman, Pierric Cistac, Thibault Goehringer, Victor Mustar, François Lagunas, Alexander Rush, and Thomas Wolf. Datasets: A Community Library for Natural Language Processing. In Heike Adel and Shuming Shi, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 175–184, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021. emnlp-demo.21. URL https://aclanthology.org/2021.emnlp-demo.21/. [54] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-ofthe-Art Natural Language Processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38–45, Online, October 2020. Association for Computational Linguistics. URL https://aclanthology. org/2020.emnlp-demos.6/. [55] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, 14

M. Perrot, and E. Duchesnay. Scikit-learn: Machine Learning in Python. Journal of Machine Learning Research, 12:2825–2830, 2011. [56] Lukas Biewald. Experiment Tracking with Weights and Biases, 2020. URL https://www. wandb.com/. [57] Plotly Technologies Inc. Collaborative data science, 2015. URL https://plot.ly. Place: Montreal, QC. [58] Charlie Marsh. uv: An extremely fast Python package and project manager, written in Rust., 2024. URL https://pypi.org/project/uv/.

Appendix A Sparsely gated linear neuron layer A.1 Parallel channels . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.2 Time and memory complexity . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

17 17 17

B Experimental details B.1 Training details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.1.1 SlimPajama . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.1.2 TinyStories . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Scaling ladder . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

18 18 18 18 19

C Additional results C.1 Neuron usage and coactivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.2 Model interpretability . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

19 19 19

D Additional details D.1 Compute resources . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Software and libraries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

27 27 27

Sparsely gated linear neuron layer

A.1 Parallel channels The full sgatlin layer consists of C parallel channels that are summed, gc = product-top-k(Wckey W query z), sgatlin(z) =

C X X

out in⊤ gc,i wc,i wc,i z

c=1 i:gc,i ̸=0 √

2 dffw ×dkey where W query ∈ Rdkey ×dmodel is shared across layers, whereas Wckey ∈ R , gc ∈ Rdffw , and √ in out dmodel 2 dffw dffw wc,i , wc,i ∈ R are all channel-specific. product-top-k(x) : R → R is the product top-k operator introduced by Lample et al. [13] that selects the k largest entries of the product set of two sets of scalars, here obtained by slicing the input vector x into two parts of equal size, and returns a sparse vector of gating scores with k non-zero entries. Pseudocode for the layer is shown in Listing 1.

1 2 3 4 5 6

class SparselyGatedLinear ( nn . Module ): d_model : int d_ffw : int d_key : int knn : int n_channels : int

7 8 9

@nn . compact def __call__ ( self , x ):

10 11 12 13 14 15 16 17 18

# Gating network n_keys = int ( math . sqrt ( self . d_ffw )) queries = nn . Dense ( features = self . d_key , use_bias = False )( x ) similar = nn . DenseGeneral ( features =(2 , self . n_channels , n_keys ) , use_bias = False , )( queries ) scores , indices = product_key_top_k ( similar , k = self . knn )

19 20 21 22 23 24 25

# Expert computation experts = self . param ( name = ’ experts ’ , init_fn = lecun_normal , shape =(2 , self . n_channels , self . d_ffw , self . d_model ) , )

26 27 28 29 30

w_in , w_out = experts [ indices ] z = jnp . einsum ( ’btd , btckd − > btck ’ , x , w_in ) z = z ∗ scores z = jnp . einsum ( ’ btck , btckd − > btd ’ , z , w_out )

31 32

return z

Listing 1: Pseudocode for the sgatlin layer using JAX [49] and Flax-like [50] syntax. A.2 Time and memory complexity We compare the time and memory complexities of different transformer feedforward layers in Table 2. In typical dense feedforward layers, the time complexity is dominated by the product of the residual stream dimension and the feedforward dimension. Scaling the number of parameters by increasing the feedforward dimension proportionally increases computational cost. MoEs decouple the total number of feedforward neurons (i.e., the product of the number of experts and the size of each expert) from the number of active feedforward neurons by activating only a subset of experts. This allows scaling the number of feedforward parameters without proportionally increasing computational cost. However, 17

Table 2: Time and memory complexities of feedforward layers. We compare the time and memory complexities of different transformer feedforward layers considered in our experiments. dmodel denotes the residual stream dimension, dffw is the number of feedforward neurons, nexperts the total number of experts, nactive the number of active experts routed per token, nchannel the number of channels, and dkey the bottleneck dimension of the gating network. Feedforward variant

Time complexity (O)

Memory complexity (O)

MLP SwiGLU MoE PEER

dmodel dffw dmodel dffw dmodel nexperts + nactive dmodel dffw √ nchannel dkey (dmodel + dffw ) + nchannel dmodel nactive √ dkey (dmodel + nchannel dffw ) + nchannel dmodel nactive

dmodel dffw dmodel dffw nexperts dmodel dffw √ nchannel dkey (dmodel + dffw ) + dffw dmodel √ dkey (dmodel + nchannel dffw ) + nchannel dffw dmodel

sgatlin

typical MoE parameterizations are usually comparably coarse-grained, and the size of each expert is large enough for the time complexity to be dominated by the product of active experts and their size. Both PEER [11] and sgatlin shrink the experts to consist of a single neuron and only activate a small, fixed subset of neurons per token. In combination with their respective efficient router/gating network, this allows scaling the number of parameters while keeping the time complexity comparably small. Note that the difference in time and memory complexity between PEER and sgatlin is due to the gating network query projection being shared across channels in the latter but not in the former.

Experimental details

B.1 Training details In the following, we provide further training details for both the language modelling experiments on the SlimPajama 627B dataset [12] conducted in Section 3 and the model interpretability experiments based on the TinyStories dataset [23] conducted in Section 4. B.1.1 SlimPajama Tokenizer We tokenize text input using the pretrained GPT-2 byte-pair encoding tokenizer with a vocabulary size of 50 257. Model We use a causal, decoder-only transformer with an RMSNorm applied prior to the attention and feedforward block and rotary positional encoding. We follow the common procedure of fixing the attention head dimension to 64 across all model scales and determining the number of attention heads according to the residual stream dimension of the respective experiment. We use the MoE layer from GPT-OSS [6] with 16 experts and top-k = 2 routing. Both PEER and sgatlin use 16 channels, top-k = 8 gating and a gating bottleneck dimension of 128. All models are trained in bfloat16 precision with float32 upcasting on the attention’s softmax operation and the final output embedding. In addition, we use activation checkpointing (flax.linen.remat) to further reduce the memory requirements during training. Optimizer We use a global batch size of 128 and a sequence length of 2048 tokens. We optimize the models using AdamW with a peak learning rate of 1 × 10−3 and apply a weight decay of 0.1, excluding scalar parameters. The learning rate follows a warmup stable square root decay schedule [17, 18] with 1000 warmup steps and a decay fraction of 0.2. Gradients are clipped to a maximum global norm of 1.0. Learning rate, weight decay, and warmup steps were selected after running a grid search over learning rate values of {1 × 10−4 , 3 × 10−4 , 1 × 10−3 }, weight decay values of {0.01, 0.03, 0.1}, and warmup steps of {100, 1000}, training a dense SwiGLU transformer with 12 layers on a compute budget of 1 × 1018 FLOPs for each configuration. B.1.2 TinyStories For our model interpretability experiments with the TinyStories dataset, we train sgatlin-based transformers with 4 layers, a residual stream dimension of 256 and 4 attention heads of dimension 64 with a batch size of 256 sequences and a sequence length of 1024 tokens. We train a custom byte-pair encoding tokenizer on the TinyStories training data with a vocabulary size of 8192. Otherwise, we reuse the same optimizer and scheduler settings as for SlimPajama, only reducing the number of warmup steps to 100 and increasing the gradient clipping norm to 2.0. 18

B.2 Scaling ladder We use the scaling ladder shown in Listing 2 to systematically increase the depth and width of our models. In our experiments, we vary the scale integer from 2 to 8. def scaling_ladder ( scale : int , feedforward_type : str ): d_model = 128 ∗ scale n_layers = 2 ∗ scale

1 2 3 4

match feedforward_type : case ’ swiglu ’ | ’ mlp ’: d_ffw = int (8 / 3 ∗ d_model / 256) ∗ 256 case ’ moe ’: d_ffw = d_model # following GPT − OSS case ’ peer ’: d_ffw = (32 + 24 ∗ scale ) ∗∗ 2 case ’ sgatlin ’: d_ffw = (16 + 12 ∗ scale ) ∗∗ 2 case _ : raise ValueError ( f ’{ feedforward_type } ␣ unknown ’)

5 6 7 8 9 10 11 12 13 14 15 16

return ( d_model , n_layers , d_ffw )

Listing 2: Python code for scaling transformer model size with different feedforward layer types.

Additional results

C.1 Neuron usage and coactivation In Figure 8, we show neuron usage and coactivation in the sgatlin-based transformer trained on TinyStories used for our model interpretability experiments. The product top-k mechanism creates neuron subpopulations. The product top-k mechanism used in the gating network of sgatlin computes an input-dependent query and calculates the dotproduct with a neuron-specific key. It efficiently computes the dot-product for a large number of possible neurons by defining the neuron keys as the product-set of multiple subkeys. As a result, each neuron shares its two subkeys with other neurons. This is illustrated in the left panel of Figure 8. In practice, this key sharing affects the empirically observed coactivation probability of neurons, effectively creating subpopulations of neurons that preferably coactivate with neurons they share a gating key with. Available neurons are effectively used. We verify that the gating network effectively uses the available feedforward capacity and activates almost all the available feedforward neurons at least once across a batch of 128 randomly chosen sequences from the TinyStories validation set, as shown in the centre panel of Figure 8. The right panel further shows that neuron activation frequency is heterogeneous, with an average Gini coefficient of 0.58 across these sequences. C.2 Model interpretability Figure 9 shows an extended version of Figure 5 in the main text, with gating embeddings for all four layers. Table 3 shows an extended version of Figure 6 in the main text, contextualizing the feedforward circuit neighbours with excerpts from the input sequences that activated them as well as the top 5 output predictions at that respective token position.

100

0.8 0.6 0.4 0.2 0

0 0

100

Neuron j

100

Layer

Layer Cumulative share of activations

0.05

Fraction of active neurons

Neuron i

97%

0.1

98%

P(j|i)

98%

Coactivation probability

99%

Shared gating keys

Gini coeﬃcient: 0.58

0.8 0.6 0.4 0.2 0 0

0.5

Cumulative share of neurons

Neuron j

Figure 8: Neuron usage and coactivation. Left As a result of the product top-k mechanism, certain pairs of neurons share gating keys. This affects the empirically observed coactivation probability of neurons, where neurons preferably coactivate with neurons that they share a gating key with. Coactivation probability is computed over 128 random sequences of the TinyStories validation set and shown for the first 128 neurons of the penultimate transformer layer. Center Almost all neurons activate at least once over 128 random sequences of the TinyStories validation set. Right Neuron activation frequency is heterogeneous, with an average Gini coefficient of 0.58 across 128 random sequences of the TinyStories validation set. We control for the skewed input token frequency of TinyStories sequences considering each unique input token only once. Layer 1

Layer 2

Input tokens

Other

Be are

was were

Names Ben

Lily Max Timmy

−5

Tom

Pronouns he

−10

her Layer 3

Layer 4

his it she

their they

Punctuation

! !"

" , ,"

. ."

−5

?" Ċ

−10

ĠĊ 0

Figure 9: sgatlin gating weights form semantic clusters (extended version of Figure 5). We collect the gating weights of sgatlin of a four layer language model trained on TinyStories for the 100 most frequent tokens and embed them into 2D space using UMAP. The resulting embedding forms clusters that are partially explained by the semantics of the corresponding input tokens. For instance, names used in the stories cluster together and pronouns form clusters in their vicinity.

Table 3: Feedforward circuit neighbours (extended). Following the procedure shown in Figure 4 of collecting reference gating weights across input sequences, we can analyse the feedforward circuits for a new input sequence by comparing them to their nearest neighbours among the reference set. In this example, we analyse the feedforward circuits at the penultimate layer for the new input sentence excerpt "... big dog named Max. Max had a red collar that he wore every day". The bold token in the ’Neighbour input excerpt’ column denotes the token at which the neighbour circuit was activated. The ’Neighbour top predictions’ column lists the 5 top output predictions at the output layer.

Pos

Token

Neighbour distance

Neighbour input excerpt

Neighbour top predictions

big

5.96e-08 5.96e-08 5.96e-08 5.96e-08 1.51e-01 1.51e-01 2.06e-01 2.20e-01 2.23e-01 2.38e-01

time, there was a big farm. On the farm time, there was a big and happy family who loved time, there was a big teddy bear. He loved time, there was a big, strong robot made of time, there was an old man. He liked to time, there was an old lady. She liked to a queen who lays eggs and makes more ants." backyard where there was a long fence. One day, . Amy had a big yard with a purple swing. Inside the boat was a little boy called Ben. Ben

dog

4.13e-02 4.92e-02 5.57e-02 6.44e-02 1.11e-01 1.32e-01 1.38e-01 1.62e-01 1.67e-01 1.71e-01

, there was a little dog named Spot. Spot loved time, there was a dog named Max. Max was upon a time, a dog named Max went for a , there was a sweet dog named Max. Max loved , there was a little bunny named Benny. Benny had , there was a little kitten named Mittens. Mittens was , there was a graceful cat named Kitty. She loved Once there was a clever kitten. Every day it would , there was a small mouse named Timmy. Timmy was was a big, strong robot made of steel. He

named | . | called | who | gie named | . | called | who | with named | and | was | called | went named | . | called | who | and named | . | who | called | and named | . | who | called | , named | . | who | called | and . | named | who | called | , named | . | who | called | that . | named | who | called | that

named

3.18e-01 3.32e-01 3.37e-01 3.54e-01 3.97e-01

, there was a dog named Max. Max was a there was a little dog named Spot. Spot loved to there was a sweet dog named Max. Max loved to a time, a dog named Max went for a walk there was a small mouse named Timmy. Timmy was always

Pos

Token

Neighbour distance

Table 3 – Continued from previous page Neighbour input excerpt

Neighbour top predictions

4.03e-01 4.18e-01 4.48e-01 5.01e-01 5.21e-01

Mia met a talking cat named Tom. Tom was a named Lily and a boy named Tom were playing in their She had a pet cat named Olive. Lily loved to a puppy named Max who was her best friend. One this is my plant. You can’t touch it."

Max

1.13e-02 2.83e-02 2.91e-02 4.29e-02 4.66e-02 4.82e-02 5.92e-02 6.47e-02 7.08e-02 7.12e-02

there was a dog named Max. Max was a wild was a little dog named Spot. Spot loved to play was a sweet dog named Max. Max loved to jog there were two birds called Jack and Jill. Every night time, a dog named Max went for a walk in was a little girl named Amy. Amy had a big was a little kitten named Mittens. Mittens was very popular there was a girl named Mia. Mia loved her jewelry was a graceful cat named Kitty. She loved to play to a little boy named Jack. He loved the fat

. | who | and | , | the . | who | and | , | ty . | who | and | , | ! and | . | , | who | ’s was | lived | went | and | wanted . | who | and | , | a . | who | , | and | that . | who | and | a | , . | who | , | and | ty . | who | and | , | ’s

6.31e-03 7.51e-03 5.01e-02 5.99e-02 6.07e-02 6.25e-02 6.51e-02 6.71e-02 7.90e-02 9.76e-02

a little boy called Ben. Ben was just 3 years a little squirrel named Sammy. Sammy loved to play and was a dog named Max. Max was a wild dog a small mouse named Timmy. Timmy was always scared of a little bunny named Benny. Benny had soft white fur a little fish named Bubbles. Bubbles lived in a big deer named Bambi. Bambi wanted to a little boy named Jack. He loved the fat car a little kitten named Mittens. Mittens was very popular with was just 3 years old. He was sailing all by

Ben | He | Ċ | One | Ċ Sammy | He | One | Timmy | Nutty Max | He | One | His | It Timmy | He | One | Every | She Benny | He | One | Every | Today Bubbles | He | One | She | Nemo B | She | He | One | Every He | Jack | One | Every | Ċ Mittens | She | One | He | Mittens Ċ | He | Ċ | One | Every

Max

1.02e-02 1.42e-02 2.15e-02 3.63e-02 5.22e-02

a girl named Mia. Mia loved her jewelry. She a dog named Max. Max was a wild dog who little dog named Spot. Spot loved to play fetch with little girl named Amy. Amy had a big yard with sweet dog named Max. Max loved to jog in the

was | loved | had | liked | wanted was | loved | had | liked | lived loved | was | had | liked | lived loved | had | was | liked | lived loved | liked | had | was | lived

Continued on next page...

Pos

Token

Neighbour distance

Table 3 – Continued from previous page Neighbour input excerpt

Neighbour top predictions

7.19e-02 9.46e-02 9.68e-02 1.03e-01 1.14e-01

cakes. One day, she went for a walk and else. One day, Tim saw a big candy store friends. One day, he heard a little girl crying princess. One day, she found a shiny belt in good boy named Timmy. Timmy always listened to his mommy

had

1.61e-01 1.97e-01 2.11e-01 2.18e-01 2.38e-01 2.69e-01 2.79e-01 2.90e-01 2.99e-01 3.11e-01

bunny named Benny. Benny had soft white fur that was girl named Amy. Amy had a big yard with a boy named Tim. Tim had a big, thin sack girl named Lily. She had a big wardrobe with smooth girl named Lily. She had a pet cat named Olive girl named Lily. Lily had a tough day at school there was a king who had a lot of money. girl named Lily. She had a puppy named Max who dress-up. She had a unique dress that her loved her jewelry. She had a big box full of

1.25e-01 1.40e-01 1.50e-01 1.63e-01 1.82e-01 2.01e-01 2.03e-01 2.26e-01 2.26e-01 2.37e-01

named Tim. Tim had a big, thin sack that named Amy. Amy had a big yard with a purple out if the cat has a home already. Maybe it to read. He had a big book with lots of named Lily. She had a big wardrobe with smooth doors a nice man who had a red shirt. He liked named Lily. She had a puppy named Max who was named Timmy. Timmy had a red bike with a big named Timmy. Timmy had a big pocket on his shirt named Lily. Lily had a tough day at school because

red

2.19e-02 1.28e-01 1.28e-01 1.37e-01 1.47e-01

nice man who had a red shirt. He liked to Timmy. Timmy had a red bike with a big wheel box. It has a red light and a number. The different boy had a red ball, and Timmy wanted a big yard with a purple swing. She loved to

Continued on next page...

Pos

Token

Neighbour distance

Table 3 – Continued from previous page Neighbour input excerpt

Neighbour top predictions

1.49e-01 1.60e-01 1.65e-01 1.89e-01 2.25e-01

Lily. She had a big wardrobe with smooth doors. Amy. Amy had a big yard with a purple swing Timmy. Timmy had a big pocket on his shirt. Tim. Tim had a big, thin sack that he very rich and had a big castle. One day,

, | box | sister | wardrobe | red , | box | toy | red | dream , | sister | red | toy | box , | toy | box | red | sister house | castle | palace | garden | crown

collar

1.21e-01 1.28e-01 1.42e-01 1.62e-01 1.77e-01 1.82e-01 1.88e-01 1.98e-01 2.04e-01 2.05e-01

. She had a unique dress that her mom made just man who had a red shirt. He liked to wear a big wardrobe with smooth doors. She loved to play red bike with a big wheel in the front. He brave knight with a colorful shield and a sharp spear. big yard with a purple swing. She loved to play big and had a shiny label that said "brilliant small toy car with a zipper on it. He can . Benny had soft white fur that was perfect for snugg different boy had a red ball, and Timmy wanted to

that | with | and | , | made . | and | that | on | with . | and | that | inside | where . | that | on | and | in . | and | to | on | , . | that | in | and | on on | that | . | with | , . | on | and | around | that that | and | , | with | coat and | that | , | . | in

that

2.94e-01 3.04e-01 3.18e-01 3.20e-01 3.40e-01 3.77e-01 4.14e-01 4.82e-01 4.85e-01 5.12e-01

He had a white basketball that he played with every day She had a unique dress that her mom made just for a big, thin sack that he carried everywhere he went They had a small plant that they watered every day. they found a toy car that Max had forgotten to put Tom found a big shell that was shiny and had many Benny had soft white fur that was perfect for snuggling It was the same shell that she had found yesterday and . She had colorful wings that looked like a pastel drawing a big, impressive dog that could do many tricks.

1.65e-01 1.80e-01 2.16e-01 3.03e-01 4.32e-01

big, thin sack that he carried everywhere he went. was the same shell that she had found yesterday and left had a white basketball that he played with every day. had a puppy named Max who was her best friend. a different boy there who he had never seen before.

Continued on next page...

Pos

Token

Neighbour distance

Table 3 – Continued from previous page Neighbour input excerpt

Neighbour top predictions

4.83e-01 5.05e-01 5.07e-01 5.11e-01 5.11e-01

, there was a shrimp named Sammy. Sammy lived in saw a different boy there who he had never seen before "It’s a tool we use to measure things. found a toy car that Max had forgotten to put away daddy. They packed everything they needed in a big,

wore

4.29e-01 4.65e-01 5.19e-01 5.47e-01 5.54e-01 5.60e-01 5.81e-01 5.85e-01 6.11e-01 6.13e-01

there was a bigger glass of milk with a straw set big, impressive dog that could do many tricks.Ċ , impressive dog that could do many tricks.ĊĊ see a metal thing that looks like a stick. It had sent her the package of pink stuff. Clara smiled plant was hurt. Some of the leaves and flowers fell loved her jewelry. She had a big box full of find any comfortable clothes to wear. He tried on many and strong. He protects the castle from bad guys , they gave the bottle of wine to the nice lady

every

1.75e-01 1.91e-01 3.25e-01 3.32e-01 3.35e-01 3.42e-01 3.54e-01 3.90e-01 3.91e-01 3.96e-01

basketball that he played with every day. One day, small plant that they watered every day. The plant had Lily loved to wear it every day. ĊĊOne friends. They played together every day.ĊĊOne so much to explore. Every corner was filled with surprises eggs. She laid eggs every day in her cozy nest their new puppy played together every day and were the happiest and determination you can overcome any challenge.<|endoftext|> friends and he waved his little arms to say thank you in tip top shape. Every day, she was happy

day

2.26e-01 2.67e-01 2.93e-01 3.18e-01

And I will protect the castle with you and your dragon was friendly and licked her face. Lily petted him and She laughed and ate more pie.ĊĊFrom then and enjoyed their tea and cookies.ĊĊSuddenly,

from | ." | . | , | and . | to | , | with | ! . | , | until | and | with . | together | while | , | with

Continued on next page...

Pos

Token

Neighbour distance

Table 3 – Continued from previous page Neighbour input excerpt

Neighbour top predictions

3.28e-01 3.37e-01 3.59e-01 3.62e-01 3.62e-01 3.68e-01

not to play in the wardrobe. Sadly, Lily never ’m going to protect the castle from the dragons," said they needed for the tea party. She brought out the "I enjoy eating healthy snacks." ĊĊThe moral and a kiss on the cheek. From that day on two of them enjoyed their juice.<|endoftext|>

. | alone | again | without | anymore from | ," | with | !" | and . | , | and | with | for , | and | !" | like | too . | to | for | before | , together | . | and | , | in

Additional details

D.1 Compute resources We ran our experiments on two internal Slurm-based clusters equipped with nodes of 8 Nvidia H100 GPUs with 94GB of high-bandwidth memory per GPU or 8 Nvidia H200 GPUs with 141GB of high-bandwidth memory per GPU respectively. The longest of our training runs takes approximately 12 hours on a node of 8 H200 GPUs to complete. In total we spent the equivalent of around 10 000 H200 GPU hours to run our experiments. D.2 Software and libraries For the results obtained in this paper we build on free and open-source software. We implemented our experiments in Python using JAX [49, Apache License 2.0], Flax [50, Apache License 2.0], the Deepmind JAX Ecosystem [51, Apache License 2.0], Grain [52, Apache License 2.0], HuggingFace dataset [53, Apache License 2.0], HuggingFace tokenizers [54, Apache License 2.0] and Scikit-learn [55, BSD 3-Clause License]. We utilized WandB [56, MIT license] to monitor the progress and results of experiments, and Plotly [57, MIT license] for generating the plots. We use uv for Python project dependency management [58, MIT License].

Record · ID 266167 · SHA-256 c63145c6c34e5a20
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.