Building on Efficient Foundations: Effective Training of LLMs with Structured Feedforward Layers
Authors: Xiuying Wei
Skander Moalla
Razvan Pascanu
Caglar Gulcehre
Venue: NeurIPS 2024 poster (2024)
Keywords: large language model
structured matrices
efficient architecture
Primary area: deep_learning_architectures

Abstract: State-of-the-art results in large language models (LLMs) often rely on scale, which
becomes computationally expensive. This has sparked a research agenda to reduce
these models’ parameter counts and computational costs without significantly
impacting their performance. Our study focuses on transformer-based LLMs,
specifically targeting the computationally intensive feedforward networks (FFNs),
which are less studied than attention blocks. We consider three structured linear
parameterizations of the FFN using efficient low-rank and block-diagonal matrices.
In contrast to many previous works that examined these approximations, our study
i) explores these structures from a training-from-scratch perspective, ii) scales up
to 1.3B parameters, and iii) is conducted within recent Transformer-based LLMs
rather than convolutional architectures. We demonstrate that these structures can
lead to actual computational gains in various scenarios, including online decoding
when using a pre-merge technique. Additionally, we propose a novel training
regime, called self-guided training, aimed at improving the poor training dynamics
that these approximations exhibit when used from initialization. Interestingly,
the scaling performance of structured matrices is explored, revealing steeper
curves in scaling training FLOPs, along with a favorable scaling trend in the
overtraining regime. Specifically, we show that wide and structured networks
can utilize training FLOPs more efficiently, with fewer parameters and lower
loss than dense models at their optimal trade-off. Our code is available at
https://github.com/CLAIRE-Labo/StructuredFFN/tree/main.

Building on Efficient Foundations: Effectively Training LLMs with Structured Feedforward Layers
Xiuying Wei xiuying. suuying.wei@ee. wwi@@ppf.cc CLAIRE, EPFL
Skander Moalla skander.moallae www.mmollllllllccc.......o.. CLAIRE, EPFL
Razvan Pascanu razp @ @oogle. Google DeepMind
Caglar Gulcehre caglar.gulcehre @ epfl.ch CLAIRE, EPFL
Abstract
State-of-the-ar results in large language models (LLMs) often rely scale, which becomes computationally expensive. This has sparked research agenda to reduce these models' parameter counts and computational costs without significantly impacting their performance. Our study focuses on transformer-based LLMs, specifically targeting the computationally intensive feedforward networks (FFNs), which are less studied than attention blocks. We consider three structured linear parameterizations of the FFN using efficient low-rank and block-diagonal matrices. In contrast to many previous works that cxamined these approximations, our study i) explores these structures from a maaning-from-seree to 1.3B parameters, and iii) is conducted within recent Transformer-base basedd LLMs rather than convolutional architectures. We demonstrate that these structures can lead to actual computational gains in various scenarios, including online decoding when using a pre-merge technique. Additionally, we propose novel training regime, called self-guided training, aimed at improving the poor training dynamics that these approximations exhibit when used from initialization. Interestingly, the scaling performance of structured matrices is explored, revealing steeper curves in scaling training FLOPs, along with a favorable scaling trend in the overtraining regime. Specifically, can utilize training FLOPs more efficiently, with fewer parameters and lower loss than dense models at their optimal trade-off. Our code is available https:
perspective, ii) scales up
show that wide and structured networks
Table 1: Better training FLOPs utilization of the wide and structured Networks: we compare dense Transformers trained according their optimal scaling law [1], efficient Transformers (GQA) [2] with high throughput, and wide and structured networks using LowRank parameterization in the FFN module and reduced attention heads, under the same training FLOPs. TP (throughput) refers to the maximum throughput measured over generation length of 256.
Runk (22%g ITN
Lawiooa W0% EP%
Deaee oepites sedling
Validation Loss
Method
AParam
Training FLOPPs
PPL
TP (token/s)
23
1
10"
Transformer-m
335M
1,55+119
18.29
30229
Training Flaps
Transformern m (GQA)
315M
1.5e+119
18.23
$4202
Wide Strücfured
219
1335:19
17,89
011778895
Figure Steeper scaling curves of LowRank with 63% or 32% % FEN param- eters, For more results, see Sec. 4.2.
Transfommer.
729M
7.00r119
14.29
23351
Transfommer- (GQA)
729M
7.03c+19
14.40
64737
Wide
7.006
(17% f)
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

Introduction
Transformer language models [3] have gained significant attention for their performance and scalability. These models have grown from hundreds of millions of parameters [4] to hundreds of billions [5-7], increasing the need for efficient training and inference techniques. While much research focuses on attention, feed forward networks (FFNs) account for over 60% of the model's parameters and FLOPs, significantly impacting latency.' 1 Recent large-scale models [8, 9] further increase the FFN size, leading them to dominate the cost of the model compared to the attention layer.
Structured linear transformations, such as low-rank or block-diagonal matrices, are important paradigms for reducing the computational cost of feedforward layers. However, they have not yet been thoroughly explored at sufficient scale to reduce pre-training costs and latency of the inference phase in modern LLM architectures, where the main focus so far has been on improving the efficiency of the self-attention mechanism.
In this work, we investigate structured matrices for FFN blocks from the train-from--eratch aspect, first identifying their efficiency and optimization challenges and then presenting experimental results, analyzing and characterizing the behavior of models trained with structured matrices, and comparing their results. We consider three efficient linear parametrizations: LowRank. BlockShuffle (comprising two block-diagonal matrices), and BlockDense (a combination of dense and block-diagonal matrices). First, while they have demonstrated materialized computational gains, they face challenges in the practical online decoding scenario of LLM, which may process only limited input tokens at one time, leading to under-utilization of computing resources and decreased efficiency due to the additional linear projection. We address this with a pre-merge technique that restores efficiency to the original dense parametrization. Second, we observe that these parameterizations of the FFN blocks are harder to train than standard linear layers, often exhibiting poorer training dynamics like loss spikes. To counter this, we propose a flexible and fast method we refer to as self-guided training. It employs dense matrix residual component during the initial training phase, steering the training process away from suboptimal starting points gradually.
We conduct our experiments at scale on Transformers ranging from 110M to 1.3B parameters by replacing the traditional heavy FFN with structured matrices. Our experiments first show the scaling behavior of these structured linear parameterizations and then illustrate how our proposed methods address their general efficiency and optimization challenges. First, we examine scaling performance from the perspectives of training compute and model size, highlighting the potential of structured matrices. By controlling for the same training FLOPs, we find that structured FFNs show steeper loss scaling curves than traditional Transformers at optimal trade-offs (See Fig. 4 and Fig. 1). Specifically, as seen in Table 1, our wide and structured networks use training FLOPs more efficiently, needing fewer parameters (464M vs. 729M) and achieving 17% throughput boost on the -1 scale, while still maintaining slightly better perplexity compared to the efficient Transformer [2]. Beyond training com- pute scaling, we also scale model size in the overtraining regime, with Fig. 5 showing favorable scaling trends for our wide and structured models. Second, our results on efficiency show that structured FFNs, with only 32% of the FFN parameters, can boost the training speed of the 1.3B model by 1.35 x. Furthermore, self- ceff-guided training enhances the performance of all three structured matrices (e.g.. reducing the perplexity gap of LowRank to about 0.4) without affecting the inference time speed-up.
As the first work to explore structured matrices at the scale of recent L.MSs, we hope our findings and results will shed new light on the study of efficient NLP architectures. Our contributions can be categorized into the following three aspects:
1. We investigate three types of structured matrices in Transformer pretraining and demonstrate their favorable scaling behavior compared to dense models. This is revealed through the study of scaling laws for training FLOPs, as well as model size scaling in the overtraining regime, showing that wide and structured networks can be strong candidates for architecture design.
2. We conduct an efficiency study of these structured matrices across various scenarios. We propose a pre-merge technique to maintain speed in a specific case and show the effective speed-up performance of structured matrices in other scenarios.
3. We identify optimization challenges in structured matrices and introduce a method called self- guided training, which efficiently improves training dynamics and boosts the final performance for all three types of structured matrices.
For example, we find that it composes 54% of total latency in 1.3B model.

---

2 Method
Multiple techniques have been proposed to approximate linear projections from sparsity to low-rank approximations. We provide an in-depth discussion of existing literature on approximating linear lay- ers in section 3, that better contextualizes our work. We focus on structured approximations of linear projections g(x) Wa æ that have the form g(x) (V)),, where U and V are structured matrices, e.g. low-rank or plock-diagonal. We opt for this particular form because it allows us to readily exploit the computational gains existing hardware using existing libraries with minimal alteration. Such approximations have been previously studied in different contexts. Our contributions are exploring them (i) to approximate FFN layers of transformers, (ii) when applied from initialization, (iii) testing them at up to 1.3B scale, investigating their general bottlenecks and providing scaling analyses.
N-
+ M)
min
R-MM+ B
W
(a) Dense
(b) LowRank
(c) BlockSbuffle
BlockDense
Figure 2: Structured linear parametrization: We show the structured linear parametrization with input dim. of N and output dim. of M. a) The traditional dense linear parametrization. b) LowRank parametrization with bottleneck of size R where R is less than M and N. BlockShuffle with two block-diagons matrices with blocks of size B interleaved with shuffle operations that mixes information from different blocks similar ShuffleNet. BlockDense with the first matrix as a block-diagonal and the second a low-rank or dense matrix.
2.1 Structured linear parametrization
We explore three structured matrices to approximate the standard linear layer W a, maintaining its input dimension and output dimension M of weight W.
LowRank Low-rank matrices have been widely used to decompose pre-trained weights for downstream compression [10] or to construct adapters for efficient fine-tuning [11]. Researchers [12] suggest that dense layers tend naturally converge to low-rank solutions during training, making this approximation ideal. Inspired by this, we explore low-rank matrices as alternatives to traditional linear layers, imposing this structure from initialization and investigating it during pre-training.
Formally, the low-rank approximation of a linear layer is given as Wa R U'V"")) where and R < min(M,N) Note that we use the superscript to indicate that these matrices are used to create a low-rank approximation by projecting to or from a low-dimensional code, a notation that would become useful later to distinguish such components from -iikonda ones. The parameter count and MAC (Multiply-Accumull Operations) decrease from M-N to (M+N)-R.
U" E MxR V E Rx and
BlockShuffle Dao et al. [13] proposes using the Monarch decomposition of FFT, replacing the linear layer with interleaved block-diagonal and permutation matrices. An alternative motivation for such a structure can be derived from efficient convolutional designs of ShuffleNet [14] and separable convolution [15]. For simplicity, we explore the form introduced by ShuffleNet to linear layers.
The core idea of BlockShuffle to reshape the feature dimension into two dimensions and first use a linear projection that mixes along one of these fictive dimensions, followed by linear projection that mixes along the other. More precisely, we first reshape the input features æ E N into B blocks and apply the non-tied weight of
x B to each block, then flatten the intermediate feature. To achieve global perception, we regroup elements from different blocks into new blocks and apply the same transformation for each block again.
N N B
3

---

Technically, we can express the per-block transformation using block coock-diagonal matrices and formulate the above process as Wx ~ (U (UV')\))) where block-diagonal matrices V and Uh has B blocks with shapes min(N,M) X B and M the shuffle function f(-) enables global feature mixing by cycling different blocks and can be implemented by simply transposing and reshaping inner features. The inverse function f (.)
B X min(N.MM per-block. As shown in Fig. 2. B
13
permutes the outputs back to their original order.
By separating features into two dimensions, only a few elements of the features will be processed each time. The parameter count and MAC are reduced from M to min(N,M) (M+NN) where B acts as a trade-off of accuracy and efficiency.
B
BlockDense The previous parametrization incorporates additional shuffle operations, which can be slow on the device. We propose a natural intermediate candidate between LowRank and BlockShuffle, combining the biock-diagonal projection Vo with a dense low-rank projection U". Thus, we can mix the features between blocks without permuting the inner features. The formula is defined as W x ~ UV)) where V is the block-diagonal matrix with b blocks in shape ft N and U E R M Technically, the second projection does not need to be low-rank approximation because R can be larger than M. Nevertheless, in practice, we chose R < M to limit the search space of this work, and thus use the superscript r for the second matrix. The parameters of this method are determined by two variables B and R, cutting the original burden from M N to R-(M + + Note that BlockDense can recover the LowRank approximation if we B and R < min (N.).
Remark We limit our exploration of the efficient linear parametrizations within the FFN blocks. These typically have 8- H2 parameters and MAC, where H standards for hidden state dimen- sion and 4- H as the intermediate hidden size of FFN. In contrast, the proposed parametrizations have:
H
LowRank: 10- H.R
BlockShuffle:10 10
BlockDense: 5 5- H.-.R- (1+ B
B
Although BlockDense is introduced new parameterization, the aim of this paper is not to claim it as the best candidate, but rather to investigate some general properties of structured matrices from efficiency, optimization, and scaling perspectives. Given the favorable efficiency and loss performance of BlockDense, it is included alongside LowRank and BlockShuffle here to cover a broader range of potential parameterizations.
2.2 Maintaining efficiency during online decoding
Parallelism-bound FFN With reduced FLOPs and parameters, our proposed linear parametrization
can accelerate the model for compute- bound and memory-bound scenarios [16], usually during training, prefilling, and decoding with relatively big batch size. However, for online decoding with very small batch size and sequence length of 1, a practical scenario for LLM, both FFN and structured FFN can become parallelism-bound [17] with poor utilization of the GPU resources, especially on powerful devices like A A100, Because each linear projection suffers from parallelism-bonnd efficient linear parametrization may lead to worse latency performance due to doubling the number of linear projections. We propose a pre-merge technique to mitigate the issue.
Pre-merge technique Taking ad advantage of the fact that these parametrizations do not have non- linearity, we propose to combine the structured matrices into a single dense layer and keep both the structured and the dense one for online decoding. Then, we can dynamically decide which parametriza- tion to use based on the current batch size and setting. Fig. analyzes using structured or dense forms for different batch and model sizes, allowing us decide when to use the pre-merged linear projection.
2.3 Addressing the optimization challenge
Using the efficient parametrization from initialization can suffer from optimization difficulty because the deep linear parametrization introduces additional symmetries?, which is source of proliferation
²s *Symmetrics arise hecause such factorization UV is not unique; for any invertible matrix C of size R. we have UV = VV operation into block block-diagonal matrices.
BlockShuffll fe can also be included by fusing the shuffle

---

of saddle points and generally less smooth loss function as pointed out in [18, 19]. We hypothesize that this makes poorer learning dynamics of the structured parametrization. Empirically, we found that the deep linear form U ((vx) a:) is more difficult to train than the standard linear layer. For example, in Fig. 3, it can suffer from training instability and loss spikes with large learning rate, while also converging much more slowly than the dense form with a small learning rate. We further elaborate on this, highlighting how the inconsistency of gradient updates between the structured parametrization and original linear projection affect learning dynamies in Appendix B.2.
Leaelook (8-228)
Lentanl (8-178]
Denise
Thai LoGa
3000
2500
5000
500
2000
2100
3009
Training Shaps
Training Steps
(a) Loss spikes
(b) Slow convergence
Figure 3: Poor training dynamics: Training dynamics of LowRank with rank of 128 under different training configurations. Curves correspond to 4-layer Transformer with a model width of 768 on WikiText- 103. We apply self-guided training in the first half of training. Refer to Sec. B.1 for more training dynamics visualizations of the other two structured parameterizations.
Self- seff-guide training Addressing the poor training dynamics by carefully tuning the learning rate schedule and gradient clipping coefficient might be possible, but it is costly and may switch between slow convergence and training instability. We propose a less costly and simple approach that can be used with minimal re-tuning of hyperparameters.
To motivate our proposal, we start by finding that the updates on UV scales the function of the backpropagated gradients (see App. B.2), then turn into the typical training dynamics with gradients. An issue that needs to be addressed in the early stage of training is feature specialization when the learning process assigns semantics to the different hidden units of the model, sometimes also phrased as identifying the winning ticket [20]. In this process, certain weights will need to be suppressed, and symmetries in the parametrization must be resolved.
To address this problem, we propose using the dense parametrization W a crutch to efficiently make decisions about feature specialization and then transfer them to U and V through g. To this
end, we use the following parametrization
o=a o=a-Wat(1-a)-un
(1)
is the layer's output, and a decays following a cosine scheduler. As residual component, learning W is unaffected by the additional saddles and pathologies introduced by the structured parametrization, allowing units to specialize. This guides the training of U and V. which are forced slowly to take over by providing the hidden units semantics learned by W. This approach also relates to homotopy methods such as simulated annealing, where hard optimization problem is transformed into an easier to optimize form with certain desired properties, gradually transforming the problem to its original form. Here, we consider the easier optimization problem is to train with dense matrices. By decreasing alpha from 1 to O, we transform the easier-to-optimize loss into the original parametrization we want to use.
Furthermore, we initialize Wo = Vo. making the observation that by using this initialization for the dense residual branch, we can casily start the guided training at any stage of learning (c.g.. fine-tuning) without affecting the model behavior. We refer to this as self-guiding training, as the better learning dynamics from W. which initially is just UV, guide the learning of U and V through the backpropagated gradients g.
Guided by the dense weights, which do not have the symmetry problem, it becomes much easier for the structured matrices to learn a good representation. From Fig. 3, it can be observed that the self-guided training prevents training spikes and fastens the convergence process. Notably, it benefits all three structured matrices with improved training dynamics illustrated Sec. B. and better final performance shown in Sec. 4.4.
5

---

Reducing the computational cost of self-guided training: Note that while > 0, we need to perform forward and backward passes through the dense version of the weight matrix W, which could be expensive. To address this issue, we consider the following stochastic version of the above formula, which allows us to control how often we need to use the dense residual branch:
Wa (1e))()))U(V UV),))
p<a p2a.
(2)
In our practice, p is a random variable sampled independently from a uniform distribution over O to 1 in each training forward pass. With a following cosine annealing schedule, this softer version Eq. (2) reduces the expected computation to half of Eq. (1). For example, using our method for half the training time increases the FLOPs by only 25% of the original FFN. This has negligible impact on accuracy. More ablation studies of this technique are presented in Sec. B.3.
3 Related work
Efficient techniques for training LLMs Recent advancements in attention mechanisms have significant improvements in the efficiency of attention [21-23, 2, 24-291, and the focus has shifted towards improving FFNs in LLMs, which contributes to at least half of the training time. Dynamic architectures such as mixtures of experts [30-32], or optimizers with faster convergence [33, 34) have been popular in improving training efficiency. Moreover, Xia et al. [35] employs structured pruning with learned sparsity masks and a dedicated data-loading policy to reduce the training budget.
There has been a recent focus on parameter-efficient fine-tuning methods like LoRA [11] and structured approximation of the linear layers (see, [10]). LoRA uses the low-rank approximation to reduce trainable parameters during the finetuning phase, whereas Sharma et al. [10] selectively applies low-rank decomposition to well-trained weights. While these methods used low-rank approximation of the weights, they did not focus pre- pre-training.
Structured matrices in deep learning Researchers use structured matrices in the form of dense matrices with shared parameters, like Circulant and Toeplitz [36]", and structured matrices, such as low-rank and diagonal, to reduce parameters and FLOPs while optimizing CUDA kernel use. Low-rank matrices. initially used in convolutional networks [37], have shown high efficiency in training [38], achieving up to a 2.9x speed-up with similar performance. Some studies [39, 40] adapt the rank during training and suggest regularizers to maintain SVD decomposition for better accuracy. Khodak et al. [41] propose spectral initialization and aligning weight decay of matrix products with standard linear layers. However, these studies mainly focus on ResNets [42] rather than recent LLMss. There have been other studies that aim to improve the expressiveness of structured matrices. For instance, Moczulski al. [43] uses interleaved diagonal and Fourier transforms, while Dao et al. [44] proposes butterfly parametrization for various transformations. These approaches often lack efficiency due to additional layers irregular sparsity patterns. Dao et al. [13] simplified butterfly matrices to block-diagonal ones, achieving a 2 speed-up on WikiTText-10% language modeling tasks. In this work, for accuracy and efficiency, we explored lowkrank factorization of weight matrices with reduced bottleneck dimension and block-diagonal matrices to reduce parameters in our LLM training studies.
4 Experiments
In our experiments, we empirically analyzed the performance of scaling, efficiency, and self-guided training for structured parameterization in LLMs.
4.1 Settings
Model We perform the experiments on the standard Transformer architecture [45, 4] equipped with rotary positional embeddings [46] and the Llama Tokenizer [47]. Its FFN module is composed of two linear layers and a GeLU activation. Four sizes are considered, including Transformer-s (110M), Transformer-m (335M), Transformer-1 (729M), and Transformer-xl (1.3B). For our efficient
We ran preliminary experiments using Circulant, Toeplitz matrices, and convolutions to improve the efficiency of FFNs, but our initial results were negative (slow and worse performance) and we did not pursue this direction further.

---

parameterizations, we only make the FFN module structured in most experiments to simplify our study, as the attention module has been well-studied [2. 23]. We explore two sizes that retain 63% or 32% of the dense FFN parameters by adjusting the rank and number of blocks (e.g., using rank half the FFN width LowRank reduces the parameters to 63%). In particular, to provide more comparative results with dense models in the scaling study Sec. 4.2, we further design the wide and structured networks, where both the attention and FFN modules are made structured using [2] and the structured matrices. This is because allocating more parameters to the FFN compared to the attention module is more favorable, and making them both structured helps maintain the parameter ratio between them. Detailed configurations can be found in Table 10.
For implementation, we take Dao et al. [13]'s implementation for the BlockShuffle and carefully manage memory copies for BlockDense. In our experiments, we chose B as a common divisor of M and N or R. Proper initialization is also investigated in Sec. A.I.
Training We the Wef dataset [48] and randomly select 0.5B tokens for validation, reserving the rest for training. All experiments, except for the overtraining experiments on 300B tokens in Fig. 5, are based on the Chinchilla scaling law [1]. where tokens are allocated at 20 times the baseline model size. We set hyperparameters such as learning rates and global batch size according to the scaling law studies from recent papers [49, 50]. However, for the 300B token experiments, we found that more advanced hyperparameter settings are necessary. For example, we use betas of [0.9, 0.98]. Additionally, different works tend to use very different learning rates [51, 49, 52, 47] in the overtraining regime. Thus, we follow the scaling law of hyperparameters described in Bi et al. [53] to avoid extensive hyperparameter searches. Other implementations include using A100 80G GPUs, mixed precision (bf (bfloat16 and float32), and adopting fast CUDA kernels like Flash Attention [25] for all experiments. We measure training FLOPss as in Megatron [54], including all matrix multiplications. Additional details are provided in Appendix C.
4.2 Scaling analysis
We evaluate the scaling performance of structured linear parameterizations from two perspectives. The first study investigates the scaling law of training-time compute. The second study trains these models with 300B tokens and evaluates their performance on downstream tasks. The results show that our structured matrices can serve as a strong alternative to the dense FFN, utilizing training FL.OPs more efficiently (e.g., smaller model and lower loss) and performing better in the overtraining regime.
BleckDemse FTN param))
Slofl (31%
BleckDense (63% paramo)
SlockShullle (63% params)
Dere (eptimal vealing law)
Donse (spilmal sealing low)
Validation Loss
Validation Loss
3.0
3.0
2.5
2.5
10"
10"
10"
18
10 24
Training Flops
Training Flops
(a) BlockDense
(b) BlockShuffle
Figure Scaling curves of structured matrices with linear fit for better illustration. The dense model is trained its optimal trade-off while we train structured FFNs on same number of tokens and retain 63% or 32% of the original parameters. 1) Structured matrices have steeper scaling curves with much closer results at larger sizes, showing gond scaling behavior. 2) With the same training FL.OPs, these curves indicate that structured matrices can have fewer parameters and lower validation loss when the x-axis is further extended.
Scaling law study: better training FLOPs utilization Based on the Chinchilla scaling law, we train four sizes Transformer models and then use the same amount of tokens to train the structured alternatives. First, we only make the FFN module structured to build the basic understanding, and retain 63% or 32% of the original parameters. scaling points for better illustration and show that all three structured matrices have steeper scaling
Fig. 4 and Fig. 1, we apply linear fit to the
7

---

curves compared to the dense models. indicating the significant potential of highly structured large models. More importantly. by fixing the training FLOPs, they have fewer parameters and eventually achieve very close or even slightly lower loss (e.g., LowRank with 63% parameters). Given their steeper scaling curves, we can also expect noticeably lower loss and fewer parameters for structured parameterizations per FLOP when the x-axis is further extended. Detailed numbers are provided in Table in the appendix, with comparisons among the three structured parametcrizations.
Next, the attention module is also structured using GQA [2], resulting in wide and structured networks. This further optimizes the use of training FLOPs, addressing the imbalance caused by structuring only the FFN module, which increases the relative impact of the attention module on the overall architecture. We adopt LowRank an example. demonstrates superior performance compared to the other two approaches in our settings, as demonstrated in Table 6 and Fig. 4. To align the training FLOPP,, the wide and structured networks are trained on larger number of tokens. It can be observed in Table that these models achieve lower perplexity while using much fewer parameters. For instance, the parameter count can be reduced from 729M 464M without compromising perplexity. Additionally, in terms of maximum throughput, ours models achieve an 8% and 17% boost on Transformer-m and Transformer-l, respectively, compared to the fast GQA.
In conclusion, the structured matrices and the wide and structured networks demonstrate great potential optimizing training FLOP utilization, achieving lower loss with fewer parameters. Additionally, it is important to note that our scaling curves for the structured matrices are not drawn their optimal craning-comppue trade-off, while the baseline is.
Scaling model size: better downstream performance To further illustrate the potential of structured matrices, we consider the overtraining regime and use LowRank as an example. Specifically, train four sizes of the dense model on 300B tokens, and build the wide and structured networks upon the design of the dense models by applying LowRank to the FFN and reducing the number of attention heads make the entire network structured. Then, the well-trained models are evaluated on downstream tasks, including PIQA, HellaSwag, Winogrande, and ARC tasks, using lm-eval eee-evallase verallaaion--a harnnesss with the default prompt. Fig. 5 presents the results, displaying the scaling trend across the four tasks (see detailed numbers and additional tasks in Table 8). The wide and structured models demonstrate comparable or superior performance, particularly larger sizes, solidifying their benefits over dense architectures.
60
44
56
Accuracy
52
36
20
32
250
500
ARC (challenge)
ARC (easy)
HellaSwag
PIQA
Model size (M)
Figure 5: Zero-shot performance downstream tasks in the overtraining regime. The wide and structured networks built upon dense ones by applying LowRank the FPN and reducing the number of attention heads make the entire network structured.
4.3 Efficiency study
We investigate the efficiency of structured FFN and consider different numbers of tokens T to discuss different scenarios. Here, I' corresponds to the total number of tokens a batch.

---

99
Dene
Table 2: Training time of Transformer-X/ and structured counterparts with 32% and 63% FFN parameters.
Len (00% parane)
Levilaki (2% panseo
Bar Litee 180% pereee)
Bloklieee pereesl
Latency ((0)
permee
Bhddkkafl parann)
Params. (M)
Training time (h)
Model
Transormer>kl
1274
352.2
LowBank
985
302.2
12.86
Width
BlockDenze
955
298.7
12.97
BlockShurrle
985
330.6
12.98
Figure 6: Latency of structured and dense FFNs across different FFN widths. Results are evaluated on 30000 tokens. The inter- mediate size the FFN is be times the FFN width.
LowRank
744
260.2
13.55
BlockDense
728
261.2
13.74
BlockShurr1e
744
284.9
13.81
Large number of tokens Using large T, the standard linear layers and our efficient structured parametrizations become computation-bound where FLOPs become a latency bottleneck [16]. This setting mainly concerns training, the prefill phase of inference, and extensive offline decoding. In Fig. 6, we evaluate the latency performance of structured and dense FFNs across different FFN widths with 30K tokens. With parameters and FLOPs reduced 63% or 32%, the lowrank and BlockDense achieve 1.4x or a 2.5x speed-up, respectively. BlockShuffle offers modest improvements, with 1.1 and 2.0x speed-ups for the two cases. We also measure the training time of the whole model in Table 2, and observe that LowRank with 63% FFN parameters reduces the training time by about 15% with 0.4 increased perplexity, and the one with 32% FFN parameters offers 1.35x x whole training speed-up with 1.1 increased perplexity.
(a) Width-2048
(b) Width-3072
(c) Width=5120
(d) Width-6144
Lovileak
Latency (ms))
0P 40
63
27.0
Batch size
Figure 7: Latency different batch size for different widths: Decoding latency results between dense FFN and structured matrices with 32% FFN parameters across different widths and batch sizes, Note that we have sequence length of at the decoding phase; thus, equals batch size.
Small number of tokens FFN can be parallelism- bound with small T (e.g., T 16) on the A100 GPUs. Then, when 7 gets increased, FFN becomes memory-bound and will eventually be computation-boune Online and offline decoding stages may encounter small number of tokens when unrolling the model step by step. As discussed earlier, our pre-merge method can alleviate the maallelism-bound issue and maintain the same latency with dense matrices. Fig. 7 shows the latency results for three different widths, varying the batch of tokens to determine when to use efficient alterna- tives choose pre-merged dense matrices. For example, with 2048-width FFN, it is difficult to fully utilize resources on GPU with limited tokens. The performance improves significantly when using width 5120 and 6144, such speed improvements of 2.63x speed-up of LowRank with 32% FFN parameters on = 2048 and 2.81 acceleration of BlockDense with 32% parameters on T= 1536.
4.4 Self-guided training
We apply self-guided training during the first half of training to demonstrate its effectiveness. As shown in Table 3 and Table 9, our method consistently reduces loss across all efficient parametrizations, improving the perplexity by 1.2 for Transformer-s and 0.8 for Transformer-m. Then, to enable straightforward comparison under the same training FL.OPs, we adjust the training steps for self- self-guided training and repeat those tokens the end to ensure they're fully learned by structured matrices. As be seen in Table 9 and Fig. 14a, Fig. 14b, Fig. 8, this reduces the perplexity

---

gap for Transformer-xl from 1.0, 1.2. and 1.3 to 0.4, 0.5, and 0.6 for LowRank, BlockDense, and BlockShuffle, respectively, under the same training FLOPs and can still enjoy 32% model FLOPs, which can bring about 2.6x inference speed-up. Additionally, we compare our method with another advanced baseline that trains structured matrices with more tokens, showing that the self self-guided training can achieve comparable or superior results even with the same number of tokens.
Table 3: Performance the three structured parame- terizations when applying self- guided training * in
Structured Matris (more tokens)
the first half of training. This increases 25% FFN
Perplexity
Stractured Matrix (self-guided training)
training FLOPs. For more comparisons, please refer Sec. B.3.
20
Architecture
FFN
Training FLOPss
PPL
12
Transformer-m
201M
1.55c+19
18.29
8-
LowRank
101c+19
20.60
5
10
15
20
25
30
35
40
45
69M
Tokens (B)
LowRank
1.1c+199
19.90
Figure 8: Comparisons between dense and structured
BlockDense
65M
1900+199
20.85
FFNs with 32% parameters under the same training
BlockDense*
1.19c+19
20.10
FLOPs. Structured FFNs are trained either with more
BlockShuffle
69M
100c+19
21.12
tokens or through self-guided training to match training
BlockShuffle
1.21e+19
20.36
FLOPs. The circle size represents model FLOPs.
5
Conclusion
In this paper, we conducted extensive experiments investigating the use structured matrices to parameterize FFNs in Transformers, with models scaling up to 1.3B parameters on the Ref dataset. Our primary aim was not to determine which structured matrix performs best, as this be task-dependent, but explore their common properties, including scaling, efficiency, and optimization challenges. We found that all of them exhibit steeper scaling curves compared to dense models. Moreover, our proposed methods, such as self-guided training, can enhance the performance across all structured matrices (c.g.. LowRank with the novel training strategy achieves a 1.35x inference speed-up with only 0.4 increase in perplexity). To conclude, we demonstrate that structured matrices can be strong candidates to replace the dense models in architecture design by scaling studies and also reveal the challenges of applying them.
BlockDense and BlockShuffle are more complicated than LowRank. In this work, only explored limited range of hyperparameter settings of them. However, since these ap- proaches are new, we believe that further performance improvements may be possible by better tuning their hyperparameters. We primarily focused on language modeling with limited vision experiments included in the appendix. Additionally, we did not explore the optimal scaling laws for structured matrices, which may further enhance performance. We also didn't 't investigate models in this paper that are comparable to today's 's practical LLMss such as LLaMA-3. This is not only because of the limited computing resources but also because this study to start investigating structured parameterizations of linear layers in modern LLM architecture training. We hope our findings and solutions about scaling, efficiency, and optimization will push their usage on the industry side and in future work.
Limitations:
Acknowledgments and Disclosure of Funding
We are grateful to Soham De for the insightful discussions and his valuable feedback on our work. We also sincerely thank the anonymous reviewers for their meaningful reviews and suggestions to make this better. We are also grateful to the RCP and IC cluster system administrators at EPFL for their support and assistance, especially during the project deadline. We also thank nimble. ai for their generous gift to CLAIRE lab, which also helped us to fund some of this research.
References
[1] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv: 15556, 2022.
10

---

[2] Joshua Ainslie, James Lee- Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized qulti-query transformer models from multi-head checkpoints. arXiv preprini arXiv:2305. 13345,, 2023.
[3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[4] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)99, 2019.
[5] 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.
[6] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092280 2023.
[7] Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, al. Using deepspeed and megatron train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201. 2022.
[8] Meta. Llama 3.
2024.
[9] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.
[10] Pratyusha Sharma, Jordan Ash, and Dipendra Misra. The truth in there: Improving reasoning in language models with layer-selective rank reduction. arXiv preprint trXiv:2312.135550 2023.
[11] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv::206v.... 00660065000..,. 2021.
[12] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[13] Tri Dao, Beidi Chen, Nimit S 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. In International Conference on Machine Learning, pages 4690 4721. PMLR, 2022.
[14] Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient convolutional neural network for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6848-6856, 2018.
[15] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chinn Chich Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings ofe on computer vision and pattern recognition, pages 4510-4520, 2018
the IEEE conference
[16] Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications the ACM, 52(4):65-76, 2009.
[17] Ke Hong, Guohao Dai. Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, and Yu Wang. Flashdecoding+4: Faster large language model inference on gpus. arXiv preprint arXiv:231 1.01222, 2023.
[18] Andrew M. Saxe, James L.. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamies of learning in deep linear neural networks. CORR, abs/1312.6120,, 2013. URL https://ag..emaaticccccoo semanti wwrg/CorpuSI::111 17272265.
11

---

[19] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
[20] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In ICLR, 2019.
[21] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint urXiv:2006.04766 2020.
[22] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 17283-17297, 2020.
[23] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arxiv::911..2155022019 2019.
[24] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium Operating Systems Principles, pages 61 611-626, 2023.
[25] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-effiiien exact attention with io-awareness. Advances in Neural Information Processing Systems, 00666334444999 2022.
[26] Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-deociding, 2023. URL http::://cff.ssttiie// httpe/////////////////eeeeeeeee
.html.
[27] James Lee-Thorp, Joshua Ainslie, llya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint 1xxiv:2005008824 2021.
[28] Krzysztof Choromanski, Han Lin, Haoxian Chen, Tianyi Zhang, Arijit Sehanobish, Valerii Likhosherstov, Jack Parker-Holder Tamas Sarlos, Adrian Weller, and Thomas Weingarten. From block-toeplitz matrices to differential equations on graphs: towards general theory for scalable masked transformers. In International Conference on Machine Learning, pages 3962-3983. PMLR, 2022.
[29] Shengjie Luo, Shanda Li. Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie- Yin Liu. Stable, fast and accurate: Kemeliied attention with relative positional encoding. Advances in Neural Information Processing Systems, 34:22795- 2021.
[30] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of-exppet layer. arXiv preprint 1eXiiv1701.06:: 2017.
[31] 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(1):5232-5270, 2022.
[32] Peter Belcak and Roger Wattenhofer. Fast feedforward networks. arXiv:2308. 14711, 2023.
arXiv preprint
[33] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv prepripe arXiv:2305. 14342, 2023.
[34] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficien 1Im training by gradient low-rank projection. arXiv preprint arXiv:2403.035500 2024.
[35] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llaa: Accelerating language model pre-training via structured pruning. arXiv preprint
2023.
12

---

[36] Yu Cheng, Felix X Yu, Rogerio S Feris, Sanjiv Kumar. Alok Choudhary. and Shi-Fu Chang. An exploration of parameter redundancy in deep networks with circulant projections. In Proceedings of the IEEF international conference on computer vision, pages 2857-2865, 2015.
[37] Misha Denil, Babak Shakibi, Laurent Dinh, Marc' Aurelio Ranzato, and Nando De Freitas. Predicting parameters deep learning. Advances neural information processing systems, 26, 2013.
[38] Cheng Tai, Tong Xiao, Yi Zhang, Xiaogang Wang, al. Convolutional neural networks with low-rank regularization. arXiv preprint 151.:55111000... 2015.
[39] Huanrui Yang, Minxue Tang, Wei Wen, Feng Yan, Daniel Hu, Ang Li, Hai Li, and Yiran Chen. Learning low-rank deep neural networks via singular vector orthogonality regularization and singular value sparsification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pages 678-679, 2020.
[40] Yuhui Xu, Yuxi Li, Shuai Zhang, Wei Wen, Botao Wang, Yingyong Qi, Yiran Chen, Weiyao Lin, and Hongkai Xiong. Trp: Trained rank pruning for efficient deep neural networks. arXiv preprint arXiv:2004. 14566, 2020.
[41] Mikhail Khodak, Neil Tenenholtz, Lester Mackey, and Nicolo Fusi. Initialization and regularization of factorized neural layers. arXiv preprint 2r 2021.
[42] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR), June 2016.
[43] Marcin Moczulski, Misha Denil, Jeremy Appleyard, and Nando de Freitas. Acde: A structured efficient linear layer. arXiv preprint arXiv:1511.0594 2015.
[44] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference machine learning, pages 1517-1527. PMLR, 2019.
[45] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[46] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 127033, 2024.
[47] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprini arXiv:2302.
2023.
[48] Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon Ilm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.011 011660 2023.
[49] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint 2rXiv:2312.00775 2023.
[50] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01066 2022.
[51] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint 2405::2405.04517, 2024.
[52] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13008 2023.
13

---

[53] Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.0200 2024.
[54] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale scale language model training gpu clusters using megatron-lm. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
[55] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. pmlr, 2015.
[56] NVIDIA. Transformerengine. http://gthhccoom///////////tttao 2022.
14

---

Appendix
A Structured matrices
A.1 Design choice
For BlockDense, we also investigate the reverse order of two projections, placing the low-rank dense matrix first, followed by the block-diagonal matrix. However, this change surprisingly yields worse performance. For instance, on the WefinedWeb dataset, perplexity increases from 29.17 29.65 with Transformer-s and 2.2B training tokens. In the case of BlockShuffle, unlike Dao et al. [13]. Zhang et al. [14] does not include second shuffle operation to restore the original order. Taking this into account, also experimented with removing the second shuffle operation and found almost no impact on performance. For example. with Transformer-s and -m using 32% FFN parameters, BlockShuffle without the second shuffle achieves perplexities of 29.89 and 21.19, respectively. compared to 29.95 and 21.12 12 for our adopted version. Nonetheless, we maintain the design from Dao et [13] for consistency.
For initialization, we follow the spectral initialization for LowRank, as suggested by prior works [41]. For BlockDense and BlockShuffle, motivated by Eq. (4), we propose using orthonormal initialization, setting the singular values of U,U, and ViV, to at the start. Experimentally. this stabilizes training dynamics and improves the perplexity performance (Table 4). For weight decay, we tried the Frobenius decay proposed by Khodak et al. [41]; however, it did not have a clear benefit our experiments and increased training FLOPs slightly. Hence, we adopted standard weight decay for all the structured FFNs.
Table 5: Ablation study of self gelldgguied training on LowRank trained on RefinedWeh.
Table 4: Different initialization of BlockShuffle and BlockDense, where random indicates random Gaussian ini- tialization and orthonormal indicates orthonormal initializa- tion. Data points are measured on the 4-layer Transformer and WikiText- 103 with learning rate of 1.0e-3.
Method
PPL
*$ ssize
25.97
Direct decomposition
28.56
Method
Initialization
PPL
Progressive decreasing rank
29.35
Transformer (4-layers)
random
23.24
Self-guided training
28.02
Self-guided training (slower)
27.90
BlockShuffle (B=4)
random
27.24
BlockShuffle (B=4)
orthonormal
25.33
-m size
18.29
BlockDense (B=2, R=128)
random
28.25
Self-guided training
19.90
BlockDense (B=2, R=128)
orthonormal
26.63
Self-guide training (slower)
19.81
A.2 Results on Refinedweb dataset
For the scaling points in Fig. 4, we provide detailed results in Table 6 for easier comparison. First, all efficient parameterizations approach the baseline as model size increases. For instance, LowRank with 32% of the parameters has a loss gap of 0.08 to Transformer-xl, whereas the gap is 0.12 at the scale of Transformer-s. Moreover, LowRank and BlockDense, with 63% the parameters in FFN, increase the loss by only 0.02 to 0.04 while reducing total training time by approximately 15% on Transformer-At and Transformer-L. Additionally, they accelcrate training by 1.35x with only 1-point increase in perplexity on Transformer-Xi with 32% of the parameters.
Although the main focus of the paper showcase their general properties-includingg scaling, efficiency, and optimization- wwe still provide comparisons in the appendix. From Table by controlling the training FLOPs and model size to be the same, LowRank and BlockDense demonstrate better performance than BlockShuffle in our main experiments, showing a 0.8 lower perplexity on Transformer-s and a 0.4 lower perplexity on Transformer-m. We think that for FFNs in language models, BlockShuffle may not be the optimal choice. However, we further compare these approaches on CIFAR- 10 dataset, showing that block-diagonal matrices serve good inductive bias in vision tasks Sec. A.3.
to compare different structured matrices but to
15

---

A.3 Results on CIFAR-10 dataset
Although in main experiments, the BlockShuffle performs the worst with two block-diagonal matrices, we provide experiments on the CIFARIO dataset here, showing that when locality is highly preferred, dickonia matrices may perform better than low-rank matrices.
In Table 7, experiments are conducted on 5-layer MLP (MultiLayer Perceptron) and VIT models. The 5-layer MLP consists of linear layer, batch normalization [55], and the ReLU activation function with a hidden dimension of 384. It is trained for 500 epochs with a learning rate of 1.0e-3 and a batch size of 128. For the VIT models, 12 layers with hidden dimension of 384 are used, and they are trained for 750 epochs with a learning rate of 6.0e-4 and a batch size of 512.
Since the first layer in vision tasks typically prefers from the locality, especially in 5-layer MLP where the image pixels are directly flattened into the input, we conducted experiments with and without structuring the first layer as LowRank and BlockShuffle. Both sets of controls, particularly when structuring the first linear layer, demonstrate that block-diagonal matrices can be beneficial for vision tasks. Specifically, replacing the first layer of the 5-layer MLP model with block-diagona matrix even yields better performance, as the block structure effectively groups neighboring pixels, compensating for the MLP's lack of locality. However, applying structured FFNs to the first layer of ViT can lead to significant accuracy degradation, reinforcing our decision not to use structured FFNs in the first layer in the main experiments.
Table 6: Performance of two sizes of different structured matrices with 63% and 32% of the original FFN module's parameters. We report model size, total FFN size, training tokens, training FLOPs, and training time. Note that the total structured FFN is not exactly 63% of the original because we don't I replace the first FFN module. Also, the BlockDense is slightly smaller for -m and -xl models to ensure the rank is multiple of 256 when matching parameters. Loss and perplexity evaluated 0.5B token validation set.
Model
FFN
Training
Architecture
Loss
PPL
Size (M)
Size (M)
Tokens (B)
FLOPs
Time (h)
Transformer-s
110
56.62
2.2
1.69c+18
4.0
3.2569
25.97
LowRank (R384)
90.17
37.16
2.2
1,44++18
3,3017
27.16
63%
BlockDense (B2R512)
90.17
37.16
2.2
1,44++18
3.8
3,3034
27.20
BlockShuffle (B2)
90.17
37.16
2.2
1.4e+188
4.2
3,3191
27.63
LovRank (R192)
73.95
20.94
2.2
1.22c+18
3.3748
29.22
32%
BlockDense (R2R256)
73.95
20.94
2.2
1.22e+18
3.5
3.3731
29.17
BlockShuffle (B4)
73.95
20.94
1,22c+18
3,3994
29.95
Transformer-m
335.08
201.33
6.7
1.55c+19
32.5
2.9062
18.29
LowRank (R512)
262.73
128.9 97
6.7
1.26++19
29.6
2,9508
19.12
63%
BlockDense (B4R768)
295.19
121.44
6.7
1.23c+19
29.9
2.9581
19.26
BlockShuffle fle (B2)
262.73
128.97
6.7
1.26c+19
33.1
2.9622
19.34
LowRank (R256)
202.43
68.68
1.0c+119
3,0251
32%
BlockDense (B4R384)
198.67
64.91
1.00c+19
3.0371
BlockShuffle (B4)
202.43
68.68
6.7
1.0lc+19
30.0
3.0501
21.12
Transformer-l
729.11
452.98
14.6
7.03e+19
130.5
2.6594
14.29
LowRank (R768)
566.32
290.19
14.6
5.61e+19
113.6
2,6957
14.82
63%
BlockDense (B2R1024)
566.32
290.19
14.6
5.6le+19
114.3
2.7038
14.94
BlockShuff fle (B2)
566.32
290.19
14.6
5.6le+19
124.3
2.7021
14.91
LowRank (R384)
430.66
154.53
14.6
4.42c+19
100
2.7527
32%
BlockDense (B2R512)
430.66
154.53
14.6
4,42c+19
100.9
2.7570
15.75
BlockShuffle (B4)
430.66
154.53
14.6
4.42e+19
110.3
2.7735
16.01
Transformer-xl
1274.14
805.31
25.5
2.10++20
352.2
2.5226
12.46
LovBank (R 1024)
984.73
515,90
25.5
1.66c+20
302.2
2.5541
12.86
63%
BlockDense (B4R1536)
954.59
485.75
25.5
1.61c+20
298.7
2.5628
12.97
BlockShuffle (B2)
984.73
515.99
25.5
1.66c+20
330.6
2.5633
12.98
LowRank (R512)
743.56
274.73
25.5
1.29e+20
260.2
2.6062
13.55
32%
BlockDense (B4R768)
728.49
259.65
25.5
1.27c+20
261.2
2.6204
13.74
BlockShuffle (B4)
743.56
274.73
25.5
1.29c+20
284.9
2.6254
13.81
16

---

Table 7: Experiments CIFAR 10 and vision models, where the locality is highly preferred. The first layer column the tahle indicates whether apply structured matrices to the first FFN.
Structured first FFN
Dense first FFN
Method
Model size (M)
Accuracy
Model (M)
Accuracy
5-layer MLP (H=768)
4.14
66.99
4.14
66,99
LouRank (R=192)
1.63
64.04
3.26
65,42
BlockShuffle (B=4)
1.63
67.08
3.26
65.67
VIT (H=384)
21.34
92.49
21.34
92.49
LouRank (R=24)
8.29
89.56
BlockShurfle (B=16)
8.29
90.42
9.38
92.49
A.4 Results on downstream tasks
Table Performance on downstream tasks under the zero-shot setting. We report perplexity performance of the validation set of RefinedWeh. For all downstream tasks except LAMBADA, we report accuracy results. For LAMBADA, we present the results in an accuracy/perplexity format. Implementation details are put in Appendix C.
Model
Model Size
RefinedWeb
ARC (challeng
(casy))))
HellaSwag
LAMBADA
PIQA
size
Dense
110.0
16,02
18.69
44.65
31.79
36.35/28 86
65.83
17.30
18.17
30.83
34.23/35.37
DDnse
353.I
12.34
Σ2.10
30.80
38.60
46.65713.92
70.75
Wide and Structured
219.4
13.38
20.90
49.12
36.32
42,46/17.60
69.75
ssize
Dense
729.I
10.76
25705
33..1
4353
$2.30993
13.of
Wide and Structured
464.4
11.61
23.55
$4.76
40.53
48.79/11.67
72.14
Supplementary to Fig. 5, Table presents the detailed results of wide and structured networks compared to dense models on downstream tasks. These models were trained on 300B tokens and implementation details can be found in Appendix C.
B Self-guided training
B.1 Visualization of training dynamics
We provide additional training dynamics curves, including BlockDense and BlockShuffle, in Fig. 10 and Fig. 11, illustrating that these structures are more challenging to train compared to standard linear layers. Specifically, they exhibit greater sensitivity to learning rates and more prone to loss spikes. To mitigate this, we apply self-guides training in the most challenging cases, which results in improved training dynamies, such faster convergence without loss spikes.
Furthermore, we report the loss spikes observed in the large Transformer-xl model trained on RefinedWeb dataset in Fig. 12.
B.2 Explanation of the poor training dynamics
We observed that loss spikes occur along with large gradient norms. This motivates to analyze gradient updates during the backward pass of the linear UV. Considering g as the gradient of output and the input, the standard linear layer W gradient update is gæ For the structured parametrization, the gradicnts of and V are gæ V and U gæ respectivell. Then, for W as being the updated parameters, will have the updates for W be:
AW W W = -r.gx.
(3)
And the one for UV:
-UUV=-Ir-y ga +gx V V +O)
(4)
17

---

Thus, it can be seen from Eq. (4) that the projections UT and V V can disrupt the gradient gæ T. If the norms of U and V are small, the new update vanishes faster than the original update, and in reverse, if their norm is large, the update blows up, leading to unstable training.
To be specific, we calculate the spectral norm of UUT and V V and use this indicate the maximum scale the matrix can stretch a vector. Fig. 13 shows that the largest singular value can vary significantly, being either much greater than or less than depending on the shape of the weight (input dimension, rank, number of blocks) and magnitude. Interestingly, very structured FFNs with small ranks or many blocks tend to have smaller spectral norms, while others have larger ones. This corresponds to the phenomenon in Fig. 10 and Fig. 11, where smaller FFN are prone to slower convergence and larger ones to loss spikes.
An alternative intuitive perspective of this issue to realize that learning with structured matrices has additional redundant degrees of freedom brought by symmetry. For example, increase the norm of a feature, can increase either the corresponding weights U or V. Given that gradient descent makes this decision independently, will overshoot and make learning less well-behaved.
SAN Kork (SEBO)
Komkark (K086)
Thain Loss
Teain LOKE
2
200 300
400
500
600
700
500
900
1000
200
300
400
500
600
700
800
500
2000
200
300
400
500
600
700
800
900
1000
Steps
(a) by 1.56- 3
(b) (r 1.0c 3
(c) by 5.0c -4
Lowilanee
(201
Lowitonk! 120
Lowfonns (10120) Dease
Donne
Decas
Tha Loss
Thain LORK
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2300
3000
500
1000
1500
2000
2500
3000
(d) (r -5
(c) (c) 1.06-3
(f) 500e=.00e-
Figure 9: (a-c): Training dynamics of LowRank with different ranks and the dense model under different hyper- parameters. Data points are measured on a 4-yer Transformer with model width 768 and WikiText- 103. We zoom into the beginning of training for clearer observations. (d-f): Training dynamics of LowRank and the self- gulddd training. The self-guided training overcomes the loss spikes and makes the training faster. We show the whole training curve to indicate its success. R indicates the rank of low-rank matrices.
B.3 Design choice
Experiments in this part conducted on Transformer-s and Transformer-m with ranks 192 and 256 in Table 5, respectively. We apply self-guided training during the first half of the training process.
First, we compare stochastic self-guided training with the static version. The stochastic and faster version in both model sizes brings about a 0.1 perplexity increase while reducing computation by half. Second, other techniques compared. Dense layer decomposition, which decomposes the weight directly at the midpoint training, is examined. This approach can lead to abrupt loss increases in training curves, resulting in worse performance. Strategies incrementally reducing rank require a feasible and complex change strategy and fail to address the inconsistent gradient update problem, thus still suffering from poor results Table 5.
Generally speaking, our method stands out due to its flexibility, simplicity, and efficiency. Eq. (1) makes it adaptable to any efficient linear parametrization without special constraints, while progressive rank reduction and direct decomposition require feasible solution evolve. Our guided initialization allows its usage in various stages of training without the need for well-trained
18

---

Becoden 2283281 BiooDaee 18289321 OWEES
Bockbonee (8381200
Bectierre (8081200
BlccDDenn (839512)
Mixd Dermn (838512)
Denea
Danea
Train LaRe
200
400
G60
1000
500 600
700 200
200
200 900 1000 Training Steps
700
Training Steps
Training Steps
(a) 1.5c- 3
(b) Ir 1.0c 3
(c) Ir 5c
Bockkersn (8c2, 8-1200
Boccbbere 18s2, 8-120
BeckDaase 28-0 R-128)
Trai
500
1000
1500
2000
2500
2000
1000
1500
2000
2500
500
1000
1500
2500
2000
Training Steps
Training Stope
Training Steps
(d) Ir
1.5e- 3
(e) br= 1.0c
l
Figure 10: (a-c): Training dynamics of BlockDense with 2 blocks and different ranks and the dense model under different hyper-parameters (d-f): For BlockDense (B 22 128), training dynamics of self guided training indicated by Other settings follow Fig.
Bredshufff 1881
BockSkoffe (80)
BlockShuffe
Wasdhdallte Jest
(MC)
Blockhhlll
Denae
Denoe
Train LONE
Train
Training
Training Stops
Steps
(a) 1y 1.5e 3
(b) 1r- 1.0c- 3
(c) 5c
Weakthumf 18w0) meckthumte (8-4)
Blckkhhfffe (8-0)
SbciShaffe (D+1) BbokSbalte (8-1)
Mackthulfee (8-4)
20000
Thairing Stepo
Training
Training Steps
(d) 1r 1.5e 3
(e) 1: 1.0e- 3
(I) [y 5e 4
Figure 11: (a-c): Training dynamics of BlockShuffle with different numbers of blocks and the dense model under various hyper-parameters (d-f): For BlockShuffle (B=8), 8), loss curves of self- self-guided training indicated by & Other settings follow Fig. 9. Note that the training dynamics of different structured matrices are not comparable here because their sizes are not controlled to be the same.
teacher. It is simple because Eq. (1) provides only one smooth transition. It is efficient due to the layer-specific definition in Eq. (1) and stochastic computation in Eq. (2).
B.4 Results on Refinedweb Dataset
Table 9 provides more comprehensive results of self-guided training, supplementing Sec. 4.4. For Transformer-s and -m, we first present the results of applying self-guided training to the first half of training the second row of each table block, demonstrating its effectiveness across all three structured matrices. Additionally, by controlling for equal training FLOPs as described in Sec. 4.4,
19

---

Train
2000 4000
6000
8000 10000 12000 14000
2000 4000 6000
8000
10000 12000 14000
Training Steps
Training Steps
(a) LowRank on Transformer-xl
(b) BlockDense on Transformer-xl
Figure 12: Loss curves of Transformer-xl. Structured FFN with 32% parameters exhibits slower convergence. Also, there exist loss spikes.
Kan
Rack
Dack
(a) N=768, std=0.02
(b) N=1536, std=0.02
(c) N=1536, std=0.05
(d) N=1536, std=0.05
Lupor Oreaen Vaee
(e) N=768, std=0.02 N=1336, std=0.02
(g) N=1536, std=0.05
(h) N=1536, std=0.05
Figure 13: Spectral norm of the matrix V V. where in (a-c). V is the low-rank matrix and in (d-f), V is the block-diagone matrix. N presents the imput dimension of the weight. Std indicates the standard deviation value of the normal distribution from which we sample weight elements.
we show that with self-guided training, all three structured FFNs with 32% parameters incur only 0.4-0.6 increase perplexity compared the baseline, while benefiting from smaller memory footprint and faster speed at inference time. To further illustrate this, we plot these points with the same training FLOPs Fig. 14, and specifically for LowRank in Fig. 8, highlighting that self-guided training achieves comparable performance to training on more tokens.
28
28
Dence
Dease
24
BlockDense tokens)
tokens)
BlockDense gelf-guided training!
BlockStuufle (self guided training!
Perplesity
Perplexity
20
20
16
16
12
15
25
30
40
20
35
Tokens (B)
Tokens (B)
(a) BlockDense
(b) BlockShuffle
Figure 14: Performance between Transformer, Structured FFN (BlockDense and BlockShuffle, LowRank in Fig. 8) with 32% parameters either trained with self-guided training or more tokens across four sizes. Circle size indicates model FLOPs.
FLOPs. To enable straightforward comparison, we controlled their training FLOPs to be the same.
20

---

Table Performance of self-guided training indicated by + on Structured FFN with 32% parameters under the same training FLOPs. We also include the structured FFN trained more different tokens highly advanced baseline. Model FLOPPs are calculated sample with 1024 sequence length.
Model
FFN
Model
Training
Architecture
Loss
PPL
Size (M)
Size (M)
FLOPs (G)
Tokens (B)
FLOPs
Transformer-
110.0
56.6
262.9
2.2
1.69+188
3.2569
25.97
LouRank
1.22e+18
3.3748
29.22
LowRank
2.2
1.39++18
3.3329
28.02
74.0
20.9
189.8
LowRank
3.0
1.69c+18
3.2928
26.92
LouRank
2.2
1.69++18
3,2866
26.75
BlockDense
2.2
1.22c+18
3.3731
29.17
BlockDense*
2.2
1.39c+18
3.3338
28.04
74.0
20.9
189.8
BlockDense
1.69e+18
3.2982
27.06
BlockDense*
2.2
1.69e+18
3.2856
26.73
BlockShuffle
2,2
1.22e+18
3.3994
29.95
BlockShuffle
1.39e+18
3.3583
28.74
74.0
20.9
189.8
BlockShuffle
3.0
1.69++18
3.3218
27.71
BlockShuffle
2.2
1.69++18
3.3011
27.14
Transformer-m
335.1
201.3
788.7
6.7
1.55c+19
2.9062
18.29
LowRank
6.7
1.0l+199
3.0251
20.60
LowRank
6.7
1.21++99
2.9907
19.90
202.4
68.7
517.0
LouRank
10.2
1.54e+19
2.9359
18.84
LowRank
6.7
1.55c+19
2.9310
18.75
BlockDense
6.7
1.00+199
3,0371
20.85
BlockDense
198.7
64.9
509.3
6.7
1.19++19
3.0008
20.10
BlockDense
10.4
1.55e+19
2.9491
19.09
BlockDense
1.55e+19
2.9420
18.95
BlockShuffle
100c+19
3,0501
21.12
BlockShuffle*
6.7
1.2le+19
3,0135
20.36
202.4
68.7
$17.0
BlockShuffle
10.2
1.54++19
2.9627
19.35
BlockShuffle
6.7
1,55e+19
2.9525
19.15
Transformer-l
729.1
453.0
1646.9
14.6
7.03c+19
2.6594
14.29
LowRank LouRank
4.42e+19
2.7527
15.69
430.7
154.5
1035.6
23.3
7.03e+19
2.6917
14.76
LowRank
14.6
7.01e+19
2.6850
14.66
BlockDense BlockDense
14.6
4.42c+19
2.7570
15.75
430.7
154.5
1035.6
23.3
7.03e+19
2.6946
14.80
BlockDense
14.6
7.01e+19
2.6941
14.79
BlockShuffle
14.6
4,42c+19
2.7735
16.01
BlockShuffle
430.7
154.5
1035.6
23.3
7.03e+19
2.7053
14.96
BlockShuffle
14.6
7.01c+19
2.7104
15.04
Transformer-ll
1274.1
805.3
2814.3
25.5
2.10e+20
2.5226
12.46
LowRank LowRank
25.5
1.29e+20
2.6062
13.55
743.6
274.7
1727.7
41.5
2.10c+20
2.5464
12.76
LouRank
25.5
2.10c+20
2.5539
12.86
BlockDense BlockDense
25.5
1.27e+20
2.6204
13.74
728.5
259.7
1696.8
42.2
2.10e+20
2.5590
12.92
BlockDense*
25.5
2.10e+20
2.5637
12.98
BlockShuffle
25.5
1.29++20
2.6254
13.81
BlockShuffle
743.6
274.7
1727.7
41.5
2.10e+20
2.5623
12.97
BlockShuffle
25.5
2.10e+20
2.5678
13.03
21

---

C Implementation details
Model Architecture details are provided in Table 10. We consider four baseline Transformer sizes, ranging from 110M 1.3B parameters, with widths from 768 to 2048. For structured models, we first adopt two configurations as described in the main paper, reducing the FFN module to 63% and 32% of its original parameters by adjusting the rank R. and number of blocks B. Only R and B values specifically associated with structured matrices are modified, as seen in Table Then, for more comparable results, we consider wide and structured networks, where the attention module is also structured by reducing the attention heads. We also present the transformer with GQA version [2] here, configured with 256 dimensions for the KVCache [9] and an enlarged FFN intermediate dimension following [8]. Based on this GQA version, we apply LowRank matrices the FFN module with a rank half of the model FFN width and use a smaller attention inner dimension to further reduce the parameters of the attention module. This allows us to maintain the parameter ratio between the attention and FFN modules.
Note that all experiments, we do not apply structured matrices to the first FFN module, as doing can lead non-negligible performance loss in models shallow networks. For dense models, use Gaussian random weight initialization with a standard deviation of 0.02. For structured matrices, spectral initialization is applied for LowRank, and orthonormal initialization for the other two, based on initial experiments.
Table 10: Detailed configurations of the baseline Transformers, along with those using GQA [2] and wide and structured networks. The latter which have and 32% the original FFN parameters, we adjust only the rank and number blocks for each method and put the configuration directly in Table 6. Width denotes the model width or the input and output dimensions of the attention and FFN modules. Intermediate dim. refers to the intermediate dimension of the FFN. Attention dim. specifies the dimension used in scaled-dot product attention. KV dim. represents the dimension used for KVCache, selected according to Team al. [9].
employed in the scaling study in Sec. 4.2. For other structured models
Model
Size (M)
Layers
Width
Intermediate dim.
Attention dim.
KV dim.
Transformer-s
110.0
12
768
3072
768
768
Transformer-s (GQA)
110.0
12
768
3584
768
256
Wide and Structured (R=384)
81.1
12
768
3584
512
256
Transformer-m
335.1
1024
4096
1024
1024
Transformer-m (GQA)
335.1
1024
4864
1024
256
Wide Structured (R=512)
219.4
1024
4864
512
256
Transformer-I
729.1
24
1536
6144
1536
1536
Transformer-1 (GQA)
729.1
24
1536
7424
1536
256
Wide and Structured (R=768)
464.4
24
1536
7424
768
256
Transformer-xl
1274.1
2048
8192
2048
2048
Transformer-xl (GQA)
1274.1
2048
9984
2048
256
Wide and Structured (R=1024)
799.6
2048
9984
1024
256
Dataset We use the RefinedWeb dataset [48], carefully curated subset of CommonCrawl, opti- mized for filtering and deduplication, providing 600B tokens for public use. Due to its large size, we shuffle, extract, and tokenize in advance. To manage memory efficiently, token IDs are stored using np. memmap, preventing the need to load all data into CPU memory simultaneously. The maximum sequence length set to 1024. Following scaling laws [1], we allocate tokens 20 times the number of parameters for each bascline model in all experiments, except for the 300B token training.
22

---

Table 12: Training configuration for 300B token training. Different studies [49, 52, 47, 51] employ very different learning rates in this setting, which also differ from taaning-comppl avoid extensive tuning, we follow the hyperparameter scaling rule of Trranformee proposed by Bi et al. [53), determining batch size and learning rate based on training FLOPs.
Table 11: Basic training configuration used in all experiments except for the overtraining regime with 300B tokens. Note that we apply the same global batch size (Batch) and the same peak learning rates (LR) to both dense and structured models to avoid hyperparam- eter search. The hyperparameter values selected based on Zhang et al. [50), Gu and Dao [49].
scaling studies [1]. To
Model
Tokens
Batch
LR
Model
Size (M)
Batch
LR
size Dense Structured
size
2.2B
512
6.0e-4
Dense
110.0
1280
9.1e-4
Structured
81.1
1024
9.6c-4
-m size Dense
-m size
6.7B
512
3.0e-4
Dense
335.1
1792
7.8c-4
Structured
Structured
219.4
1536
8.4e-4
-1 size Dense
size
Dense
729.1
2304
e.
Structured
14.6B
512
2.5e-4
Structured
464.4
2048
7.6e-4
-x size
-X size
Dense
Dense
1274.1
2816
6.6e-4
Structured
25.5B
512
2.0c-4
Structured
799.6
2304
7.1c-4
Training We use A100 80G GPUs for training and evaluation, employing mixed precision (bfloat16 and float32) with torch. cuda. amp to accelerate training. Training FLOPs are calculated following Megatron [54], including all matrix multiplications.
Different hyperparameters are used for 300B token training and other experiments. For basic training in training FLOPs scaling and self-guided training studies, configurations are listed in Table 11. Hyperparameter values are selected based on the scaling law studies of Zhang et al. [50], Gu and Dao [49], where we use the same learning rates and global batch size for both dense and structured models. Additional details include the AdamW optimizer with 0. weight decay, betas of [0.9, 0.9991, cosine annealing learning rate scheduler with 10% linear warm-up, and 0.1 x minimum value. Dropout is set to 0.0, and gradient clipping to 1.0.
In the overtraining regime where the training duration is super long, however, smaller betas [0.9, 0.98]
are required for stable training, even for baseline Transformers. Previous studies [49, 52, 47, 51] adopt very different learning rates in this setting, differing from training- coomringecouueeeeeeee scaling studies [1]. To avoid extensive searching, we follow the hyperparameter scaling rule of Transformer proposed by Bi et al. [53], determining the global batch size and learning rate based on training FLOPs. Specifically, batch size is defined by 0.3118 x (training FL0P5-0.125) ), and learning rate by
0. 0.2920 x (training FLOP (,2711,, giving the results in Table 12. It can be seen that our wide and
structured models trained on 300B tokens will use a slightly higher learning rate and smaller batch size compared to the larger Transformer.
Efficiency To enhance training and inference efficiency, our code is based on PyTorch but incorporates optimized CUDA kernels. We leverage Flash Attention [25], fast LayenNorm, and rotary embeddings from TransformerEngine [56]. along with fused operations including bias and GeLU, For inference speed testing, we use bfloat16. These techniques are consistently applied to all models to ensure fair latency and throughput comparisons.
D Broader Impacts
Enhancing the efficiency of Large Language Models (LLMs) can significantly reduce computational resources and energy consumption, benefiting the environment and democratizing access to advanced AI technologies. However, increased efficiency could also lead to greater dissemination of disinformation and the creation of deepfakes, posing risks to public trust and security and potentially reinforcing existing biases that impact specific groups unfairly. This research aims to promote the responsible development and deployment of L.Mss, maximizing societal benefits while acknowledging potential harms.
23

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We describe our claims and contributions in Section
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an 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
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Question:
Answer: [NA]

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix 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 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: We describe our experimental setup and implementation details in Section Appendix C.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer 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 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 novel architecture, describing the architecture fully might suffice,, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general, releasing code and data often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing 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 new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution should either be a way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset or instructions for how construct the dataset).
new model (e.g., large language model), then there
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
25

---

Answer: [Yes]
Justification:
We open-source our code. The dataset we use is publicly available.
Guidelines:
The answer NA means that the paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https:
/ //nips.co //nips.cc/publl
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https:/nnips cc/ wwcppuulcc/aacececcceeeeeeeceale httgg/gde/ooeuuubnee celbbmissiion ssisnsolicyyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only 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 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, hyperpa- rameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe our experimental setup Section and give details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Training large language models requires significant time and compute, making it infeasible to run multiple seeds. Nonetheless, results from single training runs highly reliable and transferrable, as demonstrated by neural scaling laws. Presenting results without error bars in LLM research is standard practice.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results 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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
26

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
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?
Question:
Answer: [Yes]
Justification: Computing resources are discussed throughout the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 the paper (e.g.. preliminary failed experiments that didn't t make it into the paper).
9. Code Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics pettps:/nneurs tttps://eurripp.ccppblcccc
Answer: [Yes]
Justification:
The research conducted in the paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there is a special consideration due to laws 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: Societal impacts are discussed in Appendix D.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact 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 technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
27

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not release data or models. not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions 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 best faith effort.
12. Lieenses 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 the assets used have been cited.
Guidelines:
The answer NA means that the paper does not 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 URL
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
28

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide documented code.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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 details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
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 the NeurIPS Code of Ethies, 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: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 NeurlPS 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.
29


Reviews:
[{"review_id": "ZGbyv4rPXl", "reviewer": "Reviewer_7Gik", "summary": "The paper studies efficient Transformer variants. Unlike most existing works on efficient attention, this work proposes methods to enhance the efficiency by focusing on feedforward networks (FFNs). It explores several efficient linear layer designs, and proposes techniques to address the training issues and decoding efficiency for practice. The experiments show that these structured FFNs not only reduce computational costs but also show promising scaling behavior on language modeling.", "questions": "Is there any special technique to initialize the matrices U and V? What is the scale/variance of the initialization?", "limitations": "The authors discuss the limitation of this paper in Sec. 5. I believe that another limitation is that the paper lacks evaluations on downstream tasks.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Paper is clearly written and well motivated.\n\n- While there has been voluminous literature on efficient attention using structured matrices, making FFNs efficient using structured matrices is new to the community.\n\n- This paper discuss interesting techniques to address optimization difficulty in training Transformers with structured matrices in FFNs.\n\n- The experiments covers models with different sizes and demonstrate the scaling behavior. The experiments also clearly demonstrate the efficiency gains in practice.", "weaknesses": "- The experimental study can be made more solid if the authors can additionally provide model quality analysis on downstream tasks (e.g., SuperGLUE) and the finetuning regime.\n\n- The second part of the related work section could be enhanced by including additional studies on structured matrices in Transformers, emphasizing that the majority of existing efforts focus primarily on the attention module. To name a few:\n\n  - Lee-Thorp, James, et al. \"Fnet: Mixing tokens with fourier transforms.\" arXiv preprint arXiv:2105.03824 (2021).\n\n  - Luo, Shengjie, et al. \"Stable, fast and accurate: Kernelized attention with relative positional encoding.\" Advances in Neural Information Processing Systems 34 (2021): 22795-22807.\n\n  - Choromanski, Krzysztof, et al. \"From block-Toeplitz matrices to differential equations on graphs: towards a general theory for scalable masked Transformers.\" International Conference on Machine Learning. PMLR, 2022.\n\n- Minor issues.\n  - Lines 4 & 71: transformer $\\to$ Transformer\n  - Line 168: alpha $\\to$ $\\alpha$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the valuable suggestions and careful reading of this paper. We will fix the minor issues in the revision and list the detailed responses to the questions below. Hope our reply can address the concerns.\n\n* **Q1**: The experimental study can be made more solid if the authors can additionally provide model quality analysis on downstream tasks (e.g., SuperGLUE) and the finetuning regime.\n  \n    **A1**: Thank you for your valuable suggestion. We have added the results of downstream performance in Table 1 and Table 2 in the general response PDF, as well as the fine-tuning performance in the table below. \n    * We use the lm-evaluation-harness repository for downstream tasks like PIQA and HellaSwag. To achieve good downstream performance, we train small-sized Transformers on 100B tokens. Table 1 in the general response PDF shows the consistent performance of these zero-shot tasks with the validation perplexity of training. All structured parameterizations with 32% FFN parameters have results close to the dense models (e.g., a 0.9-1.4-point averaged accuracy decrease). Additionally, we also provide the results of Transformer-XL trained based on the optimal scaling law in Table 2 in the response PDF, showing good results for structured matrices and consistent improvement with self-guided training.\n    * We applied fine-tuning to our trained models (small-sized Transformers on 100B tokens) using the famous transformer repository. As GLUE is better supported in this codebase than SuperGLUE, we evaluated two tasks : QQP and SST-2 in the GLUE benchmark within this short rebuttal period. As can be seen below, the structured parameterizations show very comparable accuracy performance to the dense model (e.g., 0.3 lower accuracy of BlockDense on QQP task) with only 32% parameters of the FFN module.\n\n       |     Method              | Validation PPL | QQP (acc) | SST-2 (acc) |\n       |----------------------|----------------|-----------|-------------|\n       | Transformer-s (110M) | 16.68 | 90.0      | 92.0  |\n       | 32% FFN params.      |   |  ||\n       |     LowRank          | 18.65 | 89.6 | **92.2** |\n       |     BlockDense       | **18.61**      | **89.7**  | 91.7  |\n       |     BlockShuffle     | 18.94          | 89.2      | 91.5 |\n \n      To conclude, the downstream and fine-tuning performance consistently shows the strong potential of these structured matrices.\n\n* **Q2**: The second part of the related work section could be enhanced by including additional studies on structured matrices in Transformers, emphasizing that the majority of existing efforts focus primarily on the attention module.\n\n  **A2**: Thanks for the good suggestion. We'll add these papers to the related works. Yes, they are all related and very interesting papers that apply structured matrices, like Toeplitz, over the sequence dimension, thus making the attention module very efficient. In contrast, in our paper, we investigated structured parameterizations in the FFN module and found that Toeplitz is not very suitable for mixing information in the hidden states dimension for NLP tasks. We also focus on their scaling, optimization, and efficiency aspects in modern LLM architecture training.\n\n* **Q3**: Is there any special technique to initialize the matrices U and V? What is the scale/variance of the initialization?\n\n  **A3** : We put the discussion about initialization in Section B.1 of the submission, where we mentioned that we used spectral initialization for LowRank and orthogonal initialization for BlockDense and BlockShuffle. Specifically, we use random Gaussian initialization with a variance of 0.02, similar to GPT-2, to initialize the original weight W. Then, as suggested by [1], U and V in LowRank parameterization are initialized by using the SVD decomposition of W. For the other two methods, we apply orthogonal initialization as recommended in the paper on deep linear networks [2]. Experiments in Section B.1 on the small dataset WikiText-103 further validated our choice.\n    \n    [1]. Initialization and regularization of factorized neural layers. ICLR'21\n   \n    [2]. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. Andrew M. Saxe, et al. 2014.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HlsGhGfEfI", "reviewer": "Reviewer_5PEQ", "summary": "This paper mainly focues on using structured matrices to substitute dense matrices in FFNs for training from scratch tasks. The authors propose BlockDense as a combination of low-ranked dense and block diagonal matrices (Figure 2), and to address the loss spike issues in low-ranked training process, the authors propose the \"self-guided training\" where we combine the efficient training with parameterized linear dense layer as $o = \\alpha W x + (1 - \\alpha) U (V x)$ to regularize / alleviate the optimization difficulties from saddle points and pathologies. \n\nThe authors conduct experiments on transformers with sizes from 110M to 1.3B and RefinedWeb dataset. The authors focus on the latency vs. batch size, width, etc. for online decoding and validation loss vs. FLOPs. The baselines are low-ranked matrices, block shuffle (Monarch decomposition) matrices, and dense matrices for FFNs training, and the authors show Block Dense in general achieve lower PPL than Block Shuffle on the same training budget (training FLOPs).", "questions": "Could you compare self-guided training with other regularization methods?\n\nCould you justify why BlockDense is a better efficient training method than low-ranked training?\n\nCould you provide more quantitative analyses on the scaling laws of block dense, low-ranked matrices, block shuffle, and dense matrices (actual slope, confidence interval, etc.)?", "limitations": "There are no other limitations. All weaknesses have been listed above.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper use structured matrices as a hardware-friendly efficient training method. Comprehensive latency benchmarking results are quite informative for deploying the BlockDense (or other methods) in practice. \n\nThis paper also provide comprehensive training-from-scratch experiments with medium-sized transformers.\n\nThe writing is clear to follow, and Figure 2 is quite helpful for understanding BlockDense parameterization.", "weaknesses": "Self-guided training is effectively adding an regularization term that combines $U$ and $V$, and we need to compare it against other regularization approaches (e.g. spectral regularization) to position the effectiveness of self-guided training with related works. Such comparison is missing from the paper (missing on both Figure 3 and Table 3) and it is hard to assess the effectiveness of self-guided training as a regularization method.  \n\nIn addition, I don't find the advantage of BlockDense over low ranked matrices. It seems that low-ranked matrices also have low latency (Figure 4) and low validation loss / perplexity on the same training budget (Table 3). The optimization challenge of low-ranked matrices can also be alleviated by self-guided training according to Figure 3. \n\nIt is also hard to identify whether the scaling law of BlockDense is necessarliy better than the dense matrices in Figure 1a. The orange / green line appear to be nearly parallel to the purple line in both $10^{18}-10^{19}$ and $\\sim 10^{20}$ segments. More quantitative analyses are needed to justify the statement on line 52 \"Interestingly, scaling curves in Fig. 1a indicate that structured matrices yield steeper loss scaling curves than the traditional Transformer at its optimal trade-off\". \n\nI currently consider the first 2 weaknesses as major (insufficient comparison and justification) and I vote for borderline reject. I am happy to raise my scores if the above concerns are addressed during the rebuttal period.\n\nEDIT: upon reading the rebuttal, I raise my score to the borderline accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for providing a constructive review and detailed comments.  \n\nBefore responding in detail, firstly we would like to clarify our paper's focus. We investigate the performance of three structured matrices in modern LLM training from efficiency, optimization, and scaling aspects and provide concrete ablations to improve our scientific understanding of these methods. As a result, 1) we explore BlockDense to cover a broader space of possible parametrization, but not claiming that it is the best candidate, and we found that it underperforms when compared to LowRank in some experimental settings; 2) the challenges and proposed solutions (e.g., self-guided training, pre-merge, and scaling curves) are found or proposed for all the parameterizations.  \n\n* **Q1**: In addition, I don't find the advantage of BlockDense over low ranked matrices. It seems that low-ranked matrices also have low latency (Figure 4) and low validation loss .... The optimization challenge of low-ranked matrices can also be alleviated by self-guided training.\n\n   **A1**: As stated earlier, this paper is not claiming that BlockDense is the best candidate but aims to investigate the performance of three structured matrices from several aspects within the FFN module of Transformers. Our findings and solutions like self-guided training also work for all of them.\n   \n  We propose BlockDense for several reasons:\n    * It is a natural intermediate between LowRank and BlockShuffle, combining low-rank projection with a block-diagonal matrix.\n    * It shows similar results to LowRank in our various latency tests and has good validation loss, which helps it to cover a broader space of possible parameterizations.\n    * This parameterization might be more beneficial in other domains or architectures, as we later found that BlockShuffle works better in vision tasks.\n* **Q2**: Self-guided training is effectively adding a regularization term that combines U and V, and we need to compare it against other regularization approaches (e.g. spectral regularization) to position the effectiveness of self-guided training with related works. Such comparison is missing from the paper (missing on both Figure 3 and Table 3) and it is hard to assess the effectiveness of self-guided training as a regularization method.\n  \n  **A2** : We provide the clarification about self-guided training and comparison with other regularization techniques below. \n    * We would like to first clarify that self-guided training is not a typical regularization technique. It does not explicitly constrain or normalize U and V but leverages the dense matrix W to shape the representation learned by structured matrices. Specifically, learning W is unaffected by the additional saddles and pathologies introduced by the structured parametrization, allowing it to learn faster by discovering good specialized features. Then, by decaying the residual contribution of the dense matrix, W can guide the training and transfer the learned hidden state semantics to U and V gradually.\n    * We initially considered weight normalization to stabilize training but found it limited in this context, and layer normalization increases latency and slows down the training convergence. In the table below, we also provide a comparison with spectral normalization [1] (we turn to spectral normalization because we found spectral regularization is for generalizability rather than stability) and orthogonal regularization [2]. We hypothesize that these techniques are less effective because they are designed to ensure the backpropagated signal does not vanish or explode, and two of them constrain the weight much. However, the challenge of structured parameterization is not only about signal propagation but also the capacity bottleneck of learning a good representation, which the typical regularization techniques can't solve. For the constraints, the spectral normalization directly scales the trained weight by its largest singular value, which may hurt the performance here. The additional term in the orthogonal regularization also distorts the spectrum.\n\n      | LowRank (32% FFN params.)  | Perplexity |\n      |-|-|\n      | Baseline| 29.22|\n      |Self-guided training| **28.02** |\n      |Weight normalization | 29.12 |\n      |Spectral normalization | 31.93 |\n      |Orthogonal regularization | 29.18|\n     \n    In the papers' appendix, we compared self-guided training with other techniques. We will also emphasize this discussion in the revision.\n\n   [1]. Spectral Normalization for Generative Adversarial Networks.\n\n   [2]. Can We Gain More from Orthogonality Regularizations in Training Deep CNNs?\n\n* **Q3** : It is also hard to identify whether the scaling law of BlockDense is necessarily better than the dense matrices in Figure 1a. ... Could you provide more quantitative analyses on the scaling laws (actual slope, confidence interval, etc.)?\n  \n  **A3** : Thanks for your good suggestion. To make the illustration clearer, we applied the linear fit to the scaling results and provided them in the general response PDF. Based on Figure 1 in the PDF, we can calculated the slope and obtained the results outlined in the table below:\n   | Method | slope|\n   |-|-|\n   | Dense (optimal scaling law) | -0.3549 |\n   | 63% FFN params.  |  |\n   | LowRank | -0.3672 |\n   | BlockDense | -0.3673 |\n   | BlockShuffle | -0.3718 |\n   | 32% FFN params.  | |\n   |LowRank | -0.3852 | \n   |BlockDense| -0.3827 |\n   |BlockShuffle| -0.3881 |\n  \n  The table shows that these structured parameterizations all have larger absolute slopes than the dense model at its optimal trade-off. This brings a smaller loss gap in the 1.3B model compared to the 110M models, indicating that structured parameterizations scale well or even better than the dense model. Additionally, their scaling curves for our structured parameterizations can be further optimized by finding a better balance between model size and training tokens.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AKVSkXJKfk", "reviewer": "Reviewer_SNmM", "summary": "In order to improve the efficiency of Large Language Models, the authors explore the use of three structured approximations in the FFN blocks of the Transformer: LowRank, BlockShuffle, and BlockDense. They consider both pre-training and decoding, which have distinct requirements and bottlenecks and a range of sizes from 110M to 1.3B. Furthermore, they introduce self-guided training, which uses a dense matrix as a residual component, which is then annealed. The proposed method achieves a 2.5x speed-up at a 0.4 PPL increase under the same training FLOPs.", "questions": "- The scaling law seems to indicate that the proposed approach leads to worse performance per FLOPs. If so, what is the benefit of the proposed approach?\n- Can you conduct one experiment at 300B tokens to show that your methods remain competitive in terms of performance at (or close to) convergence? e.g., the best-performing structured approximation with self-guided training against a vanilla vaseline.\n- Can you add the baseline where the FFN is not expanded as well as some evaluation on downstream tasks?", "limitations": "I would like the authors to acknowledge the limitations of the model size and number of tokens in comparison to state-of-the-art language models (such as LLAMA3), as their results may not be applicable at scale.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The research topic of improving the efficiency of LLMs and, therefore, making them more affordable is crucial and timely. The authors chose to study the FFN bottleneck, which gets worse as models scale, making the research even more relevant as LLMs become larger.\n- The authors conducted extensive experiments and ablations, notably on the model size, pre-training tokens, FFN width, batch size, and learning rate. Additionally, they provided scaling laws.\n- The experimental setup is common and modern. Notably, they used RefineWeb, RoPE, GeLU, FlashAttention, GQA, and tuned betas.\n- The paper is well-written and easy to follow.", "weaknesses": "- The models were trained between 2B and 25B tokens, amounting to a few thousand steps, which is not enough for a language model to converge even at a smaller scale. Therefore, one cannot conclude whether the proposed method is competitive or degrades the performance compared to the baseline at (or close to) convergence.\n- There is no downstream task evaluation of the models. Only the validation perplexity is reported.\n- The simple baseline where the FFN is not expanded (`d_model -> d_model -> d_model`) is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the thoughtful review and insights on this paper. We have added the necessary experiments as suggested and provided a detailed response below. We hope our response addresses the reviewer’s questions. \n\n* **Q1**: The models were trained between 2B and 25B tokens, which is not enough for a language model to converge even at a smaller scale. Therefore, one cannot conclude whether the proposed method is competitive or degrades the performance compared to the baseline at (or close to) convergence. Can you conduct one experiment at 300B tokens ...?\n\n   **A1** : Thanks for your good suggestion. On the one hand, we follow the Chinchilla scaling law paper to allocate training tokens. This allows us to compare the scaling curves of these parameterizations with the optimal curve of the dense model and indicate their effectiveness at larger scales. On the other hand, to directly demonstrate that they maintain performance at the overtraining regime, we trained Transformers (110M) on 100B tokens within this short rebuttal period. \n    \n  | Method| FFN size (M) | 2.2B Tokens (optimal dense scaling law) | 100B Tokens (overtraining regime)|\n  |-|-|-|-|\n  | | |Loss/ Perplexity | Loss/Perplexity|\n  |Transformer-s| 57|3.2569/25.67| 2.8143/16.68|\n  |LowRank | 21|3.3748/29.22|2.9256/18.65|\n  |BlockDense| 21| 3.3731/29.17|2.9239/18.61|\n  |BlockShuffle | 21| 3.3994/29.95|2.9413/18.94|\n  \n By comparing the performance between models trained with 2.2B and 100B tokens, it can be seen that these structured parameterizations maintain or slightly reduce the loss gap on 100B tokens, indicating that they’re still competitive at convergence.\n\n* **Q2** : I would like the authors to acknowledge the limitations of the model size and number of tokens in comparison to state-of-the-art language models (such as LLAMA3), as their results may not be applicable at scale.\n\n  **A2** : Thanks for pointing this out. We will definitely add the limitation in the revision that we didn’t investigate models in this paper that are comparable to today’s practical LLMs, such as LLaMA-3. This is not only because of the limited computing resources but also because this study is to start investigating structured parameterizations of linear layers in modern LLM architecture training. We hope our findings and solutions about scaling, efficiency, and optimization will push their usage on the industry side and in future work. \n\n* **Q3**: There is no downstream task evaluation of the models.\n\n  **A3**: We train the Transformer-s models on 100B tokens and evaluate their zero-shot downstream performance. Table 1 in the response PDF shows that the results are consistent with the validation perplexity of pre-training. Structured parameterizations with 32% FFN parameters incur only about 0.9-1.4 accuracy loss on the downstream evaluations.  Besides, we also provided similar good results of Transformer-xl trained on 26B tokens in Table 2 of the response PDF and the self-guided training consistently improves downstream performance by reducing training loss.\n\n* **Q4**: The scaling law seems to indicate that the proposed approach leads to worse performance per FLOPs. If so, what is the benefit of the proposed approach?\n\n  **A4**: Compared to the dense model (optimal trade-off), structured matrices have the advantage of utilizing training FLOPs more effectively, potentially reaching lower loss with fewer parameters (see Figure 1 and Table 2 in the original paper). We will clarify this point more clearly in the revision. Specifically,\n    * In Figure 1 of the general response PDF, we apply a linear fit to the scaling points for better illustration. We train the dense and structured models on the same amount of tokens. By fixing the training FLOPs, structured matrices have fewer parameters and eventually achieve very close or even slightly better (e.g., LowRank) loss in Figure 1. Given their steeper scaling curves, we can also expect noticeably lower loss and fewer parameters for structured parameterizations per FLOP when the x-axis is further extended.\n    * In the \"Wide and Structured Network\" section of the paper, we also apply existing efficient techniques to the attention module to further optimize the use of training FLOPs. In most experiments, we only structure the FFN module to simplify the study, which negatively increases the attention module's impact on the overall architecture. By making the whole network wide and structured, we demonstrate that even with a medium-sized Transformer (335M), we achieved 0.5 lower perplexity with only 252M parameters under the same training FLOPs.\n\n  Moreover, the scaling curves of the structured matrices can be further optimized by finding a better trade-off. The good scaling behavior makes them strong candidates for future architecture design.\n\n* **Q5**: The simple baseline where the FFN is not expanded (d_model -> d_model -> d_model) is missing.\n\n  **A5** : We understand that this baseline is to reduce the intermediate size of the FFN, resulting in a smaller FFN as well. However, we chose not to pursue this approach because our investigation focuses on structured linear transformations (e.g., low-rank or block-diagonal matrices), as stated in the introduction.\n     * First, these structured matrices serve as plug-and-play replacements for the dense layer. This allows us to extend our findings to dense layers in other architectures easily. Also, maintaining the input and output dimensions doesn't affect the design of other components.\n    * Second, techniques like pruning cannot function as plug-and-play transformations, as they may require complex dimension changes and impact other components in the network. Moreover, they fall outside the scope of our current context and are left for future work.\n\n   In summary, we excluded FFN changes to simplify our explorations. We believe these changes are also orthogonal and can be combined with the structured linear transformations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all the reviewers for their constructive feedback. Before replying to the comments one by one, we would like to highlight our contributions and clarify common questions in this general response:

In this paper, we investigate the performance of three structured parameterizations within the FFN modules in modern LLM architecture training, focusing on efficiency, optimization, and scaling aspects. In detail,

  * We conducted a comprehensive efficiency study in various scenarios, using our proposed pre-merge technique, to show their good latency performance.
  * We identified the optimization challenges of the structured parameterizations and proposed an effective method called self-guided training to boost performance for all of them.
  * We showed that these structured matrices have steeper scaling curves and can utilize training FLOPs more effectively than the dense model. To further validate this point, we trained a wide and structured network of medium size to have fewer parameters and lower perplexity.

For common concerns, we added the downstream performance in Table 2 and Table 3 in the attached PDF, showing consistent performance with validation perplexity of pre-training, and close results to the dense model. Additionally, we applied a linear fit to the scaling curves to make the illustration clearer in Figure 1 of the attached PDF.

Meta-review: The paper studies structured parameterizations (LowRank, BlockShuffle, and BlockDense) for the FFN or MLP layers to improve the efficiency of the transformers. Unlike many existing works that focus on fine-tuning, this work exploits these structured weight matrices during the pretraining stage. To address the training challenges associated with structured parameterizations, the authors introduce self-guided training, which uses a dense matrix as a residual component that is gradually annealed. Experiments show that the proposed method achieves a 2.5x speed-up with only a 0.4 PPL increase under the same training FLOPs.

While the use of structured weight matrices is not new and has been widely studied before for other architectures (such as ResNet), their exploration in the context of LLMs, particularly for pretraining, is timely and useful for improving the efficiency of LLMs. Reviewers pointed out several issues in the experiments, most of which were addressed by additional experiments in the rebuttal. The authors could include these in the final version to make the paper more robust.
Decision: Accept (poster)