A Visual Case Study of the Training Dynamics in Neural Networks

We thank the Reviewer 1hd3 for their valuable comments. We are happy that the reviewer acknowledged the clarity of our setup and problem motivation.

We address the reviewer's concerns below.

Experimental investigation and theoretical insights: Our visual sandbox provides insights into the training dynamics of neural networks. Hence, it mostly consists of an experimental yet rigorous study. We however note that in our additional experiments in Appendix A, we study phenomena that are of interest both theoretically and empirically in the literature (loss spikes, gradient norms, activation sparsity) [1, 2, 3, 4, 5].

Beyond low-dimensional embeddings: In Appendix A, we conduct additional experiments on loss spikes, gradient norms, and activation sparsity and also extend them to the case $d>2$. We note however that the goal of our study is to remain in a controlled setup (transformer model and modular addition problem) to better understand the training dynamics as is often the case in mechanistic interpretability papers. However, the insights our sandbox enables us to obtain can be generalized to other setups and architectures. In particular, as explained in the revised version of the paper, the experiments of Sections 4.1, 4.2, 4.3, and Appendix A are related to phenomena studied in the literature [1, 2, 3, 4, 5] on large-scale transformers and for various tasks (time series forecasting, computer vision, LLMs).

We thank the reviewer for their time. We hope we addressed their concerns and we would be happy to answer any additional questions.

[1] Li et al. The lazy neuron phenomenon: On the emergence of activation sparsity in transformers. ICLR 2023.

[2] Mirzadeh et al. ReLU strikes back: Exploiting activation sparsity in large language models. ICLR 2024.

[3] Foret et al. Sharpness-aware minimization for efficiently improving generalization. ICLR 2021.

[4] Ilbert et al. SAMformer: Unlocking the potential of transformers in time series forecasting with sharpness-aware minimization and channel-wise attention. ICML 2024

[5] Zhang et al. Why transformers need Adam: A hessian perspective, NeurIPS 2024.

We thank the reviewer XhQn for their valuable comments. We are happy to read that the reviewer found our findings interesting and potentially important.

We address the reviewer's concerns below.

Motivation for our controlled setup: The motivation for focusing on $d=2$ was to be able to visualize all the training processes in $R^2$. While such visualization could be extended to $d=3$, albeit to the cost of simplicity and clarity, it cannot be done for higher $d$ unless involving dimensionality reduction. This explains why we focused on 2D embeddings although our insights extend to more practical settings.

Insights into more practical settings: The behaviors identified thanks to our sandbox on the training dynamics can be generalized to a more practical setup. In particular, the efficiency of transfer learning is known in the literature. More importantly, controlling the loss spikes phenomenon is very important to stabilize the training of large models such as LLMs. To address the reviewer's concerns, we extended our study of loss profiles to phenomena studied in the literature on gradient norms and activation sparsity [1, 2, 3, 4, 5] in Appendix A.1 and A.2 and we extended these experiments to the case $d>2$ in Appendix A.3.

We thank the reviewer for their valuable comments. We hope we have addressed their concerns and we would be happy to answer any additional questions.

[1] Li et al. The lazy neuron phenomenon: On the emergence of activation sparsity in transformers. ICLR 2023.

[2] Mirzadeh et al. ReLU strikes back: Exploiting activation sparsity in large language models. ICLR 2024.

[3] Foret et al. Sharpness-aware minimization for efficiently improving generalization. ICLR 2021.

[4] Ilbert et al. SAMformer: Unlocking the potential of transformers in time series forecasting with sharpness-aware minimization and channel-wise attention. ICML 2024

[5] Zhang et al. Why transformers need Adam: A hessian perspective, NeurIPS 2024.

We thank the reviewer qxou for their detailed feedback. We appreciate that the reviewer found our work novel and valuable.

We address the reviewer's concerns point by point below.

1. My first concern is that, while the visualization tool is insightful, its restriction to small-scale transformers with $d=2$ limits its applicability to larger and more complex models commonly used in practice. It remains unclear whether insights from this sandbox will generalize to high-dimensional transformers. Further discussions on the higher dimensional embedding should be provided.

We thank the reviewer for their suggestion. To address the reviewer's concerns, we added two experiments in Appendix A.1 and A.2 on the connection between gradient norms and loss spikes and the connection between activation sparsity and models' performance. These experiments are connected to the insights presented in our submission and are agnostic to the embedding dimension. As a matter of fact, we extended these two experiments to the case $d>2$ in Appendix A.3.

2. Could the authors discuss the scalability of their sandbox tool to higher-dimensional transformer models, for example, when $d=3$? How does the computational complexity of this sandbox grow in terms of $d$?

We note that the motivation for focusing on $d=2$ was to be able to visualize all the training processes in $R^2$. While such visualization could be extended to $d=3$, albeit to the cost of simplicity and clarity, it cannot be done for higher $d$ (unless involving dimensionality reduction). However, the training behaviors identified thanks to our visual sandbox can be observed on large-scale models. We especially study phenomena of interest in large-scale models (gradient norms, activation sparsity, loss spikes) [1, 2, 3, 4, 5] in Appendix A.1 and A.2 and went beyond the case $d>2$ in Appendix A.3.

3. Are there quantitative metrics for evaluating the effectiveness of the proposed loss spike mitigation strategies?

Loss spikes occur during the training of models and hence are hard to quantitatively measure or monitor. We show in Appendix A.1 and A.3 (Figures 15, 16, 18, 19) that loss spikes and high gradient norms occur in tandem. Hence, one could think of quantitative metrics on gradient norms to better assess the quality of a loss spikes mitigation strategy. Recently, there has been some proposal of new normalization schemes to stabilize the training of neural networks training [6, 7, 8]. We believe such strategies (or similar ones) could be efficient in controlling the gradients of the layers and hence the loss spikes.

4. What types of theoretical developments/problems in training dynamics can be recognized or conjectured from the visualization results of this sandbox?

We acknowledge that the benefits of our sandbox for theoretical problems were not clearly stated. To address the reviewer's concern, we conducted additional experiments motivated by our visual sandbox in Appendix A on phenomena, such as loss spikes, gradient norms, and activation sparsity that are of interest both theoretically and experimentally in the literature on large-scale transformers [1, 2, 3, 4, 5]. We believe our sandbox can help better understand and elucidate such problems using dynamical observations on a small scale before tackling them on bigger models.

We thank the reviewer for their time and feedback. We hope to have addressed their concerns and would be happy to answer any additional questions.

[1] Li et al. The lazy neuron phenomenon: On the emergence of activation sparsity in transformers. ICLR 2023.

[2] Mirzadeh et al. ReLU strikes back: Exploiting activation sparsity in large language models. ICLR 2024.

[3] Foret et al. Sharpness-aware minimization for efficiently improving generalization. ICLR 2021.

[4] Ilbert et al. SAMformer: Unlocking the potential of transformers in time series forecasting with sharpness-aware minimization and channel-wise attention. ICML 2024

[5] Zhang et al. Why transformers need Adam: A hessian perspective, NeurIPS 2024.

Dear Reviewer, Do you mind letting the authors know if their rebuttal has addressed your concerns and questions? Thanks! -AC

Thank you to the authors for their detailed reply. Makes sense to me.