Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices

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.

Strengths

- The paper provides a novel point of view for understanding the structured matrices via Einsum operations. - The paper is clearly written with sufficient details. - 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. - 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.

Questions

- Is the scaling law of Einsum generalizable to the larger models (e.g., $\ge$7B parameters)? - 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? - Could the taxonomy parameters have any real numbers other than rational numbers?

Rating

7

Confidence

5

Soundness

4

Presentation

2

Contribution

4

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.

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.

Strengths

This paper is clearly written and easy to understand

Weaknesses

- The novelty of this work is not clear enough from my perspective - Empirical results may need further improvements to better support the proposed method

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. - 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. - 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. - 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. - 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. ## Post rebuttal After 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.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

This paper does not have direct negativie societal impact from my perspective.

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?

Rating

7

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

N/A