Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Hierarchy

We appreciate the reviewer's feedback, and hope we can clarify some of points raised here.

Re: the lack of value projection: first of all, please note that in both our proposed Algorithms 1-3 as well as all of our experiments we have indeed incorporated value projections. Furthermore, in order to justify the incorporation of value projections within our theoretical derivation, we have added a new appendix (Appendix C in the new draft) explaining how value projections theoretically emerge in our statistical mechanical interpretation of self-attention. (TLDR: value projections emerge as learnable step-size in the gradient update in Eq. (6) realized as gradient projections.)

Re: "key methodological steps are missing" We would appreciate a more concrete feedback here as to what exact steps are missing so we can further explain them in the next version of the draft. We have done our best to clearly illustrate our methodological contributions, mathematical derivations and the empirical study in the main paper as well as the appendices. But we'd like to know how we can further improve the presentation.

Re: the scale of datasets We have added a new set of experiments in Appendix I in the new draft with larger datasets and models, and the results are still consistent with what was reported in Sec 5.1 and 5.2.

Re: whether these dataset have hierarchical structure or not Please note that written text in general have a semantically meaningful hierarchical structure (words, sentences, paragraphs, sections, etc.) which can be seen as abstraction of semantics at different levels and therefore be utilized using the HSA framework regardless of the dataset it's coming from.

Re: "The multi-modal experiment only has marginal improvements" The improvement in the multi-modal experiment is indeed statistically significant.

Re: "drop in accuracy on the zero-shot transfer learning." First of all, please note that the experiments reported in Sec 5.3 are not transfer learning; we are not changing the task here. Second, these experiments are designed to examine the practical implications of Theorem 1 by measuring the amount of accuracy drop in case the standard attention matrix is replaced by the closest hierarchical attention matrix to it (in terms of the KL-divergence) in a completely zero-shot manner. And while the performance gap is more pronounced for certain tasks/layers, for others it's surprisingly quite narrow which indicates an inherent redundancy in attention calculation and the HSA's ability in utilizing that, even in 100% zero-shot setting.

Re: novelty of deriving self-attention compared to Hopfield networks The emergence of self-attention within Hopfield networks and our entropy minimization framework are indeed quite similar, as both are different forms of energy minimization techniques. However, please note that derivation of the standard Softmax attention is not the main contribution point of our work. We have simply used the energy minimization technique for deriving standard Softmax attention to extend Softmax attention to nested signals, which we have further shown to be the optimal derivation in terms of the KL-divergence. To the best of our knowledge this theoretical approach to deriving hierarchical self-attention is novel.

Re: runtime vs FLOPs The HSA algorithm presented in Algorithms 1-3 is the theoretical depiction of our dynamic programming methodology. In practice, however, we have implemented the parallelized version of the algorithm in our codebase which we are intending to release as well. Nevertheless, there's a constant overhead to HSA which hinders the runtime gains for small and moderate scale problems reported in the paper. However, as we have already observed for auto-regressive generation tasks with much larger models, the runtime difference as well as the memory footprint gap is quite pronounced. We postpone these results to future work where we will comprehensively delve into hierarchical auto-regressive generation.

Re: "Why are the differences between the gradient update in the entropy minimisation and self-attention not important?" They are, but those differences are somewhat orthogonal to the HSA operation itself. As we have shown in Sec 5.1, 5.2 vs Sec 5.3, HSA can be applied in both architectural settings.

Re: "clarify clearly and simply the exact steps required for HSA" The exact steps to perform HSA are detailed in Algorithm 1-3 and Eq. 15 in Appendix D.

We thank the reviewer for their insightful feedback.

Re: comparison to other hierarchical methods We agree with the reviewer that in general it'd be highly desirable to compare the proposed HSA framework to other hierarchical self-attention approaches. The main reason we did not perform such a comparison was that some of these methods have used the notion of "hierarchy" in somewhat different sense than what we have used within HSA, while those with similar interpretations did not provide their codebase. Regardless of that, we are intending to add such comparisons to our draft anyways. Nevertheless, it should also be emphasized that in this work, our main goal is not to propose the empirically best hierarchical self-attention algorithm out there. Instead, we aim at proposing theoretical constructs and rigorous derivations that seamlessly generalize the notion of Softmax self-attention to multi-scale and multi-modal settings.

Re: hierarchical self-attention and inductive bias. We agree with the reviewer that self-attention was in part introduced to overcome the restrictions of other neural architectures dictated by their inherent inductive biases. However, some domains such as language or multi-modal problems semantically have scale separation inductive biases embedded in them which do benefit the statistical complexity of the model (as our experiments show) if they can be somehow injected into it. But the standard self-attention mechanism cannot naturally incorporate this knowledge and that's where HSA comes into the equation. Furthermore, even when such scale separation inductive biases are not semantically grounded in a problem, by imposing hierarchical structure on the problem, HSA can effectively act as a mechanism to arrive at a better trade-off between accuracy and efficiency depending on the use-case.

Re: comparison with other self-attention operators. Please note that in this work our goal is not to develop a universally "best" attention mechanism that beats every other attention operation in the literature on all benchmarks. Instead, we are proposing a rigorous, mathematical construct (i.e. the nested signal) that can represent multi-scale, potentially multi-geometry problems within the transformer framework. Furthermore, we have extended the classical Softmax attention mechanism to this construct using our statistical mechanical derivation of self-attention, which we have further proved to produce the closest attention weight matrix to that of the classical Softmax attention theoretically. In that sense, comparing our methodology against other generic alternative attention mechanisms is not quite relevant to what we'd like to achieve in this paper.

Re: how conceptually HSA differs from previous hierarchical self-attention First, please note that in this paper, we do not propose any contribution regarding the loss function or the backward pass in general; all the contributions of this work are solely for the forward pass architecture of the model. As for the conceptual differences of HSA and other hierarchical transformer frameworks, there are two fundamental notions that differentiate HSA from previous work: (1) HSA is derived for our proposed notion of nested signal which is a very versatile tool to represent multi-scale, multi-modal data; this is in contrast to previous work that often operate under single-modality settings and/or with fixed, rigid hierarchical structures. In that sense, our framework can be seamlessly applied to a vast and diverse set of multi-scale and multi-modal problems, way beyond what we have presented in the paper. (2) The way we calculate the attention weights in HSA is theoretically grounded and as we proved in Theorem 1 can be seen as the direct generalization of the Softmax attention to hierarchical, multi-modal settings. In that sense, our method differs from other previous work that often define hierarchical attention heuristically. Due to its theoretical optimality, as we have shown in Sec. 5.3, HSA can replace classical Softmax attention operator post-training in completely zero-shot manner to reduce FLOPs with minimal accuracy degradation.

Thank you for the clarification. I do agree that providing a theoretical construction based on the notion of nested signal and the arising constrained attention construct is interesting and useful.

What I am less clear about is, empirically, which is the optimal structure that can be learned and have good empirical trade-offs. There have been works that induce hierarchical constraints. Previously, there have been similar attempts at introducing sparsity constraints - For instance, the work, 'Generating Long Sequences with Sparse Transformers', by Rewon Child et al introduced block-diagonal sparsity constraints that they showed scaled better than flat attention networks. Similarly, this was then extended to a more general learned sparsity constraint in 'Adaptively Sparse Transformers' by Correia et al. EMNLP 2019. Given the significant past exploration, a new work that provides a theoretical construct needs to relate the various constraints induced and provide comparisons to the variants in order to validate the contribution. I believe, the present work does not yet achieve that aim. I would be happy to be corrected.

We appreciate the reviewer's constructive feedback and insightful comments.

Re: Adding the inductive bias of hierarchy We agree that injecting an inductive bias to a model can be generally restrictive. However, some domains such as language or multi-modal problems semantically have scale separation inductive biases embedded in them which do benefit the statistical (on the top of computational) complexity of the model (as our experiments show) if they are injected into it. In that sense, for such domains, we'd argue it makes sense to do so. Nevertheless, as the reviewer correctly pointed out, arriving at such semantically meaningful hierarchical structure is not straightforward in every problem. In such cases, the application of HSA with an "imposed" hierarchical structure acts as a flexible mechanism to change the trade-off between the accuracy and efficiency (e.g., improving efficiency at the cost of minimal accuracy degradation).

Re: the algorithm not being parallelizable As we have mentioned in Appendix D.4. the proposed dynamic programming presented in Algorithm 1-3 is not parallel as it's basically a tree-traversal algorithm. However, as we explained next, it is indeed parallelizable as we have illustrated how to do so in the same Appendix. Furthermore, our implementation does not require any dedicated hardware for acceleration and uses the same standard operations in Pytorch. We are planning to release our code, so hopefully that'd be helpful to other researchers seeking to implement HSA in parallel manner.

Re: comparisons with other hierarchical methods We agree with the reviewer that in general it'd be highly desirable to compare the proposed HSA framework to other hierarchical self-attention approaches. The main reason we did not perform such a comparison was that some of these methods have used the notion of "hierarchy" in somewhat different sense than what we have used within HSA, while those with similar interpretations did not provide their codebase. Regardless of that, we are intending to add such comparisons to our draft anyways. Nevertheless, it should also be emphasized that in this work, our main goal is not to propose the empirically best hierarchical self-attention algorithm out there. Instead, we aim at proposing theoretical constructs and rigorous derivations that seamlessly generalize the notion of Softmax self-attention to multi-scale and multi-modal settings.

Re: replacing standard self-attention with HSA and fine-tuning Please note that in the experiments in Sec 5.3, we did not fine-tune the models after replacement; the reported results were obtained in a completely zero-shot manner without any re-training or fine-tuning. Furthermore, the hierarchical structures used in Sec 5.3 are not based on semantic structure of the text but are merely fixed, imposed hierarchies. That is, we are not really injecting any semantically meaningful inductive bias into the model using these hierarchies. Given these two experimental settings, accuracy degradations are indeed expected. But the reason we run experiments in Sec 5.3 under these adverse settings is because we wanted to examine the practical implications of Theorem 1. In particular, we'd like to see how much overall accuracy degradation occurs if we simply replace the standard attention weight matrices in a pre-trained transformer with (imposed) hierarchical ones that have the minimal KL-divergence to the original matrices. And interestingly, the results indeed confirmed the practical importance of the theoretical result in Theorem 1.

Re: incorporating other tricks during training Our proposed contribution in this paper are solely concerned with the forward pass architecture of the model, so most of the tricks used during the backward pass in standard transformers can be equally applied to a HSA-based transformer model as well. The main reason we have not incorporated any of such techniques in our empirical study was that we were not mainly after beating the SOTA but purely measuring the impact of introducing HSA while excluding any other potential factor impacting the overall performance of the models.

We appreciate the reviewer's constructive comments and feedback to further improve our work.

Re: the restrictions of the proposed formulation

1. The (lack of) value projection: first of all, please note that in both our proposed Algorithms 1-3 as well as all of our experiments we have indeed incorporated value projections. Furthermore, in order to justify the incorporation of value projections within our theoretical derivation, we have added a new appendix (Appendix C in the new draft) explaining how value projections theoretically emerge in our statistical mechanical interpretation of self-attention. (TLDR: value projections emerge as learnable step-size in the gradient update in Eq. (6) realized as gradient projections.)

2. Pre-norm vs. post-norm: The reviewer is indeed right that in larger-scale problems post-norm may become restrictive. Note that the introduction of post-norm is merely a technical trick to arrive at standard self-attention formulation in Proposition 1. In practice, however, we can simply use pre-norm and since the outputs would be bounded anyways, Proposition 1 would still hold. In fact, in our empirical study in Sec 5.3, we have used pre-norm. We have further added new experiments in the new draft that use the standard pre-norm (see the comment below).

Re: comparison to "current" transformers The FSA results in Sec 5.1 and 5.2 are obtained under the same architectural constraints as those of the HSA. However, we have added a new experimental study in Appendix I where we compare HSA against regular transformers where neither of the two models bears the restrictions mentioned in the main text anymore. The results of these experiments show that HSA still does better than regular transformers while at the same time significantly reduces the number of FLOPs required to perform self-attention.

Re: lacking finer attention weights between two modalities This is indeed an acute observation by the reviewer. Generally speaking, the scale separation inductive bias incorporated within our framework dictates the block constraint between any two non-overlapping sub-trees of a nested signal. And that's where the computational (and often the statistical) gain comes from. However, as the reviewer pointed out, in certain use cases, finer attention weights between sub-parts of different modalities are highly desirable. To accommodate such cases within our proposed framework, one can simply put the sub-parts of the two interacting sub-trees under one parent node, which is equivalent to placing the sub-parts on a common geometrical domain. This would force the algorithm to calculate separate attention weights among the sub-parts. Finally, if such common geometrical domain is not easily attainable for the two modalities (e.g. in the case of text and image), then one can resort to "no geometry" which is equivalent to unordered set geometry in our framework.

Re: replacing regular self-attention with HSA in Sec 5.3 The HSA-RoBERTa model introduced in Sec 5.3 is obtained by replacing the attention calculation operator with HSA while keeping everything else (including pre-norms, architecture of each layer and the learned weights) intact. The results in Sec 5.3 are obtained after replacement without any fine-tuning and hence the degradation of the accuracy. However, as we showed in the newly added experiments in Appendix I, by training from scratch, we can indeed beat the performance of the baseline while reducing the number FLOPs.

Re: insight into scaling up to much larger models As stated in the paper, HSA reduces the number of attention weights from O(M^2. b^2) to O(m.b^2), where M is the number of internal nodes and b is the average branching factor of the hierarchy. Since the overhead of HSA is O(1), this reduction will become more visible in terms of the memory footprint of the model as the model and context/input sizes become larger, which means that we'd be able to fit larger models in the restricted GPU memory for training.

Re: the dimensionality of query and keys We thank the reviewer for noticing the typo. Indeed the attention dimensionality doesn't need to be the same as the input dimensionality as it's the case in all of our experiments. We have fixed the typo in the new draft.

We apologize for the late response. We were preparing the new experimental results to address some of the reviewer's questions while at the same time experiencing loss of power due to a cyclone in our area during the last few days. We hope our answers can clarify the questions raised by the reviewers.