Search Spaces for Neural Model Training

1 INTRODUCTION . In the area of deep learning , increasing the size of neural models has led to dramatic advances ( Hestness et al. , 2017 ; Kaplan et al. , 2020 ; Henighan et al. , 2020 ) , motivating training with hundreds of billions of parameters ( Brown et al. , 2020 ; Fedus et al. , 2021 ) . However , larger models incur higher memory costs and runtimes , which limits training and inference tasks that can run on existing hardware ( Thompson et al. , 2020 ) . Nevertheless , larger models have been shown to train faster ( Li et al. , 2020 ) and better ( Kaplan et al. , 2020 ) , which drives their adoption . An area of research that seeks to compensate the increasing costs of model size is sparsity ( Hoefler et al. , 2021 ) , which moves weights to zero so they can be discarded from storage or computations . Sparsity emerges as a promising option to reduce costs of inference ( Luo et al. , 2017 ; Allen-Zhu et al. , 2019 ) , as neural models are capable of performing tasks on sizes smaller than they were trained for ( Narang et al. , 2017 ; Renda et al. , 2020 ) . However , since overparameterization remains critical for training , sparsity observes limited success there ( Lee et al. , 2019 ; Frankle & Carbin , 2019 ; Wang et al. , 2020 ; Tanaka et al. , 2020 ; Frankle et al. , 2021 ; Bellec et al. , 2018 ; Mocanu et al. , 2018 ; Dettmers & Zettlemoyer , 2019 ; Evci et al. , 2020a ; Jayakumar et al. , 2020 ) . The ability to run inference but not training with fewer weights points to a problem with search ( or how well can models navigate parameter spaces during training ) rather than model capacity ( LeCun et al. , 1990 ) . Literature shows adding more weights to training creates extra degrees of freedom that form new paths for optimization , rendering neural model training more effective ( Evci et al. , 2020b ) . While it is widely believed that having more weights facilitates escape from critical points during training , the question remains whether these weights are always needed ( e.g. , learn meaningful representations ) or could be eventually discarded if they only help with the optimization process . In this paper we explore how to remove weights from training without restricting its ability to effectively traverse parameter spaces towards good solutions . Surprisingly , we discover that weights pruned at inference time are transient and do not learn representations in the long-term . As such , these weights do not need to be stored throughout all of training , in contrast to weights that learn meaningful information for inference and thus need to be stored . The pruned weights however do learn short-term representations that provide the model extra degrees to escape regions of bad saddle points ( Kawaguchi , 2016 ) and high error plateaus ( Dauphin et al. , 2014 ) that can slow down learning , after which they can be discarded . With this understanding we propose new recommendations for training sparse models that consolidate recent advances in sparsity literature : ( 1 ) rewire weights to expand the parameter space , ( 2 ) update gradients of all weights to encourage alternate paths of optimization , and ( 3 ) discard intermediate representations to reduce noise from gradient accumulations . Following these recommendations , we show that sparse neural models achieve competitive results on dozens of deep learning workloads , even when satisfying constraints needed to accelerate training and inference using Sparse Tensor Cores ( Mishra et al. , 2021 ) in NVIDIA GPUs . The paper is organized as follows . In Section 2 we introduce a metric for learning and investigate the duration which weights retain meaningful information throughout training . We then devise recommendations for removing weights that do not learn long-term representations when training sparse models in Section 3 . Section 4 summarizes our methodology for sparse training . Section 5 shows sparse models constructed in this fashion perform competitively across a plethora of deep learning workloads and can target current hardware accelerators ( Mishra et al. , 2021 ) for training and inference . Section 6 relates to prior research . In Section 7 we conclude with directions for future work . 2 METRIC FOR LEARNING . The crux of this paper is to understand whether weights that are being added for training are learning meaningful representations or solely being used to facilitate the optimization process . For this purpose , we use correlations to measure the duration that weights in a neural model retain ( or learn ) information that is relevant for inference . Correlations determine how similar are two sets of data . When applied to weights learned over time , correlations measure how weights in future training time steps are similar to past ones . We postulate correlations represent the degree with which weights in a neural model learn . Weights that learn meaningful representations for inference are repeatedly reinforced throughout training , and their values form distinct temporal patterns that are correlated to past values . In contrast , weights that do not learn ( and can be removed during inference ) exhibit random behavior and have no dependence over time . We measure correlations between a pair of time series x and y using the Pearson coefficient ( Rodgers & Nicewander , 1988 ) : ρx , y ≡ E [ ( x − µx ) ( y − µy ) ] /σxσy , where µ is the mean , σ the standard deviation , and E [ · ] the expected value . Coefficients of value +1 denote two series have identical trends , 0 indicates the series are random , and −1 represents series with opposite behavior . We treat each individual weight as a time series w ∈ { w1 , w2 , . . . , wt } and compute Pearson coefficients between windows x ∈ { w1 , w2 , . . . , wt−τ } and y ∈ { wτ , wτ+1 , . . . , wt } , representing the similarity between the weights and their values at some future time τ . Correlations naturally decay with τ since temporal patterns are less likely to persist over long durations . An important metric we adopt is the time ∆ after which weights no longer correlate to their future values , namely when ρ = 0 . Figure 1 plots the time ( ∆ ) after which weights become completely uncorrelated with their past values . We observe that the correlation length grows with weight magnitudes obtained after training . While weights needed for inference exhibit correlations that persist across a significant fraction of training , weights that are not needed have short-term correlations and behave randomly over short periods . As a result , long-term correlations of large weights are characteristic of learning , signifying repeated reinforcement along the weight direction . On the other hand , the near-random motion of small weights suggests they don ’ t learn useful representations and can be removed during inference . Because weights removable during inference are absent of correlations over long periods of training , we conjecture they do not need to be stored all the time . However , we still need pruned weights during training due to their short-term correlations . More specifically , their short-term interactions are critical for training to take different paths for optimization , facilitating escape from bad saddle points ( Kawaguchi , 2016 ) or high error plateaus ( Dauphin et al. , 2014 ) that can slow down learning . Combining this understanding , in later sections we discover their learned values can be periodically destroyed throughout training ( once correlations fall to zero ) with no impact on accuracy , which paves the road for sparse training . 3 RECOMMENDATIONS FOR SPARSE TRAINING . This section outlines recommendations for training sparse neural models . While many of these ideas have been circling in sparsity literature , they remain misunderstood which hinders adoption . We use the understanding from the previous section to provide a deeper intuition in these cases . To summarize , we uncover the following steps to train sparse neural models : ( 1 ) Frequently rewire weights that participate in training . ( 2 ) Perform gradient updates for weights that do not participate in training . ( 3 ) Induce exploitation ( e.g. , stop rewiring after some amount of time , reset non-participating weights to zero , or regularize them ) . 3.1 REWIRING OF NEURAL WEIGHTS . Sparse models are commonly trained by rewiring ( or sampling ) a subset of weights from a neural model according to some criteria ( magnitude , sign , gradients , etc ) , allowing them to learn which weights are needed for inference . Some works rewire weights every training iteration ( Bellec et al. , 2018 ; Wortsman et al. , 2019 ; Raihan & Aamodt , 2020 ; Jayakumar et al. , 2020 ; Zhou et al. , 2021 ) , while others rewire every hundreds of training steps ( Evci et al. , 2020a ) or after an entire pass through the data ( Mocanu et al. , 2018 ; Dettmers & Zettlemoyer , 2019 ) . Since there are no clear reasons behind these choices , an important question becomes how often must weights be rewired so neural models can learn effectively . We study the effects rewiring rates have on accuracy when training sparse models of various sizes d , where d denotes the fraction of total weights being used in the sparse model . Figure 2 ( left ) plots the task error ( or accuracy difference between dense and sparse models ) as a function of the number of training steps r taken between rewirings . We find rewiring frequency does not matter when sparsity is low or moderate ( d ∼ 0.5 ) , but errors increase with r as models get sparser ( d 0.5 ) . Therefore , weights should be rewired frequently ( within a few training steps ) to avoid losing accuracy . Correlations can help us explain why training improves when we rewire sparse neural models . While weights that do not learn long-term representations can be discarded during training , they have shortterm correlations that create extra degrees of freedom for escaping critical points ( Evci et al. , 2020b ) . But this escape only occurs when weights are rewired . For example , neural models always operate on reduced parameter spaces if weights being used remain unchanged throughout training ( r →∞ ) , and expand exactly once after going through the entire data when they are rewired every epoch . By swapping between weights that participate ( are used in forward and backward propagations ) and do not participate ( are fixed to zero during training ) , training can take different paths for optimization using the short-term correlations , which helps neural models explore the parameter space . 3.2 UPDATES FOR WEIGHTS NOT PARTICIPATING IN TRAINING . Previously , we found neural weights not needed for inference form short-term interactions that make training more effective . Since sparse models lack weights that have similar roles , we can approximate this behavior using gradient updates for non-participating weights . Because non-participating weights do not contribute to the loss , their gradients determine whether another optimization path is better suited to traverse the region being explored in the loss landscape . Repeating gradient updates registers the importance of these paths over some period of training and can trigger rewiring when weights exceed a threshold , which allows training to explore different regions of the parameter space while operating on ( or forward and backward propagating ) a smaller set of weights . We evaluate the importance of gradients updates on non-participating weights by reducing their contribution with a scale factor s. Figure 2 ( middle ) shows the task error as a function of the scale factor s using various model sizes d. We observe error increases as s → 0 ( no gradients contribute when s = 0 ) , which can be attributed to premature convergence when training lacks expressive power to explore different paths . Dampening the gradient updates affect sparser models ( d 0.5 ) , which are more sensitive to critical points , more than larger models ( d ∼ 0.5 ) , which may still have sufficient weights to train effectively . As a result , non-participating weights should be updated regularly to retain accuracy . Gradient updates for weights that do not participate in training could mean they also learn representations , such that more capacity rather than search improves accuracy . We verify this by resetting non-participating weights to zero after every z training steps , thus removing any representations they might have learned . Figure 2 ( right ) shows the task error as a function of z . Errors are the highest when values are reset every iteration ( z = 1 ) , which is equivalent when training without any rewiring . Conversely , error decreases as weights are reset less frequently and saturates after sufficient training steps ( z ∼ 1k ) . Interestingly , this also represents the time it takes for added weights to decorrelate , as shown in previous sections , supporting our conjecture that correlations are indicative of learning , and weights can be discarded once they cease to exist . The ability to remove information from non-participating weights reinforces the idea that they make training more effective rather than augment model capacity . While most of literature trains sparse models using fewer weights and gradients ( Bellec et al. , 2018 ; Mocanu et al. , 2018 ; Dettmers & Zettlemoyer , 2019 ) , more recent success has been found by updates gradients for weights that do not participate in training ( Wortsman et al. , 2019 ; Liu et al. , 2020 ; Zhou et al. , 2021 ; Hubara et al. , 2021 ) as described above . However , so far there has been little understanding of why this helps improve sparse training , which we attribute to the use of short-term correlations to escape critical points .

Summary: This paper seeks to explore and understand the improved performance of Neural Networks trained with additional parameters even though a majority of the weights can be pruned during inference. They analyze the magnitude, correlation, and movement of weights during training and propose a hypothesis that the reason is that the weights which can be pruned during inference provide degrees of freedom in the loss landscape which allows better optimization of the non-prunable weights. Based on this, they explore and recommend a methodology for sparse training as a method to efficiently approximate a higher dimensional weight space that maintains these good optimization properties. They ablate the effects of the different recommendations and explore the hyperparameter space to have a better understanding of the dynamics of the sparse training.