Paper: Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems (ICLR-2024)

Reviewer comment (Novelty):
In terms of discovery and task design, the paper points out an intriguing issue with the softmax function (leading to slow and unstable convergence), an issue that has been identified and addressed by the previously proposed NormSoftmax. Moreover, the task designed within this study does not align with practical application needs and seems to lack significant utility, serving more as a platform for theoretical experimentation. It might be more beneficial to design an experiment around a multi-step task that is more natural and more applicable to visual applications.

Author response:
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.

Reviewer comment (Experiments):
The tasks designed in the paper are based on simpler datasets like MNIST and FashionMNIST. There is a lack of experimentation with more diverse and realistic high-resolution image datasets. For example, constructing new tasks using datasets like ImageNet and CelebA could increase the credibility and universality of the research results. We expect that the conclusions drawn still be as applicable to more complex real-world data.

Author response:
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.

Reviewer comment (Experiments):
The problem setting of the paper, the experiments, and the analytical insights are all based on the synthetic MNIST - Fashion MNIST multi-step task. The significance of the whole paper and all the results strongly depends on the assumption that the insights from the experiments with the synthetic data translate to multi-step tasks on data distributions encountered in real-world scenarios. Despite the pivotal significance of this assumption, there is no proof for the same. The authors state in Section 1 that they have provided some indication of the validity of the assumption, but this doesn’t seem to be present in the paper. Therefore, they should provide a sound background that demonstrates how the insights from the experiments with synthetic data scale to real-world data.

Author response:
**(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.

Reviewer comment (Experiments):
It is possible that the Eureka moments and the failure of ViTs in multi-step tasks are affected by the nature of the relationship between the intermediate tasks. The multi-step classification task used in the paper involves two unrelated tasks, which may affect the insights drawn from the experiments. For example, the majority of the analyses are carried out on a multi-step task based on MNIST and Fashion MNIST, which are not directly related. While the authors consider the worst-case scenario where the tasks are unrelated, they should attempt to explore the effect of the relationship between the tasks on Eureka moments.

Author response:
**(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.

Reviewer comment (Experiments):
While the authors consider the important and practical setting of multi-step classification, the nature of the multi-step task explored in the paper might not be indicative of the settings often encountered in real-world use cases. I agree with the authors that there are several challenges involved in attempting such a study on real-world data, as they have outlined in Section 1. However, the task with 2x2 inputs for the multi-step task on small-scale datasets such as MNIST and CIFAR might not provide meaningful insights that extend to large-scale settings.

Author response:
**(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.

Reviewer comment (Novelty):
Hierarchical classification can be considered as a specific case of multi-step classification where the intermediate tasks are highly related, i.e., the first task involves coarse categories while the second task involves fine categories under each coarse class. The authors should discuss the relationship between hierarchical classification and their proposed multi-step classification setting.

Author response:
**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.

Reviewer comment (Experiments):
Table 3 shows that the experiments are conducted on ViT architectures that are similar in scale to ViT-S. However, there are no experiments with larger models, which in part, is limited by the choice of dataset for the multi-step classification task. How do the insights from the smaller models translate to larger architectures such as ViT-B or ViT-L? What is the trend of Eureka moments observed in the larger models?

Author response:
**(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.

Reviewer comment (Experiments):
The paper deals with multi-step classification where the intermediate tasks are unrelated. Do the insights on Eureka moments and the optimization challenges presented in the paper hold when the tasks are related? For example, the authors can conduct experiments using the hierarchy of CIFAR-100 or ImageNet (or a subset of ImageNet).

Author response:
**(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.

Reviewer comment (Novelty):
How does the proposed multi-step setup relate to hierarchical classification? Can hierarchical classification be considered a specific variant? How would the insights from the current setup translate to a hierarchical setup, if the latter is a specific variant of the former?

Author response:
**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.

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

Reviewer comment (Experiments):
How can one study Eureka moments for larger ViT models, since it is likely that they might overfit the proposed MNIST/CIFAR-based classification tasks?

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

Reviewer comment (Experiments):
It will be interesting to see the robustness of the proposed solution, i.e., how it behaves when faced with more than two multi-step decision tasks. Also, if we extend the model beyond just the vision domain, for instance, in a multimodal setting, do the same assumptions and proposed solutions still hold?

Author response:
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.

Reviewer comment (Reproducibility):
Reproducing the results may be challenging. I kindly request the author to please provide the code for replication.

Author response:
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.

Reviewer comment (Experiments):
The authors do not sufficiently identify the source of the problem. In particular, I expected a deeper investigation into the role of initialization; missing Eureka moments seems like it could be driven by bad initialization which is then fixed with higher entropy.

Author response:
**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.

Reviewer comment (Writing):
Unusual to see artefact instead of artifact.

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

Reviewer comment (Evaluation):
How do the L1 Gradient Norms look for ViT without temperature in an extended horizon through reaching eureka moment? How does this compare to the norms for a run of ViT without Eureka moment?

Author response:
**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.

Reviewer comment (Presentation):
For roberta in Figure 7B, is the loss curve training or validation? How do the training and validation curves compare for this task?

Author response:
**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.