Residual Matrix Transformers: Scaling the Size of the Residual Stream
Authors: Brian Mak
Jeffrey Flanigan
Venue: ICML 2025 poster (2025)
Keywords: transformer models
deep learning
neural network architectures
efficient language modeling
Primary area: deep_learning->large_language_models
TL;DR: We present an improved transformer variant whose residual steam is replaced with an outer product memory matrix.

Abstract: The residual stream acts as a memory bus where transformer layers both store and access features (Elhage et al., 2021). We consider changing the mechanism for retrieving and storing information in the residual stream, and replace the residual stream of the transformer with an outer product memory matrix (Kohonen, 1972, Anderson, 1972). We call this model the Residual Matrix Transformer (RMT). We find that the RMT enjoys a number of attractive properties: 1) the size of the residual stream can be scaled independently of compute and model size, improving performance, 2) the RMT can achieve the same loss as the transformer with 58% fewer FLOPS, 25% fewer parameters, and 41% fewer training tokens tokens, and 3) the RMT outperforms the transformer on downstream evaluations. We theoretically analyze the transformer and the RMT, and show that the RMT allows for more efficient scaling of the residual stream, as well as improved variance propagation properties.

Residual Matrix Transformers: Scaling the Size of the Residual Stream
Brian Mak
Jeffrey Flanigan
Abstract
X10
The residual stream acts memory bus where transformer layers both store and access features (Elhage et al., 2021). We consider changing the mechanism for retrieving and storing information in the residual stream, and replace the residual stream the transformer with an outer prod- uet memory matrix (Kohonen, 1972, Anderson, 1972). We call this model the Residual Matrix Transformer (RMT). We find that the RMT enjoys number of attractive properties: 1) the size of the residual stream can be scaled independently compute and model size, improving performance, 2) the RMT can achieve the same loss as the trans- former with 58% fewer FLOPS, 25% fewer pa- rameters, and 41% fewer training tokens tokens, and 3) the RMT outperforms the transformer on downstream evaluations. We theoretically ana- lyze the transformer and the RMT, and show that the RMT allows for more efficient scaling the residual stream, as well as improved variance propagation properties. Code for this project be found at https://gihuu..oo/mmaa///////// matrix-trane transorrrr.s.....r...
1.4
Transformer
RMT
FLOPS
1.0
0.8
0.6
0.4
1000
1500
2000
2500
3000
3500
4000
residual stream size
x10
1.6
Transformer
1.4
RMT
modee size
1.2
1.0
0.6
0.4
1000
1500
2000
2500
3000
3500
4000
residual stream size
Figure 1. Model size and per-example FLOPS versus residual stream size for the transformer and RMT. See $3.1 for details.
1. Introduction
Kaplan et al. (2020) presented way forward for the field of NLP and AI, showing that simply scaling model size, data size, and computational budget is sufficient to advance model capabilities at an unprecedented rate. Their work lead to the the creation of many of the frontier models today (Achiam et al., 2023, Hoffmann al., 2024, Dubey et al., 2024). Mcanwhile, data, compute, and model sizes continue grow exponentially each year (AI, 2024). In fact, are quickly scaling towards the limits how much data and energy is available (Muennighoff et al., 2023, Zachary Zim-
merman Gramlich, 2023). This presents a need for more data and compute efficient models.
When it comes to large scale LLM training, few architectural modifications rival the impact that Mixture of Expert Models (Fedus et al., 2022) have had. Their work inspired many later cutting-edge models (Du et al., 2022, Jiang et al., 2024). Their Switch Transformer model had the ability to scale along the sparse parameter axis. This meant that model size could be scaled without affecting per-example compute. With this, they were able to scale LLM training in ways that had previously been unexplored. In this new search space, they found training configuration that was 7 times more efficient than standard transformer training.
'Department of Computer Science, University of California Santa Cruz, Santa Cruz, United States. Correspondence to: Brian Mak <bmak2 @ ed.s..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
In this paper, we follow a similar formula and present a new transformer-type model, the Residual Matrix Transformer

---

Residual Matrix Transformers
(RMT), that can scale the residual stream size while keeping model size and per-example compute fixed. The residual stream is defined as the sum of all the outputs of the pre- vious layers and acts as a communication channel between layers (Elhage et al., 2021). The size (i.e. dimension) of the residual stream determines how many features can store, so scaling its size can be thought of as scaling its bandwidth. In our experiments, we find that scaling the size of the residual stream while maintaining the constant model size and per-example compute improves both data and compute efficiency ($4.4).
matrix S E RVXN whose columns are the one-hot-vectors
of some tokenized string. The first layer encodes the input.
X E(S) + P(N)
(1)
Here, N E RRNN are the one-hot-vectors representing each token's position. The output of this layer is an embed- ding matrix X E DR and it contains the first instance of each token's residual stream. The next 2L layers are de- fined recursively and describe the successive application of attention and feed-forward layers to the residual stream.
Scaling the residual stream in a standard transformer resizes all the parameter matrices, which linearly influences the parameter and FI.OP counts (see Fig. 1). To address this, we replace the residual stream representation with an outer product memory matrix (Kohonen, 1972, Anderson, 1972), and find the relationship between the size of the residual stream and parameter and FLOP count becomes nearly con- stant for typical transformer dimensions. Additionally, we find that the our proposed model is more efficient in terms of data, FLOOPS, and parameters than the standard transformer.
x = (X + ((-11)
X = FF(LN(X (0) + x
(3)
where LN is LayerNorm (Xu et al., 2019). The final layer reads out the logits of the next predicted token from the residual stream. The transformer T is thus represented as:
T(S) = (2L))
(4)
The embedding layer encodes the one- hot-vector matrices using the learned embeddings WF RXV and WPB DXN
Embedding Layer
We summarize our contributions as follows:
We introduce a new transformer-type model, the Resid- ual Matrix Transformer (RMT), that replaces the resid- ual stream with an outer product memory matrix ($2).
B(S) WES
(5)
P(N) = WPEN
(6)
We experimentally find that the RMT is more efficient in terms of data, FLOPS, and parameters than the trans- former. The RMT trains with fewer resources and achieves better performance than the transformer and variants ($4.2, $4.3, and $4.5).
Attention Layer For our attention layer definition, we group the computation of each attention head separately.
We present theory that shows the RMT exhibits effi- cient scaling of the residual stream and has improved variance propagation properties ($3).
MHA(X) Σ w "HA(QUM) (
v(h)
(7)
QT K
We experimentally explore the effects of scaling the residual stream size while keeping model size and per- example compute fixed with the RMT. We find that a larger residual stream improves performance ($4.4).
SHA(Q, K, V) Softmax
V
(8)
VDn
(b) = w X K k(O)
WiTh (N)
X
V(h) = w
x
K
W W (h) K
W (h)
E R
P
X and (h)
(h)
E R
Dx Dr.
2. Model Architecture
FeedFdd Layer We define our feed-forward net- work in the standard way, using Gelu activation. Wi R D and W2 E DxDrF.
In this section we introduce the Residual Matrix Trans- former, a transformer variant whose residual stream is re- placed with an outer product memory mechanism. We first review the transformer architecture ($2. I), and then present the RMT formulation ($2.2).
FF(X)= W5Gelu(W;X)
(10)
Unembedding Layer The unembedding layer multiplies the final instance of the residual stream with unembedding weights Wy E RVD.
2. Transformer
Transformer Overview We begin by giving the equations that describe the transformer architecture (Vaswani et al., 2017). In our formulation, the input to the transformer is a
U(X) Wo X
(11)
2

---

Residual Matrix Transformers
2.2. Residual Matrix Transformer
Embedding Layer Eqs. 18 and 19 show how the initial residual matrices are formed.
Here we give relevant technical background on outer product memories and show how this mechanism can be applied to the transformer.
ft
E(S)
= Σ
E
X
E
S
(18)
h=1
Outer Product Memories We propose to replace the residual stream with an outer product memory store. An outer product memory store is created by summing the outer products of a set of key-value vector pairs (Koho- nen, 1972, Anderson, 1972, Gmitro et al., 1989). For some set of key vectors [q(P) E RDx p== and data vectors {a<() E Ph Jp=1, an outer product store M E DDxDD can be constructed using the following equation:
R
P(N) = Σ h=1
WE &
w
(h)
W N
(19)
(h) Here, w E RD.VV and w parameter. Notice the form of Eq. 18 and 19 mirror that of Eqs. 12. The result of the embedding layer is a tensor of shape D& x Dr N that contains an outer product memory matrix for every token.
E Dt and R is a hyper-
N M = Norm Σ q (P) @ (l))))
Attention Layer The attention equations for the RMT (12)
strongly resemble Eqs. 7-9 with a few key differences.
p=1
Here u 2 uv_. Norm is normalization function, which in the case of RMT is LayerNorm. One can retrieve particular data vector (=1)) using its corresponding key vector (t))
MHA(X) Σ & SHA(Q!"), K(l,, v(h)
(h) w
(20)
Q" K VD,
SHA(Q, K, V) SSofmax
V
(21)
x
q (r) '1 M
(13)
Here 1 denotes tensor contraction over the first dimension of q (r) and M.
(h) 2 = rrQ -1 X K K
(h)
.1 X
V (h)
(h) V
X
(22)
In the RMT, features are retrieved from the residual matrix using key vectors. Comparing equations Eqs. 22 to Eqs. 9, we see that the matrices w-(h) wr(h) and (h(h) are
RMT Overview The RMT is similar to the Transformer, with the residual stream vectors replaced with outer product memory stores.! We refer to each token's 's memory store as its residual matrix. Whereas in the standard transformer linear transformations are used to store and retrieve features from the residual stream, our model uses key vectors to store and retrieve data vectors from the residual matrix.
replaced by the key vectors TQ ,.
(h) (h) (k)
K .V E DR. Addi-
tionally, the operations have changed from matrix multiplies to tensor contractions over the first tensor dimension, which is exactly the retrieval operation from Eq. 13.
Eqs. 22 produce the attention inputs ((h), ((),, ( R RD.NN These matrices have the exact same dimensions as they did in the standard transformer, so the SHA operation remains the same.
The RMT model identical to Eqs. 1- 4, except now the batched residual stream instances (() E IR D..NN are tensors instead of matrices. This reflects the fact that in our model, for every token position we have residual matrices sized Dj x D. instead of residual stream vectors sized D.
In Eq. 20 we see that each attention head output acts as a data vector that is associated to its corresponding key vector (h)
x E(S) + PPNN
(14)
O E Ds before being added to the residual matrix.
x = MHA(LN(X + 1(--
Foedrd Layer To understand how the RMT feed- (15) forward layer is computed, first notice that Eq. 24 is identi- cal to Eq. 10.
x(+1)= = =FF(LN(X() +X()
(16)
R
FF(X)
pp & unvec, (FF(XFF))n::
(h)
(23)
T(S)= 2L)
(17)
h=1
'oore precisely, the residual stream vectors are replaced with unnormalized outer product memory stores and the Nre-LayerNorm operations (Xiong et al., 2020) in Eqs. 15-17 provide the missing normalization from Eq. 12.
FF
W2Gelu(W,.X
(24)
RRM''s D. equals transformer's Dr for our experiments
3

---

Residual Matrix Transformers
F = Concat (h) 1gh<R ("FF X)
(25)
how the model size and per-example compute scales with the residual stream size for both the transformer and the RMT. We follow Hoffmann et al. (2024)'s 's formulas for computing model size and per-example compute in the transformer, and derive the corresponding formulas for the RMT (Appendix SA, Table 7). We fix all model dimensions to those of GPT2-medium' (Radford et al., 2019) and vary D for the transformer and D2 for the RMT (where D, is fixed at 64). We see that the model size and per-example compute for the transformer scales proportionally to the residual stream size. In the RMT, however, there is practically no parameter or compute cost. In fact, if we increase the residual stream size of the transformer by 100%, we will also see a 100% increase in parameters and a ~ 94% increase in FLOPS. For the RMT, if we increase the residual stream size by 100% we will see a < 1% increase in both model size and FLOP count. We note that, as in the regular transformer, scaling the residual stream of the RMT will impact memory usage during training since residual stream activations need to be saved for gradient checkpointing
That is, the core operation of the layers the same. The only differences are Eqs. 23 and 25, which are adapters between the residual matrix and the core feed-forward operation. The idea behind Eq. 25 is that for every token position, (h)
we retrieve data vectors T'F X E RDxxN from the residual matrix. These data vectors are concatenated along their first dimensions resulting in the feed-forward input XpF E RRD,xXN In Eq. 24 the standard feed-forward operation is applied. Then in Eq. 23, we see that unvec; is applied to FF(XpF) E RRD.xN, As result, the first dimension of (FFF)) will be reshaped from RD2 to R x D.. Intuitively, for every token, unvecj splits the output of the forw-forward operation into R data vectors. Then, each of these data vectors is associated with key vector (h) w WFF E Rx and added to the residual stream.
Unlike in the attention formulation, matrices W1 and W2 of the feed-forward layer were not replaced with key vectors. Instead, we added key vector adapters between the linear transformations and the residual matrix. The reason we decided not to replace these matrices is that there is strong evidence that the feedforward weights store factual informa- tion in the transformer (Geva et al., 2021, Meng et al., 2022), and therefore perform a function beyond simply reading-in and writing-out features.
3.2. Moment Propagation Analysis
Here we verify that our proposed architectural modifica- tions have superior moment propagation properties to their transformer counterparts. their landmark work, Glorot Bengio (2010) showed that proper propagation of mean and variance through linear transformations at initialization is critical for deep learning. Here, we perform a similar analy- sis to Glorot & Bengio (2010) and Kedia et al. (2024) to find how the mean and variance propagates through the RMT's storage and retrieval operations. We compare their moment propagation properties to the operations they replaced.
Unembedding Layer In Eq. 27, the unembedding layer retrieves R data vectors x (h) E IR P.NN from the residual matrix using the key vectors Tt E Rdd, These data vectors are then multiplied by the unembedding weights w!( E RR V xD. to get logit predictions for the next token.
(A)
In Table 1 we present the closed-form expression for mean and variance propagation through the RMT's storage and (26) retrieval operations at initialization. We find that, with linear transformations, if weights are initialized indepen- dently with zero mean then the propagated mean will be zero. In this respect, the mean propagation properties of these components are the same as the transformer's. s.
Σ U(X)
W
(h)
X
(h)
U
h=1
(b) x U = TU
(h)
X
(27)
3. Theoretical Properties
To understand the variance propagation properties, in Ta- ble we plug in the model dimensions for GPT2-medium into the closed-form variance expressions in Table 1. Table 2 shows the results of this analysis. We follow Glorot &
In $3.1 we show that the RMT allows for more efficient ex- pansion of the residual stream compared to the transformer. Then in $3.2 we show that the elements we modified in the RMT have superior gradient and activation propagation properties than their transformer counterparts.
Bengio (2010) and consider L and Fim 1 to be km.
optimal since it prevents exploding or vanishing gradients or activations. We see that in all cases except for the atten- tion layer's storage operation, the RMT's operations show favorable variance propagation.
3.1. Resource Scaling
Here we verify that we can scale the size of the residual stream more efficiently with the RMT than the transformer in terms of model size and FLOPS consumed. Fig. shows
*We choose GPT2-medium for this example but the general scaling trends will hold true for any model shapes.
RMT's RD, equals transformer's 's D for our experiments.
4

---

Residual Matrix Transformers
Table Closed-form expressions for the propagation of and variance for storage and retrieval for the RMT. The expressions for the transformer are taken from (Kedia et al., 2024). The derivation the RMT equations can be found in Appendix şB.
COMPONENT
MODEL
OPERATION
Parus
Sout
pgin
Sin
RMT TRANSFORMER
Xout Σm= whth) @ Win
Ro 2
+ M
deo i (o + M gogee doat a (o Houue Ro (oagu + Hague
STORAGE
Cout Winn
0
din Tu (
+ H
RMT TRANSFORMER
Dout w. Xin
d 1x
RETRIEVAL
Bout W Zin
0
din (T am + H
dout (a Gout Maoue
Table Theoretical calculation of propagation of variance through storage and retrieval operations on forward backward passes. We assume Xavier Initialization and GPT2-medium model shapes. Boldface indicates where one model better than the other.
Transforner RMT
4 x 10
loss
3.5 x 109
2
LAYER
OPERATION
MODEL
Ruvu
RMT
0.4
1.6
STORAGE
TRANSFORMER
ATTN
RETRIEVAL
RMT
1.14
0.86
3 x 109
TRANSFORMER
0.5
1.5
10
1017
10 18
STORAGE
RMT
FLOPS
TRANSFORMER
1.6
0.4
FF
RMT
1
RETRIEVAL
TRANSFORMER
0.4
1.6
x 10
RMT
Transformer
loss
3.5 x 109
4. Experiments
In $4.2 and $4.3 we compare the RMT to the transformer and other relevant transformer variants and show that the RMT offers more efficient pretraining in terms of parameter, data, and compute efficiency. Then $4.4 show that the new scaling axis introduce (scaling residual stream size) is an effective way scale LLM training. In $4.5 we evaluate the downstream performance of the RMT and show that it outperforms a transformer that is 33% larger.
3 109
107
10a
model parameters
Figure 2. Scaling law curves of RMT transformer. Model sizes vary from 46M to 405M.
4.1. Pretraining Details
4.2. RMT vs Transformer
All models were trained on the WebneetText dataset (Gokaslan et al., 2019) in the infinite data regime. Full details for all experiments can be found in Appendix &C. To overcome the "Parameter Lottery" (Dey al., 2024), we performed hyperparameter tuning on both the RMT and transformer using pParam Transfer (Yang et al., 2024). We the hyperparameters found for all pretraining experi- ments except our comparison against transformer variants experiments in $4.3, since it was too computationally de- manding to run hyperparameter tuning for all variants. For the experiments found in $4.3, we used standard transformer hyperparameters for all models.
In this section we experimentally show that the RMT ex- hibits better parameter, data, FLOOP, and inference time mem- ory efficiency than the transformer. We verify these trends by examining the sealing properties of both architectures and analyzing the training curves of the largest models.
We train series of four models for both architectures. The size of the models vary from 46M parameters to 405M pa- rameters. The model dimensions were chosen such that for each model in the transformer series, there corre- sponding model in the RMT series whose model dimensions mirror all aspects except the size of the residual stream.
5

---

Residual Matrix Transformers
10
RMT
10
RMT
Transformer
Transformer
loss
6 10
loss
6 10
x 10
4 x 10
100
3 x
0.0
0.2
0.4
0.6
0.6
1.0
1.2
50000 100000
150000
200000 250000 300000
FLOPS
X10
time (6)
RMT
Figure 4. Train loss for the transformer and RMT a time basis.
Transformer
x 16
Table 4. Training statistics after training on 6B tokens.
RESOURCE
TRANSFORMER
RMT
% DIFF
10
FLOPS
1.18 x 10
8.45 x 10
-28% +43%
x 100
TIME
1.86 10
3.30 10
TRAIN LOSS
3.03
2.91
5
DEV LOsS
3.02
2.91
tokens
x10
Figure 3. Train loss curves for the transformer and RMT on per- token and per-flop basis.
reports the resources consumed reach the minimum loss achieved by the transformer baseline while Table 4 shows the resources consumed train on the full 6B tokens.
Table 3. Resources used to reach train loss 3.03 (min loss achieved by transformer).
Parameter Efficiency From the bottom plot of Fig. 2 we see that the RMT exhibits more efficient parameter scaling than the transformer. Table 3 and Table quantify this result for largest runs. We see that the RMT is able to achieve lower loss than a transformer model that is 33% larger.
RESOURCE
TRANSFORMER
RMT
DIFF
PARAMETERS
405M
305M
-25%
FLOPS
1,18 10t9
4.97 x 18
-58%
TOKENS
6,04 X 10
3.55 10
-41%
TIME
1.86 X 10
1.94
10
+4%
The reason why the RMT uses fewer parameters than the transformer is that the Q w!th) w and W-(h) weight matrices are replaced by vectors in the RMT mod- els. As a result, the RMT models have 33% fewer non- embedding parameters than their transformer counterparts. Furthermore, because of its resource scaling propertics from $3.1, we can make the RMT models have much larger residual stream than their transformer counterparts for prac- tically parameter penalty.
The residual stream size of the RMT models is be 2.5 4 times larger than their transformer counterparts. Each model is trained for its Chinchilla optimal number of train- ing tokens number of non-embedding parameters). Fig. 2 shows the results of these experiments. Due to resource constraints, we were unable to explore how these trends continue at larger model sizes.
Fig. and Fig. compare the training runs of the largest models both series. Both models mirror GPT2-medium's dimensions, except that the residual stream size of the RMT model is expanded by factor of 4. For direct comparison, we extended the number of tokens the RMT model was trained on to match the transformer baseline. The results in Fig. still show the RMT's training statistics stopped respective Chinchilla optimal train tokens. Table
Data Efficiency Fig. shows that the RMT exhibits much faster convergence on token basis than the transformer. Table 3 quantifies this improvement, showing that the RMT is 41 % more token efficient than the transformer. This type of data efficiency is particularly important in the current cli- mate of LLM training, where we are quickly scaling toward the limit of text available to train models on (Villalobos et al., 2024, Muennighoff al., 2023).
6

---

Residual Matrix Transformers
batch size found was 256.² From this infer that the ex- space needed to store the gradient checkpoints for the expanded residual stream is offset by reduced model size.
x 10
Dephwise 1.TMM
Hierarchical Aggregation
x 10
Highway Transfrmerr
At inference time, since there is no need to keep residual stream instances for gradient checkpointing, the memory savings from the reduced model size outweigh the expanded residual stream. This presents opportunities for more pow- erful models to be run with smaller devices.
RMT
Transformer
x 10
Runtime Fig. shows that on a per-time basis, our model is about even with the transformer. Table 3 shows that the RMT is 4% slower than the transformer even though it is more FLOP and data efficient. This discrepancy is due to the fact that the RMT took 7.17s per train step, while the trans- former only took 4.06s. We expect that a hardware- aware implementation of our model would significantly improve its runtime, but we consider this is out of the scope of this work. Currently, we find that runtime is the biggest limita- tion of our model. See Appendix $E for further discussion.
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
1.6
FLOPS
X10
3.500
Depthwise LSTM
3,500
Hierarchical Aggregation
3.475
3.450
Transformer
3.425
Highway Transformer
3.400
4.3. RMT vs Transformer Variants
3.376
RMT
Here we compare the RMT to transformer variants that can be viewed as non-orthogonal to our work.? We find that this class of architectures includes transformer variants that make architectural modification to the residual stream (Dou et al., 2018, Chai et al., 2020, Xu et al., 2024). We show that our architecture offers more efficient pretraining than all prior variants (Fig. 5).
1.4
1.6
1.8
2.0
2.2
model parameters
x11
Figure 5. RMT versus other transformer variants that modify the residual stream (Vaswani et al., 2017, Dou et al., 2018, Chai et al., 2020, Xu et al., 2024).
We train collection of transformer variants whose model dimensions mirror GPT2-small where applicable. Fig. shows the training curves and final train loss plotted against model size. From the results we see that the RMT is the most compute and parameter efficient variant of this collec- tion, achieving lower loss while using fewer FLOPS and parameters. We note that our approach is distinguished from prior residual stream modifications in that prior architec- tures use more parameters and per-example compute than the transformer, while ours uses fewer. Thus, from ef- ficiency standpoint, the improvements made by the prior variants are outweighed by their cost in our replications.
FLOP Efficiency From the top plot of Fig. 2 we see that the RMT architecture exhibits more efficient FLOP scaling than the transformer. Fig. 3 corroborates this by showing that on our largest runs the RMT converged faster than the transformer
per-FLOP basis. Table 3 shows that the RMT is 58% more FLOP efficient when training to the same loss, and Table 3 shows that the RMT is 28% more FLOP efficient when training to the same number of tokens.
We assert that this improved FLOP efficiency presents huge opportunities for energy savings in LLM training. Energy efficiency is another important factor to consider, because like data consumption, energy consumption is also scaling at an untenable rate (Strubell et al., 2019, Luccioni et al., 2023, Wu et al., 2022).
4.4. Effect of Scaling the Residual Stream
In $1 we proposed new axis to scale LLM training along In this section we show that this is an effective axis for scal- ing. In particular, we show that scaling model's residual stream while keeping model size, dataset size, and computa-
Memory Usage We find that during training, the memory usage of both models is about equal. At inference time, however, the RMT is more memory efficient.
'The max batch size search was constrained to powers of "Only a single residual matrix needs to instantiated per sequence for autoregressive decoding. This cost negligible compared to the size of the replaced attention parameters.
We deduced that the train time memory usage of both mod- els is about equal by performing a batch size search on our largest models. We found that, both cases, the maximum
"Residual stream modifications be combined with ours, and so are orthogonal to our work. See related work $5.

---

Residual Matrix Transformers
tional budget fixed increases model performance. We note that this type of experiment is only possible using the RMT and not the transformer ($3.1).
Table 5. Downstream perplexity results (Gokaslan al., 2019, Paperno et al., 2016, Gao et al., 2020, Merity et al., 2016). Lower is better.
We train a series of GPT2-small sized RMT models with residual stream sizes 384, 768, 1536, 2048, 3072, and 4096. All models are identical except for their residual key di- mensions D:- We choose to vary D because it does not affect the model's core computations, allowwng us directly observe the effects of scaling the residual stream. The dif- ference in total parameter and FLOP counts between the largest and smallest models is < 1%.
DATA SET
TRANSFORMER
RMT
LAMBADA OPENAI
61.34
25.21
OPENWEBTEXT
20.64
18.36
PILE 10K WIKITEXT
30.23
23.54
76.35
55.94
Table 6. Downstream accuracy results (Clark et al., 2018, Zellers et al., 2019, Bisk et al., 2020, Sakaguchi al., 2019). Higher better.
The results are shown in Fig. 6. We see that increasing the residual stream size monotonically decreases dev loss, showing that expanding the residual stream improves model performance. If we compare the model with residual size 4096 the model with residual size 768 (the residual stream size of GPT2-small), we see that the model with the ex- panded residual stream was able to train to the same final loss with 23% fewer FLOPS and 25% fewer tokens.
DATA SET
TRANSFORMER
RMT
ARC C
19.7
20.6
ARC E HELLASWAG
44.3
46.1
28.7
30.5
PIQA
61.9
63.7
WINOGRANDE
49.5
52.5
We can compare an RMT with a transformer with the same residual stream size, since the experiments performed in this section were trained with the same experimental settings as $4.2. We see that when the residual stream size is the same, the RMT model achieves final train loss of 3.42 while the GPT2 model achieves 3.43. This verifies that if the residual stream not expanded, the performance of the RMT about the same as the transformer.
than the transformer on all downstream tasks despite having 25% fewer parameters and using 28% fewer FLOPS during training.
5. Related Work
Layer aggregation models bear some resemblance to our model, since they allow later layers to access earlier lay- ers more easily. Indeed, our model can be interpreted as type of layer aggregation model one were to multiply out all key vectors and remove LayerNorm. Huang al. (2017) introduced dense connections and showed that they can improve the performance of convolutional neural net- works. Godin et al. (2017) extended dense connections to RNNs and showed that they can help with language model- ing. Shen et al. (2018), Dou et al. (2018) 1)), and EINokrashy et al. (2024) introduced dynamic layer aggregation through pepphwiss-attentinn and showed improvements on NMT and classification. Dou al. (2018) also explored different types of static layer aggregation and showed improved results in NMT. Dou et al. (2019) then showed that layer aggregation with capsules and routing by agreement can further improve results on NMT. We note that for most layer aggregation models it unrealistic to train multi-gpu settings as they would incur a huge penalty in terms of device-to-device communication. Dou al. (2018)'s model is excep- tion this rule, and include it in our comparison of transformer variants ($4.3).
10
3.55 x 109
3.5 100
loss
3.45 10"
3.4 x 109
100
T
384
1536
2048
3072
residual stream size
Figure 6. Dev loss vs residual stream size. Model size, dataset size, computational budget are held constant all runs. Residual stream size is calculated as Da x D..
4.5. Downstream Evaluation
In this section we present zero-shot downstream results of the RMT and transformer. We use our largest pretrained models from $4.2 trained on all 6B tokens. Table and Table show that the RMT performs significantly better
The class of models that we find most similar to our own are those that manage their residual stream memory. These models offer multi-gpu scalability, and, in theory, they also make important features produced by earlier layers more
8

---

Residual Matrix Transformers
available later layers. Srivastava et al. (2015) introduce Highway Networks that control which layer's s outputs are carried down the residual stream, Chai et al. (2020) iterated on Highway Networks and applied them to transformers. Xu et al. (2024) used an LSTM to manage the residual stream. Our work is different from these as we consider the scaling ramifications of expanding the residual stream instead of managing its memory.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
AI, E. Data on notable ai models, 2024. URL https://epoch. ai/data/notable--tt ai- ai-models. Accessed: 2024-12-14.
We find that modifications that involve residual stream scal- ing, normalization, and weight initialization to improve moment propagation (Kedia et al., 2024, Wang et al., 2024, Shleifer et al., 2021, Huang et al., 2020, Bachlechner et al., 2020) are orthogonal to our approach because they could be extended and applied to our model.
Anderson, J. A. A simple neural network generating an in- teractive memory. Mathematical Biosciences, 14(3):197-
220, 1972. ISSN 0025-5564, doi: http://6bo..rr/1001100 1016/ 0025-5564(72)900075-2 URL https:/wwww.sceencddreett comlsiioneeatti22/i/i05555555290
We consider transformer variants that modify components other than the residual stream (Dao & Gu, 2024, Peng et al., 2023, Sun et al., 2023, Fedus et al., 2022) to be orthogonal to our work as well since our approach can be easily integrated with these architectures.
Bachlechner, T., Majumder, B. P., Mao, H. H., Cottrell, G. W., and McAuley, J. Rezero is all you need: Fast convergence at large depth. CoRR, abs/2003. 044887, 2020. URL https:/aariv.org/bbb2000..3
Bisk, Y., Zellers, R., Bras, R. L.., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourrhh AAAI Conference on Artificial Intelligence, 2020.
6. Conclusion
We present novel architecture that replaces the residual stream with outer product memory mechanism, resulting in transformer variant with an expandable residual stream. We theoretically show that our modifications provide effi- cient scaling and favorable moment propagation properties. Our architecture allows for scaling along a new scaling law axis that, to our knowledge, was previously unexplored. We experimentally showed that scaling the size of the residual stream leads to more efficient training in terms of compute, parameter, and data efficiency.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Minh S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL httpggiihub.co/jaamnnaiie
Chai, Y., Jin, S., and Hou, X. Highway transformer: Self-gating enhanced self-attentive networks. In Juraf- sky, D., Chai, J., Schluter, N., and Tetreault, (eds.), Proceedings of the 58th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 6887-6900, Online, July 2020. Association for Computational Lin- guistics. doi: 10. (026655//12220.... aww.://0..cmmmi...iii. 616. URL http://cclathhoogy.og20200000 http:///22222..mmm.mm11111111 main.616.
Acknowledgments
We are thankful for the computing resources provided by the Pacific Research Platform's 's Nautilus cluster, supported in part by National Science Foundation (NSF) awards CNS- 1730158, ACI-1540112 ACI-1541349, OAC-1826967 OA1-21 CNS-2100237, CNS-2120019, the Univer- sity of California Office of the President, and the University of California San Diego's 's California Institute for Telecom- munications and Information Technology/uualcomm Insti- tute, and CENIC for the 100Gbps networks. We also thank the anonymous ICML reviewers and the members of JLab for their proofreading and valuable feedback.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, 1803005035 2018.
Dao, T. and Gu, A. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 10041 10077. PMLR, 21-27 Jul 2024. URL https://proceeeiings mlr, press//es//235/ dao24a.html
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

---

Residual Matrix Transformers
Dey, N., Anthony, Q.. and Hestness, J. The prac- titioner's guide to the maximal update parameteri- zation. https://www.cerebbtt http://werrebra.../looo/eeeeeeee practitioners- guide-to- Londe-to-tie the- maximal-update parameterization, Septem- ber 2024.
and Evaluation (LREC-COLING 2024). pp. 4665-4674, Torino, Italia, May 2024. ELRA and ICCL. URL https: Iaclanthology.orr/22224 2rep22222..rmmmmmi4111111111
Fedus, W. Zoph, B., and Shazeer, N. Switch transformers: scaling to trillion parameter models with simple and effi- cient sparsity. J. Mach. Learn. Res., 23(1), January 2022. ISSN 1532-4435.
Dou, Z.-Y. Tu, Z., Wang, X., Shi, S., and Zhang, T. Exploit- ing deep representations for neural machine translation. In Riloff, E., Chiang. D., Hockenmaier, J., and Tsujü, (eds.). Proceedings of the 2018 Conference on Empiri- cal Methods in Natural Language Processing, pp. 4253- 4262, Brussels, Belgium, October-November 2018, Asso- ciation for Computational Linguistics. doi: 10.1 18653/v1/ 11881457. URL https///clammhoolog.oro/// 11444777.
Gao, L.,, Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 2020.
Dou, Z.-Y., Tu, Z., Wang, X., Wang, L., Shi, S., and Zhang, T. Dynamic layer aggregation for neural machine trans- lation with
Gao, L., Tow, J., Abbasi, B.,, Biderman, S.,, Black, S., DiPofi, A., Foster, C., Golding, L.., Hsu, J., Le Noac"h, h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang. J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang. E., Thite, A., Wang, B.. Wang, K., and Zou, A. A framework for few-shot language model evaluation, 12 2023. URL http::///eendoooorg/rccoossss 1025/63626666
Proceedings of the AAAI Conference on Artificial Intelligence, :3 93, Jul. 2019. doi: 10. 1609/669........ http//aaa..v33330 v33001.33010311.3333333333333331 URL https://ojs. aaai.aiiiiaaao... arnglineexhh/AAAAVaaiidvvvvv///
Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y.. Krikun, M.. Zhou, Y., Yu. A. W., Firat, O., Zoph, B., Fedus, L.,, Bosma, M. P., Zhou, Z., Wang, T., Wang, E., Webster, K., Pellat, M., Robinson, K., Mecier-Hellstem., K., Duke, T., Dixon, L., Zhang. K., Le, Q. Wu, Y., Chen, Z., and Cui, C. GLaM: Effi- cient scaling of language models with mixture-of-expprts. In Chaudhuri, K., Jegelka, S., Song. L., Szepesvari, C., Niu, G., and Sabato. S. (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Re- search, pp- 5547-5569. PMLR, 17-23 Jul 2022. URI. http::/proceedingg..lrr mlr www.l..rees/vvle.. htm..u22c.html.
Geva, M., Schuster, R., Berant, J., and Levy, O. Trans- former feed-forward layers are key-value memories. In Moens, M.-F., Huang, X., Specia, L., and Yih, S. W.- t. (eds.), Proceedings of the 2021 Conference on Em- pirical Methods in Natural Language Processing, pp. 5484-5495, Online and Punta Cana, Dominican Repub- lic, November 2021. Association for Computational Lin- guistics. doi: 10.18653/ 1/20211 wew.plee.mlppmmee.ee main.446. URL https://cclattholoy..oo/// wet.:022eemmpmmmmm mannlp-aain.446...
Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In Teh, Y. W. and Titterington, M. (eds.), Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pp. 249-256, Chia Laguna Resort, Sardinia. Italy, 13-15 May 2010. PMLR. URL http://prroc......gg.. mlr wwwpprees/v/glloottolll....
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A. et al. The Ilama herd of models. arXiv preprint arXiv:2407.2178. 2024.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. https:/transformer- Srcuisspbb/222//tamwwwot//nn
Gmitro, A. F., Keller, P. E., and Gindi, G. R. Statistical performance of outer-product associative memory models. May 1989. doi: 10.
Appl. Opt., A0.28.001940. URL https://opg.optica.or////////// cfm?URI==o- 28- 1940.
Godin, F., Dambre, J., and De Neve, W. Improving lan- guage modeling using densely connected recurrent neural networks. In Blunsom, P., Bordes, A., Cho, K., Cohen, S., Dyer, C., Grefenstette, E.,, Hermann, K. M., Rimell, I.., Weston, J., and Yih, S. (eds.), Proceedings of the 2nd Workshop on Representation Learning for NLP, pp. 186- 190, Vancouver, Canada, August 2017. Association for LiniratioonaLLnnuusttiiiesss. doi: 10. 10.18655/w1/W17-22222 URL https://bbaathoogyygg//1112222
EINokrashy, M. AlKhamissi, B., and Diab, M. Depth- wise attention (DWAtt): A layer fusion method for data- efficient classification. In Calzolari, N. Kan, M.-Y., Hoste, V., Lenci, A., Sakti, S., and Xue, N. (eds.), Proceedings of the 2024 Joint International Confer- ence on Computational Linguistics, Language Resources
Computational Linguistics.
10

---

Residual Matrix Transformers
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Openwehtext corpus. http://Skylion007ggihhe githhuuub OpenWebTexiCorpus, 2019.
Kohonen, T. Correlation matrix memories. IEEE Trans- actions Computers, C-21(4):353-359, 1972. doi: 10.1109/TC.1972.50089777
Hall, D., Zhou, I., Adams, V., Wang. J., and Liang, P. Levan- ter, 2024. URL https://gihub..co//sattffff comd/taaforr--cfi///evteerr crfm/levanter.
Luccioni, A. S., Viguier, S., and Ligozat, A.-L. Estimating the carbon footprint of bloom, 176b parameter lan- guage model. Journal of Machine Learning Research, 24(253):1-15, 2023. URL http://imi..org/pppess//k/ 0069.html.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy. A., Osindero, S., Simonyan, K., Elsen, E.,, Vinyals, O., Rae, J. W., and Sifre, L. Training compute-optimal large lan- guage models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 9781713871088
Meng, K., Bau, D., Andonian, A., and Belinkov, Y. Locating and editing factual associations in gpt. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 17359-17372. Curran Associates, Inc., 2022. URL http:///poceeddngggennriippopppp httpppppepreeleepppp/2222222 file/ h1o6ff14444.55222277766666622233 Conference.pdf:
Huang, G., Liu, Z., Van Der Maaten, I.., and Weinberger, K. Q. Densely connected convolutional networks. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2261-2266 2017. doi: 10.1109/ CVPR.2017.243.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models, 2016.
Muennighoff, N., Rush, A., Barak, B., Le Scao, T. Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C. A. Scaling data-constrained language models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 50358-50377 Curran Associates, Inc., 2023. http://cccecedees.uuurrrpcoppppp files/paper/ URL BD25/heesds9440006eeeeeeel/aa//2 Conference.pdi
Huang, X. S., Perez, F., Ba, J., and Volkovs, M. Improv- ing transformer optimization through better initializa- tion. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learn- ing, volume 119 of Proceedings of Machine Learning Research, pp. 4475-44588 PMLR, 13-18 Jul 2020. URL https:/procceedings..o press/v 1 19/huang2ffhtmm.
Jiang, A. Q. Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C. Chaplot, D. S.. de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.-A,, Stock, P., Subramanian, S., Yang. S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024. URL https:/arxxvvorr/ abs/2401.04088.
Paper-
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q. N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset, Aug 2016.
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Der- czynski, I.., Du, X., Grella, M., Gv, K., He, X., Hou, H., Kazienko, P., Kocon, J., Kong. J., Koptyra, B., Lau, H., Lin, J., Mantri, K. S. I., Mom, F., Saito, A., Song, G., Tang, X., Wind, J., Wožniak, S., Zhang, Z., Zhou, Q., Zhu, J., and Zhu, R.-J. RWKV: Reinventing RNNs for the transformer era. In Bouamor, H., Pino, J., and Bali, K. (eds.). Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 14048-14077, Singapore, December 2023. Association for Computational Linguis- tics. doi: 10. httt:///////...innnnnnnnmmlppppp https///clanthooogx.or/2222: www.iingg-ennnn.ii......
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. 2020.
arXiv preprint arXiv:2001.0880
Kedia, A., Zaidi, M. A., Khyalia, S., Jung, J., Goka, H., and Lee, H. Transformers get stable: an end-to-end signal propagation theory for language models. In Proceed- ings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024.
URL
Kidger, P. and Garcia, C. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differ- entiable Programming workshop at Neural Information Processing Systems 2021, 1221,00 2021.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. 2019.

---

Residual Matrix Transformers
Sakaguchi, K., Bras, R. L.., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale, arXiv preprint 1rXiv:1907.16644 2019.
Wu, C.-J, Raghavendra, R., Gupta, U., Acun, B., Ardalani,
N., Maeng, K., Chang, G., Aga, F,, Huang, J., Bai, C., Gschwind, M., Gupta, A., Ott, M., Melnikov, A., Candido, S., Brooks, D., Chauhan, G., Lee, B., Lee, H.-H., Akyildiz, B., Balandat, M., Spisak, J., Jain, R., Rabbat, M., and Hazelwood, K. Sus- tainable ai: Environmental implications, challenges and opportunities. In Marculescu, D., Chi, Y., and Wu, C. (eds.), Proceedings of Machine Learning and Systems, volume 4, pp. 795-813, 2022. URL http:///moeedding..myyys..////pp files//paerr/002// fiee6222111020000856666555551111 Paper.pdf.
Shen, Y., Tan, X., He, D., Qin, T., and Liu, T.-Y. Dense in- formation flow for neural machine translation. In Walker, M., Ji, H., and Stent, A. (eds.), Proceedings of the 2018 Conference of the North American Chupter of the Associ- ation for Computational Linguistics: Human Language Technologies, Volume I (Long Papers), pp. 1294-1303, New Orleans, Louisiana, June 2018. Association for Com- putationnl Linguistics. doi: 18.186 1//53//1/N18----11 1117. URL https//accanthologg.org// 1117.
Xiong, R., Yang. Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang. H., Lan, Y., Wang, L., and Liu, T.-Y. On layer normalization in the transformer architecture. In Proceed- ings of the 37th International Conference on Machine Learning, ICML'20. JMLR.org, 2020.
Shleifer, S., Weston, J., and Ott, M. Normformer: Improved transformer pretraining with extra normalization. CoRR, abs/2110.09456, 2021. URL https://arxiv.rr//bb//1111 09456.
Xu, H., Song. Y., Liu, Q. van Genabith, J., and Xiong, D. Rewiring the transformer with depth-wise LSTMs. In Calzolari, N., Kan, M.-Y., Hoste, V., Lenci, A., Sakti, S., and Xue, N. (eds.), Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation (LREC-COLING 2024), pp. 14122-14133, Torino, Italia, May 2024. ELRA and ICCL. URL https//aclanttooog..o/////2 1231.
Srivastava, R. K., Greff, K., and Schmidhuber, J. Highway networks, 2015. URL https://arxxi.orr/abbb http///ass1155...5 505.00387.
Strubell, E., Ganesh, A., and McCallum, A. Energy and policy considerations for deep learning in NLP. In Ko- rhonen, A., Traum, D., and Marquez, L. (eds.), Pro- ceedings of the 57th Annual Meeting of the Associ- ation for Computational Linguistics, pp. 3645-3650, Florence, Italy, July 2019. Association for Computa- tional Linguistics. doi: 10.18853/v1/119-1335 URL https://elannholley.og///99999
main.
Xu, J. Sun, X., Zhang, Z., Zhao, G., and Lin, J. Understanding and improving layer normalization. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché-Buc, F., Fox, E., and Garnet,, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL. http://eeeeeaeeeeeuus.s.cccpppee flee2Heccar772277770066666688882
Sun, Y. Dong, L., Huang, S., Ma. S., Xia, Y., Xue, J., Wang. J., and Wei, F. Retentive network: A successor to transformer for large language models, 2023. URL https:/aarkivorr//bb/230......
Vaswani, A., Shazeer, N., Parmar; N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor programs v: tuning large neural networks via zero-shot hyperparameter transfer. In Proceedings of the 35th Inter- national Conference on Neural Information Processing Systems, NIPS '21, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN
NIPS' 17, pp. 6000-6000 6010,
Villalobos, P., Ho, A., Sevilla,
J., Besiroglu, T., Heim, L.., and Hobbhahn, M. Will we run out of data? limits of Ilm scaling based on human-generated data, 2024. URL https//aaxxivorr//bbbee 1.04325.
Zachary Zimmerman, M. G. and Gramlich, R. Ready-to- go transmission projects 2023. Technical report, Grid Strategies, 2023.
Wang, H., Ma. S., Dong, L., Huang, S., Zhang, D., and Wei, F. DeepNet: Scaling Transformers to 1,000 Layers IEEE Trunsactions on Pattern Analy- sis & Machine Intelligence, 46(10):6761-6774 October 2024. ISSN 1939-3539.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sen- tence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
doi: 10.1109/TPAM..2022.
3386927. URL
1109/TPAMI 2024.
12

---

Residual Matrix Transformers
A. Parameter and FLOP calculations
See Table for the equations used to compute FLOP and parameter counts for the transformer and RMT Fig. 1. $A.1 and $A.2 describe how the equations for the RMT were derived. The equations for the transformer follow Hoffmann et al. (2024). Our derivations for the RMT's FLOP and parameter counts also follow the approach taken by Hoffmann al. (2024). In particular, approximate the backward-pass FLOP count as twice the forward- pass count.
A.1. RMT Parameters
Embeddings
Word embeddings: RVD. De
Position embeddings: RND,
Key vectors: 2RD;
Embedding params: RVD, D. RND, + 2RDj
Attention (Single Layer)
QKV key vectors: 3RDx k.
O key vectors: RD2
Attention params: 3RD, RDv
Ferd-Forward Network (Single Layer)
Input key vectors: RDk
Core params: 2RD.Dpr
Output key vectors: RD2
FF params: RD2 + 2FD, + RDv
Unembedding
Key vectors: RD2
Unembedding weights: RV D
Unembedding params: RD2 + RV D.
Total Params embedding params L x attention params + FF params) + unembedding params + R(2D
D, (2LDFF V N))
A.2. RMT FLOPS
Embeddings
Word embeddings: 2NV RD,,
Position embeddings: 2NV RD.
Key vectors: 4N D.R
Embedding FLOPS: 2NVRD, +2NVAD,
Attention (Single Layer)
QKV key vectors:
SHA: 2NND,R 3RNN 2NND,R
O key vectors: 2N
Attention FLOPS: 6N D,D,
2NND, G.D...RR-NNNDR-3RR +2NND, 2ND.D.R
Feed-Forward Network (Single Layer)
Input key vectors: 2ND D.. R
13

---

Residual Matrix Transformers
Table 7. Parameter and forward pass FL.O count equations.
RESOURCE
TRANSFORMER DIL(AHD,, + 2DFF) +V)
RMT
PARAMETERS
4) D
FLOPS
4N(V LD(2D.H DrF) L(ND, (NH)
+ L(AND, +3N+ 4D.DFF))
Core params: ,NRD,,pp
- Output key vectors: 2N D.D,Da R.
FF FLOPS: 2ND ANRD, + 2ND
Unembedding
Key vectors: 2N
- Unembedding weights: 2NVRD.
- Unnembedding FLOPS: 2N 2NVRD,
Total forward pass FLOPS = embedding FLOPS + L x (attention FLOPS + FF FLOPS) + unembedding FLOPS
G1 L) + NR L(4ND, + 3N + 4D,D
B. Moment Propagation Derivation
Here we modify the derivations from (Kedia et al., 2024) derive the mean and variance propagation equations for the RMT's storage and retrival operations.
B.1. Storage
(N) For R input data vectors iin E R.,, and an output matrix Xout D..xD. the forward pass operation for single token is defined as follows.
R
Xout out
w
(h)
(b)
R
(h)T
w
x
in
The backward pass operation is then defined as follows.
(h)
w
(h)T Gout
gin
14

---

Residual Matrix Transformers
For the expected value of the forward pass, have,
R8
E
X out.,
E[
(b) (b)
Din,
/=11
n
E 20
(k)
iing
h=1
R
Σ
E E(w,
(1)
|E| |E|
(1)
(by ind.)
h=1
R
(K) E(w!
(h)
[
R
x
c m
(w is initialized with mean 0)
0
For the variance of the forward pass we have,
R
Var( (orou)))) Var(
u
(h) (h))
Finy
R
Var( wi
(h) (1))
Zing
(by ind. of weights from each other)
h=1
Σ (( + (o + (wx) 12 (i) (by ind. of weights and inputs)
h=1
(0 + 00
h=1
Rioz, + 2 00
For the expected value of the backward pass we have,
Dv
(1)
E
w
(b)
Goutg]
Dy
Blwsy Goutg] ]
B(w! (1) E[[Guuty] |
(by ind.)
i=1
0
15

---

Residual Matrix Transformers
And for the variance of the backwards path have,
D2
Var(g
(1)
Var
(1)
Goutes)
De Var(wo (R) Gousy)
(by ind. of weights from each other)
De
(1o + Mo + p(x) m (by ind. of weights and inputs)
i=1
Dx
Σ (o + G
Da + 2 )G
B.2. Retrieval
For R input matrix , IR DDDo and an output data vectors (h) E Pr, the forward pass operation for single token defined as follows.
(h) wilh) 1 X
Fout
!6)2 X
And for the backward pass we have,
G
w
(h) Jout
Notice that the forward equation for retrieval is the same as the backward table for storage, and vice versa. Then we can follow the exact same derivations in $B.1 to get,
a (h) 0
Var( (1))
(0xm + +xx.
E[Gin] =0
Var(Gin) Ro 2 (0² 0 your + H
C. Experimental Setup
C.1. Training Details
All of our models are trained on the OpenWebText dataset and use HuggingFace's GPT2 tokenizer. The vocab size is therefore set 50257 for all models. Learned positional embeddings are used, and unless otherwise specified, models trained with sequence length of 512 tokens. We use NoyeLay with the epsilon parameter to 1e-6 We do not any bias terms, nor do we use weight-tying for embedding and unembedding layers. In both models we include Mistral's GPT2 stability tweeks (attention upcasting and inverse layer scaling), as well as inverse layer scaling on layer outputs from GPT2. The loss function used zososs with its coefficient set to le 1e-4, When we report loss values, we report them without the loso term which is equivalent to normal cross-entropy loss. All models are trained with the Adam optimizer, with 3 = 0.9, B2 0.95, and - We use decoupled weight decay with scaling factor of 1e-4, Layernorm,
16

---

Residual Matrix Transformers
Table Transformer hyperparameters for scaling laws experiment
MODEL SIZE
DEVICE
BATCH SiZE
D
DFF
H
Dr
TRAIN TOKENS
49M
RTX 3090
64
6
384
1536
12
32
212M
160M
RTX 3090
64
12
768
3072
12
64
1.7B
260M
A100
896
3584
14
3.5B
405M
A100
256
24
1024
4096
16
64
6B
Table 9. RMT hyperparameters for scaling laws experiment
MODEL SIZE
DEVICE
BATCH SIZE
D2
De
DPF
TRAIN TOKENS
46M
RTX 3090
32
32
1536
142M
134M
RTX 3090
64
12
32
64
12
3072
1.1B
206M
A100
64
18
896
48
14
3584
2.3B
305M
A100
256
24
1024
64
16
4096
6B
embedding, and unembedding weights are excluded when weight decay is applied. We use linear warmup for 5% of the total training tokens, and cosine decay for the remaining tokens that decays to 10% of the max learning rate. In addition, gradient checkpoint is used with models.
All training experiments are implemented using JAX (Bradbury al., 2018), Equionx (Kidger & Garcia, 2021), and Haliax (Hall et al., 2024).
C.2. RMT Transformer Experiments
Table 8 and Table show the hyperparameters used for the experiments in $4.2.
C.2.1. ,PARAM
The "Parameterization Lottery" refers to how new architectures can appear to be successful or not successful based on their compatibility with existing hyperparameters. Furthermore, the typical hyperparameters used when training standard transformers are the product of years of tuning done by the machine learning community. To help mitigate biases associated with hyperparameter selection, we perform our own hyperparameter tuning on both the transformer and RMT's using j TTransfer. For each model type we perform 100 hyperparameter searches that follow suggestions from the TTrnnsfer paper. We sample from a logumiform distribution: learning rate on the interval [Le -4 Le '1, input alpha on the interval [le dle-_,1 1.]. attention alpha the interval [1e-1, 1.], output alpha the interval [1e 1.]. and initialization standard deviation the interval [le-1, 1.J. Our trail models have sequence length 128, 4 layers, 8 attention heads, head dimension 32, and 1024 feed-forward neurons. The embed size of the standard transformer is set to 256, and for the RMT the residual key and value dimensions are both set to 32, and the layer rank R is set to 8. These models are all trained on Nvida GTX 1080 Ti gpus with batch size of 32, train over 5K steps. We select the hyperparameters of the best performing run from each model type. For the standard transformer we have 6.5c input alpha 9.76, attention alpha 0.25, output alpha = 5.52, and init std 0.15.
C.3. Scaling Residual Stream Experiments
These models are all trained on Nvidia GeForce RTX 3090 gpus with a batch size of 64. models use a max learning rate of 1e and are trained over 50K steps. The models have 12 layers, residual matrix value dimension D. of 64, 3072 feed-forward neurons, and layer rank R of 12 for total of 135M parameters. Note that R is equal to the number of attention heads, R D, is input/output dimension of the core feed-forward operation (Eq. 24), and D, is the attention head dimension.
"Sme R from Eqs. $2.2
17

---

Residual Matrix Transformers
Table 10. Transformer variants statistics
ARCHITECTURE
MODEL SIZE
RESIDUAL STREAM SIZE
FINAL TRAIN LOSS
TRANSFORMER
160M
768
3.43
RMT
134M
2048
3.38
HIERARCHICAL AGGREGATION
165M
768
3.47
DEPTHWISE LSTM
233M
768
3.54
HIGHWAY TRANSFORMER
191M
768
3.41
C.4. Transformer Variants Experiments
All models were trained on Geforce-RTX 3090 gpus with batch size a 64. They were trained for 5000 train steps for total of 1.6B train tokens. The models all had 12 layers, DF of 3072, 12 attention heads, and D,. = 64. For the RMT model Dy: was set to 32 and D, was set to 64. Table 10 shows the models size and residual stream sizes for these models. The max learning rate was set to 1e- and we used Lecun initialization for all parameters.
D. Downstream Evaluation Details
All results from Table 6 plus the Lambada perplexities presented in Table 5 were generated using the EleutherAl evaluation harness (Gao et al., 2023). The remaining perplexity results in Table 5 were evaluated by hand.
E. Runtime Discussion
Although hardware-spesn -seerfcc implementations and diagnostics are outside the scope of this work, we hypothesize that much of the slowdown is due replacing highly optimized GEMM operations with unoptimized tensor contractions. The storage and retrieval operations are both implemented tensor contractions, whereas the storage and retrieval operations for the transformer are GEMMs. The contraction dimension of the RMT's tensor contraction is small compared to the embedding dimension of the transformer (the K dimension of the transformer's GEMM operations). We expect that naive implementation of this tensor contraction will have much lower throughput than an optimized GEMM kernel; however, we expect that a custom CUDA kernel that takes advantage of the small size the key vectors could dramatically improve runtime performance.
18


Reviews:
[{"review_id": "hrNtkorf3U", "reviewer": "Reviewer_fZLC", "summary": "To achieve more data-efficient and compute-efficient models, this paper introduces a new transformer-variant called the Residual Matrix Transformer (RMT), which replaces the traditional residual stream with an outer product memory matrix. The authors present theory showing that the RMT exhibits efficient scaling of the residual stream and improved variance propagation properties in some cases.\nExperimental results demonstrate that when using an RMT with a larger residual stream size, it is more efficient than traditional transformer models in terms of data, FLOPS, and parameters. Additionally, it proves to be superior in performance compared to various transformer variants. Ultimately, the paper shows that increasing the residual stream size leads to better model performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer fZLC’s thoughtful feedback and positive comments. We will address each of this reviewer’s concerns in the order they appear.\n\n> The experiments may lack some ablation studies to demonstrate the contributions of different factors. For instance, the authors need to demonstrate whether the superior performance of RMT over the Transformer stems from the outer product memory matrix structure or is simply a result of the increased residual stream size. \n\nThis is a valid point that we will make more clear in the paper.  In our experiments, we find that the performance gains are not attributable to the outer product memory matrix structure itself, but rather to the expanded residual stream that this structure enables. One of the RMT models in §4.4 has the same residual stream size as the Transformer in §4.3, and these models are trained with the exact same experimental settings. The final train loss that the RMT achieves is 3.42 while the final train loss that the Transformer achieves is 3.43, showing that when the residual stream size is the same, the observed performance is about the same. These results suggest that the observed performance gains of the RMT over the Transformer in §4.1 are due to the expanded residual stream. \n\n> In the main comparison, RMT outperforms the Transformer, but its residual stream is 2.5 to 4 times larger, which could also be a contributing factor to the observed performance gains. Beyond the analysis provided in §3.1 of the paper, thorough ablation experiments are crucial to validate these claims.\n\nAs mentioned above, the observed performance gains are most likely due to the residual stream being 2.5 to 4 times larger. We would like to clarify, however, that the ability to expand the residual stream to this size is made possible by the RMT’s outer product memory matrix structure. Expanding the residual stream of the Transformer would substantially increase the model’s parameter count and per-example compute cost. As such, an ablation where the Transformer’s residual stream size is expanded but the parameter count, FLOP count, and tokens consumed is fixed is not possible.\n\n> The paper states in lines 81-83 that \"RMT has improved variance propagation.\" However, Table 2 reveals that RMT underperforms the Transformer in the Attention retrieval case. These \"bad cases\" need further analysis to provide readers with a comprehensive understanding of the proposed method's performance and its potential application scenarios.\n\nThis is also a valid point, however we consider that performing a full analysis of the end-to-end signal propagation through the RMT to be out of the scope of this paper (for example, an entire paper was dedicated to this for the Transformer (Kedia et al., 2024)). We maintain that improving the signal propagation of 3 out of the 4 replaced components roughly shows that our model has superior signal propagation properties compared to the Transformer. We also want to note that we found an error in the submitted manuscript. In Table 2, the attention storage and retrieval numbers are swapped, i.e. the RMT outperforms the Transformer in the Attention retrieval case and underperforms the Transformer in the Attention storage case. An correction to the table is provided below.\n\n\n\\begin{array} {|r|r|}\\hline Layer & Operation & Model & \\frac{\\sigma^{2}\\_{x\\_{out}}}{\\sigma^{2}\\_{x\\_{in}}} & \\frac{\\sigma^{2}\\_{g\\_{in}}}{\\sigma^{2}\\_{g\\_{out}}} \\\\\\\\ \\hline Attn & Storage & RMT & 0.4 & 1.6 \\\\\\\\ \\hline Attn & Storage & Transformer & 1 & 1 \\\\\\\\ \\hline Attn & Retrieval & RMT & 1.14 & 0.86 \\\\\\\\ \\hline Attn & Retrieval & Transformer & 0.5 & 1.5 \\\\\\\\ \\hline  \\end{array}\n\n> Yes, there are essential references that could enhance understanding.\n\nWe have included all the references that we find enhance the understanding of our work, and we welcome any specific suggestions for additional sources that would strengthen the manuscript.", "claims_and_evidence": "Most of the claims made in the submission are supported by clear and compelling evidence, but some of them also need further justification (see Experimental Designs Or Analyses).", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-suited for the problem and application at hand.", "theoretical_claims": "I checked the proofs for the theoretical claims, and they were correct and well-justified.", "experimental_designs_or_analyses": "- The experiments may lack some ablation studies to demonstrate the contributions of different factors. For instance, the authors need to demonstrate whether the superior performance of RMT over the Transformer stems from the outer product memory matrix structure or is simply a result of the increased residual stream size. In the main comparison, RMT outperforms the Transformer, but its residual stream is 2.5 to 4 times larger, which could also be a contributing factor to the observed performance gains. Beyond the analysis provided in §3.1 of the paper, thorough ablation experiments are crucial to validate these claims.\n- The paper states in lines 81-83 that \"RMT has improved variance propagation.\" However, Table 2 reveals that RMT underperforms the Transformer in the Attention retrieval case. These \"bad cases\" need further analysis to provide readers with a comprehensive understanding of the proposed method's performance and its potential application scenarios.", "relation_to_broader_scientific_literature": "The key contributions of the paper align well with existing literature, building on prior findings and enhancing our understanding of the topic.", "essential_references_not_discussed": "Yes, there are essential references that could enhance understanding.", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "No", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P1YdK9YXei", "reviewer": "Reviewer_J7dP", "summary": "The paper introduces Residual Matrix Transformers (RMT), which increases the size of the residual stream in a transformer without incurring significant compute over memory overhead by using an outer product memory matrix. In training GPT-2 language models, RMT achieves better loss per unit of compute or parameters.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer J7dP for their thoughtful reading of the paper and appreciate their comment that this work presents an important research direction. We provide responses to many of their concerns grouped by subject.\n\nWe first address the concerns related to our application of µP transfer.\n\n> The use of µP for estimating optimal learning rates for different models is not necessarily sound, given that the training iterations are not constant across model sizes, which is an assumption in µP. \n\nWe argue that using µP to transfer hyperparameters across training iterations is reasonable given the empirical evidence provided in the µP paper. The authors empirically showed that their method works across sequence length, depth, batch size, and training time (see Table 1 of their paper). In fact, the only hyperparameter that the paper theoretically proved transferability across was model width. In their GPT-3 experiment (the closest setting to ours), the authors successfully transfered hyperparameters from a proxy model trained with significantly fewer training iterations. Given the success of their experiment, we felt that it was reasonable to take the same approach that they did.\n\n> In addition, I'm concerned that the larger learning rate used for RMT may unfairly favor it over standard transformers. Overall I think a more careful learning rate sweep is necessary to demonstrate the superiority of RMT convincingly.\n\nIn our main comparison between the RMT and the Transformer (§4.2), we included learning rate as one of the hyperparameters in our μP search and used the best performing learning rates of both the RMT and transformer models.\n\n> When using µP for estimating optimal learning rates, do you properly scale the found learning rate per layer from smaller to larger models?\n\nWe scale learning rate in accordance with Table 8 and Appendix B.1 of the µP paper (i.e. by $\\frac{1}{\\textrm{width ratio}}$ for hidden parameters).\n\n\nHere we give a response to the reviewer’s concerns about the runtime of our model. \n\nWe do not think runtime is an insurmountable obstacle for our model for two reasons. The first is that, without any hardware optimizations, the RMT only takes 4% more time to achieve the same train loss as the transformer. This shows that the performance gain needed is within reasonable reach. The second reason is that, while the residual stream size is much larger in the RMT, many of the parameter matrices are much smaller. One can use this fact to devise more efficient kernels for the RMT. For example, many of the “key” parameters are so small that they can fit entirely into SMEM. One can then imagine a GEMM kernel that only has to reload one operand matrix into SMEM, significantly decreasing the data transfer overhead. \n\nWe further assert that the main focus of our paper is to explore residual stream size as a new scaling axis. While we do discuss runtime in the interest of full transparency, we consider the discussion of specific hardware optimizations to be out of the scope of this work.\n\nFinally, we respond to the following concern:\n> RMT makes multiple modifications to the transformer architecture. These modifications do not represent the unique way of scaling the residual stream size and deserve to be motivated better or be ablated to illustrate which components actually matter for its performance.\n\nWe acknowledge the concern that RMT introduces several modifications to the transformer architecture that may seem arbitrary without proper justification. However, we would emphasize that these changes represent the minimal adjustments necessary to support the outer product memory matrix structure of its residual stream. Each modification is integral to the model's coherent functioning - removing any single element would create architectural inconsistencies. While we certainly don't claim RMT represents the only possible approach to scaling residual stream size, our work demonstrates one viable method for doing so, allowing us to investigate the effects of scaling along this axis.", "claims_and_evidence": "The GPT-2 experiments support the claim that RMT outperforms the standard transformer, but experiments on other datasets (e.g. images) would be helpful for assessing the generalizability of the finding.", "methods_and_evaluation_criteria": "The GPT-2 experiment on OpenWebText is a reasonable benchmark. But it would be useful to understand if applying RMT still improves performance when applied on top of the modern transformer architecture actually used in practice such as the one used in llama models (with rotary embedding, SwiGLU activation, RMSNorm instead of LayerNorm and no biases).", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The use of µP for estimating optimal learning rates for different models is not necessarily sound, given that the training iterations are not constant across model sizes, which is an assumption in µP. In addition, I'm concerned that the larger learning rate used for RMT may unfairly favor it over standard transformers. Overall I think a more careful learning rate sweep is necessary to demonstrate the superiority of RMT convincingly.", "relation_to_broader_scientific_literature": "Finding new axes worth scaling beyond the usual ones, like compute and parameters, is an important research direction. The finding that scaling the residual stream size alone leads to considerable performance gains is interesting and relevant to the community. On the other hand, as the authors discussed, this approach increases the memory overhead and data transfer, which can translate to slower runtimes. While RMT performs better than standard transformers when controlling for FLOPs, ultimately, we care about performance per unit of time, and using fewer FLOPs is not relevant if the hardware utilization is compromised. Indeed, operations like attention are often memory-bound rather than compute-bound, and techniques that reduce the data transfer can significantly speed up the runtimes even while increasing the FLOPs, as exemplified by FlashAttention [1].\n\n[1] Dao, Tri, et al. \"Flashattention: Fast and memory-efficient exact attention with io-awareness.\" Advances in neural information processing systems 35 (2022): 16344-16359.", "essential_references_not_discussed": "None that I'm aware of.", "other_strengths_and_weaknesses": "Strength: Scaling the residual stream size is a novel idea and appears to be a promising research direction.\n\nWeakness: RMT makes multiple modifications to the transformer architecture. These modifications do not represent the unique way of scaling the residual stream size and deserve to be motivated better or be ablated to illustrate which components actually matter for its performance.", "other_comments_or_suggestions": "While the authors do not address the problem of the slower runtime of RMT, I believe it is an important question and it is unclear whether the runtime overhead of RMT can be addressed even in principle due to the additional data transfer. I suggest that the authors analyze and discuss whether the runtime overhead of RMT can, in fact, be addressed and by what kind of approaches. For example, if most of the overhead is in the additional data transfer involving the larger residual stream matrix, that overhead may in fact be irreducible without changing the hardware.", "questions_for_authors": "1. When using µP for estimating optimal learning rates, do you properly scale the found learning rate per layer from smaller to larger models?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OmkldMawco", "reviewer": "Reviewer_ciWT", "summary": "The Residual Matrix Transformer (RMT) replaces the residual stream in transformers with an outer product memory matrix, allowing independent scaling of the residual stream size. This results in improved training efficiency and better performance on downstream tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer ciWT for their comments and provide responses to their concerns about runtime efficiency and baseline comparisons.\nFor concerns about runtime, we will copy and paste a relevant snippet of our response to reviewer J7dP. We encourage Reviewer ciWT to read our discussion with Reviewer J7dP if more context is needed.\n\n\nWe do not think runtime is an insurmountable obstacle for our model for two reasons. The first is that, without any hardware optimizations, the RMT only takes 4% more time to achieve the same train loss as the Transformer. This shows that the performance gain needed is within reasonable reach. The second reason is that, while the residual stream size is much larger in the RMT, many of the parameter matrices are much smaller. One can use this fact to devise more efficient kernels for the RMT. For example, many of the “key” parameters are so small that they can fit entirely into SMEM. One can then imagine a GEMM kernel that only has to reload one operand matrix into SMEM, significantly decreasing the data transfer overhead. \n\nWe further assert that the main focus of our paper is to explore residual stream size as a new scaling axis. While we do discuss runtime in the interest of full transparency, we consider the discussion of specific hardware optimizations to be out of the scope of this work.\n\nNext, we address concerns about baseline comparisons.\n> The transformer variants in §4.3 (Dou et al., 2018; Xu et al., 2024) are not state-of-the-art (e.g., recent works like Mamba or RWKV). Including stronger baselines would better contextualize RMT’s advancements.\n\nWe consider works like Mamba and RWKV to be orthogonal to our work because they primarily change the attention layer of the Transformer. Because these architectures use the residual stream in the same way the Transformer does, our method can be extended to be integrated with them. We agree with the reviewer that some clarifying discussion should be added to our paper.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "None", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "1. Runtime efficiency concerns: The authors acknowledge that runtime is currently the biggest limitation of their model, with RMT being 4% slower than the transformer despite being more FLOP-efficient. This suggests potential implementation inefficiencies that could limit practical adoption.\n2. Limited model size exploration: Due to resource constraints, the authors couldn't explore how efficiency trends continue at larger model sizes beyond 405M parameters, leaving questions about scalability to truly large models.\n3. Baseline Comparisons: The transformer variants in §4.3 (Dou et al., 2018; Xu et al., 2024) are not state-of-the-art (e.g., recent works like Mamba or RWKV). Including stronger baselines would better contextualize RMT’s advancements.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All reviewers argued for acceptance but only one engaged after the rebuttal period (J7dp). They had lingering concerns about whether hyperparameters introduced an bias that would favour RMT and whether the overhead of RMT could be significantly reduced with a better implementation.
Decision: Accept (poster)