Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems

We appreciate your careful review and helpful comments. We address your concerns and questions below and in the general comment to all reviewers.

One of the weaknesses mentioned by Lp2g is the weak alignment with practical application needs and significant utility of our work. Indeed, this work does not aim at contributing a method or dataset for some application. Instead, this work aims at understanding existing methods, highlighting a shortcoming, identifying a previously unknown failure mode and providing an experimental setup that can be used to study this phenomenon. We argue that improvements in the understanding of learning behavior of models may also enhance application performance further down the line.

We suspect that for many real world applications Eureka-Moments happen, but can not be easily observed. We suspect that tasks exist where the Eureka-Moment does not happen which results in severe underperformance. Knowing about Eureka-Moments will enable the community to understand training failures better and will guide future experiments to avoid such.

We would also like to highlight that we do not claim that NormSoftmax is a solution to the problem. We use it as an intervention to show that the problem is caused by the Softmax Attention. Note that this is the gold standard technique to understand cause-effect relationships commonly used also in, e.g., medical drug testing.

Regarding the lack of experiments on more diverse and realistic high-resolution image datasets, we refer to the general comment to all reviewers, which addresses this point with an additional experiment on ImageNet-100.

Thank you for your detailed evaluation and insightful feedback. We address your concerns as listed in the weaknesses below using the same numbering as in your review. We answered the questions below and in the general comment to all reviewers.

(a) We address the concerns about relevance and real world transfer in multiple ways. First, we provide additional experiments on realistic high-resolution data, as described in detail in the reply to all reviewers. In addition, we already provided the RoBERTa experiment in the submission (see Fig. 7(b)), where Eureka-Moments are observed on real data and show that the most promising mitigation i.e. NormSoftmax helps also in this case. We agree that this is not a theoretical proof that our findings translate to a real world setting, but, we provide sufficient empirical evidence for it.

Please note that we do show an indication for Eureka-Moments on real data and the transfer to real world problems in the paper, as promised in the introduction. We suspect that the RoBERTa experiment was overlooked or not identified as this indication. We admit that the RoBERTa experiment should be highlighted more and in particular these results should be referenced directly in the introduction, where we promise the indication of real world Eureka-Moments and transfer of our findings to the real world. We will modify the final version accordingly.

(b) Relationship between intermediate tasks. Thank you for this great suggestion. Indeed, the relation of the tasks is a super interesting direction. Is the task becoming harder or easier when tasks are related? What if similar features need to be used but the tasks are different? All these questions are quite interesting, but unfortunately very difficult to answer in a single paper. How can we measure the relatedness of tasks? Are the tasks we study even completely unrelated or do they not use similar features? When training on very related tasks, it opens the door to learn partial solutions to task 1. These small Eureka-Moments will be hard to spot, as the improvements get obscured by task 2 learning. All these additional factors make a clean study and interpretation of the results difficult and would introduce lots of complexity to the paper. We decided to stay in the simple setting for an initial study of Eureka-Moments and left more specific research for future work. Besides this, we already show the RoBERTa experiment in the paper, where tasks are most likely related.

Another point raised by Ddae is that the relation to hierarchical classification is not discussed in enough detail. We are happy to add a small discussion to the related works or add it in the introduction. Indeed, we mention hierarchical-classification briefly in our introduction as an example, where Eureka-Moments may also appear. However, this setting is not very well suited to study Eureka-Moments. A flat classification approach, as defined in [R1] does not necessarily have to solve the task via multi-step classification. To give an example: to distinguish cats from dogs the classifier does not necessarily have to classify both as quadrupeds and then classify the leaf nodes. Simply learning a single feature that distinguishes the two is sufficient. Since this alternative strategy exists, we do not know which one is chosen and the choice may even differ on a sample level. Because of these reasons hierarchical-classification as in [R1] is not the right task to study Eureka-Moments for the moment. Nevertheless, hierarchical-classification might benefit from the insights on Eureka-Moments.

(b) Small scale of models and dataset. The influence of model scale on Eureka-Moments is important and interesting. As already mentioned by Ddae, we studied the influence of model scale in the supplemental material (see section A.5 and Table 3) and reference this section briefly in the main manuscript. Regarding the concern of dataset size, we would like to refer to the general comment to all reviewers, where we show that large-scale datasets can be easily constructed to study Eureka-Moments in larger networks.

[R1] Miranda, Fábio M., Niklas Köhnecke, and Bernhard Y. Renard. "Hiclass: a python library for local hierarchical classification compatible with scikit-learn." Journal of Machine Learning Research 24.29 (2023): 1-17.

Questions

Question regarding hierarchical-classification: We addressed the relation to hierarchical-classification in the section above. Another important distinction is that intermediate tasks are learned implicitly in our setting, but labels for non-leaf nodes are given in general hierarchical-classification. We do not expect Eureka-Moments to appear when direct supervision in the form of labels is available for all sub-tasks.

The last question about studying Eureka moments for larger ViTs has been addressed in the general comment to all reviewers.

Thanks for the authors’ response and their efforts during the rebuttal period. Even though my concerns about the relevance of the setup and real-world transfer have been addressed, I feel that the following points still deserve attention and can significantly improve the paper:

1. I feel that the paper is still lacking with respect to the relationship between the intermediate tasks. I agree with the authors that there are numerous possibilities for the choice of intermediate tasks. However, the analysis of Eureka moments is incomplete without it. As the authors have mentioned, the Eureka moments may be observed to a lesser extent or may be nonexistent when the tasks are highly similar. If this is the case, Eureka moments are limited only to highly dissimilar tasks. I believe that the paper can be significantly improved by presenting well-thought experiments for 2-3 different levels of similarity between the intermediate tasks.
2. I agree with Reviewer Kjsr that the setting of 2x2 inputs for observing Eureka moments seems somewhat contrived. Do Eureka moments arise only in such 2x2 setups? Even though the authors have shown results on the RoBERTa task that doesn’t exactly follow this setup, the primary analyses and insights from the paper are based on the 2x2 multi-task setup. The authors should present results on a more natural multi-task setup in addition to the setup in the paper.

We would like to thank you for your time spent on this review and thoughtful suggestions and questions. We address your concerns in the general response to all reviewers and below, using the same numbering as in your review.

1. Please see the general comment to all reviewers for our response to the point raised regarding small-scale dataset and transferability to real-world scenarios.

2. We are not entirely sure we understand the request regarding the robustness. We would like to emphasize that we do not claim that NormSoftmax is a general one-size-fits-all solution. Instead, we use NormSoftmax as a minimal intervention of the Softmax function to show that the problems is related to the Attention block. Regarding robustness, we also conduct experiments on the harder "No Position Task" and find that NormSoftmax fails for 1 out of 4 training runs.

• The second part of the question addresses 3 or more step tasks. While never tested explicitly, we have no evidence that 3 step tasks behave fundamentally different. We suspect the Eureka-Ratio to be much smaller. While definitely interesting to investigate the behavior for tasks with more than 2 sub-tasks, we believe that it complicates studying the phenomenon drastically and might obscure other findings in the paper. In case you have a particular question or experiment in mind that helps in understanding the phenomenon using 3 or more-step tasks, we would be very grateful for your input.

• Regarding the last part of the question whether the phenomenon appears beyond the visual domain, we definitely think that this is the case and this is an interesting future research direction. We do already provide evidence that this happens for BERT training in the paper. Furthermore, we show with the "reasoning task" (see Fig. 7(a)), that we observe the same even for a single layer transformer with "pre-extracted features". This experiment simulates for example the setup of two frozen feature extractors (potentially multi-modal) that are "fused" using a transformer block.

3. We agree that reproducibility is extremely important and the code will be released the latest upon acceptance as mentioned in our paper. We strongly believe that providing code alongside a paper is important.

Answer to Kjsr

Thank you for your review. We highly acknowledge your thoughtful review and valuable comments. We address your point in the general comment to all reviewers and below.

Differences of gradient norms. One point raised by the reviewer Kjsr is that the difference between the gradient norms in Figure 5 is not big. We would like to point out that initially the difference between W_V gradients, and W_K and W_Q gradient norms is more than an order of magnitude for ViT. After about 1200 steps it settles at about 0.5 orders of magnitudes. For NormSoftmax, on the other hand, there is very little difference. Differences only become apparent after the Eureka-Moment (indicated by the end of the gray box in Figure 5).

Gradient norms for ViTs. The L1 gradient norms for ViT with and without Eureka-Moment can hardly be distinguished. Only briefly before the Eureka-Moment it can be observed that gradient norms increase sharply. Note, that this is not in contrast with our explanation. Gradient updates are small for W_K and W_Q for all vanilla ViTs. As correctly identified by you, random initialization also plays a big role. Depending on the random initialization of the weights the required amount of change of the weights is different. With a small gradient norm for W_K and W_Q many small steps will be needed to reach a Eureka-Moment. If the gradient norms are larger, fewer steps will be needed. This is in particular in line with our results when training for 3000 epochs instead of 300. In this setting all ViTs have a Eureka-Moment.

Role of initialization: To test the role of the initialization we implemented the initialization described in [1]. Initialization may help prevent getting into such a "bad energy landscape", where parts of the gradients become very small or initialization might just initialize the network weights closer to a Eureka-Moment. Alternatively, initialization may kick-start the learning process but could lead to the same problems eventually. To ensure a fair comparison, we followed exactly the same hyper-parameter search and training as described in the main paper. We find that the Eureka-Ratio of ViTs using the initialization described in [1] is similar to the one observed for normal ViTs. To be precise the Eureka-Ratio for our standard dataset (Table 1 left) is 7/10 with a mean accuracy of 89.31 and 0/4 Eureka-Ratio for the "No Position Task". Given that these results are very similar to the vanilla ViT we conclude that this type of initialization does not seem to resolve the problem. We suspect that despite the superior initialization the early training steps lead to equivalently bad energy landscapes with small gradients for W_Q and W_K and equally "far away". We will extend Table 1 with this baseline for the final version of the paper.

RoBERTa loss curve: In Figure 7(b) we show the training loss, as this is the relevant metric to study and also distinguishes it from grokking. The validation curve looks similar, but was not included since we ran evaluation very coarsely during training. We will repeat the training with more frequent evaluation and add training and validation loss curves to the supplemental material. We will also change the y-axis label in Figure 7(b) to training loss. Thank you for pointing this out.

Eureka-Moments in deeper ViTs: Depth is one of the parameters we tested in Table 3 (supplemental material). We still observe Eureka-Moments for deeper networks. In addition to these results, we added experiments with ViT-S on ImageNet-100. Please refer to the general comment to all reviewers for these experiments.

Thanks for pointing out the typos. We will revise the manuscript accordingly.

References

[1] Trockman, Asher, and J. Zico Kolter. "Mimetic Initialization of Self-Attention Layers." arXiv 2023.

I appreciate the additional experiments the authors provided for a more challenging task as well as trying a different initialization scheme. I will raise my score accordingly. However, I do still think the task is somewhat contrived to benefit from during training and think the work could benefit from additional examples of eureka moment.