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

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.

On the key contributions of our work: We highlight that this work provides at least 4 novel and impactful contributions.

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

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 than random.

4. 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.

On experiments with additional architectures and datasets. Focusing 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 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.

On comparing against MoE with other structures. Following your suggestion, we compare against two additional MoE architectures in Figure 4 in the rebuttal's 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.

On generalization to more than 2 factors. The 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 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

Thank 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.

[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

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.

On continuity of the parameterization. While 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.

On experiments with larger models. In Figure 1 in the rebuttal's 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.

In addition, we perform experiments with other architectures and datasets in Figure 2 and 3 in the rebuttal's pdf, including Vision Transformers on image generation and MLP on synthetic regression. The results confirm that our findings generalize to much broader settings.

On a studying per layer configurations. You 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.

We 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.

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.

On experiments with larger vocabulary and longer context. Thank you for raising this point as these new results strengthen the paper significantly. In Figure 1 in the rebuttal's 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.

On additional MoE comparisons. Initially 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: 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.

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.

Thank 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.

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