Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices
Authors: Andres Potapczynski
Shikai Qiu
Marc Anton Finzi
Christopher Ferri
Zixi Chen
Micah Goldblum
C. Bayan Bruss
Christopher De Sa
Andrew Gordon Wilson
Venue: NeurIPS 2024 poster (2024)
Keywords: Structured Matrices
Scaling Laws
Linear Layers
Optimization
Deep Learning
Einsum
Primary area: deep_learning_architectures
TL;DR: We learn hardware-efficient linear layers by constructing and searching a continuous space of structured matrices

Abstract: Dense linear layers are the dominant computational bottleneck in large neural networks, presenting a critical need for more efficient alternatives. Previous efforts to develop alternatives have focused on a small number of hand-crafted structured matrices, and have neglected to investigate whether these structures can surpass dense layers in terms of compute-optimal scaling laws when both the model size and training examples are optimally allocated. In this work, we present a unifying framework that enables searching among all linear operators expressible via an Einstein summation. This framework encompasses many previously proposed structures, such as low-rank, Kronecker, Tensor-Train, and Monarch, along with many novel structures. We develop a taxonomy of all such operators based on their computational and algebraic properties, which provides insights into their scaling laws. Combining these insights with empirical evaluation, we identify a subset of structures that achieve equal or better performance than dense layers as a function of training compute. To further improve their compute efficiency, we develop a natural extension of these performant structures that convert them into a sparse Mixture-of-Experts layer. The resulting layer significantly outperforms dense layers in compute-optimal training efficiency for GPT-2 language models.

Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices
Andres Potapezynski New York University
Shikai Qiu' New York University
Mare Finzi Carnegie Mellon University
Christopher Ferri
Zixi Chen New York University
Micah Goldblum
C. Bayan Bruss Capital One
Capital One
Columbia University
Christopher De Sa Comell University
Andrew Gordon Wilson
New York University
Abstract
Dense linear layers are the dominant computational bottleneck large neural networks, presenting a critical need for more efficient alternatives. Previous efforts focused on a small number of hand-crafted structured matrices and neglected investigate whether these structures can surpass dense layers in terms of compute- optimal scaling laws when both the model size and training examples are optimally allocated. In this work, we present a unifying framework that enables searching among all linear operators expressible via an Einstein summation. This framework encompasses many previously proposed structures, such as low-rank, Kronecker, Train, Block Tensor-Train (BTT), and Monarch, along with many novel structures. To analyze the framework, develop taxonomy of all such operators based on their computational and algebraic properties and show that differences in the compute-optimal scaling laws mostly governed by small number variables that we introduce. Namely, a small w (which measures parameter sharing) and large & (which measures the rank) reliably led to better scaling laws. Guided by the insight that full-rank structures that maximize parameters per unit of compute perform the best, propose BTT-MoE, novel Mixture-of-Eaie of--xxsrts (MoH) architecture obtained by sparsifying computation in the BTT structure. In contrast to the standard sparse MoE for each entire feed-forward network, BTT-MoE learns an MoE in every single linear layer of the model, including the projection matrices in the attention blocks. We find BTT-MoE provides a substantial compute-efficiency gain over dense layers and standard MoE.
1 Introduction
Neural networks primarily consist interleaved linear layers and simple non-linearities. In large foundation models such GPT-3 [4], these linear layers consume the vast majority of the parameters and computation [13], and are commonly represented by dense matrices. Substituting these dense matrices with structured matrices with fast matrix-vector multiplies (MVMs) has the potential to significantly improve the computational efficiency of these models.
Recent work by Dao et al. [6], Fu et al. [8], and Qiu et al. [19] demonstrated that incorporating certain structured matrices into neural network architectures, including transformers, can improve
Equual contribution. Correspondence to ap66048nyu. edu, sq21298nyu. edu, andreugw@cims wwt.:rmiiimmm..mm...mmmmmm edu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Kronecker
5/2
Scaling Laws GPT-2
B
1/2
X
0.7
0.5
B
1/2
0.4
Kronecker
Low: Low.Rank
Monarch
General Einsum
Monarch
Generic Einsum
B
3/2
X
0.2
10
10th
1045
10a7
1/2
Compute (FLOPs)
Figure 1: We use Einsums to parameterize wide range of structured matrices and search for the most efficient structure for compute-optimal training. Left:: diagrammatic representation of a general two-factor Einsum. We parameterize the space of Einsums through real-valued vector (0a, 03, 0y 0s, Oz, 0s 0p) E 1]" This space captures many well-known structures through specific values of 0. Middle: Example of well-known structures with their 0 values. Any omitted line implies the value of the entry in the vector Right: Compute-optinge optimaal scaling laws of example structures for GPT-2 Open Text when substituting its dense layers (see details in Section 4).
performance over dense models of the same size trained for equal number of epochs on problems such as ImageNet classification. However, these success cases do not reflect the current paradigm of large-scale training, where the models 1) are typically not trained for multiple epochs, making the expressiveness of dense matrices particularly appealing since the generalization gap vanishes, and 2) are heavily bottlenecked by compute cost, making it infcasible to train the models until convergence [13, 11]. unlike in image classification. These attributes of large-scale training make the compute-optimal scaling rather than scaling in model size alone more relevant.
In this work, we investigate how different structures perform in a compute- opttmall setting, which characterizes performance as function of training compute when allocated optimally between using larger models versus training on more data [13]. In language modeling and many other tasks, the compute-optimal scaling law has been shown to take the form
L.x + bC- as function of training compute C, where Lox. is the minimal achievable loss [13, 11, 10]. Quantifying the compute-optimal scaling laws various structures essential for understanding their practical value for training large-scale neural networks.
In addition investigating the scaling laws of existing structures, we expand the set of matrix structures beyond what has been previously considered. We do so by introducing a continuous parameterization the space of all possible structures whose matrix- vector-multtplicatioon (MVM) can be expressed as an Einstein summation (Einsum). This space contains many known structures such low-rank, Tensor- Train [17], Kronecker product [21, 25, 9], Monarch [6] and Block Tensor- Train [19]. but also includes many novel hardware-efficient structures. Indeed, all structures in this space are hardware-efficient in the sense that they are computed through a series of batch matrix multiplication primitives, which we implement through the Linear Operator abstractions available in CoLA [18]. Moreover, this space lends itself to an intuitive exploration as we can analyze how different parameters of the Einsum affect a structure's performance and scaling laws. We make our code available here.
We summarize our main contributions as follows:
We introduce a continuous parameterization of the space of structured matrices whose can be implemented via an Einstein summation (Einsum). This parameterization allows us to search wide range of hardware-efficies structured linear layers for neural network architectures beyond handful of well-known cases identified in prior work [6. 8, 19].
We develop a taxonomy of the space Einsums based on its computational and algebraic properties. We identify three key scalar quantities that characterize this space (w, vi v). (I) w 2 O, which reflects the extent of parameter sharing in a matrix. (2) & E 10, 1]. which characterizes to the rank of the structure (1) meaning full-rank). (3) E [0, 1), which
TTechnically, we consider everything that can be expressed via torch. einsum, which is slightly more general than the Einstein summation, which allows index to appear most twice.

---

relates to compute per dimension in an MVM, where the upper-bound v/ = 1 is achieved by dense matrices. Intuitively, measures how much a structure resembles dense.
We investigate the scaling laws of different Einsums language modeling, autoregressive image generation, and a synthetic regression task. We find that the best pessforrronnnn structures are the ones that do not share parameters (w 0), are full-rank (v) 1); while v can be varied with often negligible impact to their scaling laws. In contrast to previous findings, we demonstrate that structures shown in prior work [6, 8, 19] outperform dense matrices in non-compute-optinal settings can yield similar but not significantly better compute-optimal
scaling laws on these tasks.
Building on prior work in structure-aware learning rates [19], we show how to properly initialize generic Einsum layers and transfer learning rates from the original dense layers and across model sizes, leveraging insights from μPP [29, 26] and manifold optimization [3].
Based on the observed relation between the taxonomy variables and the scaling laws, we propose a new structured Mixture of Experts (MoE) architecture implementing sparse mixture of multiple structure matrices. This block tensor-train (BTT) MoE provides a sparse MoE in every single layer of each feedforward network (FFN) and attention project matrices, compared to standard MoE which operates over entire FFNs. We show BTT-MoE is significantly more compute-efficien than dense matrices and standard MoE for training GPT-2 language models.
2 Parameterizing the Space of Einsums
We now present a unifying framework that parameterizes all linear operators W E Rod xdm whose Wx can be expressed as an Einsum over the tensors x, A, J,B..... where A, B.... defines the operator W and contains all its learnable parameters. To simplify the presentation, throughout this paper we assume W is defined using only two factors A, B, but we show generalization to more than two factors is straightforward in Appendix E.
We consider the following general expression of such an Einsum
Yoep-
Σ agppp
BexcopArr6op
where the vectors x and are written as tensors with multiple indexes to allow them to interact differently with each other, and the factors A, B. Each index æ E {a, B, y, 8, e, o. p} ranges from 1 to dz. Given this general expression, we obtain different structures via different factorizations of din into do dgd,, and dout into dzd,dg, and separately a choice of dp. For example, for a low-rank matrix, we have d, dan, dg = d, 1, de dout; dā dø and dp r. For a Kronecker product, we have dæ dg v Vdn. ds = de ddout and d, = dø dp 1. We provide an extended list of examples in Appendix A.
The general expression and the above two examples can be more conveniently and intuitively represented as diagram shown in Figure (left), where each index a. 3,7,... corresponds to a (hyper)edge among the input, output, and weight factors: + {X, A}. 3 <> B). 7 <> {X, A, BJ. + A), E + [Y, + A, and P + [A This set edges can be written succinetly as E = (S C (X, subsets that contain X and Y simultaneously, as adding them simply produces an already included structure but repeated multiple times along one of the input and output axes. The structure of particular Einsum is fully specified by the the vector (dXA, dxB, dxAB. dy dyB, dyAB, dAB). which specifies the range of the indices a, B, 7, 8, e, o.p. p. When the range of an index is of size 1, the corresponding edge effectively disappears from the diagram and the expression simplfies.
[S] 2 and [X, YJ g S}. We exclude
As we will build models of varying sizes, it is more natural to think about how these entries scale with din and dout- We therefore assign a real-valued vector O E [0, 1]7 to each structure indicating that d, = df. for E (XA, XB, XAB), d, Li for j E (YA, YB, YAB), and dAB = min (din dout) (AAB with the restriction that 0X +xx=++xxxxx +0XAB = OY +0YB + OYAB = 1. For example, low-rank matrix whose rank scales as the dimension it operates on to the 1/2-th power represented as 0 = (1,0, 0, 1/2), and a Kronecker product of two factors of equal sizes is represented as 0 = (1/2, 1/2, 1/2, 0). We round all d, to its nearest integer when
in
3

---

Monarch
Dense
WA
Tensorr: Timin
Kronecker
1
2/3
ue
3/4
1/2
Low-rank
BTT sulepace
Full-rank BTT subspacc
Figure 2: Illustrating the Einsum taxonomy. The 3D graph represents relevant quantities of the Einsum structure such as the amount of parameter sharing (x-axis), its rank (y-axis), and its compute intensity / (z-axis). The structures on the left of the figure appear as dots on the graph based on their coordinates O. We highlight two key subspaces. (a) The BTT subspace, characterized by no parameter sharing
0, learning the maximum number of parameters per FLOP. (b) The full-rank BTT subspace where contains the most performant structures across multiple tasks.
and &
Section we show that the full-rank BTT subspace
instantiating the Finsum. Note this rounding only quantizes the the values of di, but leaves the space of meaningfully different Os continuous as we consider arbitrarily large matrices.
Taxonomy of the Space of Einsum Linear Structures
The space of all Einsums is a high-dimensional space containing a wide range of possible structures. A priori, it is difficult to reason about the properties of a particular point in this space given its coordinates E [0, 1)². In this section, we develop a taxonomy of the space of Einsums based on their computational and algebraic properties. We introduce three key scalar quantities that characterize this space: (1) v/ E [0, 1]. which is related to the rank of the structure, (2) P E [0, 1]. which is related to the compute intensity of the structure, i.e. FLOPs per MVM divided by the dimension, and (3) 2 O, which is related to the number of learnable parameters divided by the FLOPs per MVM. Each of these three quantities can be expressed in terms the entries of 0, which To simplify the presentation, we assume d dout = d in the rest of the section. Without loss of generality, we assume min(Oxx. OYB) min(OYA. BB so that it is more efficient to first multiply by A instead of B.
derive in Appendix B.
Rank Exponent, v/. For given 0, we have that rank () = O where v = min [1, 2 ABB 0XA 0YB). Thus,
0.90
0.80
1 implies full-
0.70
rank and decreasing values of a im- ply lower ranks until the limit of 0. In other words, the rank decreases when an Einsum increasingly allocates part of the input and output to factors that only connected the input out- put. The limit being a low-rank struc- ture seen in Figure 2 where 0XA 1 = 0YB and which creates a bottle- neck on (A, B). The opposite trend is exhibited by dense, which allocates the full dimension of the input and output to both factors as seen in Fig- ure Nonetheless, is still possible to achieve full-rank when allocating dimension to only single factors long as those values do not constitute most of the allocation, meaning 0 0XA= < 1/2 and 0 < 0YB 1/2, seen in Monarch, Tensor- Tensor-Train, and Kronecker
0.60
0.40
4..00..000..
10.00..0.10..)
0. (.
10
1
10
Compute (FLOPS)
Figure Compute-optimal frontier (highlighted points) of various Einsums follows power law scaling, As result, Einsums be scaled reach arbitrarily low reducible loss, each with different that estimated from small-scale experiments.

---

but not for the particular Block Tensor-Traie (BTT) structure [19] in Figure 2. In Section 4 we show that structures with di perform best when training neural networks.
Compute Intensity Exponent, v. Let F denote the FLOOPs required to perform a MVM and define the compute intensity F/d θ (d The upper bound is achieved by dense, which requires quadratic compute for an MVM and thus v = 1. In general, we have V = 1 + 1+0AB min(0x.. (0xA. OYB). Thus, in order to achieve lower compute intensity than dense, a structure has allocate dimensionality to factors that only connect to the input and output. As seen in Figure 2, the BTT example is able to achieve the lowest compute intensity by allocating a substantial part of the input dimension to the first factor and substantial part of the output dimension to the second factor. &b and v are not completely independent, e.g. & = v for low-rank matrices W., Σi= Buj. though exceptions exist such as for the Kronecker product where y can be arbitrarily low while maintaining & = 1. In Section 4, we show there exists a wide range of structures with varying v that perform as well as dense matrices.
Parameters-Sharing Exponent, w. Let N denote the number of parameters in the structure then N/F = θ (d.). Clearly, w for dense matrices where each parameter used exactly once in an MVM. In general, we can show that w min + AY 0XB + 0YB) - min(0xA, OYB). In Section 4 we find that structures that share parameters, that is w >> 0, have worse scaling laws that structures that do not (w = 0). To achieve w O, we have to avoid introducing edges that skip some factors, that is 6XB and Kronecker where there exists an edge that connects X to B while skipping A. In contrast, in Monarch and BTT, the edge connecting X with B also touches A.
fY A 0. Structures that skip factors in Figure 2 are Tensor- Train
4 Scaling Laws of Einsums
While prior works have shown that certain structured matrices such as Monarch and BTT have better scaling laws than dense matrices as function of model size on datasets such as CIFAR- 10. CIFAR-100, and ImageNet [6, 19]. their experimental setups do not reflect today's large-scale training. where the models 1) typically do not train for multiple epochs on the training set, and 2) are heavily cempute-hottlenecked such that we care primarily about performance as a function of training compute rather than model size (omitting the cost of training). These attributes of large-scale training make the compute-optimal scaling rather than scaling in model size alone more relevant.
In this section, we investigate the compute-optimal scaling laws of a wide range of Einsums - how their performance scales function of training compute. We will show that we can understand the systematic differences in the scaling laws of various Einsums by leveraging the taxonomy we have developed. While we do not find structure that achieves noticeably better scaling laws compared to dense matrices, we identify the set of common properties shared across wide range of structures that match the performance of dense matrices, based on which we will propose significantly more efficient alternative to dense layers in Section 5.
4.1 Main Experimental Setup
We train GPT-2 [20] language models on the Web dataset. To make our measurement of the scaling laws more robust and our experiments more affordable, we reduce the vocabulary of the original GPT-2 to 96 commonly used alphanumeric symbols. Using a small vocabulary limits the compute and parameters consumed by the language modeling head, which would otherwise obscure the scaling laws measured at small scales [13]. We train models of varying sizes from 120k to 76M parameters, with model dimension d [256, 4096) and number of transformer blocks I. E {3, 6). Each model is trained for 100k steps with a batch size of 65536 tokens and a sequence length of 128. All linear layers except the language modeling head are replaced with Einsums. We use the Adam optimizer with a base learning rate of 0.003 for a L = 3, d 256 dense model, and scale it using μP [27] and structure-aware learning rates [19] to larger models and models using Einsums in place of dense layers. We discuss learning rate scaling in detail in Section 6, showing it is crucial for the performance of Einsums. We weight normalization to stabilize the training of Einsums following Qiu et al. [19].
In Appendix Appendix D, we show our main conclusions derived from this simplified setup translate to the more standard GPT-2 evaluation with a longer sequence length of 512 and its original vocabulary of 50,257 tokens.
5

---

Full Space
w=
0.
1.00
0.5
1.0
1.0
0.90
0.80
0.70
0.60
Loss
0.50
0.40
0.30
Dense
Dense
Dense
0.0
0.5
0.5
10tH
10th
10th
10fr
10a4
10th
10th
10
10te
10th
10th
10
Compute (FLOPs)
Compute (FLOPs)
Compute (FLOPs)
Effect of w
Effect of g/
Effect of
Figure 4: The taxonomy parameters (w, v)) explain differences in the scaling laws. (Left): parameter sharing > 0) leads to worse scaling. (Middle): among structures without parameter sharing (w = 0), full-rank structures (vb = 1) scale better than low-rank structures (té < 1). (Right): 1) subspace, various structures have nearly indistinguishable scaling laws compared to dense matrices, suggesting that not implementing parameter sharing and being full-rank are the necessary and sufficient conditions for a eofpute-efficient linear layer for GPT-2.
in the (w
4.2 Analyzing the Compute-Optimal Scaling Laws
Einsum Performance Obeys Power Law Scaling. When replacing the standard dense layers with Einsums, we find the resulting model's loss continues to follow the usual L L. + bC-- compute-optimal scaling laws except with possibly different constants a, b. In Figure 3, we visualize the compute-optimal scaling laws of various Einsums our language modeling task, including those corresponding to previously proposed structures such as TT [17], Low-rank and BTT [19], as well as generic Einsum with all entries of 0 strictly positive. We report the reducible loss with an estimated Lx = 0.75 subtracted. This finding suggests that all Einsums can be scaled to reach arbitrarily low reducible loss, each with different rate that can be estimated from small-scal experiments.
Parameter Sharing Leads to Worse Scaling. As discussed in Section 3, the vast majority of Einsums implement some kind of parameter sharing, where the number of parameters N in the Einsum relates to its MVM FLOPs via N/P = ((d-w), for some 4. > 0. Figure (left), we show the scaling laws of wide range of Einsums (only including points the compute-optimal frontier) colored by w. We find larger values of w lead to significantly worse scaling laws. To search for compute-efficient structures, we should therefore focus on the subspace with 0.
Full-Rank Performs Best. Within the w subspace, we find that ve becomes the next most important parameter. Recall a E [0, 1] is defined such that the rank of the Einsum scales as (dd)). Einsums with di introduce information bottlenecks the model by preventing the linear layers from accessing information from all the feature dimensions. The smaller 1é is, the more severe this effect. In Figure 4 (middle), we show that small values of vi indeed lead to worse scaling laws. This observation further narrows down our search to the subspace with 0 and 1o = 1, i.e. the space of full-rank BTT matrices.
Any Full-Rank BTT Scales Similarly as Dense. The w = 0 and & = 1 subspace contains the Monarch matrices and its generalization BTT matrices', From computational perspective, a primary distinguishing factor among these structures is how close they resemble a dense matrix, which we characterize by their compute intensity E [o, 1) defined so that F/d d (dd"). is large whenever their exists large values (close to 1) in the remaining allowed entries (0Ax. 6 ABX, @YB. @YAB. AAB). In Figure 4 (right), we show that, somewhat surprisingly, has minimal effect on the scaling laws of these structures. Structures with different v have almost indistinguishable scaling laws compared to each other and dense matrices, which has v = This result shows that while dense matrices perform well compared to the vast majority of possible Einsums, their good performance does not arise from being dense, but rather from not sharing parameters and being full-rank.
Reconciling with Results from Prior Work. Our findings do not mean that structured matrices cannot outperform dense in other settings. Rather, they highlight that the relative performance between structures depends on what resource is controlled, as prior work has shown that low rank, Tensor-
In the 2- factor case, this space corresponds all Block Low-Rank [1] matrices that are full-rank.

---

Full Space
0=
0.70
0.60
0.50
0.40
Loss
0.30
3
0.20
Dense
Dense
Dense
0.10
10
1013
10
10
0
10ah
10
10
0
10ah
10
10
Compute (FLOPs)
Compute (FLOPs)
Compute (FLOPs)
Full Space
w 0
3 0, 4 1
0.5
1.0
1.0
0.03
Loss
0.01
0.003
Dense
Dense
Dense
10 10 10t4
10t5 10th
0.0
10 10
10th
10t 10
0.5
10
10th
10 10 10th
0.5
Compute (FLOPs)
Compute (FLOPs)
Compute (FLOPS)
Figure 5: Our findings about the effect of (w, tb, v) on the scaling laws generalize to other settings. (Top row) Transformers trained with cross-entropy for autoregressive pixel generation on CIFAR-5M. (Bottom row) MLP trained with mean- nean-ssquree-error loss on synthetic data generated by a large and randomly initialized MLP.
Train, Monarch, and BTT can significantly outperform dense in other settings such controlling for memory, model size, or inference compute [5, 6, 30, 17, 19, 15], rather than training compute. For example, when training dataset size instead of training compute is the primary bottleneck, such as on conventional vision datasets like CIFAR-10 and ImageNet, structured matrices have shown considerable advantage over dense function of model size and inference compute [6, 19, 14). In those settings, Qiu et al. [19] observe the benefits of structure likely arise through enabling computationally efficient wider layers.
4.3 Our Findings Generalize to Other Settings
We now test if our findings derived from the GPT-2 experiments can generalize to other settings. We evaluate on the following two additional tasks where there is sufficient data to measure the compute-optimal scaling laws without repeating training data. We provide additional experiment details in Appendix D.
Autoregressive Pixel Modeling on CIFAR-5M. We train transformers to autoregressively predict the RGB pixel values of images in the CIFAR-5M dataset [16], downsampled to 8 x 8 x 3 resolution. Figure 5 (top row) shows qualitatively the same results as our GPT-2 experiments, where w and vi have the most significant impact on the scaling laws, while varying v yield only slight variations. In this particular case, having v = 0.75 (BTT with BTT--ank scaling as d!/4) is better than having == 0.5 (Monarch matrices). One notable trend in this setup is that most Einsums, regardless of w or v, outperform dense at small scales. We hypothesize this improved performance is due to Einsums having larger embedding dimensions than dense layers for a fixed parameter budget and can thus preserve more information about the input pixels at smaller model sizes.
synthetic regression dataset where the target is a scalar-valued function defined by a large randomly initialized MLP, similar to the student-teacher setting in [2]. In Figure (bottom row), observe qualitatively the same results as in the GPT-2 experiments.
MLP Regression on Synthetic Data. We train MLPs
Together, these additional results suggest there is some degree of universality associated with our findings on the effect of w. v, and v on the compute-optimal scaling laws of neural networks that use Einsums in place of dense matrices.
7

---

4.0
0.60
Compute Mu Muliipitrr
3.5
Compute Multiplier
2.0
0.40
2.5
2.0
2.5
1.0
1014
10t5
10
10
10
101>
1014
10t
10
Compute (FLOPs)
Compute (FLOPs)
Number of Experts
Dense
BTT-MoE
Std. MoE
Figure 6: BTT Mixture-of-Expere has significantly better compute-optimal scaling laws than dense GPT-2 and its standard MoE variant. (Left): Compute-optimal frontier with 8. (Middle): 8 experts compute multiplier of BTT-MoE and standard MoE relative to dense as a function of FLOPs required by the dense model to achieve the same loss. (Right): Increasing the number of experts improves computational savings. Mean and standard deviation of the compute multiplier over all compute observations for 8 and 16 experts.
Structured Mixture of Experts
In Section 4, identified that Einsums with and / perform the best, and w 0 or d < 1 lead to worse-than-dense performance. The most impactful parameter on the scaling laws w, which measures how many parameters Einsum has compared to the FLOPs for MVM. Einsums that learn one parameter per FLOP perform significantly better than those that learn less than one parameter per FLOP. Therefore, natural question arises: can design structures that learn more than one parameter per FLOP, which we might expect will have even better scaling laws? Doing so requires that not all parameters are used an MVM, which necessitates sparse Mixture-of- Experts (MoE) like architecture [22, 7, 12]. Furthermore, we would like the structure be full-rank, i.e. v = In the following section, we introduce such a structure and demonstrate significant improvement over dense layers and the standard MoE architecture for training GPT-2.
5.1 More Parameters than FLOPs via Mixture of Experts
One natural candidate for constructing such layer via an Einsum is to turn BTT with BTT-rank E, which involves a sum over the rank index I,...,E;
Yes Σ
BaespAangppar
(2)
into a k-sparse sum:
Yes Σ 9p
B
(3)
ary
output p-th expert
where E E is k-sparse vector so that only i out of E. terms need to be computed. We interpret the number active experts and E the total number of experts. We compute g via a softmax over the top-k entries of the logits produced by (dense) linear gating function Liieer(X) E E There no need to make this gating function structured because its cost is negligible. We choose 0 A OABX 0YB @YAB 1/2 so that each expert is full-rank. We follow the common practice of using = 2. The resulting BTT-MoB layer a BTT with BTT-rank (sum of two Monarch matrices) with input- dependent parameters. It is similarly straightforward to construct structured MoE from other structures by sparsifying the sum over p with a gate. We use the load-balancing loss encourage equal utilization of all experts [22, 7, 23].
In contrast to the standard MoE architecture used transformer language models, which uses sparse MoE for each entire feed-forward network (FFN) [22, 7, 12]:
Y
W/ReLU( X) Σ
(4)
ourput expert

---

BTT-MoE learns an MoE in every single linear layer of the model (except the language modeling bead) and treats them equally, including the projection matrices WQ, wk WV, w in the attention
blocks [24]. It learns more fine-graing geee-gaaidd routing decisions among the experts, with (BE(E - 1)) 6M
possible combinations of the experts in a transformer with M blocks, compared to Eg 1)) M for the standard MoE architecture.
5.2 Compute Efficiency Gains
In Figure 6, we show GPT-2 with BTT-MoE achieves better computc-optimal scaling laws compared to the dense model as well as the standard MoE, with k = 2 and E E: {8, 16). BTT-MoE consistently outperforms the standard MoE and the dense model. We quantify and compare the compute efficiency gains of BTT-MoE and standard MoE over dense models via the compute multiplier. A model with a compute multiplier of À means with C' training FLOPs it achieves the same loss as dense model with XC training FLOPs. In Figure 6, we show BTT-MoE is significantly more compute- -eficceeff than the standard MoE for both 8 and E 16. In particular, with E 16 experts, BTT-MoE achieves a compute multiplier of À 5.3 compared to À = 4.1 10.3 for standard MoE.
5.3 Effect of Structures
In Figure 7 we show that replacing BTT-MoE with a sparse (k 2) sum of low-rank matrices (low-rank-Mofi or dense matrices (dense- MoE) also yields nontrivial compute multiplier (~ 2x) over the dense model, but is significantly less effective than BTT-MoE or even the standard MoE.
3.5
Compute Multiider
Lan-RanoME
2.5
2.0
While the poor relative performance of low-rank- is expected, this result shows that in addition to 0 and v/b = 1, V 1 is desirable property for the base structure a structured MoE. Using a dense structure with 1 means the experts are not complementary to each other since each one is able to represent the entire space of dense matrices.
2.5
0.5
9.0
Figure 7: Mean and std dev of compute multipliers for structured MoE. BTT is bet- ter than low rank or dense.
6 Scaling Optimization for Einsums
As prior work [19] has shown, the optimal initialization scales and learning rate depend heavily on the structure of the linear layers and are critical for successfully training models with structured layers. Fortunately, the theory of the Maximal Update Parameterization (μP) [26-28] and its application
G,.03, 0.g.
0.2.0.3.333300 0)
Test Erar (Naive)
60
60
50
50
40
40
40
30
100
10
1
30
10
10-3
10
110
10
10a2
Test Error (eP)
60
60
60
50
50
50
40
40
40
30
10
10
10
10
10
10
10
10"
Learning Rate (n)
Learning Rate (n)
Learning Rate (n)
64
256
1024
4096
Figure 8: Einsums trained with μP achieve lower error and share an optimal base learning rafe. We plot test error of layered MI.P models on CIFAR-10, 10, where the hidden layers are Einsums. We vary the model widths in 64, 256, 1024 and 4096. The naive approach uses a global learning rate independent of width or structure and initializes the Einsums parameters with unit variance.

---

to various structured matrices in Qiu et al. [19] provides a template on how to reason about the appropriate initialization and learning rate scales for all Einsums.
In short, jaP states that for a dense matrix W E ggdux, the optimal initialization standard deviation scales as o = /innddnn dout)/ ) /dd and its learning rate as 7) (1//dn) if using Adam. Furthermore, Qiu al. [19] shows we can apply pP structured matrices as long as we can cast the MVM as a series of batched matrix multiplications (BMM). As shown in Appendix C, we can indeed cast any Einsum as a series of BMMs and show that dm = dxa. dout = A ABAAA. d
dx ABdAB and d dYB. As result, we compute the optimal scaling of A. BB , and r,8. In particular, for Adam we have 17 (dxx)) ) and 1/ @ AxBdAAAB )- Figure 8 shows using jPP leads to stable optimal learning rate and better performance compared to naively using a constant global learning rate and unit initialization variance. This property allows us transfer the learning rate between structures and model sizes, saving substantial compute for hyperparameter tuning. For JaP, the learning rate refers to that used by a dense model with width 64, which we transfer to the Einsum models of different widths and structures via the scaling rule identified earlier (see details in Appendix C).
Finally, we discuss in Appendix C an alternative way to reason about the optimal learning rates of Einsums via Riemannian SGD (RSGD) [3]. We analyze the effective learning rate prescribed by RSGD at initialization for asymptotically large Einsums and find it often agrees with the μP prescription derived above.
7 Conclusion
Going beyond prior works that study hand-crafted structured matrices on case-by-case basis, we introduce a continuous parameterization over the space of all structured matrices expressible as Einsums. Using this parameterization, measure and compare the compute-optimal scaling laws of a wide range of known and novel structures, with the following key takeaways:
Compute-optimal scaling laws of Einsums are primarily governed by the parameter-sharing exponent and the rank exponent th. Across tasks, find all full-rank Einsums without parameter sharing (i.e. full-rank BTTs) scale similar to dense, while the remaining vast majority of Einsums consistently underperform dense as increases or decreases.
Existing structured matrices do not significantly outperform dense in the compute-optimal setting. While low rank, Tensor- Train, Monarch, and BTT have shown advantages over dense in other settings, such as controlling for memory model size, they generally perform worse or similar to dense when controlling for training compute. However, there are also instances in the compute-optimal regime where full-rank structured representation with no parameter sharing can outperform dense layers. This advantage is most likely due to the ability to make wider structured layers for the same computational budget as narrower dense layers, which can particularly benefit smaller vision models, as we show on CIFAR-5M.
μP prescribes effective initialization and learning rate scaling for Einsums. Breaking an Einsum down to sequence of batched matrix multiplications. we extend prior work structure-aware initialization and learning rate based on jaP to arbitrary Einsums.
MoE over structured matrices more efficient than standard MoE over entire FFNs. By replacing every single dense linear layer with sparse sum of structured matrices like BTT, compared to standard MoE which operates over entire FFNs, we create more efficient MoE architecture, achieving over 5x savings in compute on language modeling relative to dense. Scaling and improving the proposed structured MoE architecture are exciting directions for future work.
Acknowledgements: We thank Alan Amin, Nate Gruver, and Hoang Phan for helpful discussions. This work is supported by NSF CAREER IIS-2145492, NSF CDS&F-MSS 2134216, NSF HDR- 211831 18300, BigHat Biosciences, Capital One, and an Amazon Research Award.
10

---

References
[1] Patrick Amestoy, Cleve Ashcraft, Olivier Boitcau, Alfredo Buttari, Jean-Yves l'Excellent, and Clément Weishecker. Improving multifrontal methods by means of block low-rank representa- 2015.
tions. SIAM Journal on Scientific Computing, 37(3):A151-A1447
[2] Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06700 2021.
[3] Silvere Bonnabel. Stochastic gradient descent on Riemannian manifolds. arXiv 1111.5280, 2011.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901. 2020.
[5] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations. International Conference on Machine Learning (ICML). 2019.
[6] Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive Structured Matrices for Efficient and Accurate Training. International Conference on Machine Learning (ICML), 2022.
[7] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23 1120):1-39, 2022.
[8] Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch Mixer: A Simple Sub- Sub-Quadratic GEMM- Based Architecture. Advances in Neural Information Processing Systems (NeurIPS), 2023.
[9] Roger Grosse and James Martens. A Kronecker-Factoreed Approximate Fisher Matrix for Convolution Layers. arXiv 1602.01407, 2016.
[10] Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14700 2020.
[11] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
[12] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401 20X.022401.04088, 2024.
[13] 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. arXiv preprint arXiv:2001.0836 2020.
[14] Changwoo Lee and Hun-Seok Kim. Differentiable learning of generalized structured matrices for efficient deep neural networks. arXiv preprint arXiv:2310. 18882, 2023.
[15] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In Workshop Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023), 2023.
[16] Preetum Nakkiran, Behnam Neyshabur, and Hanic Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. arXiv preprint arXiv:2010.08127,, 2020.
[17] I. V. Oscledets. Tensor-Trane Decomposition. SIAM Journal Scientific Computing, 201 1.

---

[18] Andres Potapezynski, Mare Finzi, Geoff Pleiss, and Andrew Gordon Wilson, CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra. Advances in Neural Information Processing Systems (NeurIPS), 2023.
[19] Shikai Qiu, Andres Potapczynski, Mare Finzi, Micah Goldblum, and Andrew Gordon Wil- son. Compute better spent: Replacing dense layers with structured matrices. International Conference Machine Learning (ICML), 2024.
[20] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. OpenAl, 2019.
[21] Yunus Saatçi. Scalable inference for structured Gaussian process models. PhD thesis, Citeseer, 2012.
[22] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gates mixture-of-experts layer. arXiv preprint arXiv: /00l::/0000000000 .065388, 2017.
[23] Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching Ilama2 performance with 0.1 m dollars. arXiv preprint ar 4rX1v:2404.07413, 2024.
[24] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[25] Andrew Gordon Wilson, Elad Gilboa, Arye Nehorai, and John P. Cunningham. Fast Ker- nel Learning for Multidimensional Pattern Extrapolation. Advances in Neural Information Processing Systems (NeurIPS), 2014.
[26] Greg Yang and Edward J. Hu. Feature Learning in Infinite- Width Neural Networks. International Conference Machine Learning (ICML), 2021.
[27] Greg Yang and Etai Littwin. Tensor Programs IVb: Adaptive Optimization in the Infinite-Wii -iitt Limit. International Conference on Learning Representations (ICLR), 2023.
[28] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[29] Greg Yang, James B. Simon, and Jeremy Bernstein. Spectral Condition for Feature Learning. Preprint arXiv:2310. 17813, 2023.
[30] Qibin Zhao, Guoxu Zhou, Shengli Xie, Liqing Zhang, and Andrzej Cichocki. Tensor ring decomposition. arXiv preprint arXiv:1606.05535, 2016.
12

---

Appendix Outline
Appendix A shows specific examples of Einsums and their MVM computations.
Appendix shows the derivation for the taxonomy variables w, & and
Appendix C discusses learning rate scaling details and connections between joP and Rie-
mannian SGD.
Appendix D provides extended experiment results and experiment details.
Appendix discusses the generalization of our parameterization to Einsums of more than two factors.
Appendix shows how further improve performance by structuring the Einsum axes to be compatible with the attention head structure in transformers.
Appendix comments on the hardware used in our experiments.
A Examples of Einsums
We explicitly show the expression for some common structures expressed as Einsums. Any index not appearing in the expression has size of I and is thus omitted.
Dense A one-factor Einsum would have been enough represent dense, though way to represent it within two-factor Einsum is as:
Yo B
(5)
70
which is an over-parameterizzaion due to the Hadamard product By
Low Rank
Ye = BepApaXa
(6)
ad
Kronecker Product
Yse- Σ
Aas
(7)
Tensor-Train
Yoe = Σ
Bacpd
(8)
aßp
Monarch
Yep Σ
BreaAnoooor
(9)
Block Tensor- Train
Yes Σ Bycsp
(10)
amp
B
Taxonomy Derivation
Here we present the derivation for the taxonomy variables 1, v). helpful to first reason about the amount of compute required for calculating Einsum, which will allow to exclude some uninteresting Einsums from consideration and thereby simplify our analysis. We only consider two-factor Einsums representing square matrix, but generalization to more factors and non-square matrices is straightforward.
For reference, we show the general expression again:
Yáeo
Σ ag7p
(II)
13

---

Computational Complexity. There are three possible ways to compute Equation (11), depend- ing on whether the first contraction happens between (), (A, X). or (B. The required FLOPs are ( O (22+0)-min(9..f.) mal computational path depends on the entries of 0, and the optimal FLOPs F required is given by F = 0 (f -ax {min(9, 0, ., ).m))(
and (dd (dd2+9,mmin(e m 86)) respectively. The opti-
Removing the Exchange Redundancy. The factors A and B play an equivalent role in the Einsum expression, so each distinct structure is represented by two vectors that correspond to relabelling A and B each other. To remove this redundancy, we can require that first summing with is more computationally efficient. Thus, we require
min(Oa, 0e) 2 min(0;; 0s);
(12)
which also simplifies the FLOPs to F -(nt+0,-min((,...)) To be more exact, we have F = pd? 1 du + d.)
Degenerate Einsums. Since the overall Einsum is linear operator on R any Einsum that requires more than (dd) FLOPs to compute degenerate the sense that unnecessary computations are performed. For example, one such Einsum could correspond to factorization UVT, U E RRdr V E Bexd where >> d. For convenience, we will also define an Einsum whose cost is equal to O (d²) as degenerate since it is no more efficient than dense matrix. Given the expression for FLOPs, we conclude that non-degenerate Einsum are those where
0p < min(0... 0.).
(13)
Intuitively, this requirement means that the rank dimension dg. i.e. the range of the index p connecting A and B, cannot be set too high, since after some point it becomes more efficient to simply use dense matrix in place of the Einsum. In particular, we must have d, > 1 and d, > (except when d d=1] 1).
Compute Intensity Exponent, v. One defining characteristic of structured matrices is that the FLOPS F for performing pattixxvector-muutippiaaoooo (MVM) is sub-quadratic in d. Equivalently, the compute intensity F/d, i.e. FLOPS for an MVM normalized by the dimension, is sublinear. In case, all non-degenerate Einsum have sublinear compute intensity. More precisely, have F/d = (d"), where / = 1 + 0p min,0,, 0.) takes value in [0, 1) assuming non-degeneracy. The closer v is to 1, the more Einsum resembles dense matrix and vice versa. v corresponds to low rank matrices with @(1) rank.
Rank Exponent, tb. As show in Appendix C.I. the Einsum can be computed via two batched matrix multiplications where A acts as matrix consisting of dgd, blocks of da x dgdød, matrices. Thus, (A) = min(d, dzd,dsd = dmin(1,0z 0,, -0,+00+00+02 as dad = d. Similarly, we have that acting as a matrix consisting of dgd' blocks of dzd,dp d, matrices. Therefore we have that rank(B)= = (0 +0, Since W the product of A, B, and some reshape operations (which are full-rank), rank (W) d is the minimum of rank A) and rank(B):
1 min (1, 0 0, + 0, + 0s 0
m1n 2 0,- 0a -2.).
Note that, technically, when we write rank(M), M E (,, W). we mean the maximum pos- learned. Otherwise the equalities will become upper- bounds. For low-rank we have 0 = (1, 0, 0, 0, 1, 0, 0,) and hence & = 0,- For dense we have 0 (0, 0, 1, 0, 1, 0) and hence = 1. For Kronecker we have O (1/2, 1 0/2, 1/22 and hence x = 1. For BTT we have th = min(1, 0 + As + 0p).
sible rank of M when its parameters
We use the more familiar term FLOPs as stand-in for MAC's (Multiply- Accumulate),, even though MAC technically& 2 FLOPs.
14

---

Parameter Sharing Exponent w. The number of learnable parameters N in a Einsum is simply the sum of the elements in and B, which works out to be d d_ds + dzd. Since the
1
2n + d.) only Einsums with da = dg = 1, or equivalently 0g = 0g = 0, have parameters matching FLOPs. In general, the number of parameters per FLOP N/F ((d-w) with = min + 0g. f 02) min(O,, 0.) 2 0.
FLOPs is F = dpd
The w = 0 subspace of particular interest because any Einsum outside of this subspace has an artificially limited expressivity per FLOP because it reuses each parameter multiple times. In this sense, the w = 0 subspace is the space of maximally expressive Einsums that maximizes expressivity per FLOP. We note that this subspace precisely corresponds to the Block Tensor Train (BTT) structure proposed in Qiu et al. [19]. provided we allow a minor generalization of the original BTT structure so that axis connects the first factor back to the last factor when there more than 2 factors, exactly analogous to the generalization of the Tensor-Train structure to the Tensor Ring [30] structure.
C Scaling Optimization for Einsums
C.1 Learning Rate Scaling
In Section 6, we claimed that the Adam learning rates should be scaled as 1/ = O dx
and
Einsum sequence of batched matrix multiplies (BMMs) and showing df = dxx, d = dy A AdyYAA ABdAB, d dXBdXABdAB and di = dyB, and (2) applying results from Qiu et al. [19] on learning rate scaling for structured matrices expressible in terms of BMMs. We now justify (1).
y/B O (xxxXABBAB ). Arriving these scaling rules requires (1) expressing
The general 2-factor Finsum
Ysco
BBreopAorbbpxaat ag7p
Σ
(14)
can be computed in two steps. step 1,
Zgrōop
(15)
which is a BMM with 38 acting batch dimensions, the input dimension, and dpp the output
dimensions. So dA dxA, dA = dr a dY AB In step 2,
Yoed
BBy Zarbop
(16)
Bop
which is BMM with 8ō acting batch dimensions, 8yp input dimensions, and e the output dimension. So dB dxBdXABdAR and de dyn, wanted.
In our experiments, we apply the above scaling rule to transfer from a learning rate 7 used by a dense matrix with width do to learning rates for A, B as
do
do 2dxBdxABdAB
2dXA T
1 B
(17)
where the additional factor of two is to account for both A and B contributing updates the output of each layer, following Qiu et al. [19].
C.2 Connections between μP and Riemannian SGD
Riemannian SGD (RSGD) is optimization technique that allows us perform the equivalent of SGD on a Riemannian manifold consisting of points {q} with a metric 9q [3]. The updates of RSGD is almost identical to SGD: q (f+1) = (t Ja Iqi (")( except that the gradient is multiplied by the inverse of the metric 9q. In our case, we want to mimic training the Einsum parameters A and B as we were training W with SGD directly, even though we will never represent W explicitly. Thus, we identify (flattened) (A, B) with and specify its metric the pull-back metric of the
15

---

Euclidean metric ow = I on W, which is given by 9q - ((q)) where J(q) updates to is therefore
ow
0q The RSGD
q (++1) q 17 (J(I
")(((
Exactly computing the inverse metric would require O(P)) time, where Pis the number of parameters in A, B. It is, therefore, too expensive to run RSGD during training. However, we can efficiently approximate the inverse metric at initialization in way that exact for asymptotically large matrices. We compute JJJ through three blocks J A J and J and note that for fixed A as we increase both din and dout. we have that JAJ A where A and øв denote the initialization scales of A and B. Therefore have
Rs Mgd.oRL. J s3 dadso and JAJB 2
(JTJ)-1
o
0
I/(dado0))
We can now identify 0 ol/ as AA and (1/(dd as 7n. we further use the initialization scales prescribed by μP as given in Section 6, we have
VESGD + (
and JRSGD = O (11 dado min min(da, dododo)
&
dgd= min min(d., din(d..dgd,dd))
It is now interesting to compare these scalings the prescriptions uP SGD, For SGD, μP proposes to scale the learning rate as Einsums, we need not only to replace din, dout with d d or dB and dB but also scale down the learning rate as (1/ds) for A and ( /ds) for B. The final μP learning rates have two possibilities depending on whether we consider A or B as the first layer. we consider A as the first layer, then
/din) [29] for dense matrices. Unlike in Adam, to extend jP to
+
(1s dg da
[1 dadydp
1
and UPp
B- (
(ddssdd do dzdødp)
de
Evidently, the two approaches don't always agree. However, they are identical when da = de. dgd,dp I de, and dzdødp I da. In fact, many structures satisfy this condition. If we as- sume the structure is symmetric, meaning its transpose can be represented with the same 0, i.e.
da = de, d3 = ds, and d, = dg, then the remaining conditions simplify to only d, I L any symmetric Einsum with 0., 2 1/2 and 0, 20, -1 satisfy these conditions. Thus, for these structures, μPP SGD can be viewed as an approximation to RSGD that is valid at initialization the -iitt limit. While we cannot establish direct connection between RSGD and μP Adam, which is what we use in our experiments and broadly in large-scale training. μP Adam is similar to uP SGD in that it maximizes feature learning per layer [27]. The connection between jP SGD and RSGD therefore indirectly provides an alternative justification for jP Adam.
d Therefore,
Experiments
D.1 GPT-2 with Original Vocabulary and Longer Context
In Figure 9, we show our findings in Section 4 translate to the more standard GPT-2 evaluation with a longer sequence length of 512 and its original vocabulary of 50,257 tokens. We train models with L= 12 layers up to the GPT-2 Small [20] size by increasing width d. We use Adam with a base learning rate of 0.002 for L 3, d = 256 dense model, which is scaled to other models via jP. Since the language modeling head contains a significant fraction the parameters for models of this scale, we replace all layers, including the head, with Finsums.
Qualitatively, Figure differs from Figure 4 in two ways: 1) the scaling laws are less power law like and show some curvature on log-log scale, and 2) BTT with v > 0 seems to perform better than O. We believe 1) due the increased context length and vocabulary size, making the loss no longer follow a clean power law at the small scales we tested [13, 10]. This was an important motivation for performing experiments with smaller vocabulary size and context length in Section Similarly, we believe the increased vocabulary size and context length contributed to 2), larger implies at small scales a higher fraction of compute are in the transformer blocks rather than the language modeling head, which likely improves performance. By contrast, in our setup in Section 4, the model dimension d dominates the vocabulary size and context length, leading to less significant mizee effects.
16

---

Full Space
0.5
0
100
w 00
0.8
5.0
5,0
5.0
Loss
Locs
4.0
3.0
Densa
3,0
Denso
3.0 Dense
103h
20
1037
Compute (FLOPs)
Compute (FLOP))
Compute (FLOPs)
Figure 9: The taxonomy parameters (w, +)) explain differences in the Einsum scaling laws for 12-layer GPT-2 models with standard vocabulary (50,257 tokens) and context length of 512. Small w (no parameter sharing) and large & (full-rank) are necessary for structure to perform well, while variation in has a much smaller impact on performance.
D.2
Autoregressive Pixel Modeling CIFAR-5M
We train and layer transformers with Adam using a base learning rate of 3e-3 for a width 64 dense model. The width ranges from 32 to 512 for dense and 32 1024 for Einsums. All models trained for epochs with a batch size of 64. We subtract an estimated irreducible loss of 3.25 before reporting the loss in Figure 5 (top row).
D.3 MLP Regression Synthetic Data.
We train 3-layer MLP models with width d [64, 4096] for maximum of 10th steps and batch size of 4096 on effectively infinite synthetic dataset. The synthetic dataset is obtained by querying a valler target function on IR with inputs drawn from a Gaussian distribution. The target function is randomly initialized target MLP with 6 layers and hidden dimension of 1024. We minimize the Mean-Squared-Erre (MSE) loss. We train with Adam using a base learning rate of le-3 for a width 64 dense model. We report the raw MSE loss in Figure 5 (bottom row).
E Generalization to More than Two Factors
The generalization to more factors is easy to understand we consider the set of edges that define
the Einsum & {S [X, A, B, Y} |S| 2 2 and (X, Y) & S). For example, for three
factors, the edges that connect to X are [X, A} + i [X, B) + i2, [X, CJ + is, (X, A, B) + 4,, (X, B, C}
is A, C + is and {x, A. B, CJ + 17. The edges that connect Y (Y, AJ + JI, Y + j2, CY + ja, [Y, A, BJ + j4, (Y, B, C) + js (Y,A) A, C} + je
and (Y. A. B, CJ + jr. The edges between the factors are (}, + ri, [A, C} r2 and
{A, B. C} + r3. The expression for a three factor would be
Yj y= Σ
Angisierrjuussshhrer Big.ttiiriiiiiiiiiini Cighsoorrhassseerrrr Xi.
For more factors we follow the combinatorial procedure of listing the sets and adding an index based on each edge in E
F Exploiting the Attention Structure
In self-attention, given an embedding vector Rd compute q Qºx, v WVx and k W x. Af- ter computing each of the q, q, v E Rd vectors, they are reshaped produce one smaller vector per attention head: + Qhjs k; khj. + Uhy where h is an axis of size H, the number of attention heads, and is an axis of size d/H. When replacing Wº. WK W with BTTs, it is therefore more natural to have the BTT output axes eω coincide with hj so the MVM aware of the attention head structure rather than potentially mixing different heads. Similarly, when replacing O with BTT, it is most natural to have the BTT input axes
Scaling Laws on GPT-2
0.7
0.6
0.5
0.3
BTT
BTT.Atth
10
10 Compute (FLOPS)
10
10
Figure 10:
Exploiting attention head structure improves compute- efficiency by average of 17 %.
17

---

coincide with hjj. In Figure 10, we show doing so slightly improves compute efficiency by an average of 17% over naively replacing all attention and FFN matrices with BTT, which cor- responds to 0 = (1/2,0,1/2, 1/2. 1) Section 4.
G Hardware specifications
Our experiments in Section 4 are on A100 and H100 GPUs. The CIFAR-10 experiments in Section were run on RTX2080 Ti and RTX Titan GPUs.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper presents a continuously parameterized space of structured matrices, unifying many previously proposed structures matrices. It develops taxonomy for different structures and it explores the best performance subset of this space, in the compute optimal setting.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Though we show that with Mixture of Experts, BTT does perform better with the MoE counterparts of dense, vanilla full rank BTT performs similar with dense matrices, as mentioned in section 4.2 and discussion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
19

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The claims in the paper are proved in the Appendix, in particular, Appendix B and C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
The code is shared and the experiment setup is detailed in Section 4.1.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
20

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification:
The code is shared and the data is public dataset.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
httpg//ddee///eeeueeeeee soiisssiicinnyy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. cccpublic/guu selicc/ggidss/. /ttgggddde//oeeummiiieeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All essential experiment settings are detailed in section 4.1.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For scaling laws, runs for models of different sizes are all presented. There are multiple runs for different seeding because of expensive computes. For the crucial result in Figure 5, which demonstrating the well performance of MoE for BTT, error bars are plotted.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
21

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: It's mentioned in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https https://mmnt //tpprip//////ps//// cc cc/public/ htt.h/////hhisShieeleeeeeee
Answer: [Yes]
Justification: The paper doesn't 't involve living subjects and utilize public data. conforms with all the code of ethics listed.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
In discussion, the paper discusses the potentials to create more effective
models.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
22

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All resources used are open sources, with GNU General Public License v3.0. or other open licenses.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
23

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithhaeee thcode. www.coo//a/oo//o datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
There are clear comments in the code and clear readme document.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "I3F6CXkTDT", "reviewer": "Reviewer_Mhte", "summary": "In this paper, the authors explore the computational efficiency of various structured layers in language modeling tasks. Specifically, they propose a general parametrization of linear operators and conduct an empirical study on the conditions for scalable decomposition based on three key characteristics: rank, compute intensity, and parameters per flops. Moreover, the authors integrate Mixture-of-Experts (MoE) with structured layers, conducting experiments and comparing the results to those obtained using the standard MoE.", "questions": "1) Did you try other than BTT structured layers in the MoE experiments? Why do you only use BTT?\n2) It seems to me that rounding, e.g., $d_{in}^\\theta$ to the nearest integer interferes with maintaining the shape $d_{in}$. Is it actually the case and if so, how do you deal with it?\n3) Do you expect that all the observations about $\\omega, \\psi \\mu$ will remain the same with bigger dictionary?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "I believe that the topic of this paper is quite important and timely. With the increasing size of models, it is crucial to find effective ways to efficiently compress them. Among the various matrix and tensor factorization approaches available for compression, it becomes essential to unify them and determine which aspects make them most effective.", "weaknesses": "1) I understand that to check multiple different configurations of layers while having limited computational resources you need to somehow restrict your experiments. But what is lacking in my opinion, is verifying your findings with at least several runs in the proper setting.\n2) As far as I understand, classic MoE is applied only to FFN, while you also apply it to all the layers, including Q, K, V. This may affect comparison with MoE. See, for example, in Figure 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thoughtful feedback. Indeed, unifying existing structured approaches and performing extensive and well-controlled comparisons between them is an important contribution of this work. We now provide additional results and clarifications to your questions.\n\n**On experiments with larger vocabulary and longer context**.\nThank you for raising this point as these new results strengthen the paper significantly.\nIn Figure 1 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf), we now train models up to the size of original GPT-2 Small [1], adopting the vocabulary of 50,257 tokens and using a context length of 512. The results agree with our previous findings in Section 4 with a reduced vocabulary and a shorter context, showing that our results indeed hold in more realistic settings. Moreover, in Figure 2 and Figure 3, we evaluate on two additional datasets including image generation with Vision Transformers and synthetic regression with MLPs, further demonstrating the generality of our findings. We will include these results in the updated paper.\n\n**On additional MoE comparisons**.\nInitially we followed the standard approach of only comparing against the MoE that is only applied to the FFN layers. However, to address your concern on the lack of MoE modules in attention layers of the baseline, we now compare with two additional alternatives in Figure 4 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf): Dense-MoE and Low-Rank-MoE, which similarly replace all linear layers including those in the attention blocks with an MoE over dense or low-rank matrices. The results show that BTT still has a unique advantage over other structures for constructing structured MoEs, even when applying MoE to all layers. We will include these results in the updated paper.\n\n**On rounding to the nearest integer**. Indeed, rounding to the nearest integers as described in Section 2 can produce a matrix $W$ whose input and output dimensions slightly deviate from the original desired shape. We take the simplest approach to address this issue by padding or truncating the input and output vectors. At scale, the number $\\delta N$ of padded or truncated elements becomes vanishingly small relative to the original dimension $N$, as $\\delta N / N$ scales as $O(N^{-c})$ where $c$ is the smallest non-zero elements in $\\theta$. We will update the paper to clarify this consideration. \n\nThank you again for your constructive feedback and support. We made a significant effort to address your questions which has improved our work substantially; we would appreciate it if you would consider raising your score in light of our strong response and the significance of our work. We believe this paper will help provide a foundation for a nascent and immensely impactful new research area.\n\n[1] Radford et al. 2019. Language models are unsupervised multitask learners", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ktOsL4y0HP", "reviewer": "Reviewer_bZBq", "summary": "This paper proposes a general framework to cover different linear layers by continuous parametrization. The authors conduct extensive experiments to demonstrate several optimal design principles for linear layers, and further propose a novel sparse MoE architecture that improves upon existing works.", "questions": "- I am first puzzled by the novelty of this work. While the authors have conducted extensive experiments on different architectures, what are the key insights and contributions of this submission? It seems that the idea of combining BTT and (sparse) MoE is novel, but it seems straightforward and the authors may provide some more insights for it. \n- While the authors mentioned in section 2 that generalization to more than 2 factors is straightforward, it seems not really the case as the selection of different indexes ($\\alpha$, $\\beta$, … in (1)) requires manual design to ensure expressive power. As such, the authors may need to provide some examples on using more factors and see how existing methods may be covered by such more expressive frameworks. \n- Throughout this paper, I only see experiments on GPT-2, which may not be sufficient to derive general conclusions regarding the scaling laws of different Einsums. The authors may need to consider different model architectures (e.g., BERT, ViT) to better support any conclusions here. \n- Moreover, current experiments are only conducted on one data set, which is not sufficient to support such general conclusions in this submission either. The authors should also consider experiments on different data sets to derive general conclusions. \n- Also, the comparison of different MoE architectures in Figure 6 may not be enough to support the superiority of BTT-MoE. While Figure 3 and 4 indicate that BTT may be optimal for dense architectures, it may not be directly generalized to MoE architectures. The authors may need to include some other MoE architectures, possibly low-rank-MoE as low-rank performs quite close to BTT from Figure 3.\n\n## Post rebuttal \n\nAfter checking the rebuttal, I still doubt the claim that solely two parameters $\\omega$ (parameter sharing) and $\\Psi$ (closeness to full-rank) reliably led to better scaling laws, which motivates the authors to design BTT and combine it with MoE. Despite such weakness, I think the authors have made their contribution clear with sufficient support (additional architectures, datasets and other MoE structures). I have increased my score towards acceptance.", "limitations": "This paper does not have direct negativie societal impact from my perspective.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper is clearly written and easy to understand", "weaknesses": "- The novelty of this work is not clear enough from my perspective\n- Empirical results may need further improvements to better support the proposed method", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your feedback. Inspired by your comments, we have run additional experiments and provide clarifications below. We hope that you will consider these new results and clarifications in your final assessment.\n\n**On the key contributions of our work:**\nWe highlight that this work provides at least 4 novel and impactful contributions. \n\n1. We provide a unifying framework to conveniently parameterize a large and continuous space of hardware-efficient structured matrices via Einsums. We show this space contains many popular structured matrices explored in prior works, such as low-rank, Kronecker, Tensor-Train, Monarch, and Block Tensor-Train, while most structures within this space are novel. We further develop an informative taxonomy of the space of Einsums based on key properties relevant to machine learning, including the extent of parameter-sharing, matrix rank, and computational complexity.\n\n2. We perform the first extensive comparison of compute-optimal scaling laws for structured matrices in language modeling. State-of-the-art large language models (LLMs), including the recent Llama 3.1 405B, are purposely trained to be compute-optimal [1, 2], whereas prior works on training with structured matrices do not compare performance under compute optimality (e.g. by training for too many epochs) [3,4]. Our results therefore provide the more appropriate comparison for evaluating structured matrices in realistic LLM training. Indeed, our results reveal that structures such as Monarch that significantly outperform dense in other contexts at best match dense performance in compute-optimal scaling laws.\n\n3. We show that differences in the compute-optimal scaling laws across a wide range of structures are mostly governed by a small number of variables defined in our taxonomy. Small $\\omega$ (less parameter sharing) and large $\\psi$ (closer to full-rank) reliably led to better scaling laws, while $\\nu$ (how dense-like a structure is) can be varied while leaving the scaling laws almost unchanged. These insights will make future search for performant structured matrices significantly more efficient than random. \n\n4. Guided by the insight that full-rank ($\\psi=1$) structures that maximize parameters ($\\omega=0$) per unit of compute performs the best (as argued in our taxonomy and experiments), we propose BTT-MoE, a novel Mixture-of-Experts (MoE) architecture obtained by sparsifying computation in the BTT structure, proving to be 440% and 28% more compute-efficient than dense and standard MoE for training GPT-2 respectively.\n\n**On experiments with additional architectures and datasets.**\nFocusing on compute-optimal scaling laws limits the range of datasets that we can study and therefore we focused on language, as it is standard in scaling laws literature. However, to address your valid concern we have now incorporated Figure 2 and Figure 3 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf) where we run new experiments to verify that our findings regarding the relative performance of different Einsums continue to hold in other settings:  1) Vision transformers trained with cross-entropy loss for autoregressive image generation on CIFAR-5M; and 2) MLP trained with Mean-Squared-Error loss on synthetic data generated by a large and randomly initialized MLP. The results demonstrate that our findings in the GPT-2 experiments indeed generalize to other architectures and datasets. We will include these results in the updated paper.\n\n**On comparing against MoE with other structures.**\nFollowing your suggestion, we compare against two additional MoE architectures in Figure 4 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf): Dense-MoE and Low-Rank-MoE, which similarly replace all linear layers including those in the attention blocks with an MoE over dense or low-rank matrices. The results show that indeed BTT has a unique advantage over other structures for constructing structured MoEs. We will include these results in the updated paper.\n\n**On generalization to more than 2 factors.**\nThe set of indices in Equation 1 does not require manual design to ensure expressive power. Instead, we obtain Equation 1 by simply allowing all possible indices to exist. Following the discussion on Line 92, this set of indices can be directly read off from a graphical representation of the Einsum with each index corresponding to a hyperedge among the input, output, and the weight factors. As a result, generalization to more than $N > 2$ factors follows by constructing a graphical representation of the $N$ factor Einsum, assigning an index to each hyperedge, and writing down the resulting Einsum expression. It is easier to visualize this generalization and therefore we now demonstrate this process for $N=3$ in Figure 4 (c) in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf) and show how the resulting expression covers the general 3-factor case as well as BTT. We will include this new figure for $N=3$ as it visually makes the generalization to more factors apparent\n\nThank you again for your review. We made a significant effort to address your questions, which has substantially improved our work. We would appreciate it if you would consider raising your score in light of our response. \n\n[1] Dubey et al. 2024. The Llama 3 Herd of Models.\n\n[2] Hoffmann et al. 2022. Training compute-optimal large language models.\n\n[3] Dao et al. 2022. Monarch: Expressive structured matrices for efficient and accurate training\n\n[4] Qiu et al. 2024. Compute Better Spent: Replacing Dense Layers with Structured Matrices", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "tuCQe600t2", "reviewer": "Reviewer_XAQT", "summary": "The paper generalizes multiple existing structured matrices by means of Einsum. The scaling law of the structured matrices with different rank, compute intensity, and parameters-per-FLOPs is analyzed on GPT-2. Since the high-rank, non-parameter-sharing einsum operations obtain the best results, the paper proposes a sparse mixture of structured linear experts which also has high-rank, non-parameter-sharing taxonomy parameters with further generalization. Additionally, the learning rate scaling rule for the Einsum-induced structured matrices via maximal update parameterization is introduced.", "questions": "- Is the scaling law of Einsum generalizable to the larger models (e.g., $\\ge$7B parameters)?\n- Are there any insights or possible reason that the authors think why the high-rank, no parameter sharing leads to the best results whereas the compute intensity does not affect the accuracy? \n- Could the taxonomy parameters have any real numbers other than rational numbers?", "limitations": "- The framework and the analysis is confined to a macroscopic point of view--the test was conducted upon fixing the einsum configuration and the FLOPs across the layers, whereas the optimal configuration and FLOPs might vary from layer to layer.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "- The paper provides a novel point of view for understanding the structured matrices via Einsum operations.\n- The paper is clearly written with sufficient details.\n- The comprehensive analysis of the interesting aspects of the structured matrices leaves valuable insights for future research.", "weaknesses": "- Although the text is well written, it was hard to keep track of a dozen of alphabets indicating dimensions, taxonomy variables, etc. Reminding the role of each alphabet occasionally would improve the readability.\n- The continuity of the taxonomy space is questionable. It seems like the taxonomy space is discrete because rank, FLOPs, and number of parameters are all integers. Even if they are normalized, the parameters would still reside in a finite sized space of rational numbers from 0 to 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful and supportive review. We agree that the large number of variables presents readability challenges and will update the paper to provide reminders about the meanings of the indices. We now provide several clarifications and new results inspired by your comments.\n\n**On continuity of the parameterization**.\nWhile quantities including FLOPs, parameter count, and rank are discrete, the continuous parameterization of the space of Einsums with non-negative real-valued coordinates $\\theta$ is valid. First, any non-negative $\\theta$ satisfying the condition $\\theta_{{X}{A}} + \\theta_{{X}{B}} + \\theta_{{X}{A} {B}} = \\theta_{{Y}{A}} + \\theta_{{Y}{B}} + \\theta_{{Y}{A} {B}} = 1$ produces a valid structure once the resulting weight factors are rounded to have integer sizes. Including irrational entries in $\\theta$ is completely allowed. Now it suffices to show that any two such coordinates, say $\\theta$ and $\\theta + \\epsilon$ indeed represent a distinct set of structures as the models scale, for any $\\epsilon \\neq 0$. This is because for large enough dimensions $d$, any small difference of $\\epsilon$ in the coordinates will lead to a larger than 1 difference in the size of the weight factors along some axis, which will persist even after rounding to nearest integers. In other words, we have shown that the space of allowed $\\theta$ is real-valued and is a bijection with the space of unique Einsum structures.\n\n**On experiments with larger models**.\nIn Figure 1 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf), inspired by your comments, we show results from new experiments on larger GPT-2 models. We train 12 layered models up to the size of original GPT-2 Small [5], adopting the original vocabulary with 50,257 tokens and using a context length of 512. The results agree with our findings in Section 4 which used a reduced vocabulary of 96 tokens and a shorter context length of 128 to save cost. Due to computational constraints, we cannot experiment with 7B parameter models. \n\nIn addition, we perform experiments with other architectures and datasets in Figure 2 and 3 in the [rebuttal's pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf), including Vision Transformers on image generation and MLP on synthetic regression. The results confirm that our findings generalize to much broader settings.\n\n**On a studying per layer configurations.**\nYou raise a really interesting and exciting direction for future work. Indeed, a proper exploration of different structures per layer would require a significant amount of compute and, most importantly, to construct a method to explore the combinatorial search space. We believe this is an important next step and these are the type of questions that we are excited to see that our work is raising.\n\nWe value your support and thoughtful questions. We put a significant effort into our response and would appreciate it if you could consider increasing your score in light of our replies.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank all the reviewers for their feedback and questions. We provide a general response here and individual replies in separate posts below. Inspired by comments from reviewers, we include multiple new experiments encompassing new datasets, new architectures, and alternative MoE structures that significantly strengthen our findings and demonstrate their applicability in much broader settings. 

We appreciate the reviewer's recognition of this work's importance and timeliness. Scaling foundation models is primarily bottlenecked by compute in the dense linear layers, where structured matrices offer promising efficiency gains. Therefore, a comprehensive analysis of structured matrices' potential to enhance dense model scaling laws, along with identifying general properties of structures that correlate with their scaling behavior, is of significant value to the field.

**We highlight several particularly novel and impactful contributions in our work:**

1. We provide a unifying framework to conveniently parameterize a large and continuous space of hardware-efficient structured matrices via Einsums. We show this space contains many popular structured matrices explored in prior works, such as low-rank, Kronecker, Tensor-Train, Monarch, and Block Tensor-Train, while most structures within this space are novel. We further develop an informative taxonomy of the space of Einsums based on key properties relevant to machine learning, including the extent of parameter-sharing, matrix rank, and computational complexity. 

2. We perform the first extensive comparison on the compute-optimal scaling laws of structured matrices for language modeling. State-of-the-art large language models (LLMs), including the recent Llama 3.1 405B, are purposely trained to be compute-optimal [1, 2], whereas prior works on training with structured matrices do not compare performance under compute optimality (e.g. by training for too many epochs) [3,4]. Our results therefore provide the more appropriate comparison for evaluating structured matrices in realistic LLM training. Indeed, our results reveal that structures such as Monarch, which significantly outperform dense in other contexts, at best match dense performance under compute-optimal scaling laws.

3. We show that differences in the compute-optimal scaling laws across a wide range of structures are mostly governed by a small number of variables defined in our taxonomy. Small $\omega$ (less parameter sharing) and large $\psi$ (closer to full-rank) reliably led to better scaling laws, while $\nu$ (how dense-like a structure is) can be varied while leaving the scaling laws almost unchanged. These insights will make future search for performant structured matrices significantly more efficient, providing a guiding foundation for this important emerging research area.

4. Guided by the insight that full-rank structures that maximize parameters per unit of compute perform the best, we propose BTT-MoE, a novel Mixture-of-Experts (MoE) architecture obtained by sparsifying computation in the BTT structure, proving to be 440% and 28% more compute-efficient than dense and standard MoE for training GPT-2, respectively.

**We now summarize the new experiments we run inspired by the reviewers’ feedback.**  We present results and figures in the [attached pdf](https://openreview.net/attachment?id=cH4w74hFGe&name=pdf).
1. We verify that our findings regarding the relative performance of different structures continue to hold in the following additional setups: 
    - A more standard GPT-2 training setup. We train models up to the size of original GPT-2 Small [5], adopting the original vocabulary with 50,257 tokens and using a context length of 512. The results agree with our findings with in Section 4 which used a reduced vocabulary of 96 tokens and a shorter context length of 128 to save cost.

    - Vision transformers trained with cross-entropy loss for autoregressive image generation on CIFAR-5M.

    - MLP trained with Mean-Squared-Error loss on synthetic data generated by a large and randomly initialized MLP.
We chose CIFAR-5M image generation and the synthetic regression as additional tasks because they contain enough training examples required to study compute-optimal scaling laws, unlike other commonly used datasets such as CIFAR-10 or ImageNet classification.

2. We show that BTT-MoE outperforms two additional MoE architectures: Dense-MoE and Low-Rank-MoE, which similarly replace all linear layers including those in the attention blocks with an MoE over dense or low-rank matrices. The results show BTT has a unique advantage over other structures when used in the proposed structured MoE architecture.

We hope the reviewers can consider these results and clarifications, and the broader context of this work and its significance, in their final assessment.

[1] Dubey et al. 2024. The Llama 3 Herd of Models.

[2] Hoffmann et al. 2022. Training compute-optimal large language models.

[3] Dao et al. 2022. Monarch: Expressive structured matrices for efficient and accurate training

[4] Qiu et al. 2024. Compute Better Spent: Replacing Dense Layers with Structured Matrices

[5] Radford et al. 2019. Language models are unsupervised multitask learners

Meta-review: Authors propose Einstein summation as a framework to generate structured matrices instead of dense layers in transformers architectures and study the loss/ compute scaling laws. Authors identify three parameters of the structured matrix: rank, compute intensity, and parameters per flops that control the scaling laws of the resulting architecture. Authors identify a set of structured matrices  that achieve equal  on par loss as dense layers as a function of training compute. Authors combine these structured matrices with sparse mixture of experts to further improve the scaling law. 

Reviewers asked for additional experiments that considered number of factors N>2 and with larger vocabulary and context, authors provided additional results in the rebuttal that satisifed the reviewers with the exception of Reviewer bZBq who still have doubts on the the claim that solely two parameters  parameter sharing and closeness to full-rank reliably led to better scaling laws. 

We encourage the authors to revise the paper and include the additional results and to provide more analysis on the effects of the parameters that control the einstein summation.
Decision: Accept (poster)