Mixed Sparsity Training: Achieving 4$\times$ FLOP Reduction for Transformer Pretraining
Authors: Pihe Hu
Shaolong Li
Xun Wang
Longbo Huang
Venue: Accepted by TMLR ()

Abstract: Large language models (LLMs) have made significant strides in complex tasks, yet their widespread adoption is impeded by substantial computational demands. With hundreds of billion parameters, transformer-based LLMs necessitate months of pretraining across a high-end GPU cluster. However, this paper reveals a compelling finding: transformers exhibit considerable redundancy in pretraining computations, which motivates our proposed solution, Mixed Sparsity Training (MST), an efficient pretraining method that can reduce about $75$% of Floating Point Operations (FLOPs) while maintaining performance. MST integrates dynamic sparse training (DST) with Sparsity Variation (SV) and Hybrid Sparse Attention (HSA) during pretraining, involving three distinct phases: warm-up, ultra-sparsification, and restoration. The warm-up phase transforms the dense model into a sparse one, and the restoration phase reinstates connections. Throughout these phases, the model is trained with a dynamically evolving sparse topology and an HSA mechanism to maintain performance and minimize training FLOPs concurrently. Our experiment on GPT-2 showcases a FLOP reduction of $4\times$ without compromising performance.

Published in Transactions on Machine Learning Research (03/2025)
Mixed Sparsity Training: Achieving 4x FLOP Reduction for Transformer Pretraining
Pihe Hu"
kupihe com
Institute for Interdisciplinary Information Sciences Tsinghua Universily
Shaolong Li
shaolongli 60gmaiil com
Central South University
Xun Wang
wang- Ommails www.mmmltiiiiaaaam..mmmmmmmm edu..c
Institute for Interdisciplinary Information Sciences Tsinghua University
Longbo Huang
longbohuanng Stsinghua edu.cr
Institute for Interdisciplinarg Information Sciences Tsinghua University
Reviewed OpenReview: https: htttps://openreviaaw openreview. net/ ptr/formm? id= XosdLS7KVE
Abstract
Large language models (LLMs) have made significant strides in complex lasks, yet their widespread adoption is impeded by substantial computational demands. With hundreds of billion parameters, transformer-based LLMs necessitate months of pretraining across high- end GPU cluster. However, this paper reveals a compelling finding: transformers eahibit considerable redundancy prelraining compulations, which motivates our proposed solu- tion, Mixed Sparsity Training (MST), an efficient pretraining method that can reduce about 75% of Floating Point Operations (FLOPs) while maintaining performance. MST intogrates dynamic sparse training (DST) with Sparsity Variation (SV) and Hybrid Sparse Attention (HSA) during pretraining, involving three distinct phases: warm-up, ultra-sparsification, and restoration. The warm-up phase transforms the dense model into a sparse one, and the restoration phase reinstates connections. Throughout these phases, the model trained with dynamically evolving sparse topology and an HSA mechanism maintain perfor- mance and minimize training FLOPs concurrently. Our experiment on GPT-2 showcases a FLOP reduction of 4x without compromising performance.
1 Introduction
Over the past years, the field of Large Language Models (LLMs) has witnessed remarkable advancements, e.g., T5 (Raffel et al. 2020), GPT-3 (Brown et al., 2020), GLM (Du et al. 2021]. These sophisticated models, characterized by their vast scale, typically with ~ 200 billion parameters, have redefined the frontiers of language understanding, generation, and comprehension. Their prowess in tasks spanning from question-answering systems (Liu et al. (023)., programming tools (Roziere et al., 2023) vidco generation (Bi et al. 2023), has propelled them the forefront of research and applications in many fields.
However, the realization of their potential comes entwined with an imposing bottleneck: the substantial computational cost required for their pretraining. The process of priming these colossal models necessitates
"Equal contribution.
i "oorreppondigg author.

---

In response to the formidable challenge of computational overhead in transformer pretraining, this paper introduces novel method called Mixed Sparsity Training (MST). Unlike existing approaches, MST seam- lessly integrates dynamic sparse training with Sparsity Variation (SV) and Hybrid Sparse Attention (HSA) throughout the pretraining process, unfolding in three pivotal phases: warm-up, ultra-sparsification, and restoration. As shown in Figure 2] the warm-up phase transforms the dense model into an initial sparse topology, addressing uncertainties in link importance at the initial training stage. Subsequently, the deploy- ment of a novel sparse training method, fixed-Growing (MG), facilitates the simultaneous training of a sparse model while actively exploring other sparse topologies. The restoration phase then intelligently rein- states connections, effectively recovering performance loss encountered during the ultra-sparsificstion phase. layers, by incorporating Published in Transactions on Machine Learning Research (03/2025)
extensive computational resources (Chen et al. 2023a), involving protracted periods spanning several months
running on clusters comprising thousands of perffrrmree Graphics Processing Units (GPUs). For example, Narayanan et al. (2021) point out that GPT-3 175B (Brown et al. 2020) requires about 34 training days on a cluster of 1024 NVIDIA A100 GPUs, and Llama-2 70B (Touvron et al. 2023 even requires double training time on the same cluster.
The substantial computational overhead associated with pretraining transformer- based LLMs stands as a formidable barrier, restricting their widespread adoption and accessibility. While existing efforts have sought
to enhance pretraining efficiency through strategies such as training parallelism (Li et al. 2020
Bian et al,
2021 Zhao et al., 2023), hardwaro-assisted attention operators (Dao et al. 2022c Dao
2023). and mixed
precision training (Micikevicius et al. 2017 Burgess et al. 2019 Liu et al. 2022), these methods primarily
address system-executionn-lee. bottlenecks without tackling the intensive algorithm-level FLOPs associated with training. In essence, the algorithmic computation redundancy of transformer pretraining remains an understudied domain. Furthermore, traditional FLOP reduction techniques, including pruning (Ma et al.
2023 Frantar & Alistarh 2023 Syed et. al. 2023) and quantization (Dettmers al. 2022 Frantar et al. 2022: Frantar & Alistarh 2022), have demonstrated the presence of significant FLOP and parameter redundancy within transformers. However, these approaches are primarily applicable in post-training stages, rendering them impractical for addressing computational costs during the pretraining of transformers. This emphasizes the critical need for mitigating algorithmic inefficiencies inherent in the pretraining phase of transformers.
Take GPT-2-small (Brown et al. 2020 as an example. Figure I illustrates that fully-connected layers (Basha et al. 2020)) and self-attention layers Vaswani et al. 2017) contribute significantly to the FLOOPs, accounting for up to 83% and 15%, respectively. The FLOPs of fully connected layers stem from weight updates, while the ceaputattionintteniin
Fully-Conencted Layers
Self-Attention Layer
Others
self-attention operations drive the FLOPs of self-attention layers. Given the distinct characteristies of these layers, differ- ent methods are required to reduce training FLOPs. Dynamic Sparse Training (DST) (Mocanu et al. 2018 Evci et al. 2020) has emerged as an elfective solution to reduce the FLOPs of fully connected layers, demonstrating the capability to train 80%-sparse networks such as ResNet-50 (He et al.] 2016) and MobileNet HHwwadd al.
83%
2017) from scratch without performance degradation. However, direct application of DST methods in pretraining transformers, as seen in SET (Mocanu et al. 2018 for BERT (Dietrich et al. 2021), only yields modest FLOP savings, typically less than 20%, to maintain performance, owing to the extensive parameters required by transformers for complex tasks.
Figure 1: Pretraining FLOPs of GPT-2, de- tailed Appendix A.5
Another methodology to reduce training FLOPs of fully-connected layers is to utilize block sparse matrix, such Dao et al. (2022a b), yet the maximum FLOP reduction is incomparable to our MST methods shown in Table since they adopt a fixed sparse pattern whose sparsity level is limited to maintain the per- formance. Mcanwhile, sparse attention mechanisms (Child et al. 2019 Dai et al. 2019) have been proposed to reduce the FLOPss associated with attention operations. However, the application of sparse attention alone typically results in savings of non-dominant FLOPs, as depicted in Figure [ These observations prompt the following open question: Can we remove the computational redundancy in transformer pretraining without losing performance?
Additionally, MST introduces HSA to mitigate FLOPs generated by
2

---

Published in Transactions on Machine Learning Research (03/2025)
MG based Dynamic Sparse Training
Model
Size
MLP
MLP
MLP
MLP
Atten
Atten
A Lten
Atten
Training
Warm-up
Uttra-Sparsfficttion
Restoration
Step
Figure 2: The sparsity variation of MST includes three phases: warm-up, sutrasspariiiiattion and restoration. SV is combined with MG-based dynamic sparse training and HSA during the training.
sparse attention mask in the first two phases, and then transitioning to densified attention mask when the network weights are dense.
In essence, MST stands poised to revolutionize transformer pretraining, offering significant FLOP savings while providing users with a transparent and efficient method for training dense transformers. Our ex- periments with MST on GPT-2 demonstrate an exceptional FLOP reduction of 4x without compromise in performance across multiple zero-shot and few-shot downstream tasks. Moreover, MST is entirely orthogonal and can scamlessly integrate with existing system-level acceleration methods, such as training parallelism (Zhao et al. 2023) hardware-assisted attention operators (Dao) 2023), and mixed precision training {Liu et al. 2022 thus facilitating efficient transformer pretraining and achieving higher acceleration ratios. Ta- ble provides a comprehensive summary of related works aiming to reduce pretraining FLOPs, highlighting MST's state-of-the-art performance in this regard. A detailed literature review is provided in Appendix 2
Table 1: Comparison of different training techniques for saving transformer pretraining FLOPS.
Name
Methods
Models
Reduction
DynSparse (Dietrich et al.
2021
Dynamic Sparsification
BERT
2x
LIGO WWang el. al.
2023)
Layer Growth
BERT
2.2x
Monarch Dao et al.
2022b
Block Sparsification
GPT-2, BERT
2.2x
SPDF Thangarasa ct al.
2023)
Static Sparsification
GPT-2, GPT-3
2.5x
Pixelated Butterfly (Dao et al.
2022a)
Butterfly Sparsification
GPT-2
2.6x
MST (Ours)
Dynamic Sparsification
GPT-2
2
Related Work
We include most related works about dynamic sparse training and transformer pruning in this section.
3

---

Published in Transactions on Machine Learning Research (03/2025)
2.1 Dynamic Sparse Training
Dynamic sparse training (DST) is developed to reduce both computational workload and memory usage throughout the entire training process. It involves training a sparse neural network from scratch while dy- namically adjusting and updating the sparse mask during training, Sparse Evolutionary Training (SET) (Mocamu et al., 2018) removes least magnitude weights and reintroduces random weights for better ex- ploration at the end of each training epoch. SNFS (Dettmers & Zettlemoyer (219) utilizes exponentially smoothed momentum as the criterion for weight growth. RigL (Evci et al. 2020) uses the same magnitude- based pruning method while activating new weights based their gradient magnitudes through infrequent full gradient calculation, leading to improved accuracy under the same sparsity.
However, the greedy-based growth policy is highly likely to result in limited weight coverage, consequently yielding a sub-optimal sparse structure network. ITOP (Liu el al.) 2021b) delves inlo the fundamental mechanism of DST, revealing that the advantages of DST arise from exploring all potential parameters over time. Top-KAST (Jayakumar al. 2020) pruncs least magnitude weights, but it updates a supersct of aclive weights determined by backward sparsity to enable the revival of prumed weights. To achieve high- quality results, more weight gradients need to be involved in the computation throughout the training, which means there is less backward sparsity and substantial increase in computational workload. [Scwwarz et al.) (2021] proposed a tatgltt-prammterizattoon that keeps the low-magnitude parameters unaffected by learning largely, and its application contributes to enhanced accuracy within the Top-KAST framework. AC/DC Peste et al.) 2021) employs co-training to simultaneously train both sparse and dense models, allowing for a flexible transition between sparse and dense training and yielding accurate results for both sparse and dense models. MEST (Yuan et al. 2021) adopts a gradually decreasing drop and grow rate during training and introduces memory-efficient training framework designed for fast execution edge devices. Spartan (Tai et al. 2022) integrates soft masking with dual averaging-baae updates through computationally expensive process of calculating a soft top-k mask.
2.2 Transformer Pruning
Parameter pruning can certainly be applied to the compression of transformers as well by eliminating re- dundant model weights. The methods for parameter pruming in transformers fall into three main categories: structured, semi-structured and unstructured pruning.
Structured pruning. Structured pruming concentrates on removing larger structured patterns of the network such as groups of consecutive parameters or hierarchical structures like rows, columns, or sub- blocks of the transformer weight matrices, resulting in a model that does not require specific hardware or software for acceleration. Globally Unique Movement (GUM) (Santacroce et al. 2023) prunes network components based on their global movement and local uniqueness scores, aiming to maximize both sensitivity and uniqueness. LLM-Pruner (Ma et. al.) essential interconnected structures. It computes the structure importance with small amount of data and subsequently uses low-rank approximation (Hu et al. 2021 to partially recover lost knowledge after pruning. LoRAPrune (Zhang el al. 2023 combines paraneler--fficienl tuning methods with pruning to enhance performance on downstream tasks, which presents a pruming criterion based on LoRA, using the weights and gradients LoRA for importance estimation instead of relying on pre-trained weights. Sheared LLaMA (Xia et al. 2023) also focused on structured pruming and introduces dynamic batch loading, utilizing losses across different domains to dynamically adjust the components of the sampled data in each training batch. LoRAShear (Chen et al.) 2023b) proposed an innovative structure sparse optimizer called LoRA Half-Space Projected Gradient for progressive structured pruning and knowledge transfer. It employs a multi-stage knowledge recovery mechanism to effectively narrow down the performance gap between the full and compressod transformers.
2023 employs gradient information to selectively remove non-
Semi-structured pruning. Semi-structured pruning is an under-explored strategy that stays between unstructured pruning and structured pruning, exemplified by N:M sparsity (Zhou et al.] (220)) that every contiguous set of M elements precisely contains N non-zero elements in a certain layer or weight matrix. SR- STE (Zhou et al. 2020) proposed the sparse-refined straight-through estimator to enhance the induction of
4

---

Published in Transactions on Machine Learning Research (03/2025)
N:M sparsity in the network. A100 GPUs and the 2:4 fine-grained structured sparsity scheme, introduced by Nvidia (Choquette et al. 2021), which enable sparse neural networks to be accelerated on specific hardware. Semi-structured pattern retains relatively high model accuracy while facilitating efficient compression.
Unstructured pruning. Unstructured pruming centers on pruning less salient and individual parameters, wherever they are, offering greater ficxibility compared to structured pruning. SparseGPT (Frantar & Alistarh 2023) is an approach for few-shot pruning in transformers that avoids the need for retraining, being able prune up 60% of the parameters with minimal perplexity increase. It frames pruning as a sparse regression problem and addresses it using an approximate solver based on the inversion of the Hessian matrix. Syed et al. (2023) introduces an iterative pruning technique that fine-tunes the model during pruning with minimal training steps. Wanda (Sun et. al. 2023) performs weight pruning through the product of weight magnitudes and their corresponding input activations. This method prunes the network with a single forward pass without depending on sccond-order information or requiring weight update, achicving competitive performance compared to SparseGPT. [Shao et al. (2023) proposed Hessian sensitivity-awarr mixed sparsity pruning to attain at least 50% sparsity transformers without retraining.
3
Mixed Sparsity Training
This section introduces Mixed Sparsity Training (MST), our innovative approach for transformer protraining. MST seamlessly integrates Dynamic Sparse Training (DST) with Sparsity Variation (SV) and Hybrid Sparse Attention (HSA). Throughout the training, the model evolves dynamically with a sparse topology. The sparsity variation is firstly elaborated in Section 3.1 and novel topology evolution scheme, Mixed-Growing (MG), is given in Section 3.2 The hybrid sparse attention mechanism is detailed in Section 3.3 The muanced integration of dynamic sparse training with sparsity variation and hybrid sparse attention positions MST as self-consistent method, offering potent mcans to optimize training processes and enhance the pretraining efficiency of transformers.
3.1 Sparsity Variation
This subsection provides detailed explanation of the three phases in sparsity variation. The sparsity variation plays a pivotal role in dynamically allocat- ing model's parameters during the pretraining pro- cess, as shown in Figure
Sparsity
SM
Training Step
Warm-Up Phase In this phase, the dense model systematically transitions into a sparse configura- tion, ensuring an appropriate sparse initialization. Drawing inspiration from prior research (Zhu & Cupta 2018 Liu et al.) 2021a), we implement progressive pruning strategy over series of N itera- tions, gradually inducing sparsity in the dense network. The sparsity level S. for each iteration is determined
Warm Tw -up
Tw + TT Restoration
Ultra-Sparsification
Figure Sparsity variation.
by the cubic decay function:
-1- NAW
S.= -SM
(1-NAW)
for t E [0, Tw], where SM. N, and Aw represent the maximum sparsity, the number of sparsity stages, and the pruning frequency, respectively. The warm-up phase spans Tw = N NAw training steps. As mentioned by Zhu & Gupta (2018), the intuition behind this design is to quickly prune the network in the initial phase when redundant connections are abundant, and gradually reduce the number of weights pruned each time as fewer weights remain in the network, until the target sparsity reached.
Uttra-Sparsfiictioon Phase The duration of this phase is To training steps. We employ a novel topology evolution scheme termed Mixed-Growing (MG), detailed in Section 3.2 to facilitate proper sparse training.
5

---

Published in Transactions on Machine Learning Research (03/2025)
Notably, the model is predominantly trained with highly sparse topology during this phase, leveraging the observation that only a subset of parameters significantly contribute to model performance, while many parameters are redundant, as highlighted in prior works on transformer pruning (Ma et al. 2023 Frantar & Alistarh 2023
Restoration Phase In this phase, connections are strategically reinstated to enhance the model's ex- pressiveness and address potential performance loss incurred during former phases. This phase adopts a progressive approach symmetric to the warm-up phase, gradually reducing sparsity the sparse network over series of N growing iterations. The sparsity for each iteration is determined by:
t Tw - To (1---XV)) N
S. S.=
for I [Tw+ Tw+ where Ar denotes the growing frequency. The duration of the restoration phase is TR N NAR training steps. The progressive pruning and growing strategies have been empirically proven to achieve good performances in prior works (Zhu & Gupta. 2018 Liu et al. 2021a). Besides, our experiments in Section 4.2.1 also validates its superiority over other sparsity variation patterns.
Remark 3.1. The warm-up phase aims to establish a proper initial sparse topology, while the ultra- sparsification phase prinarily focuses on reducing FLOPs. Subsequently, the restoration phase diligently addresses performance degradation resulting from sparsification. The interplay of three phases in spar- sity variation ensures holistic transformation of the model, allowing it to leverage parameter redundancy while maintaining performance. The sparsity variation also forms the fundamental architecture of our pro- posed MST method.
3.2 Dynamic Sparse Training
We propose novel topology evolution scheme called Mixed-Growing (MG), specifically designed for MST methods, particularly effective in high-sparsity language models during the ultra-sparsification phase, We give a comparison of different DST schemes in Table 2 where existing methods exhibit weaknesses in transformers with massive amounts of parameters as shown in Section 4.1 By contrast, our proposed novel topology evolution scheme MG is dedicated to sparse training of transformers by introducing a more active exploration mechanism and out performss other baselines.
Table 2: Comparison of different topology evolution schemes.
SET
RigL
MEST
Mixed-Gren Groed-Groowiggg
Alg.
Mocanu et al. 2018
Evci et al.
2020
Nowaa et al. 2023
(Ours)
Prune
min(
min
min( + + vol)
min(
(1 R) max((Vol) random
Grow
random
max(/V,l)
random
Our link growing policy utilizes a hybrid growth strategy that integrates gradient. magnitude activation and random activation specific proportion broaden the exploration range of weights. This strategy roduces the likelihood of the network to converge to sub-optimal structures. As for the link pruning, we adhere to
simple criterion based on weight magnitude in consideration of the findings by Nowak et al. (2023 regarding the minor differences among various pruning criteria in previous works. The weight magnitude pruning approach avoids the introduction of extra hyperparameters and works well in the pruning steps of our MG method.
Unlike existing topology evolution schemes, MG introduces flcxibility by allowing a disparity between the number of pruned and grown links, facilitating adjustments in network sparsity based on predefined pat- terns like the one depicted in Figure 3 Specifically, MG utilizes current network sparsity Se and target

---

Published in Transactions on Machine Learning Research (03/2025)
network sparsity S to determine the sparsity distribution for each layer, employing Erdös- Rényi strat- egy Mocaauu et al. 2018). When the larget sparsity 5f differs from the current sparsity for given layer I, MG dynamically adjusts the number of growing or pruming links to match the target sparsity.
We exclude the training step index t and provide the pseudocode of MG in Algorithm [ where My signifies the binary mask outlin- ing the sparse network topology for 6, Ç represents the topology update fraction, and R denotes the ratio of connections activated through random growth. At each topology evolution step, MG eliminates subset of existing connections with the smallest abso- lute weight values, while also prioritizing the revival of a certain number of empty connections with the largest gradients. The remaining connections are randomly generated for cach layer.
Update Fraction
4 44 ccolt()))
2
0
Warm Tir
Tir +Tg
Ultra-Sparification Restoration
Figure 4: Piecewise cosine annealing.
The update fraction Ç is determined by piecewise decaying cosine annealing scheme
S1 (1 +1+00g Se 2 (1+ cos
(1,
1- (1- T TT-- ))
for E [T)-,.T), Ti), where Ç represents the update fraction magnitude, and To 0, Tw + Tu, Tw + Ty ......, TN+1 Tw + Tr + TR. This update schedule, visualized in Figure A ensures the network's robust cvolutionary ability across different sparsity levels and is better than the single cosine annealing scheme employed in prior works (Dettmers & Zettlemoyer 2019 Evci el al. 2020 An empirical comparison of different update schedules is provided in Appendix B.5
A Algorithm Topology Evolution by MG.
1: Ot, Ni, 81, s/: parameters, number of parameters, current and larget sparsity of layer I. for each layer do
N Nprume- [(1 N
N groww pruue + (6/ N
Nrand N grow R.
6: Ngrad N Ngrow Nand
7:
Iprune ArgTopK(
@ Mo, Nprune)
Igrad
Arg TrgTTpKi.0O.... Vipraao ([Vo, Ll, Ngrad)
Irand grow RandK (9 M WTlasop UI aand Mandd vrow) Nraud)
10:
Igrow Ilgrad grow U Irand grow
11: Update Mo, according to Iprune and Tgrow
12:
O: 0f Me,
13:
end for
3.3 Hybrid Sparse Attention
The majority of LLMs (Raffel et al. 2020 Brown et al. 2020 Du et al. 2021 are
built upon transformer-basee based architectures, where the sell-attention layer Vaswani ct al. 2017) plays pivotal role capturing spatial correlations within input texts. We first propose a novel unfactorized sparse strided self-attention, based on which we design the Hybrid Sparse Attention (HSA), to efficiently economize attention FLOPs while maintaining model performances.
A self-atiention layer (Child el al. 2019) transforms a matrix of input embeddings X (S, where
into an output matrix, controlled by a connectivity pattern n denotes the length of input embeddings X, and S, represents the of indices to which the i-th output vector attends. Denote Wg. We. and W, as weight matrices responsible for transforming given x; into query, key, value, respectively. The
Figure 5: Strided attention with stride length 3.
7

---

Published in Transactions on Machine Learning Research (03/2025)
self-attention output at each position is computed as the sum of values weighted by the scaled dlt-ppodduc similarity of keys and queries:
Att, = (a (1,;; (t,.)
where alz; S:) = softmax (Wez)) autoregressive models such as GPT-2 (Brown et al. 2020), full self-altention sets S; = (i = i), allowing each element attend all previous positions, including its own.
(WDD))jes j)Es, and represents the inner dimension. For
Drawing inspiration from previous works sparse self-attention (Child el al. 2019 Beltagy et al.) 2020 Dai et al. (211)). we first introduce an un- factorized sparse strided self-attention, wherein dif- ferent attention heads share the same sparse mask. While unfactorized sparse self-attention may entail more FLOPs than its factorized counterpart, it con- dense
Stride Length
Atten
sistently maintains the same performance self-attention, particularly evident in the ultra- sparsification phase in Section 4.2.3 The resultant FLOP savings of unfactorized sparse self-attention align with our targeted reduction of 4x, as demon- strated in Section AN In our unfactorized strided sparse self-attention, cach i-th output vector attends to the previous locations (depicted by orange cells in Figure 5), well every (-th location (illustrated by yellow cells in Figure 5), where denotes the chosen stride. Formally, S, {max(i - 1,0), max(i - I + 2, O)...., of n {j : (i j) mod I = 0} for 0 S i S n. This pattern is visualized in Figure Furthermore, our strided self atlf-attention surpasses fixed self-attention patterns, as validated in Section 4.2.3 under the same FLOPs budget.
1
Training
Step
Warm Ti
Tw + Ty Ultra-Sparsificttion Restoration
-up
Figure 6: Stride length of attention mask in hybrid sparse attention. Notice that when the stride length equals 1, the model utilizes a dense mask.
Based on the proposed unfactorized sparse self-attention, we are ready to propose the hybrid sparse attention efficiently economize attention FLOPs while maintaining optimal model performance. In the initial stages of training, HSA employs an unfactorized strided sparse attention mask to conserve attention FLOPs. As the model becomes denser, HSA transitions to dense attention mask, ensuring full restoration of model performance. Figure [6 visually illustrates the stride length of the attention mask in hybrid sparse attention during the training.
Remark 3.2. MST seamlessly integrates three interconnected components: the SV for temporal sparsification of MTP layers, MG for spatial sparsification of MTP layers, and HSA for optimizing self-attention layers. SV and MG work in tandem to harness the inherent parameter redundancy within transformers, dynamically shaping the model's connectivity both temporally and spatially throughout training. While SV orchestrates the temporal evolution of sparsity, MG fine-tunes spatial connectivity in real-time. Tn parallel, HSA ad- dresses the computational redundancy inherent in self-attention mechanisms by distributing the operation across multiple steps, thereby minimizing FLOPs while preserving model performance. The concurrent op- eration of SV and MG ensures a comprehensive approach sparsity management in MLP layers, with HSA enhancing computation efficiency in self-attention layers. This integrated framework enables the efficient and effective training of transformers, leveraging the synergies between temporal and spatial sparsification, and computational redundancy to achieve significant pretraining FT.PP reductions without compromising model performance.
4 Experiment
This section outlines the experinental evaluation of our proposed MST, focusing on auto-rogressive language modeling using GPT-2 (Brown el al. 2020), since GPT-2 and its variants are pivolal models in the domain of transformer- based LLMss In addition, we also take experiments on BERT (Devlin et al. 2018) to validate the offectiveness of MST. The evaluation involves benchmarking MST's performance against baseline approaches on various language modeling tasks. We also conduct an ablation study of individual components MST. The primary goal of this section is to showcase the advantages of MST in achieving 4x reduction in
8

---

Published in Transactions on Machine Learning Research (03/2025)
pretraining FLOPs while maintaining performance comparable to dense models. The results averaged four random seeds and detailed experiment configurations are deferred to Appendix A The code is available at
4.1 Performance Comparison
Table 3: Performance comparison of different pretraining methods on GPT-2.
LAMBADA (ACC)
LAMBADA
Wiki-2 (PPL)
PTB
Wiki-108 103
IBW
Method
FLOPs
(PPL)
(PPL)
(PPL)
(PPL)
Dense
847.8G
60.74
13.22
34.89
34.06
36.29
44.16
Tiny
212.7G
53.29
31.91
58.00
58.71
60.78
71.96
SS-80%
267.7G
50.59
54.77
83.98
75.98
88.00
99.44
RigL-80%
267.7G
55.75
25.81
43.00
48.14
44.82
62.31
MST (Ours)
219.4G
60.54
13.67
33.33
34.64
34.95
45.90
We conducted extensive experiments to assess the performance of our MST method against various baseline approaches across a range of zero- shot tasks from Brown et al. (2020) and few-shot. tasks from GLUE (Wang et al. 2018). Our baseline models encompass range of pretraining strate- gies, including standard dense pretraining, compact dense model (Tiny), static sparse pretraining with 80% sparsity (SS-80%). and dynamic sparse pretraining with 80% sparsity facilitated by RigL. (Evci et al. 2020) (RigI- 80%). Detailed benchmark configurations and baseline setups can be found in Appendix [A] The experiment results of zero-shot tasks are sum- marized in Table where we also present the pretraining FLOPs of dif- ferent models. The corresponding wall-clock training time is reported in Appendix B.1 The results of few-shot tasks are deferred in Appendix B.2
Validation Perplexity
60
40
20
2000
4000
6000
Training Step (K)
Figure 7: Perplexity of different
Apart from GPT-2, we also conduct experiments on BERT (Devlin et al. 2018) to validate the effectivencss of MST. The training perplexity is shown in Figure 7 where we find that MST achieves performances com-
pretraining methods on BERT.
parable to dense training, with only less than one-third of the FLOPs required for dense training.
Under our setting, the tiny dense model and MST require only 25% pretraining FLOPs of the dense model, while the two sparse training methods (SS and RigL) require slightly more FLOPs, around 35%. However, among these four FLOPPs-fficient training methods, only MST achieves performances comparable to those of the dense model. Moreover, MST outperforms the dense model in most zero-shot tasks. Notably, SS- 80% is implemented based on SPDF (Thangarasa et al. 2023), while RigL-80% is built upon DynSparse (Dietrich et al. 2021). By the way, also note that the experimental results for the remaining three basclines laaer growth Waan et al. 2023), block sparsification (Dao al.) (2022)), butterfly sparsification Dao et al. (2022c in Table can be referenced from their respective papers, which are incomparable to FLOP reduction ratio. The reason that do not reproduce the experiment results from these three baselines is that they are beyond the scope of network sparsification. Besides, a comparison of our reproduced densc model with OpenAl's official GPT-2 checkpoint. from Hugging Face is deferred to Appendix B.3
4.2 Ablation Study
To gain insights into the contributions of each component in MST, we conducted an ablation study on sparsity variation, topology evolution schema, and hybrid sparse attention from the training pipeline.
9

---

Published in Transactions on Machine Learning Research (03/2025)
4.2.1 Sparsity Variation
We evaluate the impact of different temporal sparsity variation on model performance during pretraining, Specifically, we compare several patterns: fully dense (Dense), fully sparse by MG (Sparse), sparse to dense (SD), dense to sparse to dense (DSD), gradual dense (GD), and our proposed Sparsity Variation (SV) utilizing cubic decay function, illustrated in Section 3.1] The validation perplexity, instant sparsity level, and average FLOP reduction compared to the dense model are plotted in Figure [ As expected, the figure shows that fully dynamic sparse training with high sparsity throughout fails to match dense performance. However, when connections are grown back, all other four methods achieve denso-level performance in terms of validation perplexily. This underscores the importance of the restoration phase for sparse training mel hods of transformers to recover dense cense-comparable performance.
While GD demonstrates the fastest growth, only achieves 1.5x reduction in FLOPs. In contrast, SD achieves a higher FLOP reduction of over 2x. Moreover, the inclusion of an initial dense phase, observed in DSD, outperforms SD) after 120K in terms of validation perplexity with similar FLOP efficiency, underscoring the importance of sparse topology initialization for model performance. Our SV, employing a cubic decay function for sparsity variation, features a burn-in phase that is more efficient than sudden changes in sparsity, such as DSD, while maintaining comparable model performance. Thus, our proposed MST offers an efficient temporal sparsification pattern to achieve dense-comparable performance while simultaneously maximizing FLOP savings up to 2.7x. Additionally, we provide hyperparameter recommendations for SV in Appendix A.3
-
Dasse
-
Sparse ---
GD
-9-
DSD
-D- SV
80
Avenige FLOP Reduction
Validation Porplocity locittyy
70
Sparsity (%)
60
40
30
20
20
40
60
80
100
120
140
0
40
80
100
120
140
20
40
60
80
100
120
140
Training Step (K)
Training Step (K)
Training Step (K)
Figure 8: Ablation study different sparsity variation patterns.
4.2.2 Topology Evolution Scheme
We showcase the effectiveness of proposed topology evolution scheme, Mixed-Growiin (MG), by comparing it with different topology evolution schemes in the task of training an ultra-sparse GPT-2 model with 90% sparsity. The baseline methods encompass static sparse training (SS), dynamic sparse training SET (Mocanu al. 2018), RigL (Evci al.) 2020 and MEST Yuan et al. 2021 respectively.
100
Validaton Perpleity
Parplesity
.00
+55
50
45
Traninn Step (K)
Trannin Step (K)
Figure provides an overview of the validation perplexity during the training of these meth-
Figure 9: Perplexity of different topology evolution schemes.
ods, with the right subplot offering a closer examination the highlighted region from the left subplot. SET consistently encounters loss spikes across all soeds occurring after 18K training steps, leading outlier performance. Among the remaining methods, SS demonstrates the poorest performance, underscoring the critical importance of topology evolution during sparse training. RigL. maintains its superiority over SET,
10

---

Published in Transactions on Machine Learning Research (03/2025)
highlighting the importance of gradient-based growing strategies. However, it is worth noting that MEST underperforms Righ, which diverges from observations in smaller models (Nowak et al. 2023), suggesting that sparse training for transformers differs significantly from that of smaller models. Moreover, MG outper- forms the other three methods, demonstrating its suitability for sparse training of transformers, attributed more proactive exploration mechanism, as elaborated Section
4.2.3 Sparse
We conducted a comparison of training performance among sparse models employing unfactorized strided self-attention and unfactorized fixed self-attention, as illustrated in Figure [0] with varying stride lengths. Specifically, we evaluated dense masks, fixed self-attention with stride lengths of 256 and 512, and strided self-attention with stride lengths of 128 and 256. Table 9 in Appendix [A.5] presents the training FLOPs of the different models, while Figure [1 showcases their training performance.
Notably, Fixed-256 and Strided-128 exhibit similar FLOPPs, yet they both exhibit performance gap compared to the dense mask. Conversely, Fixed-512 and Strided- 256 demonstrate comparable performance to the dense mask. We opted for Strided- 256 in our MST as it offers greater FLOP savings than Fixed-512 and achieves better performance than Strided-128.
Figure 10: Fixed at- tention mask.
4.3 Intuitive Insights
In addition to quantitative performance metrics, we present more intuitive insights into components of our proposed MST method. MST employs sparsity variation comprising three phases: warm-up, ultra- sparsification, and restoration. The concept of MST revolves around predominantly training the model during the ullra-sparsification phase to capitalize on the inherent parameter redundancy transformers, thereby contributing significantly to FLOP reduction. The other two phases serve as auxiliary stages for sup- porting the altra-sparsiiication phase to recover the full performance of dense training. The warm-up phase acts as burn-in stage for the model to establish an optimal initial sparse topology, while the restoration phase endeavors to recover the performance loss incured during the ultra-sparsification phase.
Fasidff
Srldd
100
Yalidaiion Papladk
Milteen Pareennnrooooee Prpplndy
10
15
20
30
16
16
Traliing Sigp (8)
Trising Step (K)
Figure 11: Validation perplexity of different atten- tion patterns.
To underscore the necessity mixed-growing in our ultra-sparsification phase, we initially present statistical results on parameter distribution across various scales of GPT-2 models. Subsequently, we visualize the evolution of model parameters during training validate the effectiveness of our sparsity variation. These intuitive findings offer deeper insights into the operation of MST and its implications for efficient transformer pretraining.
Why Growing We present statistical results on parameter magni- tude across various scales of GPT-2 mod-
0.20
Layerl
Standard Deviation
Layer2
0.15
Layer3
els in Figure [2 We observe that the larger model exhibits smaller stan- dard deviation in parameter magnitudes. This suggests that transformers with a substantial number of parameters tend have more homogencous magnitudes. That is why we introduce restoration phase to compensate for the performance
Layer4
0.10
0.05
0.00
Small (117M)
Medium (345M)
Large (762M)
XL (1542M)
Figure 12: Standard deviation of parameters in the first. encoder block in different scales of GPT-2 checkpoints from HuggingFace,
11

---

Published in Transactions on Machine Learning Research (03/2025)
loss in the ultra-sparsiffcation phase, as it may be harder to train a larger model with sparse neural networks throughout.
Moreover, the homogenization of parameters leads to reduced discrepancies in gradients, posing a poten- tial challenge for gradient- based topology evolution schemes. With reduced differences in gradients, these schemes may become less suitable for guiding the evolution of sparse neural architectures. To counteract this limitation, we introduce the Mixed-Growing (MG) scheme, a tailored approach designed specifically for transformers. MG is crafted to inject additional random exploration into the evolution of sparse topol- ogy, ensuring its adaptability to the unique challenges presented by transformers. The efficacy of MG over alternative topology evolution schemes is demonstrated in Section 4.2.2
Why Sparsity Variation? We visualize the model parameters through heatmaps after each phase in mixed sparsity training. For comparison, we include the visualization of regularly trained dense model. Taking the weights of the projection matrix in the attention layers depicted in Figure 13 as an example, we observe horizontal bands in the parameters of the regularly trained dense model, indicating redundancies across output dimensions. Interestingly, the model trained using our mixed sparsity training exhibits similar bands in the same dimensions the regularly trained dense model, albeit with narrower bandwidths. This suggests that our MST method effectively guides the model in learning essential parameters. Furthermore, the learned parameters are more concentrated, making them better suited for subsequent pruning in the post-training stage to facilitate faster inference of transformers, an advantage of our MST method.
.STT (Warm-up up)
MST (Ultra-Sparsiin) Spatiiiiptsiiiitiioon)
MST (Restoration)
Dense
Ourput Dimiension
2000
Ouggu Dimanssio
200
Output Dimenniom
200
Output Diraension
200
10
106
400
400
400
400
600
600
600
600-
200
400
Input Dimension
Input Dimsension
Input Dimension
Ingut Dimension
Figure 13: Heatmap of the weights of the projection matrix in the attention layer.
Additionally, the sparse topology obtained during the warm-up phase (first heatmap) displays aligned bands akin to the dense model, with minimal changes observed during the ultra-sparsification phase compared to the model after this phase (second heatmap). This underscores the effectiveness of the warm-up phase in establishing a favorable sparse initialization. However, we find that the restoration phase is also essential for restoring parameters, as the model parameters after the ultra-sparsffication phase (sccond heatmap) still notably differ from those of the original dense model.
5 Conclusion and Future Work
Our proposed MST method presents novel and effective approach for enhancing the efficiency of trans- former pretraining. By scamlessly integrating Sparsity Variation (SV), Mixed-Growing (MG), and Hybrid Sparse Attention (HSA), MST offers a comprehensive solution to address both algorithmic inefficiencies and computational demands inherent in transformer pretraining. Through experimentation GPT-2, we have demonstrated the efficacy of MST across various language modeling tasks, showcasing an exceptional FLOP reduction of 4x while maintaining the performance of dense models.
We envision that our discoveries lay the foundation for future research aimed at devising even more efficient and scalable approaches for training large-scale models. Extending the application of MST larger scales or diverse model architectures holds substantial merit and can serve corroborate its generalizability Moreover, MST is algorithmic and hardware-agnostic improvement that operates in neural connections
12

---

Published in Transactions on Machine Learning Research (03/2025)
levels. Therefore, it is entirely orthogonal and can scamlessly integrate with existing system-level acceleration
methods, such as training parallelism (Bian et al.) 2021 Zhao et al. 2023 hardware-gssisted attention
operators (Dao et al. 2022c Dao 2023), and mixed precision training (Burgess et al. 2019 Liu et al. 2022)
thus facilitating efficient transformer pretraining and achieving higher acceleration ratios.
Acknowledgments
This work was support by the National Natural Science Foundation of China Grants 52450016 and 52494974.
References
Jinze Bai, Shuai Bai, Shusheng Yang, Shijic Wang, Sinan Tan, Peng Wang, Junyang Lin, Chang Zhou, and Jingren Zhou. Qwen-vl: A frontier large vision-language model with versatile abilities. arXiv preprint or 2rxiv22308. 12966, 2023.
SH Shabbeer Basha, Shiv Ram Dubey, Viswanath Pulabaigari, and Snchasis Mukherjee. Impact of fully connected layers on performance of convolutional neural networks for image classification. Neurocomputing. 378:112 119, 2020.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 2rXiv:2004.05 05100 2020.
Luisa Bentivogli, Elena Cabrio, Ido Dagan, Danilo Giampiccolo, Medea Lo Leggio, and Bernardo Magnini. Building textual entailment specialized data sets: a methodology for isolating linguistic phenomena rele- vant to inference. In LREC. Citeseer, 2010.
Zhengda Bian, Qifan Xu, Boxiang Wang, and Yang You. Maximizing parallelism in distributed training for huge neural networks. arXiv preprint ar arXiv:2105. 14450, 2021.
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.
Neil Burgess, Jelena Milanovic, Nigel Stephens, Konstantinos Monachopoulos, and David Mansell. Bfloat16 processing for neural networks. In 2019 IEEE 26th Symposium on Computer Arithmetic (ARITH), pp. 88 91. LEEE, 2019.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robin- son. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3003 2013.
Lingjiao Chen, Matei Zaharia, and James Zou. Frugalgpt: How to use large language models while reducing cost and improving performance. arXiv preprinl or trXiv:2305.05 05176, 2023a.
Tianyi Chen, Tianyu Ding, Badal Yadav, Πya Zharkov, and Luming Liang Lorashear: Efficient large language model structured pruning and knowledge recovery. arXiv preprint arXiv:2310.18356, 2023b.
Rewon Child, Scott Gray, Alec Radford, and Πya Sutskever. Generating long sequences with sparse trans- formers. XiXio preprint arxi::1944. 100000.1000000 10509, 2019.
Jack Choquette, Wishwesh Gandhi, Olivier Giroux, Nick Stam, and Ronny Krashinsky. Nvidia a100 tensor core gpu: Performance and innovation. IEEE Micro, 41(2):29 35, 2021.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment. challenge. In Machine learning challenges workshop, pp. 177 190. Springer, 2005.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXio preprint 1901 1901002000, 2019.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXie preprint ar arXiv:2307006691, 2023.
Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International conference on learning representations, 2022a.
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; pp. 4690-4721. PMLR, 2022b.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-eficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344 16359, 2022c.
Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing perfor- mance. Xixiv preprint 1rXii:1907.044 2019.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.0753 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deop bidirectional transformers for language understanding. orXiv preprinl arXiv:1810.04880 2018.
Anastasia Dictrich, Frithjof Gressmann, Douglas Orr, Ivan Chelombiev, Daniel Justus, and Carlo Luschi. Towards structured dynamic sparse pre-training of bert. arXio preprint 2rx8:2108. 08806677, 2021.
Bill Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases. In Third International Workshop on Peraphrasing (IWP2005), 2005.
Zhengxiao Du Yujie Qian, Xiao Liu, Ming Ding, Jiczhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. arXiv preprint 1rXiv:2103.10360 2021.
Utku Evei, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making PMLR, 2020.
all tickets winners. In International Conference on Machine Learning, pp.
Elias Frantar and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quan- tization and pruming. Advances in Neural Information Processing Systems, 35:4475-4488, 2022.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. ArXiv 0bs/2201000774, 2023.
Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint Xrv::222 17323, 2022.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa. Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint or arXiu:2101.00027 2020.
Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on teztual entailment and paraphras- ing, pp. 1-9, Association for Computational Linguistics, 2007.
R Bar Haim, Tdo Dagan, Bill Dolan, Tisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Sapek- tor. The second pascal recognising textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entoilment, volume 7, pp. 783-794, 2006.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770-778, 2016.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang. Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint orXiv:1704.04861, 2017.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference Learning Represen- tations, 2021.
Siddhant Jayakumar, Razvan Pascanu, Jack Rac, Simon Osindero, and Erich Elsen. Top-kast: Top-k always sparse training. Advances in Neural Information Processing Systems, 33:20744 20754,, 2020.
David Kirk et al. Nvidia cuda software and gpu parallel computing architecture. In ISMM, volume 7, pp. 103 104, 2007.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint 2E4::00061170.. 2020.
Shiwei Liu, Tianlong Chen, Xiaohan Chen, Zahra Atashgahi, Lu Yin, Huanyu Kou, Shen, Mykola Pech- enizkiy, Zhangyang Wang, and Decebal Constantin Mocanu. Sparse training via boosting pruning plasticity with neuroregeneration. Advances in Neural Information Processing Systems, 34:9908 9922, 2021a.
Shiwei Liu, Lu Yin, Decebal Constantin Mocanu, and Mykola Pechenizkiy. Do we actually need dense wver-aarameterization?? in-time aver-parameterization in sparse training. In International Conference on Machine Learning, pp. 1089. PMLR, 2021b.
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for generic network architectures. In International Conference on Machine Learning, pp. 14139-14152. PMLR, 2022.
Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hao Hc, Antong Li, Mengshen He, Zhengliang Liu, et al. Summary of chatgpt-related research and perspective towards the future of large language models. Meta-Raa Rattoooodd, pp. 100017, 2023.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Lim-pruner: On the structural pruning of large language models. anXiv preprint #xXi:23355.1162 2023.
Mitchell P. Marcus, Beatrice Santorini, and Mary Amn Marcinkiewicz. Building large annotated cor- pus of English: The Penn Treebank. Computational Linguistics, 19(2):313-330, 1993. URL https:
//aclanthology.orgg hwwpchoooogyoorr/J990 J93-20//
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprinl Xrv:::1609 07843, 2016.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint Xrxio:1710 03740, 2017.
Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibeseu, and An- tonio Liotta. Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science. Nature communications, 9(1):3383, 2018.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa. Patwary, Vijay Kor- thikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, 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, pp. 1-15, 2021.
Aleksandra Nowak, Bram Grooten, Decebal Constantin Mocanu, and Jacek Tabor. Fantastic weights and how to find them: Where to prune in dynamic sparse training. arXiv preprint ar arXiu:2306. 12230, 2023.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaclla Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context, arXiv preprint orXie:1606.06000 2016.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison,, Luca Antiga, and Adam Lerer. Automatic diflerentiation in pytorch, 2017.
Alexandra Peste, Eugenia Iofinova, Adrian Vladu, and Dan Alistarh. Ae/de: Alternating com-
pressed/decoppresse training of deep neural networks. Advances in neural information processing sys- tems, 34:8557-8570, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, (11):4485-5557) 2020.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. Code Ilama: Open foundation models for code. arXiv preprinl ar Xiv:23008 12950, 2023.
Michael Santacroce, Zixin Wen, Yelong Shen, and Yuanzhi Li. What matters in the structured pruning of generative language models? arXiv preprint arXiv:2302.0777 03773, 2023.
Jonathan Schwarz, Siddhant Jayakumar, Razvan Pascanu, Peter E Latham, and Yee Teh. Powerpropagation: A sparsity inducing weight reparameterisation Advances in neural information processing systems, 34: 2021.
Hang Shao, Bei Liu, and Yanmin Qian. One-shot sensitivity-aware mixed sparsity pruning for large language models. arXiv preprint orXiv:2310.09499, 2023.
Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. orxiv preprint arx:2:2306. 11662, 2023.
Aaquib Syed, Phillip Huang Guo, and Vijaykaarti Sundarapandiyan. Prune and tune: Improving effi- cient pruning techniques for massive language models, 2023. URL httpp://ppanerriee.ne///oooo cKlgcx7nSZ
Kai Sheng Tai, Taipeng Tian, and Ser Nam Lim. Spartan: Differentiable sparsity via regularized transporta- tion. Advances in Neural Information Processing Systems, 42222 2022.
Vithursan Thangarasa, Abhay Gupta, William Marshall, Tianda Li, Kevin Leong, Dennis DeCoste, Sean Lie, and Shreyas Saxena. Spdf: Sparse pre-training and dense fine-tuning for large language models. In Workshop on sporsity in neural networks at international conference learning representations, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bash- lykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint orXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07401. 2018.
Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Lconid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang, and Yoon Kim. Learning to grow pretrained models for efficient transformer training. arXiv preprint arXi:230330099 2023.
16

---

Published in Transactions on Machine Learning Research (03/2025)
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. In Workshop on Advancing Newral Network Troining: Compulational Efficiency, Scalability, and Resource Optimization ( (WANTO NewrIPS 2023), 2023.
Geng Yuan, Xiaolong Ma, Wei Niu, Zhengang Li, Zhenglun Kong, Ning Liu, Yifan Gong, Zheng Zhan, Chaoyang Hc, Qing Jin, et al. Most: Accurate and fast memory-cconomic sparse training framework the edge. Advonces in Neural Information Processing Systems, 34:20838 20850, 2021.
Mingyang Zhang, Chunhua Shen, Zhen Yang. Linlin Ou, Xinyi Yu, Bohan Zhuang, et al. Pruning meets low-rank parametee--fffeient fine-tuning. arXiv preprinl or Xix:20305 18403, 2023.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint (xXix22344.12277 2023.
Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning nam fine-grained structured sparse neural networks from scratch. In International Conference on Learning Representations, 2020.
Michael H. Zhu and Suyog Gupta. To prume, not to prume: Exploring the efficacy of pruning for model compression, 2018. URL
17

---

Published in Transactions on Machine Learning Research (03/2025)
A Experiment Details
A.1 Hardware Setup
Our experiments implemented with PyTorch 2.1.1 (Paszke et al. 2017). CUDA 12.1 (Kirk al. 2007)
run on 4x NVIDIA A100 Tensor Core GPUs. Each run needs about 24 hours to train the dense model for 100K steps.
A.2 Model and Implementation Details
Considering OpenAI docs not relcase the training datasct, WebText, of GPT-2 (Brown al. 2020
the nanoGPT code base from htlps://giluub..om/kaaaaa prrrpaahhynanoGPT nanooPT7 NanoGPT is lightweight version of the GPT-2 model trained on the OpenWebText dataset (Gao et al. 2020). Our experiment implementation is derived from the small GPT-2 model. The architecture comprises 12 transformer layers and 12 attention heads, with an embedding size set to 768. The text is tokenized with the GPT-2 tokenizer (Brown et al. 2020p. We adopt the train-validation split provided by nanoGPT. The training set comprises 9 billion tokens, and the validation set contains 4.4 million tokens. During training, we optimize the cross-cntropy loss for next-token prediction. Consistent with nanoGPT, we employ GELU activations while disabling bias and Dropout. Distributed data parallelism with gradient accumulation is employed to enable batch size of 480. Training is conducted with bfloat16 precision on machines with 4 A100 GPUs.
A.3 Hyperparameter Settings of MST and baselines
We begin by detailing the hyperparameter settings for the dense model in Table sourced from NanoGPT These parameters are also utilized as the public parameters for the subscquent methods. Following this, provide the privale hyperparameter settings specific to MST, RigL, SS, and Tiny in Tables and S respectively.
Table 4: Public hyperparameters.
Hyperparameter
Value
Optimizer
AdamW
Activation function
GELU
Number of gradient accumulation steps
5 x 8
Batch size
12
Input sequence length
1024
Number of layers
Number of attention heads
12
Embedding dimensionality
768
Dropout rate
0.0
Bias
Not used
Learning rale
6 x 10-1
Minimum learning rate
6 x 10
Iteration intervals for learning rate decay
140000
Total number of training iterations
140000
Weight decay coefficient.
0.1
Warmup steps
2000
Threshold value for gradient clipping
0.1
Exponential decay rate for the moving average of the gradient 31
0.9
Exponential decay rate for the moving average the squared gradient Bg
0.95
18

---

Published in Transactions on Machine Learning Research (03/2025)
Table 5: Private Hyperparameters of MST.
Calegory
Hyperparameter
MST
Maximum sparsity Sy
96%
Number of sparsity levels N
5
Pruning frequency Aw
2000
SV
Warm-up up phase duration Tw
10000
Ultrassparsffication phase duration Tu
100000
Growing frequency AR
2000
Growing training phase duration TR
10000
Initial topology update fraction Çi
0.3
DST
Topology update interval Á
100
Ratio of connections activated through random growth
0.25
HSA
Stride of attention mask
256
Table 6: Private Hyperparameters of RigL.
Category
Hyperparameter
RigL
Sparsity SRigL
80
DST
Initial topology update fraction Ço
0.3
Topology update interval A.
100
Table 7: Private Hyperparameters of SS.
Category
Hyperparameter SS
DST
Sparsity Sas
80
Table 8; Private Hyperparameters of Tiny.
Category
Hyperparameter
Tiny
Embedding dimensionality
384
Model Size
Number of layers
6
A.4 Benchmark Details
In our performance evaluation, i.e. Table 12 in Section 4.1 5 datasets and few-shot evaluation 2 subtasks of GLUE. The detailed information about the datascts is elaborated in the following.
we perform zero-shot evaluation of the models
A.4.1 Zero-Shot Tasks
LAMBADA The LAMBADA dataset (Paperno et al. 2016), sourced from BookCorpus, comprises 10,022 passages, which are further divided into 4,869 development passages and 5,153 test passages. It evaluates the capability of computational models to capture long-range dependencies and text comprehension. This task involves predicting the final word of sentences, which exhibits the feature wherein human participants can predict the final word when provided with the entire passage but struggle when given only the preceding sentence. To succeed on LAMBADA, computational models cannot simply rely on local context, but must demonstrate the capacity track information from the broader discourse. The training dataset for language models evaluated LAMBADA includes totalling 203 million words.
19

---

Published in Transactions on Machine Learning Research (03/2025)
PTB The English Penn Treebank (PTB) corpus (Marcus et al. 1993) encompasses diverse collection of English text. drawn from sources such news articles, magazines, and other publications, particularly the section corresponding to Wall Street Journal articles, which stands out as one of the most prominent and widely used datasets for evaluating sequence labelling models. This task involves annotating each word with its respective Part-of-Speech lag. In the conventional split of the corpus, sections 0 to 18 serve as the training set (comprising 38,219 sentences and 912,344 tokens), sections 19 to 21 function as the validation set (consisting of 5,527 sentences and 131,768 tokens), and sections 22 to 24 serve as the test set (comprising 5,462 sentences and 129, 229,654 tokens).
WikiText The WikiText language modeling dataset (Merity et al. 2016 comprises compilation of more than 100 million tokens extracted from a selection of Good and Featured articles on Wikipedia, commonly used in varions language modeling tasks, including next-word prediction, text generation and text classi- fication. Compared to the preprocessed version of Penn Treebank (PTB), WikiText-2 more than twice the size, while WikiText-103 is over 110 times larger. Its composition of full articles makes it particularly suitable for models capable of capturing long-term dependencies.
IBW The One Billion Words (IBW) dataset (Chelba et al., 2013 originally introduced by the Google Brain team, is a substantial English language corpus for pretraining language models, which contains almost one billion words in the training data and is openly accessible for research purposes. This benchmark dataset covers various genres and topics ranging from news and technology to novels and is extensively used to assess the performance of statistical language models. Rescarchers leverage the 1BW dataset to pretrain language models, enhancing their performance across various downstream NLP tasks such as text classification, sentiment analysis, and language generation.
A.4.2 Few-shot
General Language Understanding Evaluation (GLUE) (Wang et al. 2018 benchmark comprises a suite of diverse language understanding tasks, ranging from sentence-level to discourse-leved including tasks like text classification, sentence similarity, etc. It serves as a standardized flatform for evaluating and comparing the performance of various NLP models, aiming to facilitate advancements in the field by providing unified evaluation framework. Our experimental section uses 2 subtasks of the GLUE benchmark-Ree and MRPC.
RTE The Recognizing Textual Entailment (RTE) datasets come from series of anmual textual entailment challenges. The authors of the benchmark combined the data from RTE1 (Dagan et al. 2005), RTE2 (Hamm et al. 2006), RTE3 (Giampiccolo et al. 2007), and RTE5 (Bentivogli et al. 2010). Examples are constructed based on news and Wikipedia text. The authors of the benchmark convert all datasets to a two-class split, where for three-class datasets they collapse neutral and contradiction into not entailment, for consistency. The RTE dataset is designed to assess models' understanding of textual entailment. The dataset sources texts from scries of annual textual entailment challenges. It consists of pairs of text, cach comprising premise and a hypothesis, where the task is to determine whether the premise logically entails the hypothesis.
The Microsoft Research Paraphrase Corpus (MRPC) DDoan & Brocket 2005 is a corpus of sen- tence pairs automatically extracted from online news sources. It is designed to evaluate models' performance in identifying semantic equivalence between pairs of sentences. The task involves determining whether the two sentences in each pair convey similar meanings.
MRPC
A.5 FLOP Calculation
As shown in Figure [ the major FLOPs of the GPT-2 model come from fully connected layers and self- attention layers. We illustrate the computation of these two parts in the following. We disregard other layers which are effectively irrelevant or contribute minimally to the total computation overhead such as LayerNorm and Softmax. Furthermore, we exclude the FLOPs associated with the topology evolution process, it occurs every 100 training steps, thus its influence the final result is deemed negligible.
20

---

Published in Transactions on Machine Learning Research (03/2025)
Initially, for sparse network comprising L fully connected layers and without considering bias terms, the required FLOPs for forward pass can be computed
PLOPstorward
1)0;
(1)
where is the sparsity, is the input dimensionality, and Or is the output dimensionality of the (-th layer. Note that Eq. [ is also adopted in (Evci et al. 2020 When it comes to training FLOPs, the caleulation involves multiple forward and backward passes across various networks. The FLOPs during the forward pass consist of two main components: the transformer and the final linear output layer, and the computation in the transformer primarily stems from the atlention blocks and the multi-layer perceptron block.
Based on the model architecture, we can first delve into the details of the FLOP calculation of the attention mechanism blocks. We denote as the input soquence length, Nembd the embedding dimensionality, Deea the size of each attention head, Qatten the ratio of FLOPs by sparse self-allention over the full self-attention, S the spansity pre-set. Besides, Table shows the Qatten in differren sparse self-attentio patterns.
Table 9: FLOPs of different sparse self-attention patterns.
Dense
Fixed
Fixed
Strided
Strided
Pattern
Mask
256
512
128
256
Qatten (%) 100.00
12.70
25.10
12.07
22.03
Thus, we have the following FLOPs:
1. Performing the projection to key, query, and value:
FLOPskav = (1 - 5) x L x ((2 x Nemba - 1) x x Nemba)
Calculating the allention scores:
FL.PP Sscores Qatten 2 L. L Nembd
3. Aggregating the values through the reduction:
FLOPsreduce = Qatten x 2 x Nembd X (L x L X Dhead)
4. Performing the final linear projection:
FLOPsproj = (1 - 5) x L X ((2 X Nomba - 1) x Nemba)
Therefore, we can compute the total FLOPs for attention blocks as:
FLOP FLOPskev + FLOP Secores FLOP FL.OPpprof
Now, we shift focus to calculating the FLOPs for the MLP blocks. Before proceeding, we denote Däw as the
leed-forward size (the dimensionality of the hidden layer between the two linear layers in the MLP block), and it commonly set to four times the value of Nemba- The calculation of FLOPs for the MLP involves summing the computational costs of two linear layers:
FLOPSHWI =(1 S) x L ((2 Nemba - 1) D
FLOP =(1 S) x L X ((2 x DEw 1) X Nemba)
21

---

Published in Transactions on Machine Learning Research (03/2025)
And we have
FL.PPMMPP FL.OPSHw FL.OPsff2
Therefore, the FLOPs of the transformer during the forward process is
Stranfoooer N (FLOP Sutten + FLOPSMLP)
where N denotes the number of layers (model depth), FLOP and FLOPSMPP denote the FLOPs required in the attention and MLP blocks, respectively.
So far, we have clucidated the FLOP calculation for the internal modules of the transformer. During the forward pass, there one remaining part, namely the final language model output layer, which used convert. hidden representations into probability distributions over the vocabulary. We denote Dveeal as the vocabulary size, and have:
FL = (1 - 5) L ((2 x Nembd 1) x Dvocal)
Then we obtain the total FLOPs for the forward process:
FLOPsforward FLOP FLOOPStransformee Stransformer + FLOPsmm
For the FLOPs of gradients backward propagation, denoted as FLOPStackwards we compute it as twice the computational cost. of the forward pass, which is adopted in existing literature (Evci et al.; 2020), i.e.,
Sbackward = 2 x Storwadd Finally, the total FLOPs during the training process can given by FLOP Stotal FLOP Sforward FLOPShhckward-
FLoPSbaakkwrd
B Supplementary Experiment Results
B.1 Wall-Clock Training Time
Table 10: Wall-clock rumning time of different methods.
Method
Rumning Time of 140K Training Steps
Dense
30.10h
Tiny
11.31h
SS-80%
31.53h
RigL-80%
32.63h
MST (Ours)
34.32h
B.2 Performance comparison Few-Shot Tasks
Table 11: Performance comparison on Few-Shot Tasks.
RTE (ACC)
MRPC (ACC)
Method
FLOPs
Dense
847.8G
52.98
71.26
Tiny
212.7G
55.87
71.63
SS-80%
267.7G
50.00
68.50
RigL-80%
267.7G
48.83
67.22
MST (Ours)
219.4G
52.62
70.96
22

---

Published in Transactions on Machine Learning Research (03/2025)
B.3 Comparison with OpenAl's Checkpoint
A comparison between reproduced dense model and NanoGPT sourced from the official GPT-2 check- point by OpenAI via HuggingFace, is presented in Table 12 Upon review of Table 12 we observe that our model exhibits comparable performance to the OpenAl model across most tasks. However, in certain instances, such as on WikiText2 and WikiText103, our model demonstrates inferior performance. This dis- crepancy may be attributed to slight variations model architecture and the utilization of different training datascts. While we employ the OpenWebTest dataset, OpenAl utilizes the WebText datasct, which is not publicly available.
Table 12: Performance comparison of our reproduced GPT with with OpenAT's official GPT-2 checkpoint from HuggingFace
LAMBADA (ACC)
LAMBADA
WikiText2
PTB
10
1BW (PPL)
Method
(PPL)
(PPL)
(PPL)
(PPL)
Dense (OpenAI)
60.43
12.04
25.19
30.46
26.38
43.39
Dense (Ours)
60.74
13.22
34.89
34.06
36.29
44.16
B.4 Training Curves of Comparative Evaluation in Section 4.1
The training curves of various methods evaluated in Section 4.1
Upon examination of Figure 14 it becomes apparent that both Tiny and RigL fail to achieve performance comparable to the dense model. Additionally, the static sparse method encounters issues, evidenced by loss spikes across all four seeds. Notably, MST emerges as the sole method capable of reaching dense performance. Furthermore, the training curve of MST exhibits a spike at 100K training steps. This occurs due to the transition of the self- atlention mask a dense configuration at this moment within the hybrid sparse attention scheme. Despite this spike, MST rapidly converges back normal training trajectory autonomously, without requiring manual intervention for restoration.
depicted in Figure 14
Dosee
Tany Denac
85
Rigl
MST
80
Validation Peplexity
70
60
50
40
30
20
40
100
120
140
Training Step (K)
Figure 14: Training Curves of Comparative Evaluation in Section
B.5 Different Update Schedules in MST
We delved deeper into the impact of update schodules by training models with various schemes, including 1-cosine, 2-cosine, N-cosine, and Decay N-cosine. The experiment outcomes, depicted in Figure 15 reveal distinct performances across these schedules. The results indicate that the single cosine scheme struggles to maintain performance during the uitra-sparsffication phase, while the 2-cosine scheme exhibits suboptimal performance during the restoration phase. Conversely, both N-cosine and Decay N-cosine schemes consis- tently yield good performance across both stages. This consistency can be attributed to the stepwise nature
23

---

Published in Transactions on Machine Learning Research (03/2025)
of network sparsity in these stages, necessitating reset of the update schedule to prevent convergence is- sues and subsequent performance degradation. Furthermore, the Decay N-cosine scheme outperforms the N-cosine scheme, as the decay of the update fraction helps stabilize the training process, thereby enhancing performance.
Perplesty (Sooothe))
3,7
Peplociy (Smoothee)
3,575
Upslate Fraetion
3,580
0.2
3.525
3.3
3,500
Decoy-Nowii
3,475
0.0
20
40
60
80
100
140
60
80
100
120
10
20
30
40
50
60
70
80
90
100
110
120
150 150
Training
Triining Step
Figure 15: Ablation study on different mixed update schedules.
B.6 Visualization of Model Parameters
We visualize the model parameters through healmaps after each phase in mixed sparsity training, In addition to Figure 13 in Section 4.3 we present heatmaps weights from various fulll connected layers in Figures [6] and Figure exhibits similar horizon bands as seen in Figure 13
indicating consistent patterns
across layers. Conversely, parameter distributions in Figures and 18 these differences, model consistently learns parameters closely aligned with those of the dense model, underscoring the effectiveness of our MST method.
appear more uniform. Despite
MST (Pruning)
MST (Sparse Training)
MST (Growing)
Dense
200
Output Dimession
2.00
Output Dimerasóm
200
Output Dinvaasin
200
16
400
400
400
400
Omtpt
101
600
600
600
600
200
400
600
200
400
600
400
600
400
600
Input Dimension
Input Dimension
Input Dimension
Input Dimension
Figure 16: Heatmap of the weights of weight matrix in the second attention layer.
MST (Pruning)
MST (Sparse Training)
MST (Growing)
Dense
0
Ouppu Dimension
200
Output Dimensioo
200
Oulput Dimension
200
Oulput Dinceusion
200
-00
600
600
600
600
100
Input Dim ensionn
Input Dimeeusion
Ingut Dimension
Input Dimenssoo
Figure 17: Heatmap of the weights of weight matrix in the first MLP layer.
24

---

Published in Transactions on Machine Learning Research (03/2025)
MST (Pruning)
MST (Sparse Traninn))
MST (Growwng)
Dense
Output Dimeension
200
Output Dimensson
2000
Ousput Dimmmsion
200
Output Diavension
200
400
400
400
400
12
200
400
600
200
400
600
0
200
400
600
200
400
600
Input Dimension
Input
Input
Input
Figure 18: Heatmap of the weights of weight matrix in the second MLP layer.
25


Reviews:
[{"review_id": "AQ60t07aaD", "reviewer": "Reviewer_HPPd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n1. Theoretically, the proposed method MST achieves a 4x FLOP reduction\n2. Experiments highlight the effectiveness of Sparsity Variation (SV), Mixed-Growing (MG), and Hybrid Sparse Attention (HSA) through comprehensive ablations\n3. This paper claim that MST is compatible with other acceleration techniques, such as training parallelism, hardware-assisted attention operators, and mixed precision training, enabling broader adoption.\n\n\n\n\n\n\n\n\n\n\n**Weaknesses**\n1. Dynamic Sparse Training (DST) is challenging to implement in real-world scenarios as memory usage decreases gradually during training, requiring additional engineering efforts to handle memory optimizations effectively. More explanation needed for this point.\n2. The claimed FLOP reductions may be theoretical and difficult to achieve in practical settings without experimental validation of real training speed improvements. Adding training time as the x-axis in Figures 7, 8, 9, and 11 could provide better evidence of the effectiveness of the proposed method.\n3. This paper only presents the experimental results on GPT-2; its performance on larger or other types of LLMs, like Llama, is not explored. It would be better if experiments on more LLM architectures could be conducted.\n4. This paper claims that MST is orthogonal to other system-level accelerations, such as training parallelism, hardware-assisted attention operators, and mixed precision training. I am interested in this point. For example, how can Hybrid Sparse Attention be combined with FlashAttention to increase FLOPs/s? Claims of MST’s compatibility with other accelerations, such as training parallelism and hardware-assisted methods, lack demonstration. Specific evaluations, such as integrating Hybrid Sparse Attention with FlashAttention, could strengthen this claim.\n5. Not strong support for “transformers exhibit considerable redundancy in pretraining computations”. It would be better to provide have empirical support for the claim that transformers exhibit redundancy in pretraining computations.\n6.  The proposed method MST primarily combines existing methods, such as Dynamic Sparse Training and Hybrid Sparse Attention, which may limit its perceived novelty.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces Mixed Sparsity Training (MST), a method for reducing the computational overhead of pretraining transformers by dynamically integrating sparsity techniques. The proposed method, MST, leverages Sparsity Variation (SV), Mixed-Growing (MG), and Hybrid Sparse Attention (HSA) through three phases: warm-up, ultra-sparsification, and restoration. It achieves a 4x reduction in FLOPs without compromising performance, as demonstrated using GPT-2 across various tasks. MST is compatible with existing acceleration techniques such as training parallelism, mixed precision, and hardware-assisted methods.\n\n\nThe key claimed contributions of this paper are:\n1. proposes a Mixed Sparsity Training (MST) framework combining multiple dynamic sparsity techniques.\n2. demonstrates up to a 4x reduction in FLOPs while retaining comparable performance to dense models.\n3. experiment to validates MST through experiments on GPT-2, showing effectiveness across zero-shot and few-shot tasks.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the **Weaknesses**.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "bYagZeCd9K", "reviewer": "Reviewer_Uasj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The proposed method achieves a 4× reduction in computational cost during pretraining without compromising model performance, a significant saving given the increasing computational demands of AI models.  \n- Dynamic Sparse Training (DST), Sparsity Variation (SV), and Hybrid Sparse Attention (HSA) are presented as innovative and promising techniques.  \n- The method is compatible with other system-level optimizations, such as mixed-precision training and distributed parallelism, enabling additional efficiency gains.  \n\n\nWeaknesses:\n\n- While the method effectively demonstrates a reduction in FLOPs, it remains unclear if this translates into tangible end-to-end speedups. Including wall clock time comparisons would strengthen the analysis.  \n- The efficiency of implementing the proposed scheme on hardware is uncertain. Although the authors suggest seamless integration with methods like flash attention, the extent of actual speedup achievable on hardware is not adequately demonstrated.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Mixed Sparsity Training (MST), a novel approach to significantly reduce the computational cost of pretraining large transformer-based models like GPT-2 while maintaining performance. MST integrates three techniques—Dynamic Sparse Training (DST), Sparsity Variation (SV), and Hybrid Sparse Attention (HSA)—to achieve a 4× reduction in FLOPs. The training process is divided into three phases: warm-up, which transitions the dense model to a sparse topology; ultra-sparsification, where the model is predominantly trained in a sparse configuration; and restoration, which recovers potential performance loss by reinstating connections. By dynamically adjusting sparse connections and using a sparse attention mask that transitions to dense over time, MST addresses algorithmic inefficiencies inherent in pretraining.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Consider including wall-clock time comparisons and analyses in the experiments to provide a clearer understanding of the practical speedup.  \n- Add a discussion on hardware implementation feasibility and its implications to better contextualize the method's real-world applicability.  \n- Expand the experiments to include models beyond BERT and GPT-2 to demonstrate the broader applicability of the proposed method.", "broader_impact_concerns": "There is no outstanding ethical issues from this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YNBmJYLrTG", "reviewer": "Reviewer_hWaJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The motivation is clear and the proposed method is well-defined.\n* The problem tackled can contribute to the development of the LLM foundation model.\n* Experiments show the improvement and demonstrate the benefit of the method.\n* Some interesting insights are provided and discussed, which would be beneficial for the readers.\n\nWeaknesses:\n* The design of the sparsity determination formula for each iteration in the warm-up phase and restoration phase is not clear. Why do you choose the specific formula? Is the method sensitive to the formula design? Any ablation for this design?\n* What are the additional FLOPs entailed from the unfactorized sparse self-attention compared to its factorized counterpart? How to balance the tradeoff between the two? I feel it to some extent similar to the relationship between structured and unstructured sparse training. But as we know, although unstructured sparse training theoretically also has good FLOPs reduction, it has strict requirements for the hardware, like GPU type. Does unfactorized sparse self-attention also have similar limitations?\n* As mentioned in the previous comment, I was wondering if MST has a requirement for the hardware. Does it achieve real running time acceleration, instead of theoretical FLOPs?\n* Experiments are only conducted on relatively small LLMs, like GPT-2.\n* What is the sparsity rate for MST in Table 1? Would the trend be different when sparsity is further increased like 90% or reduced like 60%?\n* You mention that the method is orthogonal to training parallelism. But I feel the ultra-sparsification phase as well as the model update during the phase would be influenced by the training parallelism.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new pre-training method for LLMs called mixed sparsity training (MST), which not only maintains LLMs' performance but also minimizes training FLOPs. To better deal with the tradeoff between performance and training FLOPs, MST finds a way to integrate dynamic sparse training with Sparsity Variation and Hybrid Sparse Attention during the pretraining process and consists of three phases in the process, including warm-up, ultra-sparsification, and restoration. Specifically, the warm-up phase converts the dense neural network into a sparse one, the ultra-sparsification update model weights and sparsity topology with high sparsity and less training FLOPs, and the restoration phase reinstates connections for the final output neural network. Experiments demonstrate the benefits of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* It would be good to add more explanation on the design of the sparsity determination formula for each iteration in the warm-up phase and restoration phase. \n* It would be good to also add some real running time comparison, instead of the theoretical FLOPs comparison.\n* It would be good to run more experiments on larger LLM models, like Llama.\n* It would be good to include the standard deviation of the table values.\n* It would be good if you could provide more discussion about why the proposed MST is orthogonal to training parallelism.\n* The design of MST seems to also work for smaller DNNs. Could you provide some discussions about which part is tailored for LLM pre-training?\n* You mentioned that there is a reduced discrepancy in gradients within the larger model. Could you provide more evidence for it?", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As above.
Decision: Accept (poster)