DASH: Warm-Starting Neural Network Training Without Loss of Plasticity Under Stationarity

Authors: baekrok  shin
junsoo  oh
hanseul  cho
chulhee  yun
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-27 10:26:47

Abstract: Warm-starting neural networks by initializing them with previously learned weights is appealing, as practical neural networks are often deployed under a continuous influx of new data. However, it often leads to *loss of plasticity*, where the network loses its ability to learn new information, resulting in worse generalization compared to training from scratch. This occurs even under stationary data distributions, and its underlying mechanism is poorly understood. We develop a framework emulating real-world neural network training and identify noise memorization as the primary cause of plasticity loss when warm-starting on stationary data. Motivated by this, we propose **Direction-Aware SHrinking (DASH)**, a method aiming to mitigate plasticity loss by selectively forgetting memorized noise while preserving learned features. We validate our approach on vision tasks, demonstrating improvements in test accuracy and training efficiency.

Main paper: 
 introduction: 
 When training a neural network on a gradually changing dataset, the model tends to lose its plasticity, which refers to the model's ability to adapt to new information (Lyle et al., 2023b; Dohare et al., 2021; Nikishin et al., 2022) . This phenomenon is particularly relevant in scenarios with nonstationary data distributions, such as reinforcement learning (Igl et al., 2020; Nikishin et al., 2022) and continual learning (Wu et al., 2021; Chen et al., 2023; Kumar et al., 2023) . A common explanation is that previously learned information becomes less relevant over time, necessitating models to overwrite the existing knowledge (Lyle et al., 2023b) . Under this viewpoint, various efforts have been made to mit-Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024).igate the loss of plasticity, such as resetting layers (Nikishin et al., 2022) , regularizing weights (Kumar et al., 2023) , and modifying architectures (Nikishin et al., 2023; Dohare et al., 2021; Sokar et al., 2023; Lyle et al., 2023a; Lee et al., 2023) .Perhaps surprisingly, a similar phenomenon occurs in supervised learning settings, even where new data points sampled from a stationary data distribution keep being introduced to the dataset during training. It is counterintuitive, as one would expect advantages in both generalization performance and computational efficiency when we warm-start from a model pre-trained on data points of the same distribution. For a particular example, when a model is pre-trained using a portion of a dataset and then we resume the training with the whole dataset, the generalization performance is often worse than a model trained from scratch (i.e., cold-start), despite achieving similar training accuracy (Ash & Adams, 2020; Berariu et al., 2021; Igl et al., 2020) . Liu et al. (2020) report a similar observation: training neural networks with random labels leads to a spurious local minimum which is challenging to escape from, even when retraining with a correctly labeled dataset. Interestingly, Igl et al. (2020) found that pre-training with random labels followed by the corrected dataset yields better generalization performance than pre-training with a small portion of the (correctly labeled) dataset and then training with the full, unaltered dataset. It is striking that warm-starting leads to such a severe loss of performance, even worse than that of a cold-started model or a restarted model from a pre-trained one with random labels, despite the stationarity of the data distribution.These counterintuitive results prompt us to investigate the underlying reasons for them. While some studies have attempted to explain the loss of plasticity in deep neural networks (DNNs) under non-stationarity (Lyle et al., 2023b; Sokar et al., 2023; Lewandowski et al., 2023) , their empirical explanations rely on various factors, such as model architecture, datasets, and other variables, making it difficult to generalize the findings (Lewandowski et al., 2023; Lyle et al., 2023a) . Moreover, there is limited research that explores why warm-starting is problematic in stationary settings, highlighting the lack of a fundamental understanding of the loss of plasticity phenomenon in both stationary and non-stationary data distributions. In each experiment, a constant number of data points are added to the existing training dataset. Models are trained until achieving 99.9% train accuracy before proceeding to the next experiment; the plot on the right reports the number of update steps executed in each experiment. Results are averaged over three random seeds. "Cold" refers to cold-starting and "Warm" refers to warm-starting.The Shrink & Perturb (S&P) method involves shrinking the model weights by a constant factor and adding noise (Ash & Adams, 2020) . Notably, DASH, our proposed method, achieves better generalization performance compared to both training from scratch and S&P, while requiring fewer steps to converge.In this work, we aim to explain why warm-starting leads to worse generalization compared to cold-starting, focusing on the stationary case. We propose an abstract framework that combines the popular feature learning framework initiated by Allen-Zhu & Li (2020) with a recent approach by Jiang et al. (2024) that studies feature learning in a combinatorial and abstract manner. Our analysis suggests that warmstarting leads to overfitting by memorizing noise present in the newly introduced data rather than learning new features.Inspired by this finding, we propose Direction-Aware SHrinking (DASH), which aims to encourage the model to forget memorized noise without affecting previously learned features. This enables the model to learn features that cannot be acquired through warm-starting alone, enhancing the model's generalization ability. We validate DASH using the online learning setting from Ash & Adams (2020) , employing various models, datasets, and optimizers. As an example, Figure 1 shows the promising results in terms of both test accuracy and training time.Loss of Plasticity. Research has aimed to understand and mitigate loss of plasticity in non-stationary data distributions. Lewandowski et al. (2023) explain that loss of plasticity cooccurs with a reduction in the Hessian rank of the training objective, while Sokar et al. (2023) attribute it to an increasing number of inactive neurons during training. Lyle et al. (2023b) find that changes in the loss landscape curvature caused by non-stationarity lead to loss of plasticity. Methods addressing this issue in non-stationary settings include recycling dormant neurons (Sokar et al., 2023) , regularizing weights towards initial values (Kumar et al., 2023) , and combining techniques (Lee et al., 2023) like layer normalization (Ba et al., 2016) , Sharpness-Aware Minimization (SAM) (Foret et al., 2020) , resetting layers (Nikishin et al., 2022) , and Concatenated ReLU activation (Shang et al., 2016) .However, these explanations and methods diverge from the behavior observed in stationary data distributions. Techniques aimed at mitigating loss of plasticity under nonstationarity are ineffective under stationary distributions, as shown in Appendix A.1, in line with the observations in Lee et al. (2023) . While some works study the warm-starting problem in stationary settings, they rely on empirical observations without theoretical analysis (Ash & Adams, 2020; Berariu et al., 2021; Achille et al., 2018) . The most relevant work by Ash & Adams (2020) introduces the Shrink & Perturb (S&P) method, which mitigates loss of plasticity in stationary settings to some extent by shrinking all weight vectors by a constant factor and adding noise. However, they do not explain why this phenomenon occurs or why S&P is effective. We develop and analyze a theoretical framework to explain why warm-starting suffers even under stationary data distribution. Based on our findings, we propose a method that shrinks the weight vector in a direction-aware manner to maintain properly learned features.Feature Learning in Neural Networks. Recent studies have investigated how training methods and network architectures influence generalization performance, focusing on data distributions with label-dependent features and labelindependent noise (Allen-Zhu & Li, 2020; Cao et al., 2022; Jelassi & Li, 2022; Zou et al., 2023; Deng et al., 2023; Oh & Yun, 2024) . In particular, Shen et al. (2022) examine a data distribution consisting of varying frequencies of features and large strengths of noise, emphasizing the significance of feature frequencies in learning dynamics. Jiang et al. (2024) propose a novel feature learning framework based on their observations in real-world scenarios, which also involves features with different frequencies but considers the learning process as a discrete sampling process. Our proposed framework extends these ideas by incorporating features with varying frequencies, noise components, and the discrete learning process while introducing a more intricate learning process that captures the key aspects of feature learning dynamics in gradually expanding datasets. 
a framework of feature learning: 
 Shen et al. ( 2022) consider a neural network trained on data with features of different frequencies and noise components stronger than the features. The gradient of the loss for each single data point aligns more with the noise than the features due to the larger scale of noise, making the model more likely to memorize noise rather than learn features. However, an identical feature appears in many data points, while noise appears only once and do not overlap across data points. Thus, if a feature appears in sufficiently high frequency in the dataset, the model can learn the feature. Thus, the model's learning of features or noise depends on the frequency of features and the strength of noise.Inspired by Shen et al. (2022) , we propose a novel discrete feature learning framework. This section introduces a framework describing a single experiment, while Section 3 analyzes expanding datasets scenarios. As our focus is on gradually expanding datasets, carrying out the (S)GD analysis over many experiments as in Shen et al. (2022) is highly challenging. Instead, we adopt a discrete learning process similar to Jiang et al. (2024) but propose a more intricate process reflecting key ideas from Shen et al. (2022) . In doing so, we generalize the concept of plasticity loss and analyze it without assuming any particular hypothesis class for a more comprehensive understanding, whereas existing works are limited to specific architectures.We consider a classification problem with C classes, and data are represented as (x, y) ∈ X × [C], where X denotes the input space. A data point is associated with a combination of class-dependent features V(x) ⊂ S y where S c = {v c,1 , v c,2 , . . . , v c,K } isT = {(x i , y i )} i∈[m]denotes a training set. The set N , representing the data points with non-zero gradients, will be defined below. The frequency of a feature v in data points belonging to N is denoted byg(v; T , N ) = 1 |T | (x,y)∈N 1(v ∈ V(x)),where 1(•) is the indicator function, which equals 1 if the condition inside the parentheses is true and 0 otherwise. At each step of training, if L and N are given, the model chooses the most frequent feature among the features not yet learned, i.e., arbitrarily choose v ∈ arg max u∈S\L g (u; T , N ).The model decides whether to learn a selected feature v by comparing its signal strength, represented by |T | • g(v; T , N ), with the signal strength of noise, given by γ, which reflects the key ideas of Shen et al. (2022) . If the frequency of the selected feature v is no less than the threshold γ/|T |, i.e., g(v; T , N ) ≥ γ/|T |, the model learns v and adds it to its set of learned features L. The feature learning process continues until the model reaches a point where the selected feature v has g(v; T , N ) < γ/|T |, indicating that the signal strength of every remaining feature is weaker than that of noise. At this point, the feature learning process ends.We consider a data point x to be well-classified if the model f has learned at least τ features from V(x), i.e., |L ∩ V(x)| ≥ τ , where τ < K. In this case, we consider x to have a zero gradient, meaning it cannot further contribute to the learning process. Throughout the feature learning process, the set N of data points with non-zero gradients is dynamically updated as new features are learned. At each step, when the model successfully learns a new feature, we update N by removing the data points that satisfy |L ∩ V(x)| ≥ τ , as they become well-classified due to the newly learned feature.If the feature learning process ends and the model has learned as many features as it can, the remaining data points that have non-zero gradients will be memorized by fitting the random noise present in them and will be considered to have zero gradients. This step concludes the training process. A detailed algorithm of the learning process can be found in Algorithm 2 in Appendix D.In our framework, the model selects features based on their frequency in the set of unclassified data points N . The intuition behind this approach is that features appearing more frequently in the set of data points will have larger gradients, leading to larger updates, and we treat g(v; T , N ) as a proxy of the gradient for a particular feature v. As a result, the model prioritizes learning these high-frequency features in a sequential manner. However, if the frequency g(v; T , N ) of a particular feature v is not sufficiently large, such that the total occurrence of v is less than the strength of the noise, i.e., |T | • g(v; T , N ) < γ, the model may struggle to learn that feature. Consequently, the model will prioritize learning the noise over the informative features. When this situation arises, the learning procedure becomes sub-optimal because the model fails to capture the true underlying features of the data and instead memorizes the noise.The threshold τ determines when a data point is considered well-classified and acts as a proxy for the dataset's complexity. A higher τ requires the model to learn more features for correct predictions, while a lower τ allows accurate predictions with fewer learned features. Experiments in Appendix B Figure 13 and 14 support this interpretation.The model predicts on unseen data points by comparing the learned features with features present in a given data point x. If the overlap between the learned feature set L and the features in x, denoted as V(x), is at least τ , i.e., |V(x) ∩ L| ≥ τ , the model correctly classifies the data point. Otherwise, the model resorts to random guessing.Accurately measuring training time within our discrete learning framework is challenging. To address this, we introduce an alternative for training time based on empirical observations. We consider the number of training data points with non-zero gradients at the initial stage of training as a proxy for training time, which represents the amount of "learning" required for the model to classify all data points correctly. To verify this, we conducted experiments with CIFAR-10 on ResNet-18. We used the gradient norm as a proxy for the number of data points with non-zero gradients and investigated its correlation with the number of training steps required to achieve 99.9% training accuracy.  
experiments with expanding dataset: 
 In this section, we establish the setup where the dataset grows after each j-th experiment, with j ∈ N to compare warm-start and cold-start within our proposed learning framework. To better understand the loss of plasticity under stationary data distribution, we consider an extreme form of stationarity where the frequency of each feature combination remains constant in each chunk. We investigate if the loss of plasticity can manifest even under this strong stationarity. The detailed description of the dataset across the entire experiment is as follows: Assumption 3.1. In each j-th experiment, we are provided with a training dataset T j := {(x i,j , y i,j )} i∈ [n] with n samples. For each class c ∈ [C] and each possible feature combination A ⊂ S c , we assume that T j contains exactly n A ≥ 1 data points with associated feature set A, where the values of n A are independent of j. Note that c∈[C],A⊂Sc n A = n.When training in the j-th experiment, we use the cumulative dataset T 1:j := l∈[j] T l , the union of all training datasets up to the j-th experiment. Remark 3.2. In each experiment, the feature combinations remain the same across the dataset, but the individual data points differ. This is because each data point is associated with its own specific noise, which varies across samples. Although the underlying features are the same, the noise component of each data point is unique. This approach ensures that the model is exposed to a diverse set of samples.h(v; L) ≜ 1 n c∈[C],A⊂Sc n A • 1 (v ∈ A ∧ |A ∩ L| < τ )which denotes the portion of data points containing v that cannot be well-classified by a learned feature set L. This leads to technical assumptions: Assumption 3.3. For any learned feature set L ⊂ S, if v 1 , v 2 ∈ S c for some class c ∈ [C] and h(v 1 ; L) = h(v 2 ; L), then v 1 = v 2 . Also, we assume for any class c ∈ [C] there exists some τ distinct features v 1 , . . . , v τ ∈ S c such that g(v 1 ; T j , T j ), . . . , g(v τ -1 ; T j , T j ) ≥ γ/n and g(v τ ; T j , T j ) < γ/n.This assumption leads to Lemma C.2, stating that the order in which features are learned within a class is deterministic. This is just for simplicity of presentation and can be relaxed. The last assumption is justified by the moderate number of data points in each chunk T j , ensuring the existence of both τ -1 learnable features and a non-learnable feature within a class. Throughout the following discussion, we will proceed under above assumptions unless otherwise specified.Now we analyze the warm-start and cold-start initialization methods within our framework, focusing on test accuracy and training time. In our learning framework, we denote a model at step s of the j-th experiment as f (j,s) . We denote the set of learned features and the set of memorized data for the model f (j,s) as L (j,s) and M (j,s) , respectively. We also denote the set of data points that have non-zero gradients at step s of the j-th experiment as N (j,s) . We define two respective versions of these sets and the model, one for warm and cold, denoted by the subscripts (e.g., f (j,s) warm and f (j,s) cold ). We note that, by definition, L , where s j denotes the last step of j-th experiment. Besides, we use a shorthand notation for step s j of the experiment j that we drop s if s = s j (e.g., L (j) := L (j,sj ) ). For the detailed algorithms based on our learning framework, see Algorithms 3 and 4 in Appendix D.In the test data, a feature combination A ⊂ S c of data point with class c ∈ [C] appears with probability n A /n along with data-specific noise. By Section 2.3, test accuracy for a learned set L and training time are defined as:ACC(L) ≜ 1 -C-1 C • 1 n c∈[C],A⊂Sc n A • 1 (|A ∩ L| < τ ) T (J) warm ≜ j∈[J] N (j,0) warm , T (J) cold ≜ j∈[J] N (j,0) coldBased on these definitions, the following theorem holds:Theorem 3.4. There exists nonempty G ⊊ S such that we always obtain L(1) warm = L(1) cold = G. For all J ≥ 2, the following inequalities hold:ACC L (J) warm ≤ ACC L (J) cold T (J) warm < T (J) cold Furthermore, ACC(L (J) warm ) < ACC(L (J) cold ) holds when J > γ δn where δ ≜ max v∈S\G h(v; G) > 0.Proof Idea. After the first experiment, the data points in T 1 cannot further contribute to the learning process of the warmstarted model. Consequently, even when a new data chunk is provided in subsequent experiments, the feature frequencies are too small, resulting in a weak signal strength of features that cannot overcome the noise signal strength. As a result, the model memorizes individual noise components of the new data points. This procedure is repeated with every experiment, causing the learned feature set to remain the same as at the end of the first experiment. In contrast, when receiving T 1:j at once (cold-starting), the signal strength of features is large enough to overcome the noise signal strength, allowing the model to learn many more features.Theorem 3.4 highlights a trade-off between cold-starting and warm-starting. Regarding test accuracy, the theorem concludes that cold-starting can achieve strictly higher accuracy than warm-starting. However, warm-starting requires strictly shorter training time compared to cold-starting.Detailed proof is provided in Appendix C. Theorem 3.4 suggests that the loss of plasticity in the incremental setting under the stationary assumption can be attributed to the noise memorization process. Actually, a similar observation is made in real-world neural network training. It is widely believed that during the early stages of training, neural networks primarily focus on learning features from the dataset, and after learning these features, the model starts to memorize data points that it fails to classify correctly using the learned features. To investigate this phenomenon, we conducted an experiment where CIFAR-10 was divided into two chunks, each containing 50% of the training dataset. The model was pre-trained on one chunk and then further trained on the full dataset for 300 epochs. We used three-layer MLP and ResNet-18 with 10 random seeds.Figure 2 shows the change in the model's performance based on the duration of pre-training. When pre-training is stopped at a certain epoch and the model is then trained on the full dataset, test accuracy is maintained. However, if pre-training continues beyond a specific threshold (approximately 50% pre-training accuracy in this case), warmstarting significantly impairs the model's performance as it increasingly memorizes training data points. We attribute this phenomenon to the neural network's memorization process after learning features. This is consistent with reports of a critical learning period where neural networks learn useful features in the early phase of learning (Achille et al., 2018; Frankle et al., 2020; Kleinman et al., 2024) , and with findings that neural networks tend to learn features followed by memorizing noises (Arpit et al., 2017; Jiang et al., 2020) .Remark 3.5. Igl et al. (2020) find that training a model on random labels followed by corrected labels results in better generalization compared to pre-training on a subset of correctly labeled data and then further training on the full dataset with the same distribution. Achille et al. (2018) also observe that pre-training with slightly blurred images followed by original images yields worse test accuracy than pre-training with random label or random noise images. These findings align with our observations: re-training with corrected labels after random label learning "revives" gradi-ents for most memorized data points, enabling new feature learning. Conversely, with static distributions, gradients for memorized data points remain suppressed, leading to learning from only a few data points with active gradients, causing memorization. 
proposed method: 
 4.1. Motivation: An Idealized MethodIn Section 3, we observed a trade-off between warm-starting and cold-starting. Cold-starting often achieves better test accuracy compared to warm-starting, while warm-starting requires less time to converge. The results suggest that neither retaining all learned information nor discarding all learned information is ideal. To address this trade-off and get the best of both worlds, we consider an idealized algorithm where we retain all learned features while forgetting all memorized data points. For any experiment J ≥ 2, if we consider the ideal initialization, learned features L (J-1) ideal are retained, and memorized data pointsM (J-1)ideal are reset to an empty set. Pseudo-code for this method is given in Algorithm 5, which can be found in Appendix D. We define T(J) ideal ≜ j∈[J] N (j,0)ideal as the training time with the idealized method, where N (j,0) ideal represents the set of data points having a non-zero gradient at the initial step of the j-th experiment. Then, we have the following theorem: Theorem 4.1. For any experiment J ≥ 2, the following holds:ACC L (J) cold = ACC L (J) ideal T (J) warm <T (J) ideal < T (J) coldThe detailed proof is provided in Appendix C. The idealized algorithm addresses the trade-off between cold-starting and warm-starting. We conducted an experiment to investigate the performance gap between these initialization methods.Synthetic Experiment. To see if our theoretical investigation applies more robustly to more realistic scenarios within our framework, we conducted an experiment that more closely resembles real-world settings. Instead of fixing the frequency of each feature set, we sampled each feature's existence from a Bernoulli distribution to construct V(x). This ensures that the experiment is more representative of real-world scenarios. Specifically, for each data point (x, y), we uniformly sampled y ∈ {0, 1}. From the feature set S y corresponding to the sampled class y, we sampled features where each feature's existence follows a Bernoulli distribution,1 (v y,k ∈ V(x)) ∼ Ber(p k ), for all v y,k ∈ S y .This approach allows us to model the variability in feature occurrence that is commonly observed in real-world datasets while still maintaining the core principles of our learning framework. We set the number of features, K = 50, where each p k is sampled from a uniform distribution, U(0, 0.2).In each chunk, we sampled 1000 data points and set the total number of experiments to 50, with γ = 50 and τ = 3. We sampled 10000 test data from the same distribution.As shown in Figure 3 , the results align with the above theorems. Random initialization, i.e. cold-starting, and ideal initialization achieve almost identical generalization performance, outperforming warm initialization. However, with warm initialization, model converges faster, as evidenced by the number of non-zero gradient data points, which serves as a proxy for training time. Ideal initialization requires less time compared to cold-starting, which is also consistent with Theorem 4.1. Due to the sampling process in our experiment, we observe a gradual increase in the number of learned features and test accuracy in warm-starting, mirroring real-world observations. We verified that these results are consistent across a wide range of hyperparameter values (see Figure 14 -16 in Appendix B).The aforementioned ideal method recycles memorized training samples by forgetting noise while retaining learned features. This raises the question of whether such an idealized method can be implemented in real-world neural network training. To address this, we propose our algorithm, Direction-Aware SHrinking (DASH), which intuitively captures this idea in practical training scenarios. The outlined behavior is illustrated in Figure 4 . In simple terms, DASH shrinks each weight vector proportionally to the cosine similarity between the weight vector and the negative gradient of the loss calculated with train data, with more emphasis on newer data. If the degree of alignment is small (i.e., the cosine similarity is close to or below 0), we consider that the weight vector has not learned a proper feature and shrink it significantly to make it "forget" learned infor-Feature Noise θ 1 θ 2 θ1 θ2 -∇ θ 1 L(Θ; 𝒯 1:j ) -∇ θ 2 L(Θ; 𝒯 1:j )Figure 4 . Illustration of DASH. We compute the loss L with training data T1:j and obtain the negative gradient. Then, we shrink the weights proportionally to the cosine similarity between the current weight θ and ∇ θ L, resulting in θ.mation. This allows weights to forget memorized noises and easily change its direction. Conversely, if the weight vector and negative gradient are well-aligned (i.e., the cosine similarity is close to 1), we shrink the weight vector to a lesser degree to maintain the learned information. This intuitive method aligns with the aforementioned idea of the algorithm, as it allows us to shrink weights that have not learned proper information while retaining weights that have learned commonly observed features.The shrinking is done per neuron, where the incoming weights are grouped into a weight vector denoted as θ. For convolutional filters, the height and width of the kernel are flattened to form a single weight vector θ for each pair of input and output filters. DASH has two hyperparameters: λ and α. Hyperparameter λ is the minimum shrinkage threshold since each weight vector is shrunk by max{λ, cos sim}. α is the exponential moving average coefficient for each chunk's gradient. Lower α values emphasize previous gradients, indicating less shrinkage, suitable for less complex datasets where retaining learned features is crucial, and vice versa. The algorithm is presented in Algorithm 1.To validate whether our intuition aligns well with DASH, we plotted the accuracy on previously learned data points in Figure 12 , Appendix B. As experiments progress, when applying DASH, the training accuracy of the previous dataset recovers faster than other methods. We argue that this is due to the nature of our algorithm, which forgets memorized noise while preserving learned features. As the number of experiments increases, the number of learned features also grows. Consequently, we can retain more features compared to previous experiments, resulting in an increase in training accuracy across experiments after applying DASH. 
experimental details: 
 Our online learning setup is similar to the one described in Ash & Adams (2020) . We divided the training dataset into Algorithm 1 Direction-Aware SHrinking (DASH) Require:• Model f Θ with list of parameters Θ • Training data points T 1:j • Averaging coefficient 0 < α ≤ 1 • Threshold λ > 0 Initialize: G (0) θ ← 0, ∀θ in Θ for i in 1 : j do ℓ ← Loss(f Θ , T i ) U Θ ← Gradient of loss ℓ for θ in Θ do G (i) θ ← (1 -α) × G (i-1) θ + α × U θ end for end for for θ in Θ do s θ ← CosineSimilarity -G (j)θ , θ θ ← θ ⊙ max{λ, s θ } end for return: f Θ 50 chunks, and at the beginning of each experiment, new data is combined with the existing training data. Models were considered converged and each experiment was terminated when training accuracy reached 99.9%, aligning with our learning framework. Results were averaged over three random seeds. We conducted experiments with vanilla training, i.e., without any data augmentations, weight decay, learning rate scheduling, etc. Additionally, we performed experiments with state-of-the-art (SoTA) settings that include aforementioned techniques. We evaluated DASH on Tiny-ImageNet, CIFAR-10, CIFAR-100, and SVHN using ResNet-18 (He et al., 2016) , VGG-16 (Simonyan & Zisserman, 2014) , and three-layer MLP architectures with batch normalization layer. Models were trained using Stochastic Gradient Descent (SGD) and SGD based Sharpness-Aware Minimization (SAM), both with momentum.DASH was compared against baselines (cold-starting, warmstarting, and S&P (Ash & Adams, 2020) ) and methods addressing plasticity loss under non-stationarity (L2 INIT (Kumar et al., 2023) and Reset (Nakkiran et al., 2021) ). Layer normalization (Ba et al., 2016) and SAM (Foret et al., 2020) , known to mitigate plasticity loss in reinforcement learning (Lee et al., 2023) , were applied to both warm and cold-starting. Consistent hyperparameters were used across all methods, with details provided in Appendix A.4. S&P, Reset, and DASH were applied whenever new data was introduced. We report two metrics for both test accuracy and steps required for convergence: the value from the final experiment and the average across all experiments. We first experimented with CIFAR-10 on ResNet-18 to determine if methods from previous works for mitigating plasticity based on non-stationarity can be a solution to our incremental setting with stationarity. Appendix A.1 shows that L2 INIT, Reset, layer normalization, SAM, and reviving dead neurons, are not effective in our setting. Thus, we conducted the remaining experiments without these methods.  
discussion and conclusion: 
 In this work, we defined an abstract framework for feature learning and discovered that warm-starting benefits from reduced training time compared to random initialization but can hurt the generalization performance of neural networks due to the memorization of noise. Motivated by these observations, we proposed Direction-Aware SHrinking (DASH), which shrinks weights that learned data-specific noise while retaining weights that learned commonly appearing features. We validated DASH in real-world model training, achieving promising results for both test accuracy and training time.Loss of plasticity is problematic in situations where new data is continuously added daily, which is the case in many real-world application scenarios. Our research aimed to interpret and resolve this issue, preventing substantial wastes in energy, time, and the environment. By elucidating the loss of plasticity phenomenon in stationary data distributions, we have taken a crucial step towards addressing challenges that may emerge in real-world AI, where the continuous influx of additional data is inevitable. We believe our findings contribute to the development of more efficient and sustainable AI that can adapt to ever-increasing real-world data.We hope our fundamental analysis of the loss of plasticity phenomenon sheds light on understanding this issue as well as providing a remedy. To generalize our findings to any neural network architecture, we treated the learning process as a discrete abstract procedure and did not assume any hypothesis class. Future research could focus on understanding the loss of plasticity phenomenon via optimization or theoretically analyzing it in non-stationary data distributions, such as in reinforcement learning. 
 
 Appendix: This section provides omitted experiments and hyperparameters used in these experiments. Section A.1 presents the results using non-stationary solutions. The following section compares the experiment results with other baselines, except for those discussed in Appendix A.1. Additionally, we include results conducted on the state-of-the-art (SoTA) setting with CIFAR-10, CIFAR-100 on ResNet-18, in Appendix A.3. In the Appendix A.4, we include hyperparmeters used in these experiments.The same experimental settings as in Section 5.1 are used. We conducted an experiment with an incremental training dataset comprised of 50 chunks. At the start of each experiment, a new chunk is provided and added to the existing training dataset. Before proceeding to the next experiment, the model is trained until achieving 99.9% train accuracy. Including several baselines, we compared DASH with warm-starting while resetting the momentum at the start of every experiment. We found that this approach cannot be a solution either, as shown in Tables 3 4 5 6 .All experiments are conducted using NVIDA-A6000 (48GB VRAM) GPU with two AMD EPYC 7763 64-Core Processor. Number of steps required to converge are provided in each table.In this subsection, we describe solutions that aim to mitigate plasticity loss under non-stationarity, which cannot remedy the loss of plasticity in an incremental setting with a stationary data distribution. Table 2 shows L2 INIT (Kumar et al., 2023) and Reset (Nikishin et al., 2022 ) cannot be a solution in our setting. Furthermore, applying layer normalization cannot close the gap between cold-starting and warm-starting; rather, the gap increases, as shown in Figure 5 . Also, Nikishin et al. (2022) and Sokar et al. (2023) state that loss of plasticity in non-stationary data distributions arises from inactive neurons in the model. However, this is not the case in our setting, as demonstrated in Figure 6 . This finding suggests that techniques designed to revive dead neurons in non-stationary data distributions cannot effectively address the plasticity loss observed in the incremental learning setting with stationary data, which is the primary focus of our study.In this subsection, we provide the previously omitted experimental results in Tables 3 4 5 6 , where the experimental settings are described in Section 5.1. As mentioned earlier, we also compared DASH with warm-starting while resetting the momentum at the start of every experiment, denoted as "Warm ReM" in the following tables. In the state-of-the-art (SoTA) setting, we employed weight decay and standard data augmentation techniques, such as horizontal flipping and random cropping. We also used a learning rate scheduler that reduces the learning rate step-wise by a factor of 0.2 at 60, 120, and 200 epochs. By applying the learning rate scheduler, there is no need to compare training time since training is completed at roughly the same epoch across all experiments. The weight decay was set to 0.0005, and the initial learning rate was set to 0.1. All other settings remain the same as mentioned above. We tested this setup on CIFAR-10 and CIFAR-100 using the ResNet-18 architecture.The results in Table 7 show that DASH performs similarly to or slightly worse than starting from random initialization. It appears that this is partly because all hyperparameters are tuned to maximize the performance of cold-starting, to achieve the (close-to-)SoTA test accuracy numbers. Due to the lack of computational resources, we were unable to tune hyperparameters specifically for DASH.Furthermore, we believe this aligns more closely with our theoretical anylsis in Theorem 4.1, as the hyperparameters are tuned to allow the model to learn as many features as possible, making it difficult for DASH to outperform cold-starting.Moreover, we observe that S&P cannot be used in these SoTA settings. We believe this is due to the nature of S&P, which shrinks all weights, while the SoTA setting is likely designed to avoid learning unuseful features, unlike the previous setting. Consequently, it is plausible that retaining learned features is more important than forgetting them, making S&P unsuitable for SoTA settings. Although DASH performs slightly worse than cold-starting, it is conceivable that it is better at retaining features compared to S&P and other warm-starting methods, resulting in better overall performance.The gap between warm-starting and cold-starting has been significantly reduced, likely due to data augmentation techniques and the increase in learning rate when new data is introduced. Data augmentation techniques increase the amount of feature information, allowing warm-starting to learn features that vanilla training (without augmentation) cannot (Shen et al., 2022) . Furthermore, as the learning rate is set to a higher value at the beginning of each new experiment, the model can forget previously memorized data points and escape spurious minima that were difficult to escape from, which is consistent with the findings of Berariu et al. (2021) . Despite these improvements, a gap still exists between warm-starting and cold-starting. In this section, we will provide the hyperparameters utilized in our experiments. Additionally, we present heatmaps illustrating the results for a wide range of two hyperparameters, α and λ, in DASH. The heatmaps in Figures 7 and 8 suggest that DASH exhibits robustness to hyperparameter variations, indicating that its performance is less affected by the choice of hyperparameter values. We fixed the momentum to 0.9 and the batch size to 128. The learning rate is set to 0.001 for training ResNet-18, and for other models, a learning rate of 0.01 is used. The value of ρ for SAM is chosen based on the performance of cold-starting.The default value of α = 0.3 is used, and we did not change this value frequently. The perturbation parameter σ used in the Shrink & Perturb (S&P) procedure is set to 0.01, as this value is considered optimal for perturbation, as described in Ash & Adams (2020) . Initially, we tested σ = 0.1 as the perturbation parameter, since Ash & Adams (2020) reported slightly better test accuracy compared to σ = 0.01 in some cases. However, we experienced significantly poorer generalization performance with σ = 0.1 compared to σ = 0.01, as shown in Figure 9 . The hyperparameters used in our experiments are described in Table 8 . To validate whether previously learned/memorized data points do not have large gradients when further trained with a combined dataset (existing + newly introduced data) using warm-starting, we plotted the train accuracy on the previous dataset for the first few epochs using CIFAR-10 trained on ResNet-18 (Figure 12 ). We observe that the newly introduced data does not substantially affect the previously learned/memorized data points, even when further trained with warm-starting, supporting the main idea of Theorem 3.4.To verify whether DASH truly captures our intuitions from the ideal algorithms, we conducted an experiment using CIFAR-10 trained on ResNet-18, with the same experimental settings. Figure 12 demonstrates that when applying DASH, the train accuracy on previous datasets rapidly increases after a few epochs compared to other methods. We argue that this behavior stems from our algorithm's ability to forget memorized noise while preserving learned features. As the number of experiments increases, the number of learned features also grows. For a fair comparison, we used λ = 0.05 for DASH, and when performing S&P and shrink, we shrank each weight by multiplying 0.05. In the case of S&P, after shrinking, we added noise sampled from N (0, 0.01 2 ). As stated in Section 2, we posited that τ could serve as a proxy for dataset complexity. Figure 14 shows that as τ increases, the threshold for considering a data point well-classified also increases, making it more difficult to correctly predict unseen data points. This difficulty is particularly pronounced for warm-starting, leading to a widening gap between the random initialization and warm initialization methods. Additionally, this phenomenon is observed in real-world neural network training, as depicted in Figure 13 . For datasets with higher complexity (from left to right), the gap between the two initialization methods widens, exhibiting the same trend as an increasing τ . We conducted synthetic experiments across a wide range of hyperparameters. Figure 14 uses the same setup as Section 3 but varies τ . It means that u must have been already learned in the (J * -1)-th experiment and it is contradictory.Thus, we have L To show that strict inequality for J > γ δn , it suffices to show that L  


Decision: All reviewers champion the acceptance of this manuscript. The AC encourages the authors to take the suggestions of Reviewer sYbB into account and consider experiments on larger models.