Title: Building on Efficient Foundations: Effective Training of LLMs with Structured Feedforward Layers
Venue: NeurIPS 2024
Accepted: yes
Keywords: large language model

structured matrices

efficient architecture
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

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 .

Razvan Pascanu

1 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 network s (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 a 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-scratch 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 a dense matrix as a 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 a 17% throughput boost on the -l scale, while still maintaining slightly better perplexity compared to the efficient Transformer [2]. Beyond training compute 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 × . Furthermore, self-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 LLMs, 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:

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.

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.

We identify optimization challenges in structured matrices and introduce a method called selfguided training, which efficiently improves training dynamics and boosts the final performance for all three types of structured matrices.

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 layers in section 3, that better contextualizes our work. We focus on structured approximations of linear projections g ( x )= Wx that have the form g ( x )= U ( V x ) , where U and V are structured matrices, e.g. low-rank or block-diagonal. We opt for this particular form because it allows us to readily exploit the computational gains on 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.

Params/MAC

N·M

(N + M) · R

min(N,M).

N + M

B

R·(M+

N

B

>M

M

>M

>M

Feature vector

x

Dense matrix

W

→XKX→

matrix

vb

diagonal

Low-rankmatrix

Vr

R

Bblocks

N

B

>

R

R-B

→

R

→xK炎→

Bblocks

→

N

N

Bblocks

IB

Bblocks

>

IB

N

N

N

→

N

(a)Dense

(b)LowRank

(c)BlockShuffle

(d)BlockDense

2.1 Structured linear parametrization

We explore three structured matrices to approximate the standard linear layer Wx , maintaining its input dimension N 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 to 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 Wx ≈ U r ( V r x ) where U r ∈ R M × R , V r ∈ R R × N and R < min( M,N ) . Note that we use the superscript r 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 on to distinguish such components from block-diagonal ones. The parameter count and MAC (Multiply-Accumulate Operations) decrease from M · N to ( M + N ) · R .

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 is to reshape the feature dimension into two dimensions and first use a linear projection that mixes along one of these fictive dimensions, followed by a linear projection that mixes along the other. More precisely, we first reshape the input features x ∈ R N into B blocks and apply the non-tied weight of N B × N B to each block, then flatten the intermediate feature. To achieve global perception, we regroup elements from different blocks into B new blocks and apply the same transformation for each block again.

Technically, we can express the per-block transformation using block-diagonal matrices and formulate the above process as Wx ≈ f -1 ( U b f ( V b x )) , where block-diagonal matrices V b and U b has B blocks with shapes min( N,M ) B × N B and M B × min( N,M ) B per-block. As shown in Fig. 2, 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 -1 ( · ) 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 · N to min( N,M ) · ( M + N ) B , where B acts as a trade-off of accuracy and efficiency.

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 block-diagonal projection V b with a dense or low-rank projection U r . Thus, we can mix the features between blocks without permuting the inner features. The formula is defined as Wx ≈ U r ( V b x ) , where V b is the block-diagonal matrix with b blocks in shape R B × N B , and U r ∈ R M × R . Technically, the second projection does not need to be a 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 + N B ) . Note that BlockDense can recover the LowRank approximation if we set B =1 and R< min( M,N ) .

Remark We limit our exploration of the efficient linear parametrizations within the FFN blocks.These typically have 8 · H 2 parameters and MAC, where H standards for hidden state dimension and 4 · H as the intermediate hidden size of FFN. In contrast, the proposed parametrizations have:

Although BlockDense is introduced as a 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 a relatively big batch size. However, for online decoding with a 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 A100. Because each linear projection suffers from parallelism-bound, 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 advantage of the fact that these parametrizations do not have nonlinearity, 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 parametrization to use based on the current batch size and setting. Fig. 7 analyzes using structured or dense forms for different batch and model sizes, allowing us to 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 2 , which is a source of proliferation

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 ( V x ) 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 a 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 dynamics in Appendix B.2.

(a) Loss spikes

(b) Slow convergence

LowRank (R=128, self-guided training)

Dense

Train Loss

500

1000

1500

2000

2500

3000

TrainingSteps

LowRank (R=128, self-guided training)

Dense

TrainLoss

22

500

1000

1500

2000

2500

3000

TrainingSteps

Self-guided 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 g (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 as 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 is the layer's output, and α decays following a cosine scheduler. As a 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 a 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 0 , we transform the easier-to-optimize loss into the original parametrization we want to use.

Furthermore, we initialize W 0 = U 0 V 0 , making the observation that by using this initialization for the dense residual branch, we can easily start the guided training at any stage of learning (e.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 in Sec. B.1 and better final performance shown in Sec. 4.4.

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:

In our practice, p is a random variable sampled independently from a uniform distribution over 0 to 1 in each training forward pass. With α following a 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 a 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-29], 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 on 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] 3 , 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.9 × 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 LLMs. There have been other studies that aim to improve the expressiveness of structured matrices. For instance, Moczulski et 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 or irregular sparsity patterns. Dao et al. [13] simplified butterfly matrices to block-diagonal ones, achieving a 2 × speed-up on WikiText-103 language modeling tasks. In this work, for accuracy and efficiency, we explored low-rank 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-l (729M), and Transformer-xl (1.3B). For our efficient

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 a rank half the FFN width in 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.1.

Training We use the RefinedWeb 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 ( bfloat16 and float32 ), and adopting fast CUDA kernels like Flash Attention [25] for all experiments. We measure training FLOPs 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 FLOPs more efficiently (e.g., smaller model and lower loss) and performing better in the overtraining regime.

(a)

BlockDense

(b)

BlockShuffle

BlockDense (32% FFN params)

BlockShuffle (32% FFN params)

BlockDense (63% FFN params)

BlockShuffle (63% FFN params)

Dense(optimal scaling law)

nLoss

Dense (optimal scaling law)

Validation

3.0

Validation

3.0

2.5

2.5

10

18

10

19

1020

10

18

10

19

10

20

Training Flops

Training Flops

Scaling law study: better training FLOPs utilization Based on the Chinchilla scaling law, we train four sizes of 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. In Fig. 4 and Fig. 1, we apply a linear fit to the scaling points for better illustration and show that all three structured matrices have steeper scaling

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 6 in the appendix, with comparisons among the three structured parameterizations.

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 as an example, as it demonstrates superior performance compared to the other two approaches in our settings, as demonstrated in Table 6 and Fig. 4. To align the training FLOPs, the wide and structured networks are trained on a larger number of tokens. It can be observed in Table 1 that these models achieve lower perplexity while using much fewer parameters. For instance, the parameter count can be reduced from 729M to 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 in 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 at their optimal training-compute 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, we 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 to make the entire network structured. Then, the well-trained models are evaluated on downstream tasks, including PIQA, HellaSwag, Winogrande, and ARC tasks, using lm-evaluation-harness 4 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 at larger sizes, solidifying their benefits over dense architectures.

26

60

44

75

24

56

72

40

22

52

69

36

A

20

48

66

32

18

44

Dense

Wide and Structured

0

250

500

750

0

250

500

750

0

250

500

750

0

250

500

750

ARC (challenge)

ARC (easy)

PIQA

Model size (M)

4.3 Efficiency study

We investigate the efficiency of structured FFN and consider different numbers of tokens T to discuss different scenarios. Here, T corresponds to the total number of tokens in a batch.

50

50

Dense

Dense

LowRank (63% params)

LowRank (32% params)

40

BlockDense (63% params)

1.1x

40

BlockDense (32% params)

w

S

30

BlockShuffle (63%params)

1.4x

口

30

BlockShuffle (32% params)

Latency

2.0x

20

20

2.6x

10

10

0

0

1024

2048

3072

4096

5120

1024

2048

3072

4096

5120

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 to 63% or 32%, the lowrank and BlockDense achieve a 1.4 × or a 2.5 × speed-up, respectively. BlockShuffle offers modest improvements, with 1.1 × and 2.0 × 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.35 × whole training speed-up with 1.1 increased perplexity.

(a)Width=2048

(b)Width=3072

(c)Width=5120

(d)Width=6144

3.0

Dense

0.6

1.2

LowRank

2.5

BlockDense

(sw)

0.5

1.0

BlockShuffle

2.0

Latency

0.8

0.4

0.6

1.5

0.3

0.4

1.0

0.2

0.2

0.5

12856

512

1024

1536

2048

128,56

512

1024

1536

2048

128,56

512

1024

1536

2048

12856

512

1024

1536

2048

Batch size

Small number of tokens FFN can be parallelism-bound with small T (e.g., T = 16 ) on the A100 GPUs. Then, when T gets increased, FFN becomes memory-bound and will eventually be computation-bound. Online and offline decoding stages may encounter a small number of tokens when unrolling the model step by step. As discussed earlier, our pre-merge method can alleviate the parallelism-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 alternatives or choose pre-merged dense matrices. For example, with a 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 as speed improvements of 2.63 × speed-up of LowRank with 32% FFN parameters on T =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 a straightforward comparison under the same training FLOPs, we adjust the training steps for self-guided training and repeat those tokens at the end to ensure they're fully learned by structured matrices. As can 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.6 × inference speed-up. Additionally, we compare our method with another advanced baseline that trains structured matrices with more tokens, showing that the self-guided training can achieve comparable or superior results even with the same number of tokens.

5 Conclusion

In this paper, we conducted extensive experiments investigating the use of structured matrices to parameterize FFNs in Transformers, with models scaling up to 1.3B parameters on the RefinedWeb dataset. Our primary aim was not to determine which structured matrix performs best, as this can be task-dependent, but to 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 (e.g., LowRank with the novel training strategy achieves a 1.35× inference speed-up with only a 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.

Limitations : BlockDense and BlockShuffle are more complicated than LowRank . In this work, we only explored a limited range of hyperparameter settings of them. However, since these approaches 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 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.

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

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

28

Dense

24

StructuredMatrix(moretokens)

Structured Matrix(self-guidedtraining)

Perplexity

20

+

C

16

12

8

0

5

10

15

202530

¥35

40

45

Tokens (B)

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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.

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.09288 , 2023.

Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990 , 2022.

Meta. Llama 3. https://llama.meta.com/llama3/, 2024.

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 on gemini research and technology. arXiv preprint arXiv:2403.08295 , 2024.

Pratyusha Sharma, Jordan T Ash, and Dipendra Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. arXiv preprint arXiv:2312.13558 , 2023.

Edward J 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:2106.09685 , 2021.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems , 32, 2019.

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.

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.

Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4510-4520, 2018.

Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM , 52(4):65-76, 2009.

Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, and Yu Wang. Flashdecoding++: Faster large language model inference on gpus. arXiv preprint arXiv:2311.01282 , 2023.

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. CoRR , abs/1312.6120, 2013. URL https://api.semanticscholar.org/CorpusID:17272965 .

Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks , 2(1):53-58, 1989.

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In ICLR , 2019.

Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems , 33: 17283-17297, 2020.

Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 , 2019.

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 on Operating Systems Principles , pages 611-626, 2023.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-deociding, 2023. URL https://crfm.stanford.edu/2023/10/12/flashdecoding.html .

James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.03824 , 2021.

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 a general theory for scalable masked transformers. In International Conference on Machine Learning , pages 3962-3983. PMLR, 2022.

Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie-Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems , 34:22795-22807, 2021.

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

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

Peter Belcak and Roger Wattenhofer. Fast feedforward networks. arXiv preprint arXiv:2308.14711 , 2023.

Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342 , 2023.

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507 , 2024.

Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694 , 2023.

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 IEEE international conference on computer vision , pages 2857-2865, 2015.

Misha Denil, Babak Shakibi, Laurent Dinh, Marc'Aurelio Ranzato, and Nando De Freitas. Predicting parameters in deep learning. Advances in neural information processing systems , 26, 2013.

Cheng Tai, Tong Xiao, Yi Zhang, Xiaogang Wang, et al. Convolutional neural networks with low-rank regularization. arXiv preprint arXiv:1511.06067 , 2015.

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.

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.

Mikhail Khodak, Neil Tenenholtz, Lester Mackey, and Nicolo Fusi. Initialization and regularization of factorized neural layers. arXiv preprint arXiv:2105.01029 , 2021.

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 Pattern Recognition (CVPR) , June 2016.

Marcin Moczulski, Misha Denil, Jeremy Appleyard, and Nando de Freitas. Acdc: A structured efficient linear layer. arXiv preprint arXiv:1511.05946 , 2015.

Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference on machine learning , pages 1517-1527. PMLR, 2019.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing , 568:127063, 2024.

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 preprint arXiv:2302.13971 , 2023.

Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon llm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116 , 2023.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 , 2023.

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.01068 , 2022.

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 arXiv:2405.04517 , 2024.

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.13048 , 2023.

Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954 , 2024.

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 language model training on gpu clusters using megatron-lm. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-15, 2021.

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

NVIDIA. Transformerengine. https://github.com/NVIDIA/TransformerEngine, 2022.

Appendix

A Structured matrices

A.1 Design choice

For BlockDense , we also investigate the reverse order of two projections, placing the low-rank or dense matrix first, followed by the block-diagonal matrix. However, this change surprisingly yields worse performance. For instance, on the RefinedWeb dataset, perplexity increases from 29.17 to 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 a second shuffle operation to restore the original order. Taking this into account, we 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 for our adopted version. Nonetheless, we maintain the design from Dao et al. [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 t U ⊤ t and V t V ⊤ t to 1 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 in our experiments and increased training FLOPs slightly. Hence, we adopted standard weight decay for all the structured FFNs.

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% of the parameters in the FFN, increase the loss by only 0.02 to 0.04 while reducing total training time by approximately 15% on Transformer-xl and Transformer-l. Additionally, they accelerate training by 1.35 × with only a 1-point increase in perplexity on Transformer-xl with 32% of the parameters.

Although the main focus of the paper is not to compare different structured matrices but to showcase their general properties-including scaling, efficiency, and optimization-we still provide comparisons in the appendix. From Table 6, 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 can serve as a good inductive bias in vision tasks Sec. A.3.

A.3 Results on CIFAR-10 dataset

Although in our main experiments, the BlockShuffle performs the worst with two block-diagonal matrices, we provide experiments on the CIFAR10 dataset here, showing that when locality is highly preferred, block-diagonal 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 a 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 a 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 a 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 a block-diagonal 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 replace the first FFN module. Also, the BlockDense is slightly smaller for -m and -xl models to ensure the rank is a multiple of 256 when matching parameters. Loss and perplexity are evaluated on a 0.5B token validation set.

A.4 Results on downstream tasks

Table 8: Performance on downstream tasks under the zero-shot setting. We report the perplexity performance of the validation set of RefinedWeb. 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.

Supplementary to Fig. 5, Table 8 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 are more prone to loss spikes. To mitigate this, we apply self-guided training in the most challenging cases, which results in improved training dynamics, such as faster convergence without loss spikes.

Furthermore, we report the loss spikes observed in the large Transformer-xl model trained on the 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 us to analyze the gradient updates during the backward pass of the linear UV . Considering g as the gradient of output and x as the input, the standard linear layer W gradient update is gx ⊤ . For the structured parametrization, the gradients of U and V are gx ⊤ V ⊤ and U ⊤ gx ⊤ , respectively. Then, for W ′ as being the updated parameters, we will have the updates for W to be:

And the one for UV :

Thus, it can be seen from Eq. (4) that the projections UU ⊤ and V ⊤ V can disrupt the gradient gx ⊤ . 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 UU ⊤ and V ⊤ V and use this to 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 1, 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 is to realize that learning with structured matrices has additional redundant degrees of freedom brought by symmetry. For example, to increase the norm of a feature, it can increase either the corresponding weights U or V . Given that gradient descent makes this decision independently, it will overshoot and make learning less well-behaved.

(a)

lr

=1

.

5

e

-

3

(b)

lr

=1

.

0

e

-

3

(c)

lr

=5

.

0

e

-

4

(d)

lr

=1

.

5

e

-

3

(e)

lr

=1

.

0

e

-

3

(f)

lr

=5

.

0

e

-

4

LowRank (R128)

LowRank (R128)

LowRank (R128)

LowRank (R384)

LowRank (R384)

LowRank (R384)

LowRank (R512)

LowRank (R512)

LowRank (R512)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

0000006008000090090000000Z

2003004005006007008009001000

00010060080040090090000000Z

Training Steps

Training Steps

Training Steps

LowRank (R=128)

LowRank (R=128)

LowRank (R=128)

LowRank*(R=128)

LowRank*(R=128)

LowRank*(R=128)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

WN

WN

M

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

Training Steps

Training Steps

Training Steps

B.3 Design choice

Experiments in this part are 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 are compared. Dense layer decomposition, which decomposes the weight directly at the midpoint of 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 a feasible solution to evolve. Our guided initialization allows its usage in various stages of training without the need for a well-trained

(a)

lr

=1

.

5

e

-

3

(b)

lr

=1

.

0

e

-

3

(c)

lr

=5

e

-

4

(d)

lr

=1

.

5

e

-

3

(e)

lr

=1

.

0

e

-

3

(f)

lr

=5

e

-

4

BlockDense (B2R128)

BlockDense (B2R128)

BlockDense (B2R128)

BlockDense (B2R256)

BlockDense (B2R256)

BlockDense (B2R256)

BlockDense (B2R512)

BlockDense (B2R512)

BlockDense (B2R512)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

200300

400

5006007008009001000

2003004005006007008009001000

2003004005006007008009001000

Training Steps

Training Steps

Training Steps

BlockDense (B=2,R=128)

BlockDense (B=2, R=128)

BlockDense (B=2,R=128)

BlockDense*(B=2,R=128)

BlockDense*(B=2,R=128)

BlockDense*(B=2,R=128)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

Training Steps

Training Steps

Training Steps

(a)

lr

=1

.

5

e

-

3

(b)

lr

=1

.

0

e

-

3

(c)

lr

=5

e

-

4

(d)

lr

=1

.

5

e

-

3

(e)

lr

=1

.

0

e

-

3

(f)

lr

=5

e

-

4

BlockShuffle (B8)

BlockShuffle (B8)

BlockShuffle (B8)

BlockShuffle (B4)

BlockShuffle (B4)

BlockShuffle (B4)

BlockShuffle (B2)

BlockShuffle (B2)

BlockShuffle (B2)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

200300

4005006007008009001000

2003004005006007008009001000

2003004005006007008009001000

Training Steps

Training Steps

Training Steps

BlockShuffle (B=8)

BlockShuffle (B=8)

BlockShuffle (B=8)

BlockShuffle*(B=8)

BlockShuffle*(B=8)

BlockShuffle*(B=8)

Dense

Dense

Dense

Train Loss

Train Loss

Train Loss

22

22

22

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

500

1000

1500

2000

2500

3000

Training Steps

Training Steps

Training Steps

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 in 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,

(a)

LowRank

on Transformer-xl

(b)

BlockDense

on Transformer-xl

LowRank(63%)

BlockDense(63%)

LowRank (32%)

BlockDense(32%)

Dense

Dense

22

22

Train Loss

Train Loss

2000

4000

6000

8000100001200014000

2000

4000

6000

8000100001200014000

Training Steps

Training Steps

(a) N=768, std=0.02

(b) N=1536, std=0.02

(e) N=768, std=0.02

(f) N=1536, std=0.02

60

1.8

Largest Singular Value

0.8

Largest Singular Value

1.6

Largest Singular Value

5

Largest Singular Value

10

0.7

1.4

0.6

1.2

0.5

1.0

3

6

80

100

200

300

400

0

007

400

600

008

100

200

300

T

400

0

200

T

400

T

600

T

800

Rank

Rank

Rank

Rank

90

1.25

8

Largest Singular Value

Largest Singular Value

1.00

Largest Singular Value

Largest Singular Value

6

0.4

0.75

0.2

05°0

2

0.25

10

20

30

01

20

30

10

20

30

10

20

30

Blocks

Blocks

Blocks

Blocks

we show that with self-guided training, all three structured FFNs with 32% parameters incur only a 0.4-0.6 increase in perplexity compared to the baseline, while benefiting from a smaller memory footprint and faster speed at inference time. To further illustrate this, we plot these points with the same training FLOPs in Fig. 14, and specifically for LowRank in Fig. 8, highlighting that self-guided training achieves comparable performance to training on more tokens.

(a)

BlockDense

(b)

BlockShuffle

28

28

Dense

Dense

24

BlockDense(more tokens)

24

BlockShuffle(more tokens)

BlockDense(self-guided training)

BlockShuffle(self-guided training)

Perplexity

Perplexity

20

20

16

16

12

12

8

8

0

5

10

15

202530

35

40

45

0

5

10

15

2025

35

40

45

Tokens (B)

Tokens (B)

C Implementation details

Model Architecture details are provided in Table 10. We consider four baseline Transformer sizes, ranging from 110M to 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 6. 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 to the FFN module with a rank half of the model or 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 in all experiments, we do not apply structured matrices to the first FFN module, as doing so can lead to non-negligible performance loss in models on shallow networks. For dense models, we 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.

Dataset We use the RefinedWeb dataset [48], a carefully curated subset of CommonCrawl, optimized for filtering and deduplication, providing 600B tokens for public use. Due to its large size, we shuffle, extract, and tokenize it 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 is set to 1024. Following scaling laws [1], we allocate tokens at 20 times the number of parameters for each baseline model in all experiments, except for the 300B token training.

Model

-s size

Dense

Structured

-m size

Dense

Structured

-l size

Dense

Structured

-xl size

Dense

Structured

110.0

81.1

335.1

219.4

729.1

464.4

1274.1

799.6

Batch

1280

1024

1792

1536

2304

2048

2816

2304

LR

9.1e-4

9.6e-4

7.8e-4

8.4e-4

7.1e-4

7.6e-4

6.6e-4

7.1e-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.1 weight decay, betas of [0.9, 0.999], cosine annealing learning rate scheduler with 10% linear warm-up, and 0.1 × 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-compute 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 × ( training FLOPs -0 . 125 ) , and learning rate by 0 . 2920 × ( training FLOPs 0 . 3271 ) , 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 LayerNorm, 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 LLMs, maximizing societal benefits while acknowledging potential harms.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We describe our claims and contributions in Section 1.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

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

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

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss limitations in 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 are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [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 in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We describe our experimental setup and implementation details in Section 4 Appendix C.

Guidelines:

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

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general, releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

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

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

Open access to data and code

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

Answer: [Yes]

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

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

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

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

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We describe our experimental setup in Section 4 and give details in Appendix C.

Guidelines:

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

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Training large language models requires significant time and compute, making it infeasible to run multiple seeds. Nonetheless, results from single training runs are 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 are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

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

Answer: [Yes]

Justification: 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 or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms with the NeurIPS Code of Ethics. Guidelines:

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

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

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Societal impacts are discussed in Appendix D.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

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

Answer: [NA]

Justification: The paper does not release data or models. not pose such risks.

Guidelines:

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

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All the assets used have been cited.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. NewAssets

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: 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 to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

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

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

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

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nSelf-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.\n\nQUESTIONS:\nCould 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.)?\n\nLIMITATIONS:\nThere are no other limitations. All weaknesses have been listed above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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$\n\nQUESTIONS:\nIs there any special technique to initialize the matrices U and V? What is the scale/variance of the initialization?\n\nLIMITATIONS:\nThe authors discuss the limitation of this paper in Sec. 5. I believe that another limitation is that the paper lacks evaluations on downstream tasks.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}