A Visual Case Study of the Training Dynamics in Neural Networks

Authors: Ambroise Odonnat, Wassim Bouaziz, Vivien Cabannes
Primary area: interpretability and explainable AI
Keywords: Training dynamics, Transformer, Mechanistic Interpretability, Visualization

Abstract:
This paper introduces a visual sandbox designed to explore the training dynamics of a small-scale transformer model, with the embedding dimension constrained to $d=2$
This restriction allows for a comprehensive two-dimensional visualization of each layer's dynamics. 
Through this approach, we gain insights into training dynamics, circuit transferability, and the causes of loss spikes, including those induced by the high curvature of normalization layers. 
We propose strategies to mitigate these spikes, demonstrating how good visualization facilitates the design of innovative ideas of practical interest.
Additionally, we believe our sandbox could assist theoreticians in assessing essential training dynamics mechanisms and integrating them into future theories.

Decision: reject
Reviewer scores: 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1730670312550

thread_type:
official_review

<thread_object>
creator: reviewer_1hd3

<content>
summary:
This paper examines the training dynamics of neural networks by conducting a detailed study on a small-scale transformer model, employing a novel visual sandbox. This tool facilitates the observation and analysis of the model's internal mechanisms, offering insights into more efficient training strategies.

strengths:
- The experimental motivation and setup are clearly articulated.

weaknesses:
- The presentation of the experimental results is suboptimal. In Figure 4, many values overlap, resulting in readability issues.
- There is a lack of theoretical support for the proposed method. The paper claims, "This tool allows us to observe and analyze the model’s internal mechanisms vividly, providing both theoretical and practical insights that could lead to more efficient training strategies." However, the theoretical insights remain unclear.
- The experiments are limited. The paper lacks experiments across varied datasets, tasks, and neural network architectures.

questions:
- Could you present the experimental results more clearly?
- Could you provide theoretical support for the proposed method?
- Could you conduct additional experiments across various datasets, tasks, and neural network architectures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700151115

thread_type:
official_review

<thread_object>
creator: reviewer_xhqn

<content>
summary:
In this work, the authors presented a case report that using 2D visualization to demonstrate the training dynamics of a neural network trained to solve the sparse modular addition problem. Various aspects of the training dynamics, including the two stages of learning, the influence of initialization, process of network adaptation in sequential learning, trajectory of the loss function including its anomaly (the spikes) are visualized and analyzed qualitatively.

strengths:
By choosing a well-structured task and proper visualization, the authors have demonstrated a number of interesting and potentially important phenomena regarding network training dynamics. The framework presented may be extended to study other networks and tasks, leading to better understanding of networks’ learning process.

weaknesses:
It is unclear what part of the current work can generalize. Is it the methodology of 2D embedding and visualization? If so, what is the difficulty of doing it and what is the innovative way to overcome them in the present work? And how the proposed method can generalize to other tasks and networks? These needs to be explained and clarified in a clearer way. Or, is the conclusion of the training dynamics generalizable? But here it seems unlikely, because the current observations were based on a particular network structure trained with a specific task. Even for the current network and task, the authors need to show how robust of the observations, e.g., the two stages of training, with one for embedding and one for classification. Without either the method or the conclusions generalizable to more situations, the impact of the work might be limited.

questions:
I suggest the authors re-organize the paper to better explain and clarify what is the problem solved by this case report, and how it can help readers working in training networks in general.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730736025951

thread_type:
official_review

<thread_object>
creator: reviewer_qxou

<content>
summary:
This paper presents a novel visual sandbox designed to explore the training dynamics of a small-scale transformer model with a restricted embedding dimension of $d=2$. This setup allows for two-dimensional visualization, providing a detailed look into the model’s internal mechanisms. The sandbox helps uncover hierarchical learning processes, circuit transferability, and causes of loss spikes, particularly those due to high curvature in normalization layers. The authors propose methods to mitigate these spikes, highlighting how visualization can lead to more effective training strategies. They suggest that this tool could assist in developing theories and improving training processes for more complex neural networks.

strengths:
- The paper introduces a novel tool for visualizing the training dynamics of transformers. The visualization offers both practical and theoretical insights, making complex mechanisms more accessible.

- The authors identify a hierarchical learning process inside a pretrained transformer via their visualization tool, showing how low-level features are established before higher-level refinement. This is valuable for understanding model behavior and training stages.

- The tool also helps to recognize loss spikes caused by high curvature in the model's normalization layers. Based on this, the authors suggest practical mitigation strategies, potentially making training more stable and efficient.

weaknesses:
- 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.

- There is no comparison between a pretrained model and a fine-tuned model based on insights derived from the sandbox. Consequently, it remains unclear whether the recommendations informed by the sandbox visualizations effectively contribute to improved model performance on practical tasks.

questions:
- 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$?

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

- What types of theoretical developments/problems in training dynamics can be recognized or conjectured from the visualization results of this sandbox?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732465803272

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment

comment:
We thank all the reviewers for thoroughly and carefully reading our paper. We are grateful to them for acknowledging the **clarity** of our setup and motivations (Reviewers XhQn,  1hd3) and for finding our study **novel and valuable** to better understand training dynamics (Reviewer qxou).

Several reviewers raised a common concern regarding the takeaways of our study for more practical settings as well as the low dimensional case treated in our study ($d=2$). To clear out any possible misunderstanding regarding the scope of our work, we provide clarifications below.

- **Low-dimensional embeddings:** Our focus on a simplified transformer architecture in low dimension was motivated by the need to be able to visualize all the steps the neural network operates during training. Our visual sandbox leads to detailed and rigorous observations of the training process and thanks to the controlled setup (modular addition problem), one can better interpret the training dynamics. 

- **Beyond $d=2$**: Thanks to our visual sandbox, we studied several phenomena on the training dynamics (e.g., layers learning, transfer learning, loss spikes). As explained in our paper, these phenomena are of great interest for the training of large-scale transformer models such as LLMs. Moreover, we note that many of the studied behaviors can be analyzed **independently of the embedding dimension $d$.** To address the reviewers' concerns, we conducted three additional experiments connected to the phenomenon studied in our submission: **(1)** connection between loss profiles and saddle points via gradients norms in Appendix A.1, **(2)** connection between activation sparsity and models' performance in Appendix A.2, **(3)** extension of these 2 experiments for $d > 2$. 


- **Practical takeaways**: These experiments are strongly connected to phenomena observed in the literature on large-scale transformer models [1, 2, 3, 4, 5] which makes our study relevant in practical settings. As demonstrated by the research process described in our work, we believe our visual sandbox can enable rigorous study in small-scale scenarios to elucidate and then study them on bigger models. We believe this can be valuable to researchers and practitioners to better understand the training dynamics of neural networks.

**Updates to the paper (in blue)**

We list below the updates to the paper (in blue):
- Added Figures 15, 16, 17, 18, 19,and  20 in Appendix A
- Experiment on gradient norms in Appendix A.1
- Experiment on activation sparsity in Appendix A.2 
- Similar experiments with $d>2$ in Appendix A.3

We thank the reviewer for their time and valuable feedback that helped us improve our work. We remain at their disposal to continue this collaborative discussion until the end of the rebuttal.

---
[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. 

---
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper applies various methods for visualizing the training dynamics of a small transformer model with embedding dimension 2, trained on a sparse modular addition problem.  Various aspects of training dynamics are observed and studied including the impact of initialization, observation of two stages of learning, circuit transferability, and loss spikes during training.  

Some reviewers felt the paper presented a novel tool for visualizing training dynamics. Other reviewers noted the task and visualization strategies were well chosen to provide insight into the training dynamics.  The visualizations and analysis showed interesting phenomena in the training dynamics inducing a multi-part learning process and loss spikes caused by normalization layers.  These may lead to insights for improving training.  

Reviews generally agreed, however, that the embedding dimension was too small to draw strong conclusions about larger models, limiting the applicability to real-world models and problems.  Moreover, it is not readily clear how the method in the paper may even be applied to larger models or other task functions.  Some reviewers also pointed out the method lacked a strong theoretical basis.

Reviewers reached consensus that questions about the practical drawbacks of the method outweighed the demonstrated insights and observations.