Title: Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices
Venue: NeurIPS 2024
Accepted: yes
Keywords: Structured Matrices

Scaling Laws

Linear Layers

Optimization

Deep Learning

Einsum
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

Abstract

Dense linear layers are the dominant computational bottleneck in 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 to investigate whether these structures can surpass dense layers in terms of computeoptimal 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, Block Tensor-Train (BTT), and Monarch, along with many novel structures. To analyze the framework, we develop a taxonomy of all such operators based on their computational and algebraic properties and show that differences in the compute-optimal scaling laws are mostly governed by a small number of variables that we introduce. Namely, a small ω (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, we propose BTT-MoE, a novel Mixture-of-Experts (MoE) 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 of interleaved linear layers and simple non-linearities. In large foundation models such as 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

Y

X

B

A

1/2

1

1

Kronecker

Y

X

B

A

1/2

1/2

1/2

1/2

Monarch

Y

X

B

A

1/2

1/2

1/2

1/2

Y

X

B

A

10 14

10 15

10 16

10 17

Compute (FLOPs)

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Loss

Scaling Laws on GPT-2

Kronecker

Low-Rank

Monarch

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 infeasible 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-optimal setting, which characterizes performance as a 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 = L ∞ + bC -a as a function of training compute C, where L ∞ is the minimal achievable loss [13, 11, 10]. Quantifying the compute-optimal scaling laws of various structures is essential for understanding their practical value for training large-scale neural networks.

In addition to 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 of the space of all possible structures whose matrix-vector-multiplication (MVM) can be expressed as an Einstein summation (Einsum). 2 This space contains many known structures such as low-rank, Tensor-Train [17], Kronecker product [21, 25, 9], Monarch [6] and Block TensorTrain [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 matrix-vector-multiplication can be implemented via an Einstein summation (Einsum). This parameterization allows us to search a wide range of hardware-efficient structured linear layers for neural network architectures beyond a handful of well-known cases identified in prior work [6, 8, 19].

We develop a taxonomy of the space of Einsums based on its computational and algebraic properties. We identify three key scalar quantities that characterize this space ( ω, ψ, ν ) . (1) ω ≥ 0 , which reflects the extent of parameter sharing in a matrix. (2) ψ ∈ [0 , 1] , which characterizes to the rank of the structure ( ψ = 1 meaning full-rank). (3) ν ∈ [0 , 1) , which

relates to compute per dimension in an MVM, where the upper-bound ν = 1 is achieved by dense matrices. Intuitively, ν measures how much a structure resembles dense.

We investigate the scaling laws of different Einsums on language modeling, autoregressive image generation, and a synthetic regression task. We find that the best-performing structures are the ones that do not share parameters ( ω = 0 ), are full-rank ( ψ = 1 ), while ν 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] to outperform dense matrices in non-compute-optimal 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 µ P [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 a 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-efficient 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 ∈ R d out × d in whose matrix-vector-multiply y = Wx can be expressed as an Einsum over the tensors x , A , 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

where the vectors x and y are written as tensors with multiple indexes to allow them to interact differently with each other, and the factors A , B . Each index x ∈ { α, β, γ, δ, ϵ, ϕ, ρ } ranges from 1 to d x . Given this general expression, we obtain different structures via different factorizations of d in into d α d β d γ , and d out into d δ d ϵ d ϕ , and separately a choice of d ρ . For example, for a low-rank matrix, we have d α = d in , d β = d γ = 1 , d ϵ = d out , d δ = d ϕ = 1 and d ρ = r . For a Kronecker product, we have d α = d β = √ d in , d δ = d ϵ = √ d out and d γ = d ϕ = d ρ = 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 a diagram shown in Figure 1 (left), where each index α, β, γ, . . . corresponds to a (hyper)edge among the input, output, and weight factors: α ↔ { X , A } , β ↔ { X , B } , γ ↔ { X , A , B } , δ ↔ { Y , A } , ϵ ↔ { Y , B } , ϕ ↔ { Y , A , B } and ρ ↔ { A , B } . This set of edges can be written succinctly as E = { S ⊆ { X , A , B , Y } : | S | ≥ 2 and { X , Y } ̸⊆ S } . We exclude 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 a particular Einsum is fully specified by the the vector ( d XA , d XB , d XAB , d YA , d YB , d YAB , d AB ) , which specifies the range of the indices α, β, γ, δ, ϵ, ϕ, ρ . When the range of an index is of size 1, the corresponding edge effectively disappears from the diagram and the expression simplfies.

As we will build models of varying sizes, it is more natural to think about how these entries scale with d in and d out . We therefore assign a real-valued vector θ ∈ [0 , 1] 7 to each structure indicating that d i = d θ i in for i ∈ { XA , XB , XAB } , d j = d θ j out for j ∈ { YA , YB , YAB } , and d AB = min( d in , d out ) θ AB , with the restriction that θ XA + θ XB + θ XAB = θ YA + θ YB + θ YAB = 1 . For example, a low-rank matrix whose rank scales as the dimension it operates on to the 1 / 2 -th power is represented as θ = (1 , 0 , 0 , 0 , 1 , 0 , 1 / 2) , and a Kronecker product of two factors of equal sizes is represented as θ = (1 / 2 , 1 / 2 , 0 , 1 / 2 , 1 / 2 , 0 , 0) . We round all d i to its nearest integer when

BTT subspace

Full-rank BTT subspace

Y

X

B

A

Tensor-Train

3/4

1/3

2/3

1/4

Y

X

B

A

1/2

1/2

1/2

1/2

Monarch

Y

X

B

A

Dense

Y

X

B

A

Dense

1

1

Y

X

B

A

Low-rank

1/4

1

1

Y

X

B

A

Kronecker

1/2

1/2

1/2

1/2

4/5

4/5

Y

X

B

A

BTT

1/5

1/5

1/10

1/2

1

1

1/2

2/3

0

instantiating the Einsum. Note this rounding only quantizes the the values of d i , but leaves the space of meaningfully different θ s continuous as we consider arbitrarily large matrices.

3 A 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 θ ∈ [0 , 1] 7 . 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) ψ ∈ [0 , 1] , which is related to the rank of the structure, (2) ν ∈ [0 , 1] , which is related to the compute intensity of the structure, i.e. FLOPs per MVM divided by the dimension, and (3) ω ≥ 0 , 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 of the entries of θ , which we derive in Appendix B. To simplify the presentation, we assume d in = d out = d in the rest of the section. Without loss of generality, we assume min( θ XA , θ YB ) ≥ min( θ YA , θ BX ) , so that it is more efficient to first multiply by A instead of B .

Rank Exponent, ψ . For a given θ , we have that rank ( W ) = Θ ( d ψ ) where ψ = min(1 , 2+ θ AB -θ XA -θ YB ) . Thus, ψ = 1 implies fullrank and decreasing values of ψ imply 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 are only connected the input or output. The limit being a low-rank structure as seen in Figure 2 where θ XA = 1 = θ YB and which creates a bottleneck 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 Figure 2. Nonetheless, it is still possible to achieve a full-rank when allocating

10 15

10 16

10 17

Compute (FLOPs)

0.30

0.40

0.50

0.60

0.70

0.80

0.90

Loss

(

1

2

,

1

2

, 0,

1

2

,

1

2

, 0,

1

4

)

(1, 0, 0, 0, 1, 0,

3

4

)

(

1

2

,

1

4

,

1

4

,

1

4

,

1

2

,

1

4

,

1

4

)

(

1

3

, 0,

2

3

, 0,

2

3

,

1

3

,

1

4

)

dimension to only single factors as long as those values do not constitute most of the allocation, meaning 0 < θ XA ≤ 1 / 2 and 0 < θ YB ≤ 1 / 2 , as seen in Monarch, Tensor-Train, and Kronecker

but not for the particular Block Tensor-Train (BTT) structure [19] in Figure 2. In Section 4 we show that structures with ψ = 1 perform best when training neural networks.

Compute Intensity Exponent, ν . Let F denote the FLOPs required to perform a MVM and define the compute intensity as F/d = Θ( d ν ) . The upper bound is achieved by dense, which requires quadratic compute for an MVM and thus ν = 1 . In general, we have ν = 1+ θ AB -min( θ XA , θ YB ) . Thus, in order to achieve lower compute intensity than dense, a structure has to 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 a substantial part of the output dimension to the second factor. ψ and ν are not completely independent, e.g. ψ = ν for low-rank matrices W ij = ∑ r k =1 B ik A kj , though exceptions exist such as for the Kronecker product where ν can be arbitrarily low while maintaining ψ = 1 . In Section 4, we show there exists a wide range of structures with varying ν that perform as well as dense matrices.

Parameters-Sharing Exponent, ω . Let N denote the number of parameters in the structure then N/F = Θ( d ω ) . Clearly, ω = 1 for dense matrices where each parameter is used exactly once in an MVM. In general, we can show that ω = min( θ XA + θ YA , θ XB + θ YB ) -min( θ XA , θ YB ) . In Section 4 we find that structures that share parameters, that is ω > 0 , have worse scaling laws that structures that do not ( ω = 0 ). To achieve ω = 0 , we have to avoid introducing edges that skip some factors, that is θ XB = θ YA = 0 . Structures that skip factors in Figure 2 are Tensor-Train 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 .

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 a function of model size on datasets such as CIFAR10, 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 compute-bottlenecked 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 as a 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 a structure that achieves noticeably better scaling laws compared to dense matrices, we identify the set of common properties shared across a wide range of structures that match the performance of dense matrices, based on which we will propose a significantly more efficient alternative to dense layers in Section 5.

4.1 Main Experimental Setup

We train GPT-2 [20] language models on the OpenWebText 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 120 k to 76 M parameters, with model dimension d ∈ [256 , 4096] and number of transformer blocks L ∈ { 3 , 6 } . Each model is trained for 100 k 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 use 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.

10 14

10 15

10 16

10 17

Compute (FLOPs)

1.00

0.30

0.40

0.50

0.60

0.70

0.80

0.90

Loss

Full Space

Dense

0.0

0.5

10 14

10 15

10 16

10 17

Compute (FLOPs)

1.00

0.30

0.40

0.50

0.60

0.70

0.80

0.90

Loss

=0

Dense

0.5

1.0

10 14

10 15

10 16

10 17

Compute (FLOPs)

1.00

0.30

0.40

0.50

0.60

0.70

0.80

0.90

Loss

=0,

=1

Dense

0.5

1.0

Effect of

ω

Effect of

ψ

Effect of

ν

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 -a compute-optimal scaling laws except with possibly different constants a, b . In Figure 3, we visualize the compute-optimal scaling laws of various Einsums on our language modeling task, including those corresponding to previously proposed structures such as TT [17], Low-rank and BTT [19], as well as a generic Einsum with all entries of θ strictly positive. We report the reducible loss with an estimated L ∞ = 0 . 75 subtracted. This finding suggests that all Einsums can be scaled to reach arbitrarily low reducible loss, each with a different rate that can be estimated from small-scale 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 F via N/F = Θ( d -ω ) , for some ω > 0 . In Figure 4 (left), we show the scaling laws of a wide range of Einsums (only including points on the compute-optimal frontier) colored by ω. We find larger values of ω 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 ω = 0 subspace, we find that ψ becomes the next most important parameter. Recall ψ ∈ [0 , 1] is defined such that the rank of the Einsum scales as Θ( d ψ ) . Einsums with ψ < 1 introduce information bottlenecks in the model by preventing the linear layers from accessing information from all the feature dimensions. The smaller ψ is, the more severe this effect. In Figure 4 (middle), we show that small values of ψ indeed lead to worse scaling laws. This observation further narrows down our search to the subspace with ω = 0 and ψ = 1 , i.e. the space of full-rank BTT matrices.

Any Full-Rank BTT Scales Similarly as Dense. The ω = 0 and ψ = 1 subspace contains the Monarch matrices and its generalization BTT matrices 3 . From a computational perspective, a primary distinguishing factor among these structures is how close they resemble a dense matrix, which we characterize by their compute intensity ν ∈ [0 , 1) defined so that F/d = Θ( d ν ) . ν is large whenever their exists large values (close to 1 ) in the remaining allowed entries ( θ AX , θ ABX , θ YB , θ YAB , θ AB ) . In Figure 4 (right), we show that, somewhat surprisingly, ν has minimal effect on the scaling laws of these structures. Structures with different ν have almost indistinguishable scaling laws compared to each other and dense matrices, which has ν = 1 . 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-

10 12

10 13

10 14

10 15

Compute (FLOPs)

0.10

0.20

0.30

0.40

0.50

0.60

0.70

0.80

Loss

Full Space

Dense

0.0

0.5

10 12

10 13

10 14

10 15

Compute (FLOPs)

0.10

0.20

0.30

0.40

0.50

0.60

0.70

0.80

Loss

=0

Dense

0.5

1.0

10 12

10 13

10 14

10 15

Compute (FLOPs)

0.10

0.20

0.30

0.40

0.50

0.60

0.70

0.80

Loss

=0,

=1

Dense

0.5

1.0

10 12

10 13

10 14

10 15

10 16

Compute (FLOPs)

0.03

0.01

0.003

Loss

Full Space

Dense

0.0

0.5

10 12

10 13

10 14

10 15

10 16

Compute (FLOPs)

0.03

0.01

0.003

Loss

=0

Dense

0.5

1.0

10 12

10 13

10 14

10 15

10 16

Compute (FLOPs)

0.03

0.01

0.003

Loss

=0,

=1

Dense

0.5

1.0

Train, Monarch, and BTT can significantly outperform dense in other settings such as 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 as a 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 × 8 × 3 resolution. Figure 5 (top row) shows qualitatively the same results as our GPT-2 experiments, where ω and ψ have the most significant impact on the scaling laws, while varying ν yield only slight variations. In this particular case, having ν = 0 . 75 (BTT with BTT-rank scaling as d 1 / 4 ) is better than having ν = 0 . 5 (Monarch matrices). One notable trend in this setup is that most Einsums, regardless of ω or ψ , 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.

MLP Regression on Synthetic Data. We train MLPs on a 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 5 (bottom row), we observe qualitatively the same results as in the GPT-2 experiments.

Together, these additional results suggest there is some degree of universality associated with our findings on the effect of ω, ψ, and ν on the compute-optimal scaling laws of neural networks that use Einsums in place of dense matrices.

10 14

10 15

10 16

10 17

10 18

Compute (FLOPs)

0.20

0.30

0.40

0.50

0.60

0.70

Loss

10 15

10 16

10 17

10 18

Compute (FLOPs)

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Compute Multiplier

8

16

Number of Experts

0

1

2

3

4

5

Compute Multiplier

Dense

BTT-MoE

5 Structured Mixture of Experts

In Section 4, we identified that Einsums with ω = 0 and ψ = 1 perform the best, and ω > 0 or ψ < 1 lead to worse-than-dense performance. The most impactful parameter on the scaling laws is ω, which measures how many parameters an Einsum has compared to the FLOPs for an MVM. Einsums that learn one parameter per FLOP perform significantly better than those that learn less than one parameter per FLOP. Therefore, a natural question arises: can we 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 in an MVM, which necessitates a sparse Mixture-of-Experts (MoE) like architecture [22, 7, 12]. Furthermore, we would like the structure to be full-rank, i.e. ψ = 1 . 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 a layer via an Einsum is to turn a BTT with BTT-rank E , which involves a sum over the rank index ρ = 1 , . . . , E :

where g ∈ R E is a k -sparse vector so that only k out of E terms need to be computed. We interpret k as the number of active experts and E as the total number of experts. We compute g via a softmax over the topk entries of the logits e produced by a (dense) linear gating function e = Linear( X ) ∈ R E . There is no need to make this gating function structured because its cost is negligible. We choose θ AX = θ ABX = θ YB = θ YAB = 1 / 2 so that each expert is full-rank. We follow the common practice of using k = 2 . The resulting BTT-MoE layer is a BTT with BTT-rank 2 (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 ρ with a gate. We use the load-balancing loss to encourage equal utilization of all experts [22, 7, 23].

In contrast to the standard MoE architecture used in transformer language models, which uses a sparse MoE for each entire feed-forward network (FFN) [22, 7, 12]:

into a k -sparse sum:

BTT-MoE learns an MoE in every single linear layer of the model (except the language modeling head) and treats them equally, including the projection matrices W Q , W K , W V , W O in the attention blocks [24]. It learns more fine-grained routing decisions among the experts, with ( 1 2 E ( E -1) ) 6 M possible combinations of the experts in a transformer with M blocks, compared to ( 1 2 E ( E -1) ) M for the standard MoE architecture.

5.2 Compute Efficiency Gains

In Figure 6, we show GPT-2 with BTT-MoE achieves better compute-optimal scaling laws compared to the dense model as well as the standard MoE, with k = 2 and 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 a dense model with λC training FLOPs. In Figure 6, we show BTT-MoE is significantly more compute-efficient than the standard MoE for both E = 8 and E = 16 . In particular, with E = 16 experts, BTT-MoE achieves a compute multiplier of λ = 5 . 3 ± 0 . 3 , compared to λ = 4 . 1 ± 0 . 3 for a 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-MoE) or dense matrices (denseMoE) also yields a nontrivial compute multiplier ( ∼ 2 × ) over the dense model, but is significantly less effective than BTT-MoE or even the standard MoE.

While the poor relative performance of low-rank-MoE is expected, this result shows that in addition to ω = 0 and ψ = 1 , ν < 1 is a desirable property for the base structure in 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.

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

10

4

10

3

10

2

30

40

50

60

Test Error (Naive)

(

4

5

, 0,

1

5

, 0,

4

5

,

1

5

, 0

)

10

4

10

3

10

2

30

40

50

60

(

1

3

,

1

3

,

1

3

,

1

3

,

1

3

,

1

3

,

1

3

)

10

4

10

3

10

2

30

40

50

60

(

1

2

,

1

2

, 0,

1

3

,

1

3

,

1

3

, 0

)

10

4

10

3

10

2

Learning Rate (

)

30

40

50

60

Test Error (

P)

10

4

10

3

10

2

Learning Rate (

)

30

40

50

60

10

4

10

3

10

2

Learning Rate (

)

30

40

50

60

64

256

1024

4096

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Compute Multiplier

BTT-MoE

Dense-MoE

Low-Rank-MoE

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, µ P states that for a dense matrix W ∈ R d in × d out , the optimal initialization standard deviation scales as σ = Θ( √ min( d in , d out ) /d 2 in ) and its learning rate as η = Θ(1 /d in ) if using Adam. Furthermore, Qiu et al. [19] shows we can apply µ P to 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 d A in = d XA , d A out = d YA d YAB d AB , d B in = d XB d XAB d AB and d B out = d YB . As a result, we can compute the optimal scaling of σ A , σ B , η A , and η B . In particular, for Adam we have η A = Θ( 1 d XA ) and η B = Θ( 1 d XB d XAB d AB ) . Figure 8 shows using µ P leads to a stable optimal learning rate and better performance compared to naively using a constant global learning rate and unit initialization variance. This property allows us to transfer the learning rate between structures and model sizes, saving substantial compute for hyperparameter tuning. For µ P, 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 a case-by-case basis, we introduce a continuous parameterization over the space of all structured matrices expressible as Einsums. Using this parameterization, we 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 ψ. Across tasks, we 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 or 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 a 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 a sequence of batched matrix multiplications, we extend prior work on structure-aware initialization and learning rate based on µ P to arbitrary Einsums.

MoE over structured matrices is more efficient than standard MoE over entire FFNs. By replacing every single dense linear layer with a sparse sum of structured matrices like BTT, compared to standard MoE which operates over entire FFNs, we create a more efficient MoE architecture, achieving over 5 × 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&E-MSS 2134216, NSF HDR2118310, BigHat Biosciences, Capital One, and an Amazon Research Award.

References

Patrick Amestoy, Cleve Ashcraft, Olivier Boiteau, Alfredo Buttari, Jean-Yves l'Excellent, and Clément Weisbecker. Improving multifrontal methods by means of block low-rank representations. SIAM Journal on Scientific Computing , 37(3):A1451-A1474, 2015.

Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701 , 2021.

Silvere Bonnabel. Stochastic gradient descent on Riemannian manifolds. arXiv 1111.5280 , 2011.

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:1877-1901, 2020.

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.

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.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research , 23 (120):1-39, 2022.

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-Quadratic GEMM-Based Architecture. Advances in Neural Information Processing Systems (NeurIPS) , 2023.

Roger Grosse and James Martens. A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers. arXiv 1602.01407 , 2016.

Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701 , 2020.

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.

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.04088 , 2024.

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.08361 , 2020.

Changwoo Lee and Hun-Seok Kim. Differentiable learning of generalized structured matrices for efficient deep neural networks. arXiv preprint arXiv:2310.18882 , 2023.

Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: Highrank training through low-rank updates. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023) , 2023.

Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. arXiv preprint arXiv:2010.08127 , 2020.

I. V . Oseledets. Tensor-Train Decomposition. SIAM Journal on Scientific Computing , 2011.

Andres Potapczynski, Marc 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.

Shikai Qiu, Andres Potapczynski, Marc Finzi, Micah Goldblum, and Andrew Gordon Wilson. Compute better spent: Replacing dense layers with structured matrices. International Conference on Machine Learning (ICML) , 2024.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. OpenAI , 2019.

Yunus Saatçi. Scalable inference for structured Gaussian process models . PhD thesis, Citeseer, 2012.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 , 2017.

Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching llama2 performance with 0.1 m dollars. arXiv preprint arXiv:2404.07413 , 2024.

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.

Andrew Gordon Wilson, Elad Gilboa, Arye Nehorai, and John P. Cunningham. Fast Kernel Learning for Multidimensional Pattern Extrapolation. Advances in Neural Information Processing Systems (NeurIPS) , 2014.

Greg Yang and Edward J. Hu. Feature Learning in Infinite-Width Neural Networks. International Conference on Machine Learning (ICML) , 2021.

Greg Yang and Etai Littwin. Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limit. International Conference on Learning Representations (ICLR) , 2023.

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.

Greg Yang, James B. Simon, and Jeremy Bernstein. A Spectral Condition for Feature Learning. Preprint arXiv:2310.17813 , 2023.

Qibin Zhao, Guoxu Zhou, Shengli Xie, Liqing Zhang, and Andrzej Cichocki. Tensor ring decomposition. arXiv preprint arXiv:1606.05535 , 2016.

Appendix Outline

Appendix A shows specific examples of Einsums and their MVM computations.

Appendix B shows the derivation for the taxonomy variables ω , ψ and ν.

Appendix C discusses learning rate scaling details and connections between µ P and Riemannian SGD.

Appendix D provides extended experiment results and experiment details.

Appendix E discusses the generalization of our parameterization to Einsums of more than two factors.

Appendix F shows how to further improve performance by structuring the Einsum axes to be compatible with the attention head structure in transformers.

Appendix G 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 a size of 1 and is thus omitted.

Dense A one-factor Einsum would have been enough to represent dense, though one way to represent it within a two-factor Einsum is as:

which is an over-parameterization due to the Hadamard product B γϕ A γϕ .

Low Rank

Kronecker Product

Tensor-Train

Monarch

Block Tensor-Train

B Taxonomy Derivation

Here we present the derivation for the taxonomy variables ( ω, ψ, ν ) . It is helpful to first reason about the amount of compute required for calculating an Einsum, which will allow us to exclude some uninteresting Einsums from consideration and thereby simplify our analysis. We only consider two-factor Einsums representing a square matrix, but generalization to more factors and non-square matrices is straightforward.

For reference, we show the general expression again:

Computational Complexity. There are three possible ways to compute Equation (11), depending on whether the first contraction happens between { A , B } , { A , X } , or { B , X } . The required FLOPs † are Θ ( d 2+ θ ρ ) , Θ ( d 2+ θ ρ -min( θ α ,θ ϵ ) ) , and Θ ( d 2+ θ ρ -min( θ β ,θ δ ) ) respectively. The optimal computational path depends on the entries of θ , and the optimal FLOPs F required is given by F = Θ ( d 2+ θ ρ -max { min( θ α ,θ ϵ ) , min( θ β ,θ δ ) } ) .

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 to each other. To remove this redundancy, we can require that first summing with A is more computationally efficient. Thus, we require

which also simplifies the FLOPs to F = Θ ( d 2+ θ ρ -min( θ α ,θ ϵ ) ) . To be more exact, we have F = ρd 2 ( 1 d α + 1 d ϵ ) .

Degenerate Einsums. Since the overall Einsum is a linear operator on R d , any Einsum that requires more than Θ ( d 2 ) FLOPs to compute is degenerate in the sense that unnecessary computations are performed. For example, one such Einsum could correspond to a factorization UV ⊺ , U ∈ R d × r , V ∈ R r × d where r ≫ d. For convenience, we will also define an Einsum whose cost is equal to Θ ( d 2 ) as degenerate since it is no more efficient than a dense matrix. Given the expression for FLOPs, we conclude that non-degenerate Einsum are those where

Intuitively, this requirement means that the rank dimension d ρ , i.e. the range of the index ρ connecting A and B , cannot be set too high, since after some point it becomes more efficient to simply use a dense matrix in place of the Einsum. In particular, we must have d α > 1 and d ϵ > 1 (except when d = 1 ).

Compute Intensity Exponent, ν . One defining characteristic of structured matrices is that the FLOPs F for performing matrix-vector-multiplication (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 our case, all non-degenerate Einsum have sublinear compute intensity. More precisely, we have F/d = Θ( d ν ) , where ν = 1 + θ ρ -min( θ α , θ ϵ ) takes value in [0 , 1) assuming non-degeneracy. The closer ν is to 1 , the more a Einsum resembles a dense matrix and vice versa. ν = 0 corresponds to low rank matrices with Θ(1) rank.

Rank Exponent, ψ . As we show in Appendix C.1, the Einsum can be computed via two batched matrix multiplications where A acts as a matrix consisting of d β d γ blocks of d α × d δ d ϕ d ρ matrices. Thus, rank ( A ) = min( d, d β d γ d δ d ϕ d ρ ) = d min(1 ,θ β + θ γ + θ δ + θ ϕ + θ ρ ) as d α d β d γ = d . Similarly, we have that B acting as a matrix consisting of d δ d ϕ blocks of d β d γ d ρ × d ϵ matrices. Therefore we have that rank ( B ) = d min(1 ,θ β + θ γ + θ δ + θ ϕ + θ ρ ) . Since W is 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 ) :

Note that, technically, when we write rank ( M ) , M ∈ { A , B , W } , we mean the maximum possible rank of M when its parameters are learned. Otherwise the equalities will become upperbounds. For low-rank we have θ = (1 , 0 , 0 , 0 , 1 , 0 , θ ρ ) and hence ψ = θ ρ . For dense we have θ = (0 , 0 , 1 , 0 , 0 , 1 , 0) and hence ψ = 1 . For Kronecker we have θ = (1 / 2 , 1 / 2 , 0 , 1 / 2 , 1 / 2 , 0 , 0) and hence ψ = 1 . For BTT we have ψ = min(1 , θ γ + θ ϕ + θ ρ ) .

Parameter Sharing Exponent ω . The number of learnable parameters N in a Einsum is simply the sum of the elements in A and B , which works out to be N = d ρ d 2 ( 1 d α d δ + 1 d β d ϵ ) . Since the FLOPs is F = d ρ d 2 ( 1 d α + 1 d ϵ ) , only Einsums with d δ = d β = 1 , or equivalently θ δ = θ β = 0 , have parameters matching FLOPs. In general, the number of parameters per FLOP N/F = Θ( d -ω ) with ω = min( θ α + θ δ , θ β + θ ϵ ) -min( θ α , θ ϵ ) ≥ 0 .

The ω = 0 subspace is 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 ω = 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 an axis connects the first factor back to the last factor when there are 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 η A = Θ( 1 d XA ) and η B = Θ( 1 d XB d XAB d AB ) . Arriving at these scaling rules requires (1) expressing an Einsum as a sequence of batched matrix multiplies (BMMs) and showing d A in = d XA , d A out = d YA d YAB d AB , d B in = d XB d XAB d AB and d B out = d YB , 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).

The general 2-factor Einsum

can be computed in two steps. In step 1,

which is a BMM with βγ acting as batch dimensions, α the input dimension, and δϕρ the output dimensions. So d A in = d XA , d A out = d YA d YAB d AB . In step 2,

which is a BMM with δϕ acting batch dimensions, βγρ input dimensions, and ϵ the output dimension. So d B in = d XB d XAB d AB and d B out = d YB , as wanted.

In our experiments, we apply the above scaling rule to transfer from a learning rate η used by a dense matrix with width d 0 to learning rates for A , B as

where the additional factor of two is to account for both A and B contributing updates to the output of each layer, following Qiu et al. [19].

C.2 Connections between µ P and Riemannian SGD

Riemannian SGD (RSGD) is an optimization technique that allows us to perform the equivalent of SGD on a Riemannian manifold consisting of points { q } with a metric g q [3]. The updates of RSGD is almost identical to SGD: q ( t +1) = q ( t ) -η q g -1 q ( t ) ∇ q L ( q ( t ) ) , except that the gradient is multiplied by the inverse of the metric g q . In our case, we want to mimic training the Einsum parameters A and B as if we were training W with SGD directly, even though we will never represent W explicitly. Thus, we identify (flattened) ( A , B ) with q and specify its metric as the pull-back metric of the

Euclidean metric g W = I on W , which is given by g q = J ( q ) ⊺ J ( q ) where J ( q ) = ∂ W ∂ q . The RSGD updates to q is therefore

Exactly computing the inverse metric would require O ( P 3 ) time, where P is 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 a way that is exact for asymptotically large matrices. We compute J ⊺ J through three blocks J ⊺ A J A , J ⊺ B J B , and J ⊺ A J B and note that for a fixed θ as we increase both d in and d out , we have that J ⊺ A J A ≈ d β d ϵ σ 2 B I , J ⊺ B J B ≈ d α d δ σ 2 A I , and J ⊺ A J B ≈ 0 where σ A and σ B denote the initialization scales of A and B . Therefore we have

We can now identify Θ ( 1 / ( d β d ϵ σ 2 B ) ) as η A and Θ ( 1 / ( d α d δ σ 2 A ) ) as η B . If we further use the initialization scales prescribed by µ P as given in Section 6, we have

It is now interesting to compare these scalings to the prescriptions of µ P SGD. For SGD, µ P proposes to scale the learning rate as Θ( d out /d in ) [29] for dense matrices. Unlike in Adam, to extend µ P to Einsums, we need not only to replace d in , d out with d A in , d A out , or d B in and d B out , but also scale down the learning rate as Θ(1 /d β ) for A and Θ(1 /d δ ) for B . The final µ P learning rates have two possibilities depending on whether we consider A or B as the first layer. If we consider A as the first layer, then

Evidently, the two approaches don't always agree. However, they are identical when d α = d ϵ , d β d γ d ρ ≤ d ϵ , and d δ d ϕ d ρ ≤ d α . In fact, many structures satisfy this condition. If we as- sume the structure is symmetric, meaning its transpose can be represented with the same θ , i.e. d 2

d α = d ϵ , d β = d δ , and d γ = d ϕ , then the remaining conditions simplify to only d ρ ≤ α d . Therefore, any symmetric Einsum with θ α ≥ 1 / 2 and θ ρ ≤ 2 θ α -1 satisfy these conditions. Thus, for these structures, µ P SGD can be viewed as an approximation to RSGD that is valid at initialization in the infinite-width limit. While we cannot establish a 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 µ P SGD in that it maximizes feature learning per layer [27]. The connection between µ P SGD and RSGD therefore indirectly provides an alternative justification for µ P Adam.

D 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 a L = 3 , d = 256 dense model, which is scaled to other models via µ P. Since the language modeling head contains a significant fraction of the parameters for models of this scale, we replace all layers, including the head, with Einsums.

Qualitatively, Figure 9 differs from Figure 4 in two ways: 1) the scaling laws are less power law like and show some curvature on a log-log scale, and 2) BTT with ν > 0 seems to perform better than ν = 0 . We believe 1) is due to 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 a smaller vocabulary size and context length in Section 4. Similarly, we believe the increased vocabulary size and context length contributed to 2), as a 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 finite-size effects.

10 16

10 17

10 18

Compute (FLOPs)

3.0

4.0

5.0

Loss

Full Space

Dense

0.0

0.5

10 16

10 17

10 18

Compute (FLOPs)

3.0

4.0

5.0

Loss

=0

Dense

0.5

1.0

10 16

10 17

10 18

Compute (FLOPs)

3.0

4.0

5.0

Loss

=0,

=1

Dense

0.5

0.8

D.2 Autoregressive Pixel Modeling on CIFAR-5M

We train 2 and 3 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 to 1024 for Einsums. All models are trained for 2 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 on Synthetic Data.

We train 3-layer MLP models with width d ∈ [64 , 4096] for a maximum of 10 6 steps and a batch size of 4096 on an effectively infinite synthetic dataset. The synthetic dataset is obtained by querying a scalar-valued target function on R 8 with inputs drawn from a Gaussian distribution. The target function is a randomly initialized target MLP with 6 layers and a hidden dimension of 1024. We minimize the Mean-Squared-Error (MSE) loss. We train with Adam using a base learning rate of 1e-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 if we consider the set of edges that define the Einsum E = { S ⊆ { X , A , B , Y } : | S | ≥ 2 and { X , Y } ̸⊆ S } . For example, for three factors, the edges that connect to X are { X , A } ↔ i 1 , { X , B } ↔ i 2 , { X , C } ↔ i 3 , { X , A , B } ↔ i 4 , { X , B , C } ↔ i 5 { X , A , C } ↔ i 6 and { X , A , B , C } ↔ i 7 . The edges that connect to Y are { Y , A } ↔ j 1 , { Y , B } ↔ j 2 , { Y , C } ↔ j 3 , { Y , A , B } ↔ j 4 , { Y , B , C } ↔ j 5 { Y , A , C } ↔ j 6 and { Y , A , B , C } ↔ j 7 . The edges between the factors are { A , B } ↔ r 1 , { A , C } ↔ r 2 and { A , B , C } ↔ r 3 . The expression for a three factor would be

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 x ∈ R d , we compute q = W Q x , v = W V x and k = W K x . After computing each of the q , k , v ∈ R d vectors, they are reshaped to produce one smaller vector per attention head: q i → q hj , k i → k hj , v i → v hj where h is an axis of size H, the number of attention heads, and j is an axis of size d/H. When replacing W Q , W K , W V with BTTs, it is therefore more natural to have the BTT output axes ϵϕ coincide with hj so the MVM is aware of the attention head structure rather than potentially mixing different heads. Similarly, when replacing W O with a BTT, it is most natural to have the BTT input axes

10 14

10 15

10 16

10 17

Compute (FLOPs)

0.3

0.4

0.5

0.6

0.7

Loss

Scaling Laws on GPT-2

BTT

BTT-Attn

αγ coincide with hj. 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 corresponds to θ = (1 / 2 , 0 , 1 / 2 , 0 , 1 / 2 , 1 / 2 , 0) in Section 4.

G Hardware specifications

Our experiments in Section 4 are run on on A100 and H100 GPUs. The CIFAR-10 experiments in Section 6 were run on RTX2080 Ti and RTX Titan GPUs.

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 a 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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 no 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 a 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced.

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 to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

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 be 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 a dataset and/or model, the authors should describe the steps taken to 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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 no 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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.

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: It's mentioned in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or 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 in 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 in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper doesn't involve living subjects and utilize public data. It conforms with all the code of ethics as listed.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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.

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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 not use 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 is used and, if 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset'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 a 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is 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. Crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on 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 to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis paper is clearly written and easy to understand\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThis paper does not have direct negativie societal impact from my perspective.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

comments: 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.

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.
2) 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.

QUESTIONS:
1) Did you try other than BTT structured layers in the MoE experiments? Why do you only use BTT?
2) 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?
3) Do you expect that all the observations about $\omega, \psi \mu$ will remain the same with bigger dictionary?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-13