Structured Recurrent Mixers for Massively Parallelized Sequence Generation

# Structured Recurrent Mixers for Massively Parallelized Sequence Generation
Source: [https://arxiv.org/html/2605.08696](https://arxiv.org/html/2605.08696)
###### Abstract

Over the last two decades, language modeling has experienced a shift from predominantly recurrent architectures that process tokens sequentially during training and inference to non\-recurrent models that process sequence elements in parallel during training, which results in greater training efficiency and stability at the expense of lower inference throughput\. Here we introduce the Structured Recurrent Mixer, an architecture that allows for algebraic conversion between a sequence parallel representation at train time and a recurrent representation at inference, notably without the need for specialized kernels or device\-specific memory management\. We show experimentally that this dual representation allows for greater training efficiency, higher input information capacity, and larger inference throughput and concurrency when compared to other linear complexity models\. We postulate that recurrent models are poorly suited to extended sequence length scaling for information\-rich inputs typical of language, but are well suited to scaling in the sample \(batch\) dimension due to their constant memory per sample\. We provide Mojo/MAX inference implementations of SRMs exhibiting 12x the throughput and 170x the concurrency of similarly powerful Transformers inferenced on vLLM, increases characteristic of Pytorch implementations resulting in a 30% increase in compute\-constant GSM8k Pass@k\. We conclude by demonstrating that SRMs are effective reinforcement learning training candidates\.

00footnotetext:The author would like to that IBM for support during research and the writing of this manuscript\. Code for this paper can be found at[https://github\.com/blbadger/mrm](https://github.com/blbadger/mrm)## 1Introduction

Is model speed or accuracy more important for language tasks? The answer usually depends on the task, specifically whether or not it admits a way of quickly checking whether a sample is good or bad\. For tasks where samples cannot be quickly verified, model accuracy is typically more important than speed, but for many tasks verification is feasible: for example, unit or functional tests may be applied to code samples, outcome evaluation to agentic workflows, self\-verification to mathematical theorem proof generation via Lean\. If rapid verification is possible, it is often more efficient \(in terms of minimal compute spent per correct output\) to generate many samples from smaller and less\-capable models than fewer samples from larger, more accurate ones\(Brown et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib1); Chen et al\.,[2025](https://arxiv.org/html/2605.08696#bib.bib2)\)\.

Most functional language benchmarks do not focus on efficiency in these terms but instead on the likelihood of a correct answer or correct next token given one or more attempts, although some more recent benchmarks have begun to incorporate efficiency in terms of monetary cost per task \(for example Arc\-AGI\(Foundation,[2026](https://arxiv.org/html/2605.08696#bib.bib3)\)\)\. The shift away from recurrent\-style architectures \(recurrent neural networks\(Rumelhart et al\.,[1985](https://arxiv.org/html/2605.08696#bib.bib4)\)and LSTMs\(Hochreiter and Schmidhuber,[1997](https://arxiv.org/html/2605.08696#bib.bib5)\)\) occurred in the context of next token prediction accuracy as well as their relation to applied benchmarks in which model goodness metrics are based on a model’s accuracy per sample regardless of the compute necessary to generate that sample\. For outputs that are easily verified, a metric of accuracy per compute applied seems more appropriate, and the findings that smaller Transformer models\(Vaswani et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib6)\)can outperform larger ones in this context leads to the question: are architectures that perform more efficient inference than Transformers better suited to this paradigm of accuracy per compute?

For large\-scale language modeling one usually wants to train causal models to predict each next token in a context window in parallel via token shifting, with a change in model behavior during inference to remove this sequence parallelism for efficiency so that the model predicts one token at a time with a cache storing information from previous tokens\. Sequence parallel models, particularly those with𝒪​\(n2\)\\mathcal\{O\}\(n^\{2\}\)time and𝒪​\(n\)\\mathcal\{O\}\(n\)space complexity fornntokens, are efficient and stable to train but inefficient for sequential token generation, whereas recurrent models \(in this work defined as those with𝒪​\(n\)\\mathcal\{O\}\(n\)time and𝒪​\(1\)\\mathcal\{O\}\(1\)space complexity\) are efficient for sequence generation but typically do not admit easy ways to train efficiently and stably, as they are not intrinsically parallelizable over the sequence dimension\.

As a result of these complementary deficiencies, and because of the large latency of global device memory access relative to arithmetic operations present in GPUs, there have been numerous attempts to hybridize linear\-complexity architectures \(Mamba SSMs for Granite 4\(Mishra et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib7);[ibm,](https://arxiv.org/html/2605.08696#bib.bib8)\)and Nemotron 3\(Blakeman and et al\.,[2025](https://arxiv.org/html/2605.08696#bib.bib9)\), Gated Delta Nets for Qwen 3\.5\(Qwen Team,[2026](https://arxiv.org/html/2605.08696#bib.bib10)\)and Kimi Linear\(Team et al\.,[2025](https://arxiv.org/html/2605.08696#bib.bib11)\)\) to transformer backbones\. Hybridization, however, introduces a significant challenge: although these models gain the advantages of \(smaller\) variable\-sized caches, they also exhibit the same drawback as Transformers with respect to cache scaling, just with lower constant values\.

These observations motivate the introduction of a new architecture, the Structured Recurrent Mixer, that we find exhibits some of the desirable information retention and training efficiency characteristics of quadratic complexity models while improving on the the efficient inference computation of current recurrent models\. From an informational perspective we show that recurrent models are ill\-suited to unbounded sequence length scaling but are better applied to batch dimension scaling, and show that the SRM exhibits large throughput and concurrency increases relative to other models tested without any effort into optimizing parallelization via device\-specific kernels\. We address the question of whether designing an architecture that exhibits efficient sample throughput properties could lead to accuracy beyond that seen for a representative Transformer, with both training and inference compute held constant\. Finally we investigate the use of reinforcement learning to convert the many\-sample efficiency back to single\- or few\-sample efficiency\.

## 2Our Contribution

This paper details experiments testing two primary insights\. The first is that masked mixers should be more easily linearizable \(in the sequence dimension\) than transformers due to their data\-independent token mixing operations, allowing for an effective bridge between quadratic complexity and recurrent models\. We find that

Designing recurrent model token mixing operations to capture qualitative features of trained higher\-complexity models recovers some of those models’ training efficiency and information retention

and secondly we introduce a re\-evaluation of the benefits of constant space models when applied to information\-dense inputs such as language, proposing that

Recurrent language models are better suited to batch \(sample number\) dimension rather than sequence length scaling due to their constant\-sized, contiguous memory per sample

for which scaling the SRMs are shown to have unparalleled throughput characteristics even among recurrent models by construction, which is desirable characteristic for any task where samples may be quickly tested\. We introduce the following:

1. 1\.A qualitative feature mapping method from quadratic\-complexity to linear\-complexity models
2. 2\.A dual\-architecture sequence parallel \(at train time\) and batch parallel \(at inference\) recurrent model
3. 3\.A Mojo/MAX inference engine for SRMs with \>7x the throughput of Pytorch implementations
4. 4\.A resampling method for applying GRPO to large batches without decreasing exploration

## 3Related Work

Recurrent neural networks were among the first models applied to the problem of variable\-length sequence generation\(Rumelhart et al\.,[1985](https://arxiv.org/html/2605.08696#bib.bib4)\)\. The Long Short\-Term Memory model, a recurrent\-like model introducing multiple hidden states and associated memory streams for these sates, was introduced in order to stabilize training and better model long sequences\(Hochreiter and Schmidhuber,[1997](https://arxiv.org/html/2605.08696#bib.bib5)\)\. The greater training efficiency of Transformers compared to LSTMs both with respect to loss per parameter and loss per input context length account in part for why these architectures have been largely supplanted by Transformers for most language modeling tasks\(Kaplan et al\.,[2020](https://arxiv.org/html/2605.08696#bib.bib12)\)\.

Attempts to convert Transformers directly into recurrent\-like architectures via linearized attention have experienced somewhat limited success: direct swap\-ins of linear attention for dot\-product attention\(Katharopoulos et al\.,[2020](https://arxiv.org/html/2605.08696#bib.bib13); Shen et al\.,[2018](https://arxiv.org/html/2605.08696#bib.bib14)\)results in poor training efficiency and stability for causal modeling\(Poli et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib15)\), although progress has been made by reformulating the linear attention operation\. Along these lines, the RWKV architecture\(Peng et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib16)\)was introduced, and these models retain key and value vectors but perform memory mapping operations requiring the use of custom GPU kernels for efficient computation\.

The most well\-explored model architectures that allow for conversion between sequence parallel and recurrent representations exist in the state space framework, in which a state vector acts as the the recurrent hidden layer and is updated via linear operations\. Refined from S2/S3\(Fu et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib17)\)models, Mamba\(Gu and Dao,[2024](https://arxiv.org/html/2605.08696#bib.bib18)\)and Mamba 2\(Dao and Gu,[2024](https://arxiv.org/html/2605.08696#bib.bib19)\)SSM architectures incorporate selectivity at the token level and sophisticated memory management which allows for efficient parallel scan\-based state updates and thus conversion between sequence parallel and recurrent\-like forms\.\(Dao and Gu,[2024](https://arxiv.org/html/2605.08696#bib.bib19)\)further investigates the nature of SSMs as subsets of structured matrix models via semiseparable matrix representations, where token mixing weight matrices may be formulated in terms of lower\-rank sub\-matrices\.

## 4Method

Structured Recurrent Mixers are structured so that one is able to algebraically convert the sequence parallel representation \(which uses matrix multiplication to mix token information\) to the recurrent representation \(which must read and write to constant\-sized memory and use a fixed number of operations per token generated\)\. Training proceeds by structuring token mixing matrices by expanding weight vectors along rows or columns, applying an optional trainable decay rate, and proceeding with shifted next token prediction\.

### 4\.1Masked Mixer Parameterization

SRMs are based on the Masked Mixer, a transformer\-like architecture in which attention is swapped for masked MLPs \(matrix multiplications\)\(Badger,[2025](https://arxiv.org/html/2605.08696#bib.bib20)\)\. These models apply masked matrix multiplication \(or equivalently masked 1\-D convolution with a stride size of 1\) to the sequence dimension, which can be expressed asY=X​M\+BY=XM\+BwhereM∈ℝn×nM\\in\\mathbb\{R\}^\{n\\times n\}andX,Y,B∈ℝd×nX,Y,B\\in\\mathbb\{R\}^\{d\\times n\}\. Notably the values ofMMare triangular masked for both training and inference, and are not data\-dependent but are parameterized by trainable values in the model\. These architectures exhibit similar time and space complexities to Transformers but form fewer activations \(ddper layer\) in the token mixing operation, resulting in higher throughput at a fixed number of parameters which is balanced by their lower training efficiency on a per\-parameter basis\.

### 4\.2Trained Masked Mixer Weights: Qualitative Feature Representation

Initial investigations detailed in Table[S2](https://arxiv.org/html/2605.08696#S1.T2)demonstrated that a naive restriction of token mixing matrices to a recurrent form \(row repeat matrices\) results in relatively inefficient causal language model training\. This motivated an analysis on the structure of trained masked mixer token mixing layers, which revealed three notable qualitative features: matrices with columns containing near\-identical values, matrices with rows containing near\-identical values, and matrices with weight decay with increased distance from the main diagonal \(or some combination of these three features\) as shown in Figure[1](https://arxiv.org/html/2605.08696#S4.F1)\. As we explore in the next section, structuring a recurrent SRM to represent these qualitative features results in substantially higher training efficiency\. Curiously, not all features of trained mixer weight matrices are apparently necessary for efficient recurrent modeling: most layers exhibit lower near\-constant main diagonal weights compared to off\-diagonal weights \(Figures[1](https://arxiv.org/html/2605.08696#S4.F1),[S1](https://arxiv.org/html/2605.08696#S1.F1)\), but introducing an independent weight for the main diagonal to SRMs does not result in significantly increased training efficiency for larger models \(Table[S5](https://arxiv.org/html/2605.08696#S1.T5)\)\.

![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/mixer_weight_figure.png)Figure 1:Causal Masked Mixer token mixing matrix weights, \(a\) annotated features on whole\-matrix maps, \(b\) weights near the main diagonal\. Blue indicates negative, and red positive values\.
### 4\.3Architecture

As the token mixing layers of the SRM are restricted in order to retain a recurrent representation, we now illustrate what kind of restrictions are required\. In brief, the restrictions we explore are relatively simple: rows must contain identical elements or columns must contain identical elements, and optionally a constant factor may be multiplied to each diagonal\. Effectively we perform matrix\-vector multiplication \(with masking\), which is why the operation may be converted into a fixed\-memory recurrent operation\.

To illustrate an example row repeat mixing operation, consider inputX∈ℝd×nX\\in\\mathbb\{R\}^\{d\\times n\}withddhidden dimension andnntokens, outputY∈ℝd×nY\\in\\mathbb\{R\}^\{d\\times n\}, mixer weight matrixW∈ℝn×nW\\in\\mathbb\{R\}^\{n\\times n\}, biasB∈ℝd×nB\\in\\mathbb\{R\}^\{d\\times n\}, and trainable decay constantλ∈\(0\.9,1\]\\lambda\\in\(0\.9,1\]\. We obtainYYfrom inputsX,B,WX,B,Was shown in Equation[1](https://arxiv.org/html/2605.08696#S4.E1)for the case wheren=3n=3for simplicity\. Equation[1](https://arxiv.org/html/2605.08696#S4.E1)may be expressed in its recurrent representation in Equation[2](https://arxiv.org/html/2605.08696#S4.E2), where the recurrent hidden state \(cache\) is simply the sum term\.

\(\|\|\|Y0Y1Y2\|\|\|\)=\(\|\|\|X0X1X2\|\|\|\)​\(α0λ​α0λ2​α00bλ​α100α2\)\+\(\|\|\|β0β1β2\|\|\|\)\\begin\{pmatrix\}\|&\|&\|\\\\ Y\_\{0\}&Y\_\{1\}&Y\_\{2\}\\\\ \|&\|&\|\\end\{pmatrix\}=\\begin\{pmatrix\}\|&\|&\|\\\\ X\_\{0\}&X\_\{1\}&X\_\{2\}\\\\ \|&\|&\|\\end\{pmatrix\}\\begin\{pmatrix\}\\alpha\_\{0\}&\\lambda\\alpha\_\{0\}&\\lambda^\{2\}\\alpha\_\{0\}\\\\ 0&b&\\lambda\\alpha\_\{1\}\\\\ 0&0&\\alpha\_\{2\}\\end\{pmatrix\}\+\\begin\{pmatrix\}\|&\|&\|\\\\ \\beta\_\{0\}&\\beta\_\{1\}&\\beta\_\{2\}\\\\ \|&\|&\|\\end\{pmatrix\}\(1\)
Yn=αn​Xn\+βn\+∑m=0n−1λn−m​αm​Xm\\begin\{split\}Y\_\{n\}&=\\alpha\_\{n\}X\_\{n\}\+\\beta\_\{n\}\+\\sum\_\{m=0\}^\{n\-1\}\\lambda^\{n\-m\}\\alpha\_\{m\}X\_\{m\}\\end\{split\}\(2\)
Similarly, a column repeat matrix multiplication with per\-token biases may be reduced to the recurrent representation given in Equation[3](https://arxiv.org/html/2605.08696#S4.E3), where again the sum term is the constant\-size cache andαn\\alpha\_\{n\}is factored out upon each update\. We implement two methods for combining column and row repeat outputs: linear combinations where each head has one column and one row repeat matrix and the output is linearly combined, as well as what we term ‘mixed’ heads where heads are evenly divided into column and row repeat layers and combined as other headed models are \(via concatenation before out projection\), which we term ‘head\-parallel’ layers\. In either case next token prediction is efficient and requires at most two BLAS level 1 operation per recurrent layer at inference, although linearly combined layers with both row and column repeat layers require double the cache size as mixed or row\- or column\-only layers\.

We implement multiple heads per token mixing layer by training unique projectionsIh∈ℝdh×dI\_\{h\}\\in\\mathbb\{R\}^\{d\_\{h\}\\times d\}for each head0,1,…,h0,1,\.\.\.,h, for head dimensiondh=d/hd\_\{h\}=d/h, performing token mixing on these heads independently, and projecting the concatenated outputs viaPo​u​t∈ℝd×dP\_\{out\}\\in\\mathbb\{R\}^\{d\\times d\}as shown in Equation[4](https://arxiv.org/html/2605.08696#S4.E4)\. Equivalently, one input projectionI∈ℝd×dI\\in\\mathbb\{R\}^\{d\\times d\}may or may not be applied, heads are sliced and token mixing performed and outputs concatenated and optionally projected back\.

Y=X​\(α0λ​α1λ2​α2⋯λn−1​αn0α1λ​α2⋯λn−2​αn00α2⋯λn−3​αn⋮⋮⋮⋱⋮000⋯λ0​αn\)\+β⟹Yn=αn​Xn\+βn\+αn​∑m=0m=n−1λn−m​Xm\\begin\{split\}Y=X\\begin\{pmatrix\}\\alpha\_\{0\}&\\lambda\\alpha\_\{1\}&\\lambda^\{2\}\\alpha\_\{2\}&\\cdots&\\lambda^\{n\-1\}\\ \\alpha\_\{n\}\\\\ 0&\\alpha\_\{1\}&\\lambda\\alpha\_\{2\}&\\cdots&\\lambda^\{n\-2\}\\ \\alpha\_\{n\}\\\\ 0&0&\\alpha\_\{2\}&\\cdots&\\lambda^\{n\-3\}\\alpha\_\{n\}\\\\ \\vdots&\\vdots&\\vdots&\\ddots&\\vdots\\\\ 0&0&0&\\cdots&\\lambda^\{0\}\\alpha\_\{n\}\\end\{pmatrix\}\+\\beta\\implies Y\_\{n\}&=\\alpha\_\{n\}X\_\{n\}\+\\beta\_\{n\}\+\\alpha\_\{n\}\\sum\_\{m=0\}^\{m=n\-1\}\\lambda^\{n\-m\}X\_\{m\}\\end\{split\}\(3\)
![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/srm_token_mixing.png)Figure 2:SRM architecture overviewY=Po​u​t​\(\(I0​X\)​W0\+β0∘\(I1​X\)​W1\+β1∘⋯∘Ih−1​\(X\)​Wh−1\+βh−1\)Y=P\_\{out\}\\bigg\(\(I\_\{0\}X\)W\_\{0\}\+\\beta\_\{0\}\\circ\(I\_\{1\}X\)W\_\{1\}\+\\beta\_\{1\}\\circ\\cdots\\circ I\_\{h\-1\}\(X\)W\_\{h\-1\}\+\\beta\_\{h\-1\}\\bigg\)\(4\)
We also implement the SRM recurrent operation for non\-unitary kernels, where each token mixing layer mixes along both sequence and a limited number of hidden dimension elements, with the number of hidden dimension elements mixed equal to the kernel size\. We train separate filters for each kernel with weightsW∈ℝk×n×nW\\in\\mathbb\{R\}^\{k\\times n\\times n\}parameterized as column repeat matrices \(Equation[3](https://arxiv.org/html/2605.08696#S4.E3)\)\. In our implementation, shown in Equation[5](https://arxiv.org/html/2605.08696#S4.E5), we zero padX∈ℝd\+k×nX\\in\\mathbb\{R\}^\{d\+k\\times n\}and use a stride of size 1\.

Y=∑i=0kX\(i:d\+i,:\)​Wi\+βY=\\sum\_\{i=0\}^\{k\}X\_\{\(i:d\+i,\\;:\)\}\\;W\_\{i\}\+\\beta\(5\)
Algebraic equivalence allows the SRM to be represented by either recurrent or sequence parallel operations without any change to the underlying cache structure or model parameters: for instance, we implement an SRM as a Pytorch module can perform either parallel or recurrent forward passes on\-demand without modification to the underlying memory, which greatly simplifies pipelines in which fast, parallelized inference needs to be efficiently swapped with sequence parallel forward/backward passes for gradient formation and backpropegation \(typical of online reinforcement learning\)\.

### 4\.4Computational Complexity

The computational complexity of SRMs at inference, as for other recurrent models, is𝒪​\(n​d2\)\\mathcal\{O\}\(nd^\{2\}\)time and𝒪​\(d\)\\mathcal\{O\}\(d\)space\. It is notable that SRMs are ‘true’ constant\-space models in that contain constant factors which themselves do not scale with sequence length \(see for example\(Qin and Zhong,[2023](https://arxiv.org/html/2605.08696#bib.bib21)\)\)\.

It is worth noting that the quadratic complexity of SRMs at train time is effectively hidden by the much larger constant factors associated with memory transfer and gradient production for token hidden layers, causing these models to resemblen​d2nd^\{2\}as long asnc​t​xn\_\{ctx\}is not substantially larger thandd\. This is also true of Transformers, and accounts for the lack of increased throughput of models that train with linear complexity \(Mamba for example\) over SRMs at train time with increased context \(Tables[S9](https://arxiv.org/html/2605.08696#S1.T9)\-[S12](https://arxiv.org/html/2605.08696#S1.T12)\)\.

Global memory access on device \(GPU\) typically requires an order of magnitude more clock cycles than matrix arithmetic operations, and as such is limiting for inference throughput for Transformers\. Batching inputs allows for more efficient memory access patterns, but still does not typically result in tensor core saturation for transformer models during inference\. In the batch dimensionbb, fornntokens the memory scaling characteristics of each transformer layer is𝒪​\(n​b​d\)\\mathcal\{O\}\(nbd\), whereas for each layer of the SRM the corresponding scaling is𝒪​\(b​d\)\\mathcal\{O\}\(bd\)\.

## 5Experimental Results

### 5\.1SRM architecture optimization

Experimentation shows that multi\-headed SRMs are more efficient to causally train than non\-headed models, that using both column and row repeat matrices outperforms row repeat alone or column repeat alone, and that models with a decay term are more efficient to train than SRMs without decay terms \(Table[S2](https://arxiv.org/html/2605.08696#S1.T2)\)\. The most efficiently trainable architecture we tested consisted of ‘mixed’ heads such that half of the heads are row repeat and half column repeats, each with hidden dimensiond/hd/hsuch that the total cache size per layer isdd\. We find that use of unique projections leads to slightly more efficient training than the use of parallel heads, and that curiously the use of in and out projections results in a decrease in per\-step loss despite their redundancy, although there is an approximately 10% training throughput cost to using projections compared to not using them,\(Table[S2](https://arxiv.org/html/2605.08696#S1.T2)\-[S7](https://arxiv.org/html/2605.08696#S1.T7)\)\. We employ head projections in all models unless otherwise noted\. Pytorch\-based recurrent and sequence\-parallel SRM mixer representation implementations are confirmed to be numerically stable and yield identical outputs even when cache values and decay constants are used in half precision\.

### 5\.2Causal Training Efficiency

As one of the largest challenges to recurrent models is efficient and stable causal language model training, we measured the training efficiency of SRMs as compared to both quadratic\-complexity and other linear\-complexity models\. For Transformers, we use the Llama 2 architecture\(Touvron et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib22)\)featuring RoPE, RMS norms and improved parameter initialization which have been shown to outperform early causal Transformer architectures, and models ofdm=xd\_\{m\}=xachieve similar training throughputs of Mixers withdm=2​xd\_\{m\}=2xover the range ofdmd\_\{m\}observed here\(Badger,[2025](https://arxiv.org/html/2605.08696#bib.bib20)\)\. We observe training efficiency through the lens of real world limitations, specifically GPU throughput and memory required, which can be converted to loss achieved per FLOP applied during training\.

We find that the SRM is slightly less efficient to train than𝒪​\(n2​d2\)\\mathcal\{O\}\(n^\{2\}d^\{2\}\)complexity Transformer and Masked Mixer models of similar sample throughput across diverse language understanding \(FineWeb\-edu\) and mathematics and coding \(FineMath\) text corpora\. For these datasets, we find that SRMs are more efficient to train than Mamba \(2\) models \(Figure[3](https://arxiv.org/html/2605.08696#S5.F3), Tables[S3](https://arxiv.org/html/2605.08696#S1.T3),[S4](https://arxiv.org/html/2605.08696#S1.T4)\)\. We note that we find Mamba models to be only slightly less efficiently trainable than Transformers of similar throughput on a per\-parameter basis \(Table[S3](https://arxiv.org/html/2605.08696#S1.T3)\) as reported by\(Dao and Gu,[2024](https://arxiv.org/html/2605.08696#bib.bib19)\), but that the substantially lower throughput on device \(for H100s see Table[S6](https://arxiv.org/html/2605.08696#S1.T6),[S11](https://arxiv.org/html/2605.08696#S1.T11)\) for parameter\-equivalent models results in more efficient SRM training in practice\. On older devices \(V100s\) Mamba models experience a \>20x drop in throughput and are virtually untrainable as a result, whereas SRMs perform consistently across all devices tested\.

Early in training RWKV \(4\) models exhibit similar efficiency as SRMs but suffer from catastrophic numerical instabilities across various hyperparameter settings, compute datatypes, and architectural configurations and as such are not explored in detail\.

![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/srm_scaling_figure.png)Figure 3:SRM causal training on FineWeb\. \(a\) Throughput\- and memory\-equivalent model training:dm=1024d\_\{m\}=1024SRM,dm=512d\_\{m\}=512Transformer,dm=256d\_\{m\}=256Mamba, allnl=16n\_\{l\}=16\. \(b\) SRM sample scaling characteristics for fixed model sizes, \(c\) SRM training compute \(assuming theoretical maximum FLOPs on device as measured by V100 throughput\) and loss scaling\.
### 5\.3Benchmark Evaluations

To assess how the training efficiencies impact benchmark performance, we compare functional characteristics of linear\-complexity models by evaluating them on a number of single\-sample benchmark tasks without regard to the compute required to generate samples\. We benchmark for general question answering via Arc\-Easy\(Clark et al\.,[2018](https://arxiv.org/html/2605.08696#bib.bib23)\), sentence completion using HellaSwag\(Zellers et al\.,[2019](https://arxiv.org/html/2605.08696#bib.bib24)\)and Lambada \(OpenAI processed\)\(Paperno et al\.,[2016](https://arxiv.org/html/2605.08696#bib.bib25)\), information retention via SQuAD\(Rajpurkar et al\.,[2016](https://arxiv.org/html/2605.08696#bib.bib26)\), SQuAD 2\(Rajpurkar et al\.,[2018](https://arxiv.org/html/2605.08696#bib.bib27)\), Longbench\(Bai et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib28)\), IFEval\(Zhou et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib29)\), SWDE\(Arora et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib30)\), and xWinoGrad\(Tikhonov and Ryabinin,[2021](https://arxiv.org/html/2605.08696#bib.bib31)\)using Eleuther\-AI’s evaluation harness\(Gao et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib32)\)\. In Table[1](https://arxiv.org/html/2605.08696#S5.T1)we detail benchmarks for compute\-equivalent Transformer, SRM, and Mamba models, and we find that Transformers top the most benchmarks of all three models and that SRMs outperform Mamba models on most benchmarks observed\.

Table 1:Training compute\-matched SRM, Transformer, and Mamba benchmark scores after FineWeb\-edu training\. All benchmarks are higher is better apart from WikiText\.
### 5\.4Informational Retention and Capacity

The ability of quadratic\-complexity models to retain input information in their hidden layers is of relatively minor importance for causal language modeling for the reason that these models are exposed to a cached representation of each input token upon each prediction step, and the prediction itself typically requires the knowledge of only a small subset of previous tokens\. For linear\-complexity models, however, the model is only exposed to a fixed\-size cache of all previous tokens and cannot simply reference individual token representations separately\. SRMs in particular hold a cache size equivalent to one token’s embedding per layer for the entire input\. This motivates the measurement of the amount of information present in the hidden states of causal\-trained models \(which we term ‘information retention’\) as well as the ability of these models to store input information in one or more states \(‘information capacity’, or equivalently the upper bound of information retention\)\. These questions are especially important to SRMs because Masked Mixers trained for maximum information retention contain qualitatively different token mixing matrices than those of models trained for next token prediction \(Figure[1](https://arxiv.org/html/2605.08696#S4.F1),[S1](https://arxiv.org/html/2605.08696#S1.F1)\) and SRMs are by design only reflect those present for next token prediction\.

We measure model information retention by testing the capability of a CLM\-trained model’s embedding to regenerate input via a trained decoder\. Given the output \(i\.e\. the vector directly responsible for a single token prediction, namely the last hidden layer’s last embedding\) of token\-predicting modelℰ\\mathcal\{E\}applied to inputxx,y=ℰ​\(x,θe\)y=\\mathcal\{E\}\(x,\\theta\_\{e\}\), we seek to invertyyto recoverxx\. We do this by first freezingθe\\theta\_\{e\}and then training a decoder to recoverxxgiven an unrolled form ofyy, where the decoder yieldsz=𝒟\(U\(ℰ\(x,θe\),θd\)z=\\mathcal\{D\}\\left\(U\(\\mathcal\{E\}\(x,\\theta\_\{e\}\),\\theta\_\{d\}\\right\)and the difference inxxandzzin terms of both Cross\-Entropy Loss as well as the more interpretable entropy ratio, the fraction of input information the model retains, after\(Badger,[2026](https://arxiv.org/html/2605.08696#bib.bib33)\)\.

We find that causal SRMs exhibit higher information retention than causal Mamba models, and that these differences exist before and after training \(Table[2](https://arxiv.org/html/2605.08696#S5.T2)\); furthermore comparing our results to previous work\(Badger,[2026](https://arxiv.org/html/2605.08696#bib.bib33)\)we find that SRMs exhibit similar information retention characteristics to quadratic complexity models\. We measure information capacity, or the amount of information a model is capable of retaining from its inputs, in two ways: encoder\-decoder style input regeneration and copy task training efficiency\. The encoder\-decoder input regeneration is identical to that detailed above except the encoderθe\\theta\_\{e\}is initialized from scratch and is unfrozen during training\. The copy task is a measure of a model’s ability to learn to copy 512 tokens in relatively few \(10k instead of 200k\) training steps\. Copying may be thought of as a less information\-intensive task than regeneration, such that copy task training illustrates a model’s ability store smaller amounts of information relatively little training whereas input regeneration measures a models’ ability to store more information with more training\.

We find that SRMs exhibit higher informational capacity as measured by encoder\-decoder input regeneration \(Table[3](https://arxiv.org/html/2605.08696#S5.T3)\) and similar or slightly better copy task training efficiency when adjusted for compute but not on an equivalent embedding size basis, with variation dependent on SRM architecture \(Table[4](https://arxiv.org/html/2605.08696#S5.T4)\)\. To conclude, the qualitative features transferred to SRMs from causal Masked Mixers typically do not reduce information capacity, and SRMs exhibit superior information retention and superior information capacity to Mamba models on a per\-compute basis\.

Table 2:Causal pretrained model information retention\. All modelsdm=512,nl=16,nc​t​x=512d\_\{m\}=512,n\_\{l\}=16,n\_\{ctx\}=512Table 3:Architecture Information Capacity, measured via encoder\-decoder input reconstruction\. All modelsdm=512,nl=16,nc​t​x=512d\_\{m\}=512,n\_\{l\}=16,n\_\{ctx\}=512Table 4:Copy accuracy \(512 tokens copied at 10k training steps\)dm=512,nl=16,nc​t​x=1024d\_\{m\}=512,n\_\{l\}=16,n\_\{ctx\}=1024\. ‘k’ denotes kernel number for SRM mixer layers, and H=0 denotes a non\-headed model\. Most SRM configurations are typically less efficient copy learners perdmd\_\{m\}size but more efficient per compute than Mamba models and achieve high copy accuracy with limited training steps\.These experiments suggest that while SRMs are capable of capturing nearly all of the information present in a reasonably sized input, causal\-trained SRMs are typically not capable of doing so\. This is true for other model architectures trained for causal modeling\(Badger,[2026](https://arxiv.org/html/2605.08696#bib.bib33)\), but is particularly detrimental for recurrent models because they store fixed\-sized caches\. Recurrent models may be thought of as forming a compression of the input in their cache, and in order to retrieve this information accurately the compression must be of high fidelity and therefore scalingnc​t​xn\_\{ctx\}beyond the limit of compression achievable by the model necessarily leads to information loss and poor performance\. Compression limits are highly dataset\-dependent, and for low\-information data such as DNA sequences are much lower \(in terms of bits required per byte of input\) than they are for high\-information inputs like language\. For language, assuming a 0\.54 bits per byte \(14\.8x\) compression achieved by recent LLMs\(Deepseek,[2025](https://arxiv.org/html/2605.08696#bib.bib34)\)adm=1024d\_\{m\}=1024embedding would be expected to be able to losslessly compress up to1024​p​a​r​a​m​s∗2​b​y​t​e​sp​a​r​a​m∗0\.5​t​o​k​e​n​sb​y​t​e∗14\.8=151551024\\;params\*2\\frac\{bytes\}\{param\}\*0\.5\\frac\{tokens\}\{byte\}\*14\.8=15155tokens, although this is likely an overestimation given that it assumes that decoding the cache results in no further compression\. We test the idea that scaling recurrent models’ context beyond their compression limit necessarily results in performance degradation by observing copy fidelity as the number of copied tokens increases for a fixed training dataset and SRM model, and find a decline in copy accuracy \(from \>90% to <30%\) as the number of copied tokens increases \(Figure[S2](https://arxiv.org/html/2605.08696#S1.F2)\)\. This motivates the following investigation into scaling recurrent models in the batch, rather than the sequence, dimension\.

### 5\.5Recurrent Throughput and Concurrency

For causal language model inference, tokens must be generated sequentially \(or perhaps a few at a time\) as the true distribution is not knowna priori, such that parallelization across the sequence dimension is no longer possible\. Instead one can parallelize across the batch, or sample, dimension with the relatively safe assumption that the samples are independent\. It is clear that the recurrent representation of SRMs is more efficient to parallelize in the sample dimension than linear space models such as the Transformer, as forsssamples andlllayers a recurrent model requiress​l​dsldspace per layer and at each next token prediction transferss​l​dsldelements from memory, whereas Transformers requires​l​d​nc​t​xsldn\_\{ctx\}space and transfers​l​d​nc​t​x/2sldn\_\{ctx\}/2elements from memory on average\. If memory transfer alone was limiting, this means that one would expect to observe a throughput increase of\(nc​t​x/2\)\(n\_\{ctx\}/2\)x in SRMs compared to Transformers of equivalent size\. Asymptotically asnc​t​x→∞n\_\{ctx\}\\to\\inftythis would be true of real devices, but for our context lengths the penalty for larger batch sizes through the token feedforward layers is non\-trivial\. Even for the relatively small context length sizes used in this work, however, we find that there is an order\-of\-magnitude increase in throughput \(Table[5](https://arxiv.org/html/2605.08696#S5.T5)\)\. Likewise, we would expect for annc​t​xn\_\{ctx\}x increase in the memory required in Transformers compared to SRMs asnc​t​x→∞n\_\{ctx\}\\to\\infty, and we observe a two orders\-of\-magnitude increase in maximum concurrency \(Table[5](https://arxiv.org/html/2605.08696#S5.T5)\)\.

Table 5:Throughput \(tokens per second\) per context length, 1x V100 \(16GB\) with Pytorch implementations, SRMs aredm=1024d\_\{m\}=1024and Transformers aredm=512d\_\{m\}=512, bothnl=16,h=4n\_\{l\}=16,h=4\. Note that the decrease in batch size fornc​t​x\>1024n\_\{ctx\}\>1024for SRMs is not due to the model using more memory, but due to the memory required to store generated tokens via a non\-optimized in\-memory store\.When we compare throughput and concurrency of the SRM to Mamba models, we find that the former experiences much higher \(\>500​x\>500x\) concurrency and approximately7​x7xthe throughput of equivalently powerful models of the latter, although Mamba models are in turn observed to exhibit increased throughput compared to Transformers \(Table[6](https://arxiv.org/html/2605.08696#S5.T6)\)\. We attribute this substantial difference to tiny number of transformations SRMs perform during sequence mixing \(two vector\-scalar multiplications and a vector addition per layer\)\.

Table 6:Maximum throughput \(tokens per second\) and concurrency \(samples\) per model, 1x H100 \(96GB\) with Pytorch implementations, allnc​t​x=512,nl=8n\_\{ctx\}=512,n\_\{l\}=8\. RWKV is limited by default to 1024 batch size, and the maximum predicted batch for thedm=512d\_\{m\}=512model is around 2000\.To demonstrate the benefits of building a production\-grade inference framework for recurrent models, we implemented recurrent SRMs using Modular MAX inference framework\([Modular,](https://arxiv.org/html/2605.08696#bib.bib35)\)and verified that implementing an SRM in MAX resulted in substantially higher throughput \(7\.5x\) and concurrency \(1\.75x\) compared to that obtained by Torch implementations using identical hardware\. We then compared the throughput of SRMs to those of equivalently capable transformers applied using vLLM, a widely applied efficient inference engine built around paged attention\(Kwon et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib36)\)\. We observe similar efficiency increases in our MAX SRM implementation compared to a Transformer inferenced via vLLM as we did for Pytorch implementations of both models: the SRM exhibits 12x the throughput and 170x the concurrency of a pretraining\-matched \(26x the throughput and 225x the concurrency of a parameter\-matched\) Transformer when generating 1k tokens on a single Nvidia H100 \(Tables[7](https://arxiv.org/html/2605.08696#S5.T7),[S1](https://arxiv.org/html/2605.08696#S1.T1)\)\. On the optimized Albatross inference engine, the RWKV exhibits increased throughput compared to the Transformer but still exhibits only a fraction of the throughput of the SRM\.

Table 7:Maximum throughput and concurrency per model, 1x H100 \(96GB\) with optimized engine inferences \(SRMs via our Mojo/MAX implementation, Transformers and Mamba on vLLM with optimized kernels if appropriate, RWKV \(7\) on Albatross\)\. Non\-SRM models are alldm=512d\_\{m\}=512, and all arenl=8n\_\{l\}=8with an 8k\-size tokenizer\. Mamba models exhibit a memory leak and do not reach their theoretical limit of 20k concurrent samples, achieving only 300\.At the throughputs achieved by SRMs, the operation of finding the next tokens given the model’s logits becomes a bottleneck: for example, the throughput of thedm=512d\_\{m\}=512with argmax on the logits to compute the next token is only 63% that of the same model without applying argmax \(Table[S1](https://arxiv.org/html/2605.08696#S1.T1)\)\.

We comment that these results illustrate the versatility of the SRM architecture: unlike Mamba and RWKV which specialized device kernels for performant inference, SRMs are easily incorporated into other inference frameworks and do not require substantial rewrites for new hardware\. This is a significant problem for training as well as inference: Mamba kernels perform poorly on older hardware \(V100s\) due to differences in memory management, and both Mamba and RWKV kernels will likely need to be rewritten for new GPU architectures in order to maintain performance as L1/L2/shared cache memory management is likely to change in the future\.

### 5\.6Verifiable Output Scaling

We next investigated whether the 10\-30x throughput increase for SRMs versus Transformers results in materially higher performance on benchmarks in which outputs may be verified, and in particular we focus on GSM8k with verification in the form of an oracle that tests equivalence of the model’s last numerical output for that question with the correct answer\. The following question is addressed: given a fixed amount of compute and memory \(over both train and test time\), what is most efficient for obtaining at least one correct sample per question? Equivalently, how does Pass@k scale with increasedkkfor SRMs, and in particular does it follow a similar power law found for Transformers across many tasks\(Brown et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib1)\)? We find that the Pass@k scaling for SRMs is nearly identical to that for Transformers for GSM8k \(Figure[4](https://arxiv.org/html/2605.08696#S5.F4)\)\. In practical terms, this means that to reach a given target coverage the transformer requires around 10x the inference compute, and for a constant compute budget the SRM gives a correct answer to around 30% more GSM8k questions than the Transformer\. We find that this relationship holds regardless of preparation for the GSM8k dataset: samples from SRMs and Transformers pretrained on 13B tokens of FineMath\-4\+ followed by SFT on the GSM8k training dataset are by eye more coherent and exhibit higher Pass@k accuracy than those from models trained on FineWeb, but again SRMs exhibit around 30% more coverage per fixed compute\.

![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/sample_scaling_fig.png)Figure 4:\(a\) GSM8k Inference Scaling, in compute per SRM sample\. ‘Math’ models are pretrained on FineMath 4\+ followed by SFT on GSM8k, others are pretrained only on FineWeb\. \(b\) Compute\-constant GRPO \(50 samples per batch for SRM, 5 for Transformers\) on Math\-trained models, with balanced resampling for SRMs\. SRMs aredm=1024d\_\{m\}=1024and Transformerdm=512d\_\{m\}=512
### 5\.7Reinforcement Learning

Greater inference throughput afforded via SRMs would be expected to benefit reinforcement learning methods that make use of parallel generation with the goal of maximizing the probability that one or a few outputs are correct\. One class of this criteria is where a model is used to generate multiple token sequences \(rollouts\) independently in parallel, which are then evaluated and reinforced depending on whether they fulfill some criteria\. This class encompasses PPO\(Schulman et al\.,[2017](https://arxiv.org/html/2605.08696#bib.bib37)\)as well as GRPO\(Shao et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib38)\)and its variants\(Liu et al\.,[2025](https://arxiv.org/html/2605.08696#bib.bib39)\)\. With some simplification, GRPO proceeds by first generating a groupGGof samples from a model for a given question, verifying the goodness of each sample using reward functions, and then performing gradient descent to change the model’s parameters such that ‘good’ outputs with higher rewards are more likely to be sampled whereas bad outputs are less likely\. Parallel generation coupled with random sampling \(typically top\-p nucleus sampling\(Holtzman et al\.,[2020](https://arxiv.org/html/2605.08696#bib.bib40)\)with temperature\) results in distinct samples usually containing non\-identical token sequences\. In this work we evaluate on outcomes only, such that the rewardr=1\.0r=1\.0if evaluation is correct andr=0\.0r=0\.0if incorrect\.

𝒥=𝔼\[q∼Pq,\{Gr\}∼πθ\(\{oi\}i=1Gr\|q\)\]1Gr∑i1\|oi\|∑t=1\|oi\|\(Ai,t−β𝔻K​L\(πθ\|\|πθr​e​f\)\)\\mathcal\{J\}=\\mathbb\{E\}\[q\\sim P\_\{q\},\\\{G\_\{r\}\\\}\\sim\\pi\_\{\\theta\}\(\\\{o\_\{i\}\\\}\_\{i=1\}^\{G\_\{r\}\}\|q\)\]\\frac\{1\}\{G\_\{r\}\}\\sum\_\{i\}\\frac\{1\}\{\|o\_\{i\}\|\}\\sum\_\{t=1\}^\{\|o\_\{i\}\|\}\\left\(A\_\{i,t\}\-\\beta\\mathbb\{D\}\_\{KL\}\(\\pi\_\{\\theta\}\|\|\\pi\_\{\\theta\_\{ref\}\}\)\\right\)\(6\)
One of the primary challenges with parallelized rollout RL algorithms such as GRPO is the inherent limitation that the model must be capable of generating at least one ‘good’ output of the batch \(here wherer=1\.0r=1\.0\), or the model has no reference with which to better approximate\. GRPO is particularly susceptible to decreases in model exploration when few samples are rewarded because an implicit penalty is applied for each incorrect sample which limits exploration relative to SFT\-trained models, which manifests in lower Pass@k \(k\>1\) accuracy of GRPO\- versus SFT\-trained models\(Shao et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib38)\), which we find worsens with as the generation batch sizeGGincreases \(Figure[S3](https://arxiv.org/html/2605.08696#S1.F3)\)\. To prevent this, we introduce a balanced resampling approach where for each sampled batchGGof generations per question, given a training batch sizebbthe generations are evaluated and we resampleGGto include up tob/2b/2‘good’ samples, before resamplingGGselecting for ‘bad’ samples to make the training batch reach sizebb\. Denoting this resampledGGasGrG\_\{r\}, and simplifying for the case where one training step is taken per generation batch such thatπθo​l​d=πθ\\pi\_\{\\theta\_\{old\}\}=\\pi\_\{\\theta\}then consequentlyπθ​\(oi\|q,oi\)/πθo​l​d​\(oi\|q,oi\)=1\\pi\_\{\\theta\}\(o\_\{i\}\|q,o\_\{i\}\)/\\pi\_\{\\theta\_\{old\}\}\(o\_\{i\}\|q,o\_\{i\}\)=1, and thus we maximize the objective function given in Equation[6](https://arxiv.org/html/2605.08696#S5.E6)when batches are large \(ie for SRMs\) and use the standard GRPO objective for small\-batch training\. In practice we typically pool samples from multiple questions perGG, such that resampling occurs over these pooled batches\.

With FineMath 4\+ pretrained models, we performed SFT on GSM8k followed by GRPO and confirmed that balanced resampling substantially increases Pass@k for SRMs \(Figure[S3](https://arxiv.org/html/2605.08696#S1.F3)\)\. We then compared GRPO training of SRMs generating 50 samples per question versus Transformers with 5 samples per question, both with a single forward/reverse pass per batch and SRMs using balanced resampling\. The results of both Pass@k after GRPO, as well as the difference between models before and after GRPO are shown in Figure[4](https://arxiv.org/html/2605.08696#S5.F4)\(b\), and we observe that the single\-sample accuracy is slightly higher for SRMs than for Transformers\. The tradeoff between higher single\-sample accuracy and lower exploration is clearly visible for Transformers, but in contrast we find that SRMs with their larger, resampled batches lead to increased successful exploration\.

## 6Discussion

### 6\.1Limitations of this work

The SRM architecture exhibit a few notable limitations: firstly as implemented it is not suitable for modeling samples where the sequence dimension far exceeds the hidden dimension due to the quadratic complexity \(w\.r\.t\. token length\) at training, secondly there is no built\-in method for context length and candidate methods have yet to be tested, and finally the architecture is most easily optimized for larger parameter numbers with fewer activations compared to Mamba or Transformer models, and is unsuitable for applications requiring minimal parameter counts\. This work is limited in scale and also in scope, and does not explore a number of optimizations \(parallel scan for recurrent prefill, for example\)\. One of the main outstanding question relates to the benefits or drawbacks one receives for various classes of problems when using a Transformer compared to an SRM: is it better to use a somewhat more capable model yielding fewer samples or a somewhat less capable one that gives \>10x the throughput?

### 6\.2The Future

Many\-threaded versus single\-threaded performance has increased exponentially the last couple of decades, as has the ratio of global memory bandwidth to arithmetic operation latency \(the ‘memory wall’\) on device\(Kirk and Wen\-Mei,[2016](https://arxiv.org/html/2605.08696#bib.bib41); Gholami et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib42)\), and it is likely that these features will continue for the foreseeable future\. Although there is a limit to how parallelizable current inference scaling techniques such as Chain of Thought\(Wei et al\.,[2023](https://arxiv.org/html/2605.08696#bib.bib43)\)are, with some caveats one would expect for the performance \(i\.e\. useful sample production efficiency\) of parallel sample generation to far outstrip single\-sample sequential generation in the years to come because of the correspondence of parallel sample generation to parallel threads on device\. Of the architectures studied in this work, SRMs are the most suitable for parallel generation and furthermore are ideally situated for dealing with the memory wall, as they exhibit a higher ratio of arithmetic operations to memory accesses\. If these thread and memory characteristics hold in the future, SRMs will continue to increase their relative efficiency compared to memory\-intensive models such as Transformers\. Signs of practical aspects of this can be seen in this work: the ratio of SRM to Transformer throughput on V100s atnc​t​x=512n\_\{ctx\}=512is 9\.66x, whereas for the newer H100 is 10\.56x \(using compiled Pytorch implementations\)\.

## References

- Brown et al\. \[2024\]Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V\. Le, Christopher Ré, and Azalia Mirhoseini\.Large language monkeys: Scaling inference compute with repeated sampling, 2024\.URL[https://arxiv\.org/abs/2407\.21787](https://arxiv.org/abs/2407.21787)\.
- Chen et al\. \[2025\]Mouxiang Chen, Binyuan Hui, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Jianling Sun, Junyang Lin, and Zhongxin Liu\.Parallel scaling law for language models, 2025\.URL[https://arxiv\.org/abs/2505\.10475](https://arxiv.org/abs/2505.10475)\.
- Foundation \[2026\]ARC Prize Foundation\.Arc\-agi\-3: A new challenge for frontier agentic intelligence, 2026\.URL[https://arxiv\.org/abs/2603\.24621](https://arxiv.org/abs/2603.24621)\.
- Rumelhart et al\. \[1985\]David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams\.Learning internal representations by error propagation\.Technical report, 1985\.
- Hochreiter and Schmidhuber \[1997\]Sepp Hochreiter and Jürgen Schmidhuber\.Long short\-term memory\.*Neural Computation*, 9\(8\):1735–1780, 1997\.doi:[10\.1162/neco\.1997\.9\.8\.1735](https://doi.org/10.1162/neco.1997.9.8.1735)\.
- Vaswani et al\. \[2023\]Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N\. Gomez, Lukasz Kaiser, and Illia Polosukhin\.Attention is all you need\.2023\.URL[https://arxiv\.org/abs/1706\.03762](https://arxiv.org/abs/1706.03762)\.
- Mishra et al\. \[2024\]Mayank Mishra, Matt Stallone, Gaoyuan Zhang, Yikang Shen, Aditya Prasad, Adriana Meza Soria, Michele Merler, Parameswaran Selvam, Saptha Surendran, Shivdeep Singh, Manish Sethi, Xuan\-Hong Dang, Pengyuan Li, Kun\-Lung Wu, Syed Zawad, Andrew Coleman, Matthew White, Mark Lewis, Raju Pavuluri, Yan Koyfman, Boris Lublinsky, Maximilien de Bayser, Ibrahim Abdelaziz, Kinjal Basu, Mayank Agarwal, Yi Zhou, Chris Johnson, Aanchal Goyal, Hima Patel, Yousaf Shah, Petros Zerfos, Heiko Ludwig, Asim Munawar, Maxwell Crouse, Pavan Kapanipathi, Shweta Salaria, Bob Calio, Sophia Wen, Seetharami Seelam, Brian Belgodere, Carlos Fonseca, Amith Singhee, Nirmit Desai, David D\. Cox, Ruchir Puri, and Rameswar Panda\.Granite code models: A family of open foundation models for code intelligence, 2024\.URL[https://arxiv\.org/abs/2405\.04324](https://arxiv.org/abs/2405.04324)\.
- \[8\]Granite 4\.0 \- IBM Granite — ibm\.com\.[https://www\.ibm\.com/granite/docs/models/granite](https://www.ibm.com/granite/docs/models/granite)\.\[Accessed 16\-04\-2026\]\.
- Blakeman and et al\. \[2025\]NVIDIA: Aaron Blakeman and Aaron Grattafiori et al\.Nvidia nemotron 3: Efficient and open intelligence, 2025\.URL[https://arxiv\.org/abs/2512\.20856](https://arxiv.org/abs/2512.20856)\.
- Qwen Team \[2026\]Qwen Team\.Qwen3\.5: Towards native multimodal agents, February 2026\.URL[https://qwen\.ai/blog?id=qwen3\.5](https://qwen.ai/blog?id=qwen3.5)\.
- Team et al\. \[2025\]Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, Wentao Li, Enzhe Lu, Weizhou Liu, Yanru Chen, Weixin Xu, Longhui Yu, Yejie Wang, Yu Fan, Longguang Zhong, Enming Yuan, Dehao Zhang, Yizhi Zhang, T\. Y\. Liu, Haiming Wang, Shengjun Fang, Weiran He, Shaowei Liu, Yiwei Li, Jianlin Su, Jiezhong Qiu, Bo Pang, Junjie Yan, Zhejun Jiang, Weixiao Huang, Bohong Yin, Jiacheng You, Chu Wei, Zhengtao Wang, Chao Hong, Yutian Chen, Guanduo Chen, Yucheng Wang, Huabin Zheng, Feng Wang, Yibo Liu, Mengnan Dong, Zheng Zhang, Siyuan Pan, Wenhao Wu, Yuhao Wu, Longyu Guan, Jiawen Tao, Guohong Fu, Xinran Xu, Yuzhi Wang, Guokun Lai, Yuxin Wu, Xinyu Zhou, Zhilin Yang, and Yulun Du\.Kimi linear: An expressive, efficient attention architecture, 2025\.URL[https://arxiv\.org/abs/2510\.26692](https://arxiv.org/abs/2510.26692)\.
- Kaplan et al\. \[2020\]Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B\. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei\.Scaling laws for neural language models, 2020\.URL[https://arxiv\.org/abs/2001\.08361](https://arxiv.org/abs/2001.08361)\.
- Katharopoulos et al\. \[2020\]Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret\.Transformers are rnns: Fast autoregressive transformers with linear attention, 2020\.URL[https://arxiv\.org/abs/2006\.16236](https://arxiv.org/abs/2006.16236)\.
- Shen et al\. \[2018\]Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li\.Efficient attention: Attention with linear complexities\.*CoRR*, abs/1812\.01243, 2018\.URL[http://arxiv\.org/abs/1812\.01243](http://arxiv.org/abs/1812.01243)\.
- Poli et al\. \[2023\]Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y\. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré\.Hyena hierarchy: Towards larger convolutional language models, 2023\.URL[https://arxiv\.org/abs/2302\.10866](https://arxiv.org/abs/2302.10866)\.
- Peng et al\. \[2023\]Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S\. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui\-Jie Zhu\.Rwkv: Reinventing rnns for the transformer era, 2023\.URL[https://arxiv\.org/abs/2305\.13048](https://arxiv.org/abs/2305.13048)\.
- Fu et al\. \[2023\]Daniel Y\. Fu, Tri Dao, Khaled K\. Saab, Armin W\. Thomas, Atri Rudra, and Christopher Ré\.Hungry hungry hippos: Towards language modeling with state space models, 2023\.URL[https://arxiv\.org/abs/2212\.14052](https://arxiv.org/abs/2212.14052)\.
- Gu and Dao \[2024\]Albert Gu and Tri Dao\.Mamba: Linear\-time sequence modeling with selective state spaces, 2024\.URL[https://arxiv\.org/abs/2312\.00752](https://arxiv.org/abs/2312.00752)\.
- Dao and Gu \[2024\]Tri Dao and Albert Gu\.Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024\.URL[https://arxiv\.org/abs/2405\.21060](https://arxiv.org/abs/2405.21060)\.
- Badger \[2025\]Benjamin L\. Badger\.Masked mixers for language generation and retrieval, 2025\.URL[https://arxiv\.org/abs/2409\.01482](https://arxiv.org/abs/2409.01482)\.
- Qin and Zhong \[2023\]Zhen Qin and Yiran Zhong\.Accelerating toeplitz neural network with constant\-time inference complexity\.In Houda Bouamor, Juan Pino, and Kalika Bali, editors,*Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing*, pages 12206–12215, Singapore, December 2023\. Association for Computational Linguistics\.doi:[10\.18653/v1/2023\.emnlp\-main\.750](https://doi.org/10.18653/v1/2023.emnlp-main.750)\.URL[https://aclanthology\.org/2023\.emnlp\-main\.750/](https://aclanthology.org/2023.emnlp-main.750/)\.
- Touvron et al\. \[2023\]Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie\-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom\.Llama 2: Open foundation and fine\-tuned chat models, 2023\.URL[https://arxiv\.org/abs/2307\.09288](https://arxiv.org/abs/2307.09288)\.
- Clark et al\. \[2018\]Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord\.Think you have solved question answering? try arc, the ai2 reasoning challenge\.*ArXiv*, abs/1803\.05457, 2018\.
- Zellers et al\. \[2019\]Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi\.Hellaswag: Can a machine really finish your sentence?, 2019\.URL[https://arxiv\.org/abs/1905\.07830](https://arxiv.org/abs/1905.07830)\.
- Paperno et al\. \[2016\]Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández\.The lambada dataset: Word prediction requiring a broad discourse context, 2016\.URL[https://arxiv\.org/abs/1606\.06031](https://arxiv.org/abs/1606.06031)\.
- Rajpurkar et al\. \[2016\]Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang\.Squad: 100,000\+ questions for machine comprehension of text, 2016\.URL[https://arxiv\.org/abs/1606\.05250](https://arxiv.org/abs/1606.05250)\.
- Rajpurkar et al\. \[2018\]Pranav Rajpurkar, Robin Jia, and Percy Liang\.Know what you don’t know: Unanswerable questions for SQuAD\.In Iryna Gurevych and Yusuke Miyao, editors,*Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics \(Volume 2: Short Papers\)*, pages 784–789, Melbourne, Australia, July 2018\. Association for Computational Linguistics\.doi:[10\.18653/v1/P18\-2124](https://doi.org/10.18653/v1/P18-2124)\.URL[https://aclanthology\.org/P18\-2124](https://aclanthology.org/P18-2124)\.
- Bai et al\. \[2024\]Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li\.LongBench: A bilingual, multitask benchmark for long context understanding\.In*Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\)*, pages 3119–3137, Bangkok, Thailand, August 2024\. Association for Computational Linguistics\.doi:[10\.18653/v1/2024\.acl\-long\.172](https://doi.org/10.18653/v1/2024.acl-long.172)\.URL[https://aclanthology\.org/2024\.acl\-long\.172](https://aclanthology.org/2024.acl-long.172)\.
- Zhou et al\. \[2023\]Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou\.Instruction\-following evaluation for large language models, 2023\.URL[https://arxiv\.org/abs/2311\.07911](https://arxiv.org/abs/2311.07911)\.
- Arora et al\. \[2024\]Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré\.Simple linear attention language models balance the recall\-throughput tradeoff, 2024\.
- Tikhonov and Ryabinin \[2021\]Alexey Tikhonov and Max Ryabinin\.It’s all in the heads: Using attention heads as a baseline for cross\-lingual transfer in commonsense reasoning, 2021\.
- Gao et al\. \[2024\]Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou\.The language model evaluation harness, 07 2024\.URL[https://zenodo\.org/records/12608602](https://zenodo.org/records/12608602)\.
- Badger \[2026\]Benjamin L\. Badger\.Language model memory and memory models for language, 2026\.URL[https://arxiv\.org/abs/2602\.13466](https://arxiv.org/abs/2602.13466)\.
- Deepseek \[2025\]Deepseek\.Deepseek\-v3 technical report\.2025\.URL[https://arxiv\.org/abs/2412\.19437](https://arxiv.org/abs/2412.19437)\.
- \[35\]Modular\.MAX: A high\-performance inference framework for AI — modular\.com\.[https://www\.modular\.com/open\-source/max](https://www.modular.com/open-source/max)\.\[Accessed 20\-04\-2026\]\.
- Kwon et al\. \[2023\]Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica\.Efficient memory management for large language model serving with pagedattention\.In*Proceedings of the 29th symposium on operating systems principles*, pages 611–626, 2023\.
- Schulman et al\. \[2017\]John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov\.Proximal policy optimization algorithms, 2017\.URL[https://arxiv\.org/abs/1707\.06347](https://arxiv.org/abs/1707.06347)\.
- Shao et al\. \[2024\]Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y\. K\. Li, Y\. Wu, and Daya Guo\.Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024\.URL[https://arxiv\.org/abs/2402\.03300](https://arxiv.org/abs/2402.03300)\.
- Liu et al\. \[2025\]Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin\.Understanding r1\-zero\-like training: A critical perspective, 2025\.URL[https://arxiv\.org/abs/2503\.20783](https://arxiv.org/abs/2503.20783)\.
- Holtzman et al\. \[2020\]Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi\.The curious case of neural text degeneration, 2020\.URL[https://arxiv\.org/abs/1904\.09751](https://arxiv.org/abs/1904.09751)\.
- Kirk and Wen\-Mei \[2016\]David B Kirk and W Hwu Wen\-Mei\.*Programming massively parallel processors: a hands\-on approach*\.Morgan kaufmann, 2016\.
- Gholami et al\. \[2024\]Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W\. Mahoney, and Kurt Keutzer\.Ai and memory wall, 2024\.URL[https://arxiv\.org/abs/2403\.14123](https://arxiv.org/abs/2403.14123)\.
- Wei et al\. \[2023\]Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou\.Chain\-of\-thought prompting elicits reasoning in large language models, 2023\.URL[https://arxiv\.org/abs/2201\.11903](https://arxiv.org/abs/2201.11903)\.
- von Werra et al\. \[2020\]Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec\.TRL: Transformers Reinforcement Learning, 2020\.URL[https://github\.com/huggingface/trl](https://github.com/huggingface/trl)\.

## S1Appendix

### S1\.1Training Details

Pretraining proceed as follows: first datasets were pretokenized and padded as necessary, and models were initialized and compiled where possible and optimized via the Pytorch implementation of AdamW withβ1=0\.9,β2=0\.999\\beta\_\{1\}=0\.9,\\beta\_\{2\}=0\.999\. Most pretraining datasets use a constant token number of200,000∗512∗128=13\.1​B200,000\*512\*128=13\.1Btokens, and this is held constant for longer\-context trainings \(here we have200,000∗1024∗64=13\.1​B200,000\*1024\*64=13\.1Btokens\)\. We useη=5∗10−4\\eta=5\*10^\{\-4\}for SRMs and Mamba models andη=2∗10−4\\eta=2\*10^\{\-4\}for Transformers after optimizing this hyperparameter independently, and use a 4000\-step warm\-up with linear learning rate decay\.

We train for supervised finetuning by masking input tokens and training only on model\-generated outputs, minimizing Cross\-Entropy Loss between these and the answers provided by the GSM8k training dataset\. We train withη=1∗10−4\\eta=1\*10^\{\-4\}, gradient clipping with maximum grad norm of 0\.1, weight decay of 0\.1, warmup ratio of 0\.1, and optimize using the Pytorch implementation of AdamW withβ1=0\.9,β2=0\.999\\beta\_\{1\}=0\.9,\\beta\_\{2\}=0\.999, and selected the checkpoint corresponding to the minimum evaluation loss \(which invariably corresponded to around the third epoch with these hyperparameters regardless of the model being trained\)\.

GRPO training was performed using the ‘trl‘\[von Werra et al\.,[2020](https://arxiv.org/html/2605.08696#bib.bib44)\]GRPO trainer, and was modified for resampling and SRM\-specific cache handling for that model\. For all models we useη=2∗10−5\\eta=2\*10^\{\-5\},β=0\.04\\beta=0\.04, gradient clipping with a maximum grad norm of0\.10\.1, warm\-up ratio of 0\.1, weight decay of 0\.1, and top\-p sampling \(top\-p of 0\.9, temperature=0\.7\) and optimize via the Pytorch implementation of AdamW\.

For all training runs, Pytorch models were compiled when possible and trained via DDP with automated fp16/fp32 mixed precision, except for Mamba autoencoder models which were trained via bf16/fp32 mixed precision training due to numerical stability issues\. The main exceptions to compilation were SRMs with GRPO, where turning on and off the cache for switching between recurrent generation and sequence parallel training causes shape changes that are incompatible with older GPUs \(V100s\) although not with newer ones, and Mamba models for which we used optimized kernels that are themselves compiled and cannot be trivially re\-compiled into the Pytorch graph\.

### S1\.2MAX model details

MAX implementations are designed for demonstration only, and use a hybrid graph approach to maintain consistency with Pytorch implementations: the throughputs measured are for implementations of SRM Mojo/MAX static graphs compiled with with embedded dynamic Pytorch tensors for the cache vectors \(we provide but did not profile purely static MAX graph implementations due to hardware constraints\)\.

When we profiled MAX SRMs, we found that there existed a significant \(around 200ms\) delay while calling empty MAX graphs from python driver code, and that this delay persists when calling SRM implementations with our driver code\. Observing that forward passes and cache updates may be internally looped inside the forward pass code and thereby avoiding calling the graph multiple times per generation \(although notably MAX does not currently support compilation of default argmax operations\), we adjust for this delay in Table[7](https://arxiv.org/html/2605.08696#S5.T7), and provide the raw throughput values in Table[S1](https://arxiv.org/html/2605.08696#S1.T1)\.

Table S1:Raw H100 \(NVL\) MAX SRM throughputs, with delay from driver code calls\.We found a bug in the current \(26\.1 and 26\.2\) MAX open source repository all machines tested: there exists a memory leak in which LLVM exhibits an out\-of\-memory error that depends on the number of Mojo modules called during model instantiation before compiling, without regard to the tensor size associated with those modules or what modules are called\. For architectural consistency with Pytorch\-profiled models \(meaning that SRMs use four heads per layer, mixed row/column heads, nonparallel heads and input and output projections\) this necessitated profiling models with fewer layers, and we chose 8 rather than 16\. As the memory error is module call\-dependent, we found that parallelized four\-headed SRMs without input/output projections were able to be initialized with all 16 layers\.

### S1\.3Benchmark Details

Functional benchmarks were performed using modified versions of Eleuther AI’s LM evaluation harness\[Gao et al\.,[2024](https://arxiv.org/html/2605.08696#bib.bib32)\]\. Code relating to SRM cache instantiation, updates, and destruction interfacing was added to a fork of the evaluation harness, and is available at https://github\.com/blbadger/lm\-evaluation\-harness\. Mamba and Transformer models used Huggingface implementations in that module\. All benchmarks are zero\-shot unless otherwise noted\.

![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/more_features.png)Figure S1:Token mixing weights from trained masked mixers\. Autoencoders trained for maximum information retention \(into the last hidden layer of the last token\)\.Table S2:SRM training loss at 200k steps \(13\.1 B tokens\), FineWeb\-edu\. All modelsdm=512,nl=16d\_\{m\}=512,n\_\{l\}=16\.Table S3:FineWeb\-edu training, 13B tokens,nc​t​x=512n\_\{ctx\}=512\. SRMs are mixed headed with decay\.Table S4:FineMath 4\+ training, 13B tokens, compute\-equivalentTable S5:Constant diagonal models offer no consistent training efficiency increase\. Alldm=512,nl=16d\_\{m\}=512,n\_\{l\}=16Table S6:Training Throughput \(samples/sec\) and Memory \(GB per device\) and parameter count for SRM models \(h=4h=4,nl=16n\_\{l\}=16, mixed heads, decay, nonparallel\) on 2x H100 GPUs \(NVL\)\.Table S7:Training Throughput \(samples/sec\) and Memory \(GB per device\) for head\-parallel SRM models \(h=4, mixed heads, decay\) on 2x H100s\.Table S8:Training Throughput \(samples/sec\) and Memory \(GB per device\) for parallelized, non\-projected SRM models \(h=4, mixed heads, decay\) on 2x H100s\.Table S9:Training Throughput \(samples/sec\) and Memory \(GB per device\) for head\-parallel SRM models \(h=4,nc​t​x=512,b=128h=4,n\_\{ctx\}=512,b=128, mixed heads, decay\) on 2x H100\.Table S10:Training Throughput \(samples/sec\) and Memory \(GB per device\) for head\-parallel SRM models \(h=4,nc​t​x=1024,b=64h=4,n\_\{ctx\}=1024,b=64, mixed heads, decay\) on 2x H100\.Table S11:Training Throughput \(samples/sec\) and Memory \(GB per device\) for Mamba 2 models \(h=8,nc​t​x=512,b=128h=8,n\_\{ctx\}=512,b=128\) on 2x H100\. \*with gradient accumulationTable S12:Training Throughput \(samples/sec\) and Memory \(GB per device\) for Mamba 2 models \(h=8,nc​t​x=1024,b=64h=8,n\_\{ctx\}=1024,b=64\) on 2x H100\. \*with gradient accumulation![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/copy_scaling_figure.png)Figure S2:Copy accuracy plummets as the number of tokens to be copied increases\. SRM,dm=256,nl=16d\_\{m\}=256,n\_\{l\}=16, the total context window is double the tokens copied \(FineWeb corpus\), trained for 10k steps,b=64b=64samples per step\.![Refer to caption](https://arxiv.org/html/2605.08696v1/Figures/grpo_g_scaling.png)Figure S3:GRPO exploration on GSM8k\. Models are pretrained on 13B tokens of FineMath 4\+ followed by SFT on GSM8k training dataset\.