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: The authors proposed a new point of view to explain why deep learning models often require more weights are needed to train models rather than to run inference for tasks. One example is that sparse models usually perform well in inference, but do not perform well in training. The main content of this viewpoint is that more weights during training can expand the search space, create extra degrees of freedom for search, and form new optimization paths when model optimization falls into critical points, thereby facilitating model training. Based on this point of view, this paper proposes a series of steps to augment search spaces of sparse models during training to approximate the behavior of larger models. The experimental results show that this series of steps can effectively reduce the model error.