Energy Transformer

Summary

This paper presents a new transformer architecture that is based on minimizing energy with two networks viewed as Hopfield networks (one for the attention within the data point, and another for global features). The result energy-based model resembles (and perhaps motivates) empirical choices in the vanilla transformer, and unlocks interpretations of the learned features, as well as better performance on graph-based benchmarks. The authors present the idea that their energy transformer could be useful for other domains, such as computer vision, NLP, audio, etc.

Strengths

S1: The problem is highly relevant. Finding better transformer architectures is important in modern machine learning. The ET is a reasonable and well-motivated contribution. S2: The improvements on graph benchmarks seem to be really strong. S3: I like the visual interpretation of the intermediate states of the tokens. S4: The paper is very well written.

Weaknesses

W1: Limited evaluation (only on graphs). I suspect the authors didn't evaluate large-image representation learning because of the issue of learning global features (see question Q4 below). I am not sure why the authors didn't try NLP. I think trying NLP is a very crucial step for convincing the community that ET works. W2: Missing discussion about the computational cost of the architecture (at least I do not see it in the main text). How easy it is to implement your method and scale it? The answers to these questions should be a part of the main text. They are also integral to convincing the community about the scalability of ET.

Questions

Q1: How does your visualization from Figure 2 relate to patterns learned by CNNs? Q2: How does the MLP weights relate to the weights $\xi_{\mu j}$ in Eqation 5? Can you explain in more detail? Q3: Have you tested the crucial importance (for minimizing the energy) of the second term of the RHS in the equation of Line 177? I think you have done that in the Appendix (in the SM). Could you comment on the findings of your ablations a bit? Q4: Have you tried downstream tasks of the representations for ImageNet? It would be interesting to compare with MAE. I have a suspicion that the features are not going to be useful for transfer learning because it seems like they only learn local information (the example with the missing eye for the dog that you give in Fig. 3 is insightful). Then, my question becomes: Why is MAE (speculatively) more useful than ET in terms of obtaining features for transfer learning? How can we let ET learn more global information? Q5: Can you show reported results from previous works in Table 1?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

the authors adequately addressed the limitations

Summary

This paper proposes a new Transformer-like architecture which is is inspired by modern Hopfield networks, and falls into the category of energy-based models. The authors argue that the existence of an energy function leads to several advantages such as more principled design and theoretical interpretation. The energy transformer (ET) they design a recurrent network similar to a deep transformer architecture with shared weights, however with the additional property that the recurrent dynamics is shown to minimize an energy function, and can be seen as an instance of a Hopfield network. The authors evaluate the performance of the architecture on image completion (qualitative), graph anomaly detection and graph classification tasks, demonstrating strong and competitive performance. Please note that due to receiving 6 papers to review from NeurIPS alone, I have allocated a time budget of 4h per paper, and my review is based on that. In particular I did not study the appendix in great depth. I regret this situation and apologize for possible inaccuracies.

Strengths

- Clear presentation - The combination of attention mechanism and Hopfield networks makes this architecture interesting for theoreticians. - Strong evaluation results on the graph anomaly detection and graph classification tasks.

Weaknesses

- Insufficient discussion of limitations. - Evaluation on the masked image completion task is not quantitative, and not compared to related architectures/baselines such as ViT.

Questions

- In lines 148-152, the functional similarity between the HN sub-block and a MLP seems plausible but was not fully clear to me. Could this discussion be made more explicit in terms of equations? - The two descriptions of the difference to previous modern Hopfield networks in lines 172-175 and lines 183-184 did not seem consistent to me: The first states that the difference is that the keys and queries are dynamical variables, while the second states the main difference is the additional symmetrizing term in the attention mechanism. - in line 192, the text states "e.g. by replacing them with a learnable MASK token". The "e.g." makes it unclear in how far this is a generic example, and what was actually done in the experiment. - In the architecture used in sect. 4, the main text does not explain how the adjacency structure of the graph enters the model. From the appendix I understand that the attention operation is constrained to the 1-step environment of the node. Maybe this would be good to discuss in the main text. - In general, the ET block can be understood as a deep transformer model, however with shared weights instead of independent weights in each layer. What limitations does this imply on the type of tasks which are well suitable for the ET architecture? Is the stacked model of used in sect. 5 a result of the need to overcome the shared weight limitation?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

While some weaknesses in the empirical performance are pointed out, e.g. lines 200-203, there is no broader discussion of technical and theoretical limitations.

Summary

This paper proposes a new transformer architecture namely Energy Transformer, consists of layer norm, multi-head energy attention and Hopfield network, within a single recurrent ET block. The ET block models the dynamics of token updates by minimising the sum of energy from both the multi-head attention and the Hopfield network. The main contribution is the new designed ET architecture.

Strengths

The paper proposes a novel transformer architecture, which uses the Hopfield network, and models the dynamics between tokens with a time differential equation. The proposed architecture seems promising in the applications for CV, NLP and other domains.

Weaknesses

The experiments can not fully support the claim. - Empirical experiments on image reconstruction is hard to quantify and are also not compared to any of the SOTA algorithms. The value of the reconstruction error shown in the Appendix is hard to understand/evaluate without comparing to any other algorithms. - For instance, FID scores in most of the image generation tasks, or AUC in [1] can also be quantitative measures for the ET. - Transformers are widely used in NLP tasks, but this is almost completely ignored in the discussion. Since there is no theoretical justification for the proposed architecture, showing empirical results on various applications (including but not limited to NLP) may help to convince the reader. - In my humble opinion, quantitative experiments on graph anomaly detection and classification are not sufficient to show the full power of such generative models. [1] Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. In NeurIPS, 2019.

Questions

- Is there any theoretical justifications that show the ET has either faster/slower convergence in training compared to a vanilla transformer/energy-based models, or ET obtains a better optimum than EBM after optimisation? - Why there are no quantitative experiments on e.g. image generation? How ET performs when the input image is 100% masked(generate from noise)? Can ET be applied to image anomaly detection as in [1]?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

I don't see a detailed discussion on limitations of the work.

Summary

The paper presents a new architecture for transformer (ET). The architecture is designed by first defining an energy function that forces the alignment of the query of a patch to the keys of a few other patches and to make each patch to look like a natural image patch. A dynamics (differential equation) is induced to minimize the energy function, giving rise to a transformer layer containing the novel attention mechanism and a Hopfield network component. Experiments are performed using this new transformer to fill partially masked images and to detect graph anomalies and classify graphs. Effectiveness is demonstrated.

Strengths

The design is well motivated and clearly rationalized, nicely exploiting the recently observed connection between Hopfield networks/associative memories and transformers,. The proposed architecture is novel and interesting to this reviewer. The experimental validation looks convincing. The interpretability aspect of the model relative to ViT is particularly appealing to this reviewer. The paper is also very well written, pleasant to read. The proposed structure is well positioned in relation to the literature of attention and associative memories.

Weaknesses

Despite its well-motivated design, the proposed structure still lacks rigorous theoretical justification for its effectiveness. I have a few minor editorial comments: 1. In Eq (4), the notations g_{jB} are not explained until later in the paper. 2. You perhaps need a notation for the range of \alpha, otherwise in the unlabelled equation in page 6, sum over \gamma is unclear. — If I am not mistaken, \gamma and \alpha range in the same set.

Questions

In the introduction (beginning of page 2), you mention that large storage capacity of associative memories is important for learning applications. I feel this is a correct statement intuitively however have not seen a rigorous argument connecting the storage capacity to learning. Could you explain the reasoning behind the statement or give a reference?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

NONE

Thank you for the comments

I am keeping my rating.

Thanks, good rebuttal. The paper should be accepted.

I read the whole rebuttal and am satisfied with the authors' response.