Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems
Authors: Saeed Amizadeh
Sara Abdali
Yinheng Li
Kazuhito Koishida
Venue: NeurIPS 2025 poster ()
Keywords: Attention Mechanism
Transformers
Hierarchical Data
Multi-modal Data
Geometric Deep Learning
Primary area: deep_learning
TL;DR: We mathematically generalize the Softmax attention mechanism to hierarchical and multi-modal settings.

Abstract: Transformers and their attention mechanism have been revolutionary in the field of Machine Learning. While originally proposed for the language data, they quickly found their way to the image, video, graph, etc. data modalities with various signal geometries. Despite this versatility, generalizing the attention mechanism to scenarios where data is presented at different scales from potentially different modalities is not straightforward. The attempts to incorporate hierarchy and multi-modality within transformers are largely based on ad hoc heuristics, which are not seamlessly generalizable to similar problems with potentially different structures. To address this problem, in this paper, we take a fundamentally different approach: we first propose a mathematical construct to represent multi-modal, multi-scale data. We then mathematically derive the neural attention mechanics for the proposed construct from the first principle of entropy minimization. We show that the derived formulation is optimal in the sense of being the closest to the standard Softmax attention while incorporating the inductive biases originating from the hierarchical/geometric information of the problem. We further propose an efficient algorithm based on dynamic programming to compute our derived attention mechanism. By incorporating it within transformers, we show that the proposed hierarchical attention mechanism not only can be employed to train transformer models in hierarchical/multi-modal settings from scratch, but it can also be used to inject hierarchical information into classical, pre-trained transformer models post training, resulting in more efficient models in zero-shot manner.

Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems
Saeed Amizadeh, Sara Abdali, Yinheng Li, and Kazuhito Koishida
Microsoft
Redmond, WA 98052
{saamizad, saraabdali, yinhengli, kazukoi)@microsoffr
Abstract
Transformers and their attention mechanism have been revolutionary in the field Machine Learning. While originally proposed for the language data, they quickly found their way to the image, video, graph, etc. data modalities with various signal geometries. Despite this versatility, generalizing the attention mechanism to scenarios where data is presented at different scales from potentially different modalities is not straightforward. The attempts to incorporate hierarchy and multi- modality within transformers are largely based on ad hoc heuristics, which are not seamlessly generalizable similar problems with potentially different structures. To address this problem, in this paper, we take fundamentally different approach: we first propose mathematical construct to represent multi-modal, multi-scale data. We then mathematically derive the neural attention mechanics for the pro- posed construct from the first principle of entropy minimization. We show that the derived formulation is optimal in the sense of being the closest to the standard Softmax attention while incorporating the inductive biases originating from the perrarchcallgeommeti information of the problem. We further propose an effi- cient algorithm based on dynamic programming to compute our derived attention mechanism. By incorporating it within transformers, we show that the proposed hierarchical attention mechanism not only can be employed train transformer models in hicrarchical/mulli-mooo inject hierarchical information into classical, pre-trained transformer models post training, resulting in more efficient models in zero-shot manner.
settings from scratch, but it can also be used to
Introduction
The field of Deep Learning has recently experienced a spectacular breakthrough with the rise of Large Language Models (LLMs)). no secret that this success is largely owed the Transformer architecture [90 and its self-attention mechanism. Although they were originally proposed to work
with language 24
52
13,
18, transformers have found their way deal with images 25
100 85).
video 9
65
4 46 audio 32
91
45 16 graphs [1OI] 61 72 76), groups 42
83 manifolds
38 and point clouds 3
107 without significantly altering their basic neural attention mechanism.
This is mainly due to the fact that, unlike many other architectures, transformers incorporate data geometry not by architectural priors but by explicit, black-box, position embedding functions, which can be easily replaced from one domain to another.
Despite this versatility, information quite often comes in diffferent modalities and different scales. In terms of geometry, this means that we deal with problems where each datapoint may occupy multiple,
inconsistent geometries at potentially different scales. This is indeed challenging. even for transformers! To this end, various novel (often heuristic) neural architectures have been proposed deal with multi-modal [55
40 69 and hierarchical data [54 105
67 94
19
106
110
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Aside from the heuristic- based nature of many of these empirical architectures. they quite often suffer from more practical dilemma. On one hand, many such frameworks tend to partially discard geometrical or hierarchical information depriving the learning task from valuable domain knowledge which can significantly reduce the model's statistical complexity. On the other hand, by incorporating the full geometrical knowledge of different modalities and their hierarchical structure within these heuristic frameworks, we often end up with highly woblem-speccife architectures that are hardly generalizable to other similar problems. To address this challenge in a unified and principled way, in this work, we take a radically different approach:
Instead of coming up with yet another heuristic neural architecture right off the bat, we first propose a mathematical construct called nested signal to formally represent information. As we show, the proposed formalism enables us to coherently represent different geometrical domains at different scales while maintaining its generality across different problems.
hierarchical
In order to define mathematically-soon mechanism. First, we show that the standard Softmax self-attention [90] can be mathematically derived from the principle of entropy minimization. Then by generalizing this principle to nested signals, we derive the hierarchical self-attention (HSA) neural mechanics which is the generalization of the Softmax attention for nested signals.
neural operations on nested signals, we turn to the attention
We further show that the attention weights derived from the HSA are optimal in the sense of being the closest to flat Softmax attention weights in terms the total KL-divergence, while at the same time adhering to the hierarchical structure of the data.
Next, we propose an efficient algorithm based on dynamic programming to calculate the HSA, that is provably faster than its direct evaluation. By implementing HSA within the transformer architecture, we empirically show that we are able to train models that can seamlessly incorporate the seerarchica/mmuliimooee domain knowledge to arrive at better and more efficient transformers.
Last but not least, we show that HSA can further replace the standard Softmax self-attention operation in pre-trained transformers and significantly reduce the number of self-attentin FL.OPs while incurring minimal Accuracy drop, in entirely zero-shot manner.
2 Representing Hierarchical Data
In Geometric Deep Learning, signal is defined as the mapping x: Ω - C, where the set Q the domain of the signal and C is a vector space, typically Bdd with d being the channel dimension. For example, an RGB image is a signal where Q is the 2D grid and C = IR, i.e. the RGB color space. Similarly, text can be seen niche applications in Geometric Deep Learning (17 extend the notion of signals to the domain of graphs, gauges, manifolds, etc. by defining the appropriate structure for Ω. We refer to the set of all such possible domains as D. The elements f2 E DD are not necessarily vector spaces (e.g. 2D grid). In order to numerically handle these spaces, define special signal En Ω -> I' for each Ω D which maps the elements of each domain in D to R"; we refer to this special signal as the position embedding. Given En, each signal defined on Ω is seen generalize the notion of signal to encompass traditional tabular features.
signal with 52 being the 1D grid and C word embedding space. More
: (?) C. In Appendix C we
In this section, we introduce the notion of nested signals which is the key modeling tool to represent multi-modal, hierarchical data. To this end, we first define the set of all simple signals as the set of all possible signals defined on all possible domains; that is, S = {x Ω C Ω E D). The signals defined on different domains may have different channel dimensions; to make the channel dimension uniform across different domains, we vero-pad the lower dimensional signals to the maximum channel dimensionality d across different domains, such that each element of S has the same channel dimension d regardless of its domain.
Definition 2.1 (Nested Signal). The set of d-dimensional nested signals up to depth f, .e, is recursively defined as No. = x Ω + U Ω E D,U E (N-1,RR)) where No = Rdd Also, define N = N₂ as l + oc; each element domain Ω E D of a nested signal x, denoted by r(z).
A is then referred to as nested signal. The top-level
For example, website is a nested signal where at the top level, we have webpages de- fined on the nodes of a graph domain representing the link structure between the web- pages. Each webpage is in turn another nested signal where at its top level we have
2

---

an unordered set of textboxes and images constituting the page. Going one level further, word embeddings or pixel values to
each textbox or image is (simple) signal assigning the nodes of 1D or 2D grid domains, respectively. While in theory, the domains f2 E 'D can be infinite, in practice, we mostly deal with nested and simple signals defined on finite S!'s. In particular, a nested signal z is said to be finite if the domains !tss at all of its nesting levels are finite. Given the set of position embeddings ε = {En Ω E D), finite nested signal can be represented by signal hierarchy as follows:
Fig. Topp) depicts this example.
UnI
Definition 2.2 (Signal Hierarchy). For a finite nested sig- nal x, its signal hierarchy hg is a tree with the root node Rs associated with r(z), the top level domain of z. The
Feat
Houess
children of Rg are defined as Bzz = {he(w) | u E r(r) where x(u) is the value of signal (possibly another nested signal) at If z=(u) is a vector instead of a sig- nal, then hu((w) is simply z(u). Furthermore, each child hz(u) E chd(R., is annotated by (r(z)(u),
unbeddinge
(u), the position embedding vector dictated by its parent node.
Page
Papez
Page
Page
We denote the nodes (and equivalently their correspond- ing sub-trees) in the signal hierarchy h2 by upper-case letters. Any set of sibling nodes in fgg is referred as a fam- ily, The members of a family are nested signals (or real vectors for the leaf nodes) that reside on the same domain Q and therefore share the same position embedding fune- tion En. Furthermore, for A E hr, chd(A) A)), sib(A) and ((A) represent the set of A's children, its siblings and the index set of the leaf node descendants of A, respectively. Two nodes in h are called unrelated if neither of them is descendant of the other. For two unrelated nodes A and B, their immediate common ancestor denoted by ica(A, B), while their highest distinct ancestors are denoted by A' and B', respectively, where we have A', B' E chd(ica(AA (A B)); i.e., A' and B' are always siblings even if A and B are not. See Appendix A for the notational details well as a visual demonstration of the tree-related concepts.
Pass
Ward
Word
Ward
Word
Peets
webeedingg
embeedding anbeddingg
Graph position embedding
mored set position embeddinn 20 gnid position embedding position ambedding
Figure 1: (Top) A nested signal repre- senting a website. (Bottom) Its signal hierarchy representation. Different col- ors encode different types of position embeddings assigned to each node.
Since sibling nodes share the same position embedding function, the relative positional distance (or similarity) between them is well-defined More generally, for any two unrelated nodes A, B E hzz we can form a well-defined positional distance between them by comparing the position embeddings of A' and B' which is well-defined since A' and B' are always siblings. The implication of this construction is indeed powerful as it would enable the signal hierarchy formalism to define meaningful positional distance between any two unrelated nodes in the hierarchy regardless of their modalities or signal types. Fig. Bottom) shows the signal hierarchy representation for our earlier website example.
Hierarchical Self-Attention
The nested signal formalism and its signal hierarchy representation introduced in the previous section provide a systematic way to represent hierarchical data that can potentially span across different modalities and domains. However, the question remains what kind of neural architectures can handle such versatile data structure? To answer this question, we note that for non- -in-aierarc simple signals, the transformer architecture first introduced by [90] allows for a unified representation learning methodology that can accommodate various signal domains (as long as the position embedding is available), not to mention its remarkable success in revolutionizing deep learning. However, extending the attention mechanism to nested signals is not straightforward as the information in such signals can come with different signal domains at different scales.
To address this problem, in this section, we first propose a statistical mechanical framework that ele- gantly derives the classical Softmax attention mechanism from the principle of entropy minimization when a finite (simple) signal viewed as a physical system with N particles. By generalizing our proposed construction to nested systems, we then derive novel, theoretically-rigrrous mechanism
3

---

for calculating self-attention within nested signals, which we refer as Hierarchical Self-Attention (HSA). By its direct construction, the proposed HSA mechanism aims reducing the total entropy of the nested system, or equivalently put, increasing information within the learned representation of the nested signal. We further show that our proposed construction to derive HSA optimal in the sense of Kullback-Leibler (KL.) divergence from the Softmax attention weights if the hierarchical structure were to be ignored. This result will subsequently open the door for the application of our proposed formulation to approximate the inefficient Softmax attention pre-trained transformers using the more efficient bierarchical calculations a hierarchy exists and can be imposed in a given problem. Finally, propose efficient algorithm based dynamic programming that calculates HSA for a given signal hierarchy hg in (MM b²), where M is the number of families in h2 and b is its maximum branching factor (i.e. family size).
3.1 Softmax Attention Revisited
Let x { E Rdd 1 NJ be finite signal with N elements in Rd' with the corresponding position embeddings {e, E R 1 < i < N). To calculate self-attention over æ, one needs to define
the set of query variables Q = {4= E Rd
S S N] and key variables K = (h; R.d
N), where s and k; is are (linear) functions of z4. Then the conditional entropy of Q given K is:
H(Q K) = - ,(Q, Ko (0Q K)dQdK =Eq., log ((Q K)]
(1)
where p(Q, K) and p(Q K)) are the unknown joint and posterior distributions over Q and K. While the joint distribution can be approximated using the Monte Carlo method, the posterior can be approximated by variational distribution E(Q K), which gives rise to the variational upper-bound on the conditional entropy:
HuB (Q K) EE,, log E(Q K)] 2 H(Q (K)
We further represent the variational distribution by the Boltzmann distribution, i.e. E(Q K) Z(KJ expl- K)/r], where (Q, K), Z(K), and r are the energy function the partition function and the temperature parameter, respectively. The variational upper-bound then can be written as:
(u K) = Eq k (0Q K)/T] + Ex log Z(K)
The end goal of representation learning is to transform the input signal (i.e. the query variables Q) into a "better" representation. principled way to arrive at better representation is modify Q such that its information content is maximized, or equivalently its entropy is minimized. Since we cannot directly calculate the entropy, we can work with its variational upper-bound HyB as a proxy. Then, the entropy minimization approach amounts to gradient descent on Hon w.r.f. each q::
de + q: A Vg (Q K) = .... (V
ISiSN
(4)
where À > 0 the step size.
Proposition 3.1 (Softmax Attention). For the energy function ({Q, K), defined as:
N
sQQ. K) = N log
Σ
Σ
exp
|1::
k2 12+ of egl
(5)
N N-
2vd
j=L.
if Q and K variables are normalized using the LayerNorm function [0] then for (Nv0)-1 x 1 and sample size 1, the Eq. (4) reduces to:
N
exply kj/ efe,)
ISi5N, q: - qi +
E
ki
(6)
Ltmltyy explat he/ Vd of
N
$=1.441
essentially, this is Softmax attention via residual connection.
Proof. See Appendix H.I
Note that the energy function needs satisfy f (x)p_ co.

---

Note that (6) is similar to the original attention formulation proposed by 190, except for a few differences: (1) there is no separate value linear projection; the value projection emerges later as we incorporate learnable step-size (see Appendix D). (2) the LayerNorm is applied post-linear projection as opposed to pre-normalization in the original formulation, and (3) the residual addition is applied post-linear projection. In other words, with few minor modifications, the original Softmax attention operation can be interpreted as maximizing the information content in the representation. But the real importance of the formulation in (4) is that depending how we define the energy function, we can arrive at various types of attention mechanisms tailored to different applications. We use this feature in the next section to derive hierarchical self-attention (HSA) mechanism for nested signals.
3.2 Generalizing Attention to Nested Signals
We derive a self-attention mechanism for finite nested signals represented via signal hierarchy tree. We follow the same recipe as the previous section by defining an appropriate energy function. But first, for any two unrelated nodes A and B in hig we define the interaction energy A-+B:
AA-B= = -En("" (B') (c()" + 2vd (((A) |E(B) (E
1
ΣΣN
11 kyll
where | set cardinality, (-(-) is the position embedding dictated by ica( A, B) (i.e. S2 = ((i and A' and B' are the highest distinct ancestors of A and B. as defined in Section 2 Intuitively speaking, the interaction energy A captures the dissimilarity between the nested signals rooted A and as weighted sum of their highest non-common ancestors' position dissimilarity (the first term) and the average Euclidean distance between their leaf nodes (the second term). By calculating energy (dissimilarity) at the subtree level instead of individual leaves, we inherently encode the inductive bias that the leaf nodes of subtree (i.e. a nested signal) can be pooled into a single representative (i.e. the subtree's root) while roughly maintaining the underlying semantics. This is referred to scale separation in Geometric Deep Learning [17.. fundamental prior dealing with multi-scale physical systems, benefiting us both statistically (by taming the curse of dimensionality) and computationally (by providing efficient algorithms). Using the interaction energy definition, the energy of the signal hierarchy rooted at non-leaf node A is recursively defined:
((A)
Σ
|((B)] log
exp (- Q(B)) +
Σ CE
Ie(C)) exp ( V
(8)
IE(A))
BEChO(A)
For leaf nodes, ((A) is set to co. of (Ry) is the energy of the whole signal hierarchy ha Intuitively, (8) states that the energy of a system (a signal hicrarchy tree) is the weighted sum of the energy contribution of its subsystems (immediate subtrees) where the weights are proportional to the size of each subsystem. The contribution of each subsystem, in turn, is a non-linear combination (via the weighted log-sum-exp function, which is the addition operation in the log-space) of the energy of the subsystem itself (the recursion term) and its interactions with its sibling subsystems (the second term). It is easy see that for single-level /, (i.e. simple signals), ((R)) reduces to (5). Having defined the energy function, we can follow the recipe in (4) to calculate the HSA for ht by recursively computing the gradients Vg. (o (R.)) E Idd for each leaf node ge. E ((R2) as:
[B(B)- V Vg ((B) + ECCESD(DA)) [[C)) V B
|((B')
V (R.)(R
(9)
aBB)) + ECESit(B [(C))](Bi,C)
[(((..)
where aB*) exp (- dBB')())))))))))))))) 3(B", C) = exp ( (-B
(10)
and B' denotes the child of R2 which contains 9i as leaf. It is not difficult to show that for the quadratic interaction energy function in (7), if both Q and K variables are normalized beforehand
using a LayerNorm layer, then the recurrence in (9) can be unrolled and written in the matrix form
(see (36) in Appendix H.3
V + OK where:
o (Rz)...., V ( ((a))= K ke,,,...,
(11)
and e - (R_ (R.) is the attention matrix; that is, 8G., is the coefficient of the key variable k, for computing the attention update V v R.,) for the query variable q (.. However, θ is different from classical attention matrix in the sense that many of its entries share the same values. In
5

---

particular, for any two sibling nodes A and B in hz, the corresponding entries between the leaves of A and B form block in O with one value; that is, fi.5 = AA.B. B. Vi E ((A). ((B). In other words, the attention weight between any leaf node in A and any leaf node in B is approximated by one value AAB; we refer to this approximation between the leaves of sibling nodes in h2 as the block constraint which makes the attention matrix a hierarchical matrix [35 [6]. Fig. Left illustrates the self-attention matrix for a toy example signal hierarchy with the block constraint. The block constraint is directly administered by the form of the interaction energy function in 7 as well as the signal hierarchy energy recurrence in (8)
The block constraint effectively reduces the degrees of freedom for an attention matrix from O(I((R,)))) = O(M² b²) to (MM bb)), where (,,,),, M and are the total number of leaf nodes, the number families e. non-leaf nodes) and the maximum branching factor in h,, respectively. With- out it, we essentially go back to the standard Softmax attention mechanism where the unormalized attention weights before Softmax are calculated by evaluating the interaction energy function for every pair of leaf nodes. We refer to this process as flattening a nested signal. Fig Right) shows the self- attention matrix for the flattened version of our earlier toy example without the block constraint. Flat- tening is not only computationally costly (by being quadratic M instead of linear), may also hurt the model statistically. Note that by enforcing coarse-grained attention weights through the block con- straint, we effectively administer a form of regularization guided by the scale separation prior which is in turn induced from the prior knowledge of the hierarchical structure in the problem. By flattening nested signal, we simply discard this prior knowledge which can make the model prone to over- werr-fitting. Note that the block constraint by itself merely enforces tied values for the atten- tion weights over the leaves of sibling nodes; it does not, however, specify what those values should be. That is, there are infinitely many attention matrices that ad- here to the block constraint; our proposed formulation in ( is just one of them. How- ever, as we show next, our proposed formu- lation is optimal in the sense of being the closest approximation
the standard Soft- max attention if the nested signal were to be treated as a flat, simple signal.
Figure 2: (Left) The self-attention matrix for toy sig- nal hierarchy with the block constraint. Each contiguous tile here represents one tied value for the corresponding cells. (Right) The self-attention matrix for the flattened (or simple) signal without the block constraint.
Theorem 3.2 (The optimality of HSA).
Let both Q and K variables be normal- ized using the LayerNorm function. For
the given interaction energy function ta in
77. if Q = -Oo,,))e(K,, xx)((R is the self-attention matrix for the nested signal derived from the proposed gradient recurrence in (9) (as depicted by (TT),, then for the temperature parameter
(Ie(R=)lVd) - = 1 is stochastic matrix; that is, it is non-negative and we have Ô = 1. Moreover, Ô is the closest attention matixx with the block constraint to the classicul Softmax attention matrix for the flattened signal in terms of total Ki-ddiveggence; that is,
① = arg min EEB
Σ DKL(OL,10(..)
(12)
(EE(R=)
where B C [ is the space of all stochastic attention matrices that admit the block constraint induced by hrs and 0%, (Vi ((R,)) are the rows of the attention matrix for the flattened version of the signal:
exp(-vi-)) EkEa(R)) expl (op(-Wi -1-)))) Vi.j E ((Rz)
0.
(13)
Proof. See Appendix H.2
This result is crucial in the sense that it shows our proposed HSA mechanism for nested signals formalized by (8) and p is the closest approximation to the classical attention mechanism while at the same time adhering to the block constraint (induced by the hierarchical structure of the nested signal), which benefits the model computationally and statistically.
From the practical perspective, this result has another important implication: if we replace the interaction energy function tiizy with the original cosine similarity in transformers (where the

---

position information is simply added to the signal), our proposed methodology provides the closest hierarchical approximation of the original Softmax attention. Practically speaking, this means that if we have access to some form of hierarchical information lt in problem at inference time, we can simply replace the self-attention operation in pre-trained transformer- based models by HSA and arrive at much more efficient calculations without the need for major re-training. Note that the direct evaluation the recurrence in ( for all query variables qi still takes O M logs M). In Appendix we prove that the HSA can be computed in O(M b²) using dynamic programming algorithm. Furthermore, we propose transformer encoder architecture based on the HSA in Appendix
4 Efficient Calculation of HSA
Algorithm 2: The Top-down Pass
Algorithm 1: Hierarchical Self-Attention
Input: A fgg 12 RR Output: ((A) E IRd
Input: h
2 Output:
e ((R)) &- - (Rm))
Function ComputeAttention (A, u,
foreach E chil( A) do
((C) < +/2 exp (u
ComputeSuuffiicietttttssss/ Computeattenttion(R. u, 0)
LogSigmoid ((C) ((C)) ((C)
foreach i E ((R,) do
t u Loggimmoid ( o(C)
V (R22 ((L))
ComputeAttention (C, (CC)
8 end
end
return (V ie {(R,))
End Function
Algorithm 3: The Bottom-up Pass
Input: E hgg
// A node in the signal hierarchy
2 Output: ((A) E R. (AA) e R, 0(A) e gd Function Compite@uffiiienttttaa. (A):
if' A is a leaf then
((A) + 20
(A(A) + (A)
// //A) is the query at leaf A // k(A) is the key at leaf
(=(A) t k(A)
(A) + (AA
// v(A) is the value at leaf A
else
10
foreach C E chd( A) do ComptteSufficientttatee end
11
12
((1) +- Eccond(A) [(())
log [exp ( ((C) exp ( [C))]
13
[[(A]]
14
(A(A) <- [TAT Ecechd(aa) (C)
15
Pa(A) < [((A)] Ececha(A)) (((C)lpx(C)
(A(A)= t E(A)) E Ecechd(A) (((C)lp.(C)
17
end
18
vB E sib(A):
1// A-BB 8l ((A) +
-1 (a)) (((AA)(({(BB) Vd loggl [((B]
19
((A) - - log |BE exp(v)
20
0(A) - cxp ({)) EBESSb(A) eXp(VAA-AR) Pc
21
End Function
Even though our proposed HSA formulation in Eq. brings down the degrees of freedom for the attention matrix O(M b2), the naïve implementation of the recurrence in Eq. (9) for all query variables 4s still takes O M log, M) time. However, we note that the calculation of V q. øl Rg) and V, (R,) for any two leaf nodes i, E ((Rz) shares some common intermediate calculations corresponding to the shared segment of the two paths that connect the root node to i and j. This indeed gives rise to the notion of common optimal substructure which the hallmark of problems that can be efficiently solved by dynamic programming. To this end, we propose dynamic programming
7

---

algorithm that computes V in Eq. (IT) in O(M . b2) time by traversing the signal hierarchy tree in two passes: a bottom-up pass followed by a top-down pass. Essentially, the former computes the energy function of-) while the latter calculates the attention vectors Va ((-) for all i E { RZz). Algorithms | illustrate these steps. For the formal correctness results as well as further practical details for our proposed algorithm, see Appendix E.
5 Experimental Results
In this section, we present an empirical study aiming at two main goals: (1) showing the capability of the HSA mechanism in incorporating useful domain hierarchy knowledge into training better transformer models from scratch, and (2) demonstrating the unique capacity of HSA as post-training approximation of the Softmax attention in pre-trained transformer models in order to reduce the self-attention computation FLOPS in zero-shot manner.
5.1 Hierarchical Language
Despite its unimodality, natural language data often comes in semantically meaningful hierarchy (e.g. sections, paragraphs, sentences, etc.) which can be seen as granular abstraction of the underlying semantics in the data. Nonetheless, most transformer-based frameworks ignore this hicrarchical structure which not only discards valuable prior knowledge about the semantics of the text, but in the long context scenario, can also result in loss of information due truncation (which is common practice for long sequences in order to manage the computational complexity of the Softmax attention). HSA avoids truncation for long sequences by effectively reducing the computational complexity via incorporating the hierarchical abstraction.
For empirical assessment, we have chosen the text classification problem for the sentiment analysis task on two datasets: IMDB [57 1. and Elec [59 2 for sentiment classification in movie
reviews and Amazon electronics product reviews, respectively 3 The rationale for choosing these datasets lies in their inclusion of lengthy texts, which means they can benefit from hierarchical representation. For details, see Appendix I.]
Signal Hierarchy:
We represent each text datapoint in our datasets as a 3-level signal hierarchy: paragraphs, sentences and tokens. The position embedding each level is the 1D grid embedding materialized by random Fourier features [48]. The tokens form the leaves of each signal hierarchy and are represented via vector embeddings. We have experimented with two token-emm oekeneemeeddings in our experiments: the simple Word2Vec [60 4] and the richer, transformer-based 75 77 S.
Experimental Settings: We have used similar architectures for both the baseline and the HSA, each amounting to 1.2M trainable parameters. For a fair comparison, we have used the same training hyper-parameters for both models. See Appendix for the details of experimental settings.
Word2 embedding
T5-small embedding
Dataset
Model
Acc
F1 Score
Acc
F Score
FSA
0.6739 = 0.0004
0.6739100.004
0.757710.0024 0.757740.0024
IMDB
HSA
0.746911.0000 0.0029
0.7468110.0027
0.8129110.0010
0.8129:1000100
FSA
0.7182+40.0001 0.7182+00.001
0.8212+000014
0.8212+0.0014
Flee
HSA
0.7549-0.00005
0.754910.0005
0.8521110.0222
0.852110.0022
Table 1: The sentiment classification Accuracy/F1 score comparison for the Flat Self-Attention
(FSA), iee. the Softmax attention, and the Hierarchical Self-Attention (HSA).
HSA vs. Flat Self-Attention: Table depicts the test Accuracy and F1 Score of sentiment classifica- tion for the two models on the IMDB and Elec datasets. As these results show, HSA consistently and significantly outperforms the standard Softmax self-attention across the datasets well as the ooken-embeddiins. The superiority of HSA over the standard self-attention can be attributed to two main factors: (1) by incorporating the semantic hierarchical knowledge of the problem within the attention computation process, HSA effectively employs form of regularization based on the scale separation prior that protects it against potential overfitting, and (2) for long input sequences, unlike the standard self-attention mechanism, HSA can evade truncation of the input sequence by effectively reducing the memory and the compute footprints of the attention mechanism.
Word2Vec vs. T5 embedding: From Ta we also observe that the classification results signifi- cantly improve for both models by replacing the basic Word2Vee Vec token embedding with the richer T5
8

---

embedding. This is not surprising, but it also shows that our proposed HSA framework can be incor- porated as a (shallow) adaptor on the top of pre-trained foundational models and adapt them for a new domain. Furthermore, we can see the gap between the HSA and the standard self-attention intensifies for simpler token embeddings. In other words, where we do not have access to pre-trained embedding models, the superiority of HSA and its hierarchical inductive bias is even more significant. This points to the potential sigmificant boost we can gain by training HSA-based, multi-modes foundational models instead of the classical transformers. Due to its demanding computational requirements, we leave this empirical investigation for future work. Nonetheless, in Appendix ! we have experimented with training the HSA-based transformer from scratch (as opposed to on the top of pre-trained embedding) and showed superior generalization capability compared to the classical transformer.
5.2 Multi-modal News Classification
Model FSA
Acc
FI Score
0.7921 + 0.0036
0.7902 + 0.0003
In order to showcase the capabilities of our proposed framework in multi-modal settings, we have performed experiments for the news classification task on N24News dataset (93], where for each news article not only we have language and image modalities present, but the text itself consists of multiple sub- modamldai,sses, i.e. headline, abstract, image caption and body.
DeepSet
0.7578 + 0.0096
0.7590 + 0.0065
HSA
0.7052-t0.0155
0.8091+:0.0100
Table 2: Accuracy/FI-scomm comparison for the Flat Self-Attention (FSA), iee, the Softmax attention, DeepSet(102), and the Hierarchical Self-Attention (HSA) on N24News dataset.
Original RoBERTa
HSA-RoBERR
Dataset
Acci
Pref
Reef
FLMM)
Accf
Pre1
Rect
FL(MM.
IMDB(264)
0.9558
0.9558
0.9558
214.94
0.9494
0,9501
0.9494
4.32
AGNEWS(54)
0,9469
0,9469
0,9469
8.99
0,9422
0.9423
0.9422
0.8357
ColA((12)
0.8150
0.8348
0.8017
0.4441
0.7687
0.7608
0.7821
0.1912
SST-2(26)
0.9403
0.9404
0,9402
2.08
0.9025
0.9083
0.9014
0.4132
MRPC(55)
0.9117
0.9006
0.8938
9.33
0.8553
0.8613
0.7963
0.8481
RTE(70)
0.7833
0.7870
0.7796
15.11
0.7400
0,7400
0.7377
1.29
QNLI(38)
0.9267
0.9267
0.9268
4.45
0,5072
0.3398
0.7531
0.5643
Table 3: The FLOPSs comparisons for zero-shot HSA approximation of RoBERTa-base layers 7,9,1 11 and RoBERTa-large layers 2,,.2,2 (for IMDB). We have reported MFLOPs per impacted layers as well as Accuracy (Acc), Precision (Pre) and Recall (Rec). The FLOPs are computed based on the average seq. length (shown in parentheses) for each dataset.
Baselines: For N24News dataset, most approaches in the literature concatenate a subset of the text sub- modaliioaites and use that as the representation of the whole article. There are also a few multi-modal methods that incorporate the image modality as well, the best of which achieves 91% Accuracy and 90% Fl Score using 21 2 trainable parameters [93], not to mention incorporating other tricks such as using multiple loss functions to achieve the SOTA performance. For our experimental evaluation of HSA, however, would need to keep these other contributing factors out, and instead compare moderate size models within our computational budget that are only different in their attention mechanisms. To this end, for our bascline method, we concatenate headline, abstract and body into one text sequence and use that to train a classical transformer (realized via one-level signal hierarchy). As the second baseline, to incorporate the image modality as well as the text; see Appendix [ for details. For all baselines as well our HSA-based model, we ensure the number of trainable parameters is around 12M.
incorporate multi-moda model based on the DeepSet architecture [022]
Signal Hierarchy: For the HSA-based model, each news article is represented as a signal hierarchy where at the top level the image modality as well as the text sub-modalities are represented by the key-value signal type (see Appendix C). The headline, abstract and caption sub-trees are further divided into tokens in the next level using the ID Grid signal type; whereas, the body is divided into paragraphs (again using ID Grid signal) where each paragraph is treated as a leaf by pooling the text embedding of the whole paragraph. To embed the text components at the leaves, we use e5-base 933 6); whereas, for image leaves, we use VIT 26 77. Both of these models have shown superior performance in various benchmarks 77.
Results: Table 2 shows the test accuracy and Fl Score for the three competing methods for the N24News multi-class classification problem. From these results, we can see that our HSA methodol- ogy outperforms the baselines and the difference significant. Interestingly, despite incorporating the additional modality of image, the performance of DeepSet significantly declines compared to the

---

vanilla uni-modal, flat attention. This signifies the fact that it is not enough to only incorporate other information modalities within the model, but also how they are incorporated is equally important to boost the model's generalization. In that sense, our proposed nested signal formalism along with its hierarchical attention mechanism provide principled methodology to incorporate different information modalities within transformer model.
5.3 Zero-Shot Hierarchical Approximation
An important feature of our proposed framework is that Theorem 3.2 gives us the theoretical basis for approximating Softmax attention via HSA given an appropriate hierarchical structure. This means that HSA can seamlessly replace regular Softmax attention after training, and depending on the task and the original model, the accuracy may not experience significant drop. The main objective for such replacement post-training is reduce the number of FLOPs needed for the self-attention operation. To further examine this idea, we have adopted the classical pre-traine tr-inrrined RoBERTa model [52] and have replaced the Softmax self-attention operation in it with HSA, and then run it against some benchmark classification datasets. During this experimentation, we made a few insightful observations. First, in general, the performance drops significantly if we replace Softmax attention with HSA for all hidden layers of RoBERTa, and some amount of fine-tuning is needed to regain the original performance. However, zero-shot replacement is still feasible if only a subset of layers go through HSA replacement. In particular, earlier layers seem to be more sensitive to HSA approximation while the final layers are more amenable to it. Furthermore, we observed that by interleaving HSA layers and regular Softmax layers, we can significantly reduce the accuracy gap.
Based on these observations, we applied HSA approximation to layers 7, 9 and 11 in RoBERTa- base and 16, 18, 20, 22 and 24 in RoBERTa-large. As for the hierarchy, instead of using the sentence/paragraphlltee.. structures in text, opted to fixed hierarchies generated by non-overlapping hopping windows on the input text. In particular, we used a four level hierarchy where the layers' branching factors from top to bottom are 16, 8, 4 and 2. For more experimental results on different hierarchy structures and different HSA layer combinations, see Appendix L Table 3 compares HSA-equipped RoBERTa (henceforth HSA-RoBERTH and the original RoBERTa in terms of FLOPs as well as Accuracy on 5 GLUE benchmarks (92). IMDB benchmark (58) and AGNEWS benchmark [04]. As these results show HSA layers significantly reduce the number of FLOPs for attention computation, and depending on the task the accuracy drop can be minimal. Keep in mind these results are obtained completely zero-shot without any fine-tuning. Indeed fine-tuning can further close the accuracy gap while maintaining the performance gain by HSA. This points to another HSA's strong potential: to be used as a self-attention approximation technique for long-context problems. We leave the further exploration of this direction to future work.
6 Conclusions
In this paper, we propose HSA, novel mathematical framework for generalizing classical Softmax self-attention mechanism to hierarchical problems that not only occupy multiple scales but may be also defined multiple geometries. Unlike many existing work that approach these problems via heuristic neural architectures, we mathematically derive our formulation from the principle of entropy minimization given the (nested) data signal is seen statistical mechanical system. Given its strong theoretical and algorithmic properties, we empirically showed that HSA can be used to inject hierarchical domain knowledge into training of transformer models and hence produce models with better generalization. We further showed that HSA can be used as a self-attention approximation technique for pre-trained models to significantly reduce the FLOPs needed for self-attention the test time. This opens the door for HSA to be used in long context scenarios, even after training.
One high-impact future application of HSA is training large-scale foundational models that can naturally handle multi-modal and hierarchical inputs using the HSA formalism. On the theoretical side, HSA can be also extended to include non-Softmax attention mechanisms (See Appendix K) The other important future direction is application of HSA to transformer decoder for hierarchical auto-regressive generation. This is important specially because it has the potential boost LLMs in terms of both generalization (by incorporating hierarchical, multi-modal domain knowledge) and speed (due to the low-rank nature of HSA computation). Due its significance, we have laid the foundations of hierarchical decoding via HSA in Appendix G
10

---

References
[2]
[3 http://maaa..ffafrrorrtl////a/12
[4] httpss//ooaaa.ooggaaaaaaaaa/ccer
[5] stee://mnggingfeeeccoooo/22/1//3
[6]
77
[8] Saeed Amizadeh, Bo Thiesson, and Milos Hauskrecht. Variational dual-tree framework for large-scale transition matrix approximation. arXiv preprint arXiv: 2012.
[9] Anurag Arna,, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. Proceedings of the IEEE/CVF international conference computer vision, pages 6836-6846, 2021.
[10] Jimmy Lei Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 2016.
[11] Tadas Baltrušaitis, Chaitanya Ahuja, and Louis-Phillppe Morency. Multimodal machine learning: A survey and taxonomy. IEEE transactions pattern analysis and machine intelligence, 4132):22-443, 2018.
[12] Khaled Bayoudh, Raja Knani, Fayçal Hamdaoui, and Abdellatif Mtibaa. A survey on deep multimodal learning for computer vision: advances, trends, applications, and datasets. The Visual Computer, 38(8):2939-299 2022.
[13] Iz Beltagy, Matthew E Peters, and Anman Cohan, Longformer: The long -ocummett transformer. arXiv preprint arXiv:2004.051500 2020.
[14] Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space- attention all you need for video understanding?
[15] Mathieu Blondel, Andre Martins, and Vlad Niculae. Learning classifiers with fenchel-young losses: Generalized entropies, margins, and algorithms. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 606-615. PMLR, 2019.
[16] Zalán Borsos, Raphaël Marinier, Damien Vincent, Eugene Kharitonov, Olivier Pietquin, Matt Sharifi, Dominik Roblek, Olivier Teboul, David Grangier, Marco Tagliasacchi, al. Audiolm: language modeling approach to audio generation. IEEE/ACM Trunsactions on Audio, Speech, and Language Processing, 2023.
[17] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković, Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104. 13478, 2021.
[18] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[19] Ziang Cao, Changhong Fu, Junjie Ye, Bowen Li, and Yiming Li. Hift: Hierarchical feature transformer for aerial tracking. In Proceedings of the IEEF/CVF international conference on computer vision, pages 15457-15466, 2021.
[20] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 2019.

---

[21] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 1at.:200099979,,,,.. 2020.
[22] Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. arXiv preprint wwXiv:1909.000111 2019.
[23] Soham Deshmukh, Benjamin Elizalde, Rita Singh, and Huaming Wang. Pengi: An au- dio language model for audio tasks. Advances in Neural Information Processing Systems, 36:1 18900-88108 2023.
[24] Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova, Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[25] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11999, 2020.
[26] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale, 2021.
[27] Subhabrata Dutta, Tanya Gautam, Soumen Chakrabarti, and Tammoy Chakraborty. Redesigning the transformer architecture with insights from multi-particll dynamical systems. Advances in Neural Information Processing Systems, 34:5531-5545 2021.
[28] Xinjie Fan, Shujian Zhang, Bo Chen, and Mingyuan Zhou Bayesian attention modules. Advances in Neural Information Processing Systems, 33:16362-16376 2020.
[29] Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[30] Jing Gao, Peng Li, Zhikui Chen, and Jianing Zhang. A survey on deep learning for multimodal data fusion. Neural Computation, 52(5):829-866 2020.
[31] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
[32] Yuan Gong, Yu-An Chung, and James Glass. Ast: Audio spectrogram transformer. arXiv preprint arXiv:2104.01772 2021.
[33] Meng-Hao Guo, Jun-Xiong Cai. Zheng-Ning Liu, Tai-Jiang Mu, Ralph R Martin, and Shi-Min Hu. Pct: Point cloud transformer. Computational Visual Media, 7:187-199, 2021.
[34] Wenzhong Guo, Jianwen Wang, and Shiping Wang. Deep multimodal representation learning:
A survey. leee Access, 7:63373-63394 2019.
[35] Wolfgang Hackbusch. A sparse matrix arithmetic based on h-matrices. part i: Introduction to h-matrices. Computing, 62(2):89-108. 1999.
[36] Wolfgang Hackbusch and Boris N Khoromskij. sparse h-matrix arithmetic. part ii: Applica- tion to dinei-idmeesionall problems. Computing, 64:21-47, 47, 2000.
[37] Dongchen Han, Tianzhu Ye. Yizeng Han, Zhuofan Xia, Siyuan Pan, Pengfei Wan, Shiji Song. and Gao Huang. Agent attention: On the integration of softmax and linear attention. In European Conference on Computer Vision, pages 124 140. Springer, 2025.
[38] Lingshen He, Yiming Dong, Yisen Wang, Dacheng Tao, and Zhouchen Lin. Gauge equivariant transformer. Advances in Neural Information Processing Systems, 34:27331 2021.
12

---

[39] Yunpeng Huang, Jingwei Xu, Zixu Jiang, Junyu Lai, Zenan Li, Yuan Yao, Taolue Chen, Lijuan Yang. Zhou Xin, and Xiaoxing Ma. Advancing transformer architecture in long-context large language models: A comprehensive survey. arXiv preprini arXiv:231 1..3311.133/... 2023.
[40] Zhicheng Huang, Zhaoyang Zeng, Bei Liu, Dongmei Fu, and Jianlong Fu. Pixel-bert: Aligning image pixels with text by deep multi-modal transformers. arXiv preprint arXiv:2004. 00049, 2020.
[41] Zhongzhan Huang, Mingfu Liang, Jinghui Qin, Shanshan Zhong, and Liang Lin. Under- standing self-attention mechanism via dynamical system perspective. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 1412-1422, 2023.
[42] Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups. In International Conference on Machine Learning, pages 4533-4543, PMLR, 2021.
[43] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnn:: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5165. PMLR, 2020.
[44] Wonjae Kim, Bokyung Son, and Ildoo Kim. Vilt: Vision-and and--anddl-dad lan-uaag transformer without convolution or region supervision. In International conference on machine learning, pages 5583-5594. PMLR, 2021.
[45] Khaled Koutini, Jan Schlüter, Hamid Eghbal-Zadeh, and Gerhard Widmer. Efficient training of audio transformers with patchout. arXiv preprint arXiv:2 10.05069,, 2021.
[46] Kunchang Li, Yali Wang, Peng Gao, Guanglu Song, Yu Liu, Hongsheng Li, and Yu Qiao. Uniformer: Unified transformer for efficient spatiotemporal representation learning. arXiv preprint arXiv:2201 arX11:2201.04676, 04676,, 2022.
[47] Liunian Harold Li, Mark Yatskar, Da Yin, Cho-Jui Hsieh, and Kai- Kai-Wei Chang. Visualbert: A simple and performant baseline for vision and language. arXiv preprint arXiv: 1rX8v:1908003555 2019.
[48] Yang Li, Si Si, Gang Li, Cho-Jui Hsich, and Samy Bengio. Learnable fourier features for spatial positional encoding. Advances in Neural Information Processing Systems, 34:1 15586998,,, 2021.
[49] Zewen Li, Fan Liu, Wenjie Yang. Shouheng Peng, and Jun Zhou. A survey of convolutional neural networks: analysis, applications, and prospects. IEEE transactions on neural networks and learning systems, 33(12))6999-7011 2021.
[50] Junyang Lin, An Yang, Yichang Zhang, Jie Liu, Jingren Zhou, and Hongxia Yang. Interbert: Vision-and-language interaction for multi-modal pretraining. arXiv preprint arXiv:2003. 13198, 2020.
[51] Yang Liu and Mirella Lapata. Hierarchical transformers for multi-document summarization. arXiv preprint 120.::9905.3164 2019.
[52] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
[53] Yun Liu, Yu-Huan Wu, Guolei Sun, Le Zhang, Ajad Chhatkuli, and Lue Van Gool. Vision transformers with hierarchical attention. Machine Intelligence Research, 21(4):670-683, 2024.
[54] Ze Liu, Yutong Lin, Yue Cao. Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[55] Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining aank-tgi visiolinguistic representations for vision-and-language tasks. Advances in neural information processing systems, 32, 2019.
13

---

[56] Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie- Yan Liu. Understanding and improving transformer from multi-particle dynamic system point of view. arXiv preprint arXiv: 0a0./99966.02, 2019.
[57] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christo- pher Potts. Learning word vectors for sentiment analysis. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea, editors, Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 2011, Association for Computational Linguisties.
[58] Andrew Maas, Raymond E. Daly, Peter Pham, Dan Huang, Andrew Y. Ng. and Christo- pher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies. pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
[59] Julian McAuley and Jure Leskovec. Hidden factors and hidden topics: understanding rating dimensions with review text. Proceedings of the 7th ACM conference on Recommender systems, 2013.
[60] Tomas Mikolov, Kai Chen, Gregory S. Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. In International Conference Learning Representations, 2013.
[61] Erxue Min, Runfa Chen, Yatao Bian. Tingyang Xu, Kangfei Zhao, Wenbing Huang, Peilin Zhao, Junzhou Huang, Sophia Ananiadou, and Yu Rong. Transformer for graphs: An overview from architecture perspective. arXiv preprint arXiv:2202.08455, 2022.
[62] Niklas Muennighoff, Nouamane Tazi, Loic Magne, and Nils Reimers. Mteb: Massive text embedding benchmark, 2023.
[63] Fionn Murtagh and Pedro Contreras. Algorithms for hierarchical clustering: an overview. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery, 2(1):86 2012.
[64] Piotr Nawrot, Szymon Tworkowski, Michal Tyrolski, Lukasz Kaiser, Yuhuai Wu, Christian Szegedy, and Henryk Michalewski. Hierarchical transformers more efficient language models. arXiv preprint 1rXiv22110.33111. 2021.
[65] Daniel Neimark, Omri Bar, Maya Zohar, and Dotan Asselmann. Video transformer network. In Proceedings the IEEE/CVF international conference on computer vision, pages 3163-3172 2021.
[66] Tam Nguyen, Tan Nguyen, and Richard Baraniuk. Mitigating over-smoothing in transformers via regularized nonlocal functionals. Advances in Neural Information Processing Systems, 36, 2024.
[67] Raghavendra Pappagari, Piotr Zelasko, Jesús Villalba, Yishay Carmiel, and Najim Dehak. Hierarchical transformers for long document classification. In 20/9 JEEE automatic speech recognition and understanding workshop (ASRU), pages 838 844. IEEE, 2019.
[68] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention, 2021.
[69] Aditya Prakash, Kashyap Chitta, and Andreas Geiger. Multi-modal fusion transformer for end-to-end autonomous driving. In Proceedings of the IEEF/CVF conference computer vision and pattern recognition, pages 7077-7087, 2021.
[70] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020.
[71] Tanzila Rahman, Mengyu Yang, and Leonid Sigal. Tribert: Human-centric audio-visual representation learning. Advances in Neural Information Processing Systems, 34:9774-9787, 2021.
14

---

[72] Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501-14515 2022.
[73] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008. 0220000..0222.... 2020.
[74] Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Mapping of attention mechanisms to generalized potts model. arXiv preprint 2rX1v:2304.07235, 2023.
[75] Raanan Rohekar, Yaniv Gurwicz, and Shami Nisimov. Causal interpretation of self-attention in pre-trained transformers. Advances in Neurul Information Processing Systems, 36, 2024.
[76] Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Sele-suue graph transformer on large-scale molecular data. Advances in neural information processing systems, 33:12559- 2020.
[77] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge, 2015.
[78] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 3531-3539, 2021.
[79] Pranav Shetty and Suraj Singh. Hierarchical clustering: survey. International Journal of Applied Research, 7(4): 2021.
[80] Alexander Shim. A probabilistic interpretation of transformers.
arXiv preprint
arXiv:2205.01080 2022.
[81] Ryan Singh and Christopher I. Buckley. Attention as implicit structural inference. Advances in Neural Information Processing Systems, 01124499-244444 2023.
[82] Masahiro Suzuki and Yutaka Matsuo. A survey of multimodal deep generative models. Advanced Robotics, 36(5-6):261-278, 2022.
[83] Kai Sheng Tai, Peter Bailis, and Gregory Valiant. Equivariant transformer networks. In International Conference on Machine Learning, pages 6086-6095, PMLR, 2019.
[84] Bo Thiesson and Jingu Kim. Fast variational modc-secking. In Artificial Intelligence and Statistics, pages 1230-1242. PMLR, 2012.
[85] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
[86] Thuan Nguyen Anh Trang, Khang Nhat Ngo, Hugo Sonnery. Thieu Vo, Siamak Ravanbakhsh, and Truong Son Hy. Scalable hierarchical self-attention with learnable hierarchy for long-range interactions. Transactions on Machine Learning Research, 2024.
[87] Thanh-Dat Truong, Chi Nhan Duong, Hoang Anh Pham, Bhiksha Raj, Ngan Le, Khoa Luu, et al. The right to talk: An audio-visual transformer approach. In Proceedings of the IEFF/CVF International Conference on Computer Vision, pages 1105-1114, 2021.
[88] Yao-Hung Hubert Tsai, Shaojie Bai, Paul Pu Liang, Zico Kolter, Louis- Chiis-Phi Morency, and Ruslan Salakhutdinov. Multimodal transformer for unaligned multimodal language sequences. In Proceedings of the conference. Association for computational linguistics. Meeting, volume 2019, page 6558. NIH Public Access, 2019.
[89] Constantino Tsallis. Possible generalization of holtzmann-gibhe statisties. Journal of statistical physics, 52:479-487, 1988.
15

---

[90] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[91] Prateek Verma and Jonathan Berger. Audio transformers: Transformer architectures for large scale audio understanding. adieu convolutions. arXiv preprint 1eXiv:2105.0033:0 2021.
[92] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. In Tal Linzen, Grzegorz Chrupata, and Afra Alishahi, editors, Proceedings of the 2018 EMNLP Workshop BlackboxNLPP: Analyzing and Interpreting Neural Networks for NLP, pages 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics.
[93] Zhen Wang, Xu Shan, Xiangxie Zhang, and Jie Yang. N24news: A new dataset for multimodal news classification, 2022.
[94] Chuhan Wu, Fangzhao Wu, Tao Qi, and Yongfeng Huang. Hi-transformer: Hierarchical interactive transformer for efficient and effective long document modeling. arXiv preprint urXiv:2106.000 2021.
[95] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chenggi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1):4-24, 2020.
[96] Peng Xu, Xiatian Zhu, and David A Clifton. Multimodal learning with transformers: A survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 2023.
[97] Zichao Yang, Diyi Yang. Chris Dyer, Xiaodong He, Alex Smola, and Eduard Hovy. Hierarchi- cal attention networks for document classification. In Proceedings of the 2016 conference of the North American chapter of the association for computational linguistics: human language technologies, pages 1480-1482 2016.
[98] Zihao Ye, Qipeng Guo, Quan Gan, Xipeng Qiu, and Zheng Zhang. Bp-transformer: Modelling long-range context via binary partitioning. arXiv preprint arXiv:1911.0400 2019.
[99] Botao Yu, Peiling Lu, Rui Wang, Wei Hu, Xu Tan, Wei Ye, Shikun Zhang, Tao Qin, and Tie-Yan Liu. Museformer: Transformer with fine-and coarse-grained attention for music generation. Advances in Neural Information Processing Systems, 35:1376- 2022.
[100] Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vggan. arXiv preprint 2rXiv:2110.046)) 2021.
[101] Seongjun Yun, Minbyul Jeong, Rachyun Kim, Jaewoo Kang, and Hyunwoo Kim. Graph transformer networks. Advances in neural information processing systems, 32, 2019.
[102] Manzil Zaheer. Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fer- gus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[103] Haopeng Zhang and Jiawei Zhang. Text graph transformer for document classification. In Conference on empirical methods natural language processing (EMNLP), 2020.
[104] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-Ievel convolutional networks for text classification. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gareett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc.,
2015.
[105] Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Ô Arik, and Tomas Pfister. Nested hierarchical transformer: Towards accurate, data-efficient and interpretable visual understand- ing. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 3417-3425, 2022.
16

---

[106] Bin Zhao, Maoguo Gong, and Xuelong Li. Hierarchical multimodal transformer to summarize videos. Neurocomputing, 468:360-369, 2022.
[107] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEE/CVF international conference computer vision, pages 16259-
16268, 2021.
[108] Shihao Zhou, Duosheng Chen, Jinshan Pan, Jinglei Shi, and Jufeng Yang. Adapt or perish: Adaptive sparse transformer with attentive feature refinement for image restoration. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2952-2963, 2024.
[109] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandku- mar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in neural information processing systems, 34:17723- 2021.
[110] Zhenhai Zhu and Radu Soricut. H-transformer-lne Id: Fast one-dimensional hierarchical attention for sequences. arXiv preprint arXiv:2107.119000 2021.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have presented our claimed theoretical contributions in Sections
D
E
and
K
The claimed empirical contributions can be found in Sections and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please see Appendix M
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiaatinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer: [Yes]
Justification: Please Appendix H
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please Section 5 and Appendix !
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
19

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: While the datasets used in our empirical studies are publicly available, our code is not being released the moment until after submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Please see Section and Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have provided 95% CI for all the results in the main paper that involved variability, which is the training process in our case. However, we have not provided CI for the results presented in Appendix The rest of the results across the main paper and the appendix do not involve variability as there's ' no training involved.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
20

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please see Appendix We have also reported the GFLOPs required for half of our experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't fidn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:://eeeerrpeeccccpbbbtt///ta
Answer: [Yes]
Justification: To the best of our knowledge and judgment, do not have any deviation from the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our contribution in this work is primarily toward theoretical research without any tie any particular application.
Guidelines:
21

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work doesn't 't release any data or model.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
We have not used any asset development of our proposed framework.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
22

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not introduce any new asset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codeemoodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
23

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Our core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
24

---

Contents
Infroduction
2
Representing Hierarchical Data
3
Hierarchical Self-Aftention
Softmax Attention Revisited
3.2
Generalizing Attention to Nested Signals
Effciennt Calculation of HSA
Experimental Results
5.1
Hierarchical Language
5.2
Multi-modal News Classification
5.3
Zero-Shot Hierarchical Approxiaaton
10
6
Conclusions
10
Notations
26
Related Work
26
Generalizing The Notion of Signal
28
D
The Emergence of The Value Projection Matrix
28
Algorithmic Details
29
Correctness and Complexity
29
Black-box Attention Computation
29
E.3
GPU Implementation
30
Hierarchical Transformer Encoder
30
Hierarchical Auto-regressive Generation
31
Hierarchical Causal Masking
32
G.2
Hierarchical Caching
32
Proofs
33
H.1
Proposition 1: Softmax Attention
33
Theorem The Optimality of HSA
34
3
Theorem 2: The Correctness and The Complexity of Algorithms
38
Experimental Settings
40
Datasets
40
25

---

Model Architectures
40
1.3 Training Syper-parameters
41
J
Comparison to The Classical Transformer Architecture
41
Going Beyond Softmax Attention
42
K.I Sparse Attention as Energy Minimization
42
K.2
The Hierarchical Generalization
43
Zero-shot Approximation of Self-attention: Ablation Study
43
Experimental Setup
43
Results
44
M Limitations
45
A
Notations
Table 4 summarizes our notations in the main paper. Moreover, Fig 3 visually demonstrates some of our tree-related notations.
Th root node
Rx
The signal hierarchy hx
one-hop path
The immediate common ancestor (ica) of A and B
multi-hop path
The highest distinct ancestors of A and B
A
B'
Any two unrelated nodes and
A
B
The indices of the leaf node descendants of A, i.e. ((A)
Figure 3: The visual demonstration of some of our tree-related notations in the paper.
B Related Work
Hierarchical models: The notion hierarchy has played a key role in data representation and
clustering in Machine Learning 63 79). In the context of transformers, the idea of multi-scale
attention been mainly used to combat the long-context challenge in language [39 67 98 97
51]. but it has also made its way into vision [54] [05] [53] audio [99 and graphs [86]. Nevertheless,
26

---

Notation
Description
Bold face
A vector or a matrix when is not obvious from the context
A column vector of zeros
A column vector of ones
A simple nested signal
(u)
The value of signal æ: at position u
Ω
The signal domain
C
The vector space containing the signal range The dimensionality of C. i.e. the channel dimension The set of all signal domains present in the problem The position embedding function for domain Q The dimensionality of each position embedding The all position embedding functions for all domains in The set of all possible simple signals in the problem The set all possible nested signals up to depth The set of all possible nested signals in the problem An equivalent notation for
D
en
S
N
N
No
he
The signal hierarchy representing the finite nested signal The nodes in the signal hierarchy h2
A, B. B.C....
L
The leaf node in the signal hierarchy ha corresponding the query variable g:
R
The root node of the signal hierarchy hr
E(A)
The of indices of the leaf node descendants of node The children of node A The parent of node A
chd(A)
Pa(A)
sib(A) A)
The siblings of node A
ica( A, B)
The immediate common ancestors of the unrelated nodes A and B
A', B'
The highest distinct ancestors of the unrelated nodes A and B The number of non-leaf nodes of the signal hierarchy hg
M
b
The maximum branching factor of the signal hierarchy hg Set cardinality
H(Q K)
The conditional entropy of the query variable Q given the key variable K The (directional) interaction energy between the unrelated nodes A and B The energy of node A
AA+B
o(A)
((A)
energy of node wrt query vector & The (directional) attention weight between query d and key k;
The gradient
Oig
The attention matrix
B
The set of (hierarchical) stochastic matrices respecting the block constraint wrt lr2 Table 4: The notations used in the main paper.
most of these frameworks deal with single modality that occupies the same geometry, just different scales. Our proposed framework, in contrast, can incorporate an arbitrary number of nnutuall--incosisee
geometries within its representation of the multi-scale data. Another related line of work based on hierarchical matrices [35] 36] that have been used traditionally for clustering 84 as well as transition matrix approximation g but more recently for attention matrix [110].
Multi-modal models: Multi-modality has been vastly explored in Machine Learning [1] and more recently within various neural architectures, using various fusion techniques 30
As for multi-modal transformers [96]. most frameworks are tailored toward a fixed sct of modalities,
12
e.g. vision-language [44
40 55 [09], audio-visual [87]. audio-language [23]. graph-language
[03]. Dision-pose-une
771,, audio-vision ciciiovviioo--le vision-language [88] etc. The fusion of different modalities in
these frameworks typically takes place via heuristic operation at the embedding or the attention stages resulting in distinct architectural variants, which are typically categorized as (1) single-stream (e.g. 47). (2) multi-stream (e.g. (55)), and (3) hybrid-stream (e.g. [50]). However, most of these frameworks either ignore the geometrical (positional) information for some of the input modalities, impose artificial restrictions input geometries such as alignment.
Geometric Deep Learning: Geometric Deep Learning [17 studies the invariance and equivariance
properties of deep learning models by introducing the notion of signal and its geometry which is
27

---

explicitly modeled via the signal's 's domain. We build our framework also based on the same notion of signal and generalize further to nested signals which can represent hierarchical. multi-modal data which potentially encompass multiple domains. Also, most frameworks within Geometric Deep Learning achieve the desired equivariance properties through the model's architecture (e.g. CNNs [49
GNNs [95], and Group-equivaraint CNNs (29]). A prominent exception is the Teannsffrrrrr
42
where the desired group-equivariance is achieved by explicit modeling of the position information and its separate similarity computation (as opposed to adding it to the feature vectors). The formulation of the position information in our framework is in part inspired by the LieTransfomer.
The theoretical foundations of self-attention: Despite its revolutionary success in Deep Learning, there has been quite little effort to understand the theoretical foundations of self-attention. These
efforts provide various interpretations of self- ceff-attention, including the probabilistic view [0 28 the
causal view [75], the structural inference view [8], the dynamical system view [41
56 27 the
statistical mechanical view (74). the variational denoising view [66], the clustering view 13 and the Hopfield network view 73 this paper, we provide a statistical mechanical perspective to derive self-attention from the first principle of entropy minimization; in that sense, our interpretation is closely related to the statistical mechanical, denoising and Hopfiled network views. More importantly, our interpretation lends itself to straightforward generalization to the hierarchical self-attention mechanism which, as we show, is both theoretically optimal and efficiently computable.
C Generalizing The Notion of Signal
In standard Geometric Deep Learning, signals typically represent data structures in Computer Vision, Audio Processing, Natural Language Processing and Graph and Manifold Processing. But the notion off signal is quite versatile and can be generalized to include feature representations in classical Machine Learning. In particular, countable, discrete set with no additional structure. In this case, if the elements f2 are conceptually indistinguishable, then any signal x on l is said to be defined the position embedding Eg maps all the elements of f2 to the constant vector 0. The latter conveys that there is no positional information associated with the signal. As an example, vector set can be seen as a signal defined on an unordered set.
note the special case where the signal domain St is
unordered set and subsequently,
On the other hand, if the elements of Ω are distinguishable, we can define bijective position embedding Eq to carry that information into the position vector space. We refer to signals defined on such Ω domains key-value signals. For instance, a tabular feature vector in classical Machine Learning can be seen as a set of key-value pairs where the keys are the feature names and the values the feature values, and hence modeled as key-value signal. In this case, a text embedding model can be used to map the feature names into a vector space and regard the results as the position embeddings of those features. In other words, the notion of signal in work is quite generic and encompasses not only the signal types in Geometric Deep Learning but also the classical tabular feature vectors.
D
The Emergence of The Value Projection Matrix
The derived formulation for Softmax attention in (6) deviates from the classical Softmax attention in that it lacks separate value projections, which can be quite restrictive as it significantly reduces the model's degrees of freedom. Nevertheless, the value projections derived formulation by considering learnable step-size for the gradient update in (6). In particular, instead of setting step size to À 1, let W. where W, Rddd is trainable parameter. By doing so, (6) changes to:
be theoretically injected into our
expla //kk/VVatee j=1o7i DiELi exp(1] ky bv//))))Te)
q: 5 + Σ
W.ky, ISISN
(14)
By defining 0; = Wal, = W,W,D,,, we effectively arrive at separate value projections, where W. can be seen as the value projection matrix used in the standard Softmax attention formulation.
Note that by introducing learning step-size in the form of projection matrix, we effectively project the direction of the gradient vector into new direction. So in that sense, [14 is no longer strict
28

---

gradient ascent update. In other words, depending on the learned projection matrix W. and the value of gradient vector for point g. we may decrease or even increase the upper-bound on the conditional entropy. This extra degree of flexibility indeed enables the transformer model to best adapt to the end task. And therefore, we have adopted separate value projections our code as well all of our reported experiments, similar to the standard transformer architecture.
E
Algorithmic Details
In Section we proposed an efficient algorithm based on dynamic programming for the calculation of HSA. In this section, provide formal correctness results and complexity analysis for this algorithm. We, furthermore, introduce an extension of the proposed algorithm to more generic algorithmic template with replaceable black-box base attention computation module. Finally, we discuss some practical details on how to implement our proposed algorithm on modern GPUs.
E.1 Correctness and Complexity
First off, it is not hard to show that for the case of flat hierarchy, Algorithms 1 reduce to the standard Softmax attention calculations. other words, the standard Softmax attention calculation is special case of our proposed algorithm here. However, showing the correctness and the complexity of Algorithms theorem.
3 for the general case, is more involved, which we achieve through the following
Theorem E.1. For given signal hierarchy h if both query and key variables are normalized via the LayerNorm function, then Algorithms 2 compute {V ar (Bx)) j E {(R)) in Eq. @ in O(M b²)
based on the interaction energy function defined in Eq. 7 where b and M are the branching factor and the number of families in hg, respectively.
Proof. See Appendix
Once the attention values are computed using Algorithm]] leaf nodes can be updated via the residual connection:
I
the query vector representations at the
9. - q ((R2) V o(R.), Vie ((()))))
(15)
And that would conclude the HSA operation.
E.2 Black-box Attention Computation
is important to note that Lines 16-18 in Algorithm 3 perform the standard Softmax attention mechanism on the members of family that contains node A. In other words, our proposed HSA algorithm can be seen as divide-ande sorid--and-conqqer algorithm where the attention computation on the whole sequence (i.e. the hierarchy's leaves) is broken down into attention computation on the much smaller families in the hierarchy (aka the sub-problems) via the bottom-up part of the algorithm, and then these intermediate results (aka the sufficient statistics) are combined through the top-down part of the algorithm to produce the final self-attention output. From this perspective, if the average branching factor (i.e. the family size) in the hierarchy is b, then on average, the sub- uub-problem attention calculation takes Obb) time and memory for each node A, which makes the O(b²) complexity for the entire family. Then intuitively for the total of M families
hierarchy, the final computational complexity comes to O(M b2). As special case, for flat hierarchies where there is only M = family of size N (i.e. the sequence length), the complexity becomes (N
More importantly, from the practical perspective, the divide-and-conquee view of the proposed algorithm encapsulates the sub-problem self-attention computation (in Lines 16-18 Algorithm 3) as a black-box module that can be easily replaced by any exact approximate function that computes the standard Softmax attention. This has a significant practical implication, as it allows the HSA algorithm to invoke any efficient attention computation frameworks in the literature as its base attention calculation sub-module. For instance, the quadratic factor 12 O(M b22) can be further reduced to linear if one employs one of the many approximation techniques proposed for as the black-box sub-problem attention
efficient computation of Softmax attention 21 13 68 43
computation module in Lines 16-18 in Algorithm
29

---

E.3 GPU Implementation
The Algorithms| 2 are technically classical tree-travee traversal algorithms which are typically not fit for parallel processing on GPU. Indeed, this would introduce practical challenge for incorporation of HSA within modern Deep Learning frameworks. To address this challenge, in this section, we present two major techniques for introducing parallelization both at the node level for one signal hierarchy as well as at the batch level across multiple signal hierarchies.
First, we note that all the summations in Algorithms Sa 2 can be done in parallel for different sets of nodes in fig. In particular, if summation statement can be parallelized for K nodes of h2, it can be implemented as a (sparse) matrix by dense vector multiplication WE, where W= ( is the sparse coefficient matrix and v = [u=]sx] I contains the values of the input terms. In particular, wi, is the weight of the jth term for computing the summed quantity at the ith node (typically 1 or 0). As for the quantities Algorithms and 2 p (.). (-(()) and 17(+) can be parallelized over all the nodes in hy; that is, in order to compute each one of these quantities for all nodes of hg. only one sparse matrix-vector multiplication is needed given the appropriate coefficient matrix. The computation of ((-) and (-) is also parallelizable over the nodes belonging to the same depth in h in other words, given the appropriate coefficient matrices, we would need sparse matrix-vector multiplications to calculate each one of these quantities for all nodes in hg where D is the depth of hg. Since the coefficient matrices in this scheme are highly sparse, we have represented the coefficient matrices using sparse tensors and used the efficient implementation of sparse matrix by dense vector multiplication in Pytorch to carry out the tree-based summations in A [nd
The other fundamental aspect of parallelization in Deep Learning is batch computation, which typically boils down to matrix operations for the standard batches of fixed-size tensors. However, in our scenario, the signal hierarchies in each batch are trees with different structures as well as potentially different signal types/modalities appearing in arbitrary arrangements for each signal hierarchy in the batch. This effectively makes the classical batch computation impossible for signal hierarchies. To address this challenge, we propose a completely different technique for batch parallelization. As explained above, we already have a method to parallelize the computations within each signal hierarchy; we can further parallelize the computations across different signal hierarchies in batch by making them part of one hierarchy. In particular, we introduce a dummy root node and make each signal hierarchy in the batch a direct child of it. The position embedding for this dummy root is set to unordered embedding: that is, no position embedding. This way, end up with only one, wide signal hierarchy in our batch that is just one level deeper than the deepest signal hicrarchy in the original batch. By performing the parallel version of Algorithms [ (as described above) on this one "concatenated" signal hierarchy, we effectively compute all the targeted quantities for all signal hierarchies in the batch at the same time. We refer this batch processing technique as breadth- wise tree concatenation.
F Hierarchical Transformer Encoder
The proposed HSA mechanism does not introduce any trainable parameters on its own; it is simply an attention operation. However, similar to classical transformers, we can add trainable linear projections before performing HSA. This gives rise the hierarchical transformer encoder (HTE) architecture which is capable of operating on signal hierarchies representing finite nested signals. Similar to classical transformers, we also add multiple heads as well as point-wise linear projection of the output of HSA followed by some non- con-linea The same way the classical transformer layers do not change the query sequence length or the position embeddings of its tokens, HTE layers do not alter the structure of the hierarchy tree or its nodes' positional embeddings Figure 4 depicts our proposed architecture for each HTE layer.
Aside from HSA, HTE is different from classical transformer encoder in two ways. First, the LayerNorm operation is performed after linear projection as opposed to before it. As mentioned before, by doing so, the attention operation will minimize a proper energy function which is in turn a proxy for minimizing the entropy of the representation. Second, unlike simple signals in standard transformers, signal hierarchies can contain different modalities and signal domains within their
EEven though, the same position embeddings are fed to each layer, in our implementation, we have designed a separate linear projection per position embedding type per layer to project the position embeddings before the HSA operation.
30

---

Pointwise Feed- forward NN
Per-iiyality Linear Projections
LayerNorm
LayerNorm
Concat
Linear
HSA
o
Different heads
Figure 4: The proposed Hicrarchical Transformer Encoder (HTE) layer architecture.
different families across the hierarchy. Therefore, using the same linear projection layer for all this various types of information may be an ver-simplification. To this end, our proposed framework allows different linear projection per the type of input information. For example, the leaf vectors coming from the language and vision modalities can be transformed using their own separate linear projection layers. Note that this distinction is only allowed the linear projection layer; the HSA operation itself universal and does not treat different types of information differently. Also, is assumed the different types of information in given problem (including modalities and signal domains) are priori known and fixed, even though each signal hierarchy in the input dataset can be an arbitrary, variable-depth composition of these known types. By making this assumption, can know ahead of time how many linear projection layers are needed within each HTE layer.
The HTE layers can be cascaded to form a hierarchical transformer based on the HSA. Furthermore, different types of pooling operations can be introduced to (gradually) coarsen the hierarchical structure of the input nested signal. In particular, using local pooling operation, the Ieaf nodes of the input signal hierarchy are either merged together or completely pooled into their parents resulting a coarser representation of the underlying nested signal. Furthermore, since the channel dimensionality is constant across the hierarchy, global pooling is also well-defined which reduces the whole signal hierarchy into single, fixed-size vector of d dimensions (e.g. by taking the average). Depending on the application, global pooling can also be realized by taking a specific leaf node's query vector of the output signal hierarchy (e.g. in per-token classification tasks uni-modal, hierarchical data).
G Hierarchical Auto-regressive Generation
The HSA-based, encoder-only architecture introduced in Appendix Fis primarily suitable for classifi- cation and regression applications. However, for auto-regressive generation such as causal language modeling, we would need to have a decoder. One straightforward approach to use an encoder- decoder architecture where the encoder is HSA-based while the decoder is the standard sequential decoder. In particular, in this scheme, the hicrarchical self-attention is only incorporated for the initial prompt while for the generated text, we simply compute the standard flat attention. While simple, this solution does not take the full advantage HSA, especially if the generated text allows for the similar hierarchical structure as the prompt text. For instance, if the hierarchy is built upon the sentence and paragraph structure of the prompt text, then it fairly reasonable for the generated text have the same hierarchical construct as well. The same can be said when the hierarchy is based fixed hopping windows over the text. such cases, a HSA-based, decoder-only architecture is needed to incorporate the hierarchical structure of the generated text during auto-regressive generation.
Theoretically speaking, for HSA-based decoder during auto-regressive generation, we would need to maintain a dynamic signal hierarchy where every generated token augments the signal hierarchy with at least one new leaf node and possibly multiple non-leaf nodes. Once the signal hierarchy updated, the HSA calculations are, in principle, the same as before. Nevertheless, there are two major issues here specific auto-regressive generation. First, unlike the HSA mechanism introduced far, due to causal generation of tokens, leaf nodes are only allowed to attend to the other leaf nodes that have appeared before them; that is, we would need a hierarchical causal masking mechanism. Second, running the full HSA algorithm for every generated token is inefficient as it would re-compute some of the sufficient statistics in Algorithm 3 which clearly redundant, In the following sections, we address these two problems.
31

---

G.1 Hierarchical Causal Masking
In the standard auto-regressive generation using the self-attention mechanism, in order to prohibit tokens from attending to the future tokens, one incorporates causal mask in calculation of the attention weights via an appropriate lower-triangular mask matrix. However, this straightforward approach will not work with hierarchical self-attention mechanism because attention weights between all tokens are not computed simultaneously but rather hierarchical fashion.
Nevertheless, one can easily show that if the standard causal masking is applied at each level of the hierarchical attention calculation, the end, no leaf token will attend to future tokens (i.e. the tokens to its right) in the hierarchy. In particular, as explained in Appendix E.2 Lines 16-18 of Algorithm encapsullte black-box Softmax self-attention function that applied for each family in the hierarchy. For applying hierarchical causal masking, we can simply apply the standard causal masking within this black-box self-aatention calculation. This equivalent to replacing the sibling function sib(A) in lines 16-18 of Algorithm 3 with (A) which restricts A's siblings to the ones its left (i.e. previous tokens). This simple black-box causal masking will further propagate through the hierarchy such that at the end, the leaf nodes will only attend to other leaf nodes that are located their left. Figure illustrates this process through a toy example.
Standard Causal Masking @ Level 1
Slandard Causal Masking @ Level 2
Standard Causal
Masking @ Level
Level 1
Level 2
Level
The signal hierarchy
Hierarchical Causal
Masking
Figure 5: An illustration of the proposed Hierarchical Causal Masking scheme for hierarchical auto-regressive generation.
G.2 Hierarchical Caching
In standard auto-regressive generation, every generated token merely attends to the tokens seen so far whose projections are cached via a key-value cache. This makes the attention computation for cach token linear in the (generated + prompt) sequence length. This simple idea, however, is not directly applicable to the hierarchical case. One important distinction that we need to keep in mind that in the hierarchical case, are not only generating sequence but also hierarchy that comes with it; in other words, the generated sequence is the set the leaf nodes of a hierarchy that needs to be maintained and updated well. As such, any caching mechanism would need to maintain and update the signal hierarchy and not just its leaf nodes. Note that caching the hierarchy means maintaining its structure well its nodes' sufficient statistics pre-computed by Algorithm
Nevertheless, during generation, we do not need to keep the entire signal hierarchy. In particular, in our HSA framework not every leaf node directly attend to every other leaf node; instead, leaf nodes
32

---

that are not in the same family only attend to each other at the coarse scale through their highest distinct ancestors. This means that during generation, newly generated token (leaf node) only needs to directly attend to its previously generated leaf siblings and not other leaf nodes. Instead it will indirectly attend to other leaf nodes en masse by attending to their highest ancestor that is not an ancestor of the new token. Following this scheme, we would only need to cache sub-tree of the original hierarchy that consists of the ancestor line of the latest generated token as well as their immediate children nodes. We refer to this sub-tree right-skewed because only the right-most sibling in each family across the signal hierarchy is allowed to have children. Figure ([A) illustrates the maximal right-skewed sub-tree for the toy hierarchy in Figure 6 B).
Once the the right-skewed sub-tree of the signal hierarchy is extracted, we can simply update as new tokens are generated. However, we have to be careful as not all of the newly generated tokens are added the latest family: some new tokens may start a new family via higher level of the hierarchy. For example, if the hierarchy for language data is built based upon the sentence and paragraph structure in the text, a new token is not always going be part of the latest sentence or paragraph; it may start a new sentence even new paragraph. In such cases, more nodes need to be added to or deleted from the cache other than the new token's 's leaf node. These two cases are illustrated in Figure (D).
Finally we note that during the entire generation process the hierarchical cache remains a right- skewed tree which means that the CPU and memory complexity for calculating attention and maintaining the cache would be Ib NN where N is the length of the generated sequence so far and is the average branching factor of the hierarchy. This is in stark contrast to the classical key-value caching where the memory and computation are of O(N) complexity, and hence shows the potential computational advantage of our hierarchical scheme.
(A)
(B)
(C)
(D)
Figure 6: An illustration of the proposed hierarchical caching mechanism for hierarchical auto- regressive generation: (A) The original signal hierarchy built on the prompt text. (B) The right-skewed sub-tree of the original hierarchy. (C) The updated hierarchy after generation of new token that does not end the latest family. (D) The updated hierarchy after generation of another token that does end the latest family. The green leaf nodes depict the latest generated tokens in each step.
H Proofs
H.1
Proposition 1: Softmax Attention
Proof. Since each Q and K variables is normalized via LayerNorm layer, have || d and ||=|12 = d, v1 < < N, which would reduce the energy function in (5) to:
(Q, vd
log
exp (af (atky//aa
N
j=1.j$1
By taking the gradient of w.r.t. qi we get:
N
exp(y ki// Vdl
V., Vg (()())
Σ
k2 1 1
j= L exp kt/ Vd)
By plugging V,, ( K) into and setting T = (N (d) -1 A and the sample size to 1, we will get the Softmax attention formulation in
33

---

H.2
Theorem 1: The Optimality of HSA
First off, since each Q and K variables is normalized via a LayerNorm layer, we have | [| 12 =
1 d and
d, v1 - i - N, which would reduce the interaction energy function A B in to:
A P EB(()) Vd Vd. (e(A)) 18(B))
1
E (ER(B))EE(B)
Σ
q kj
(16)
Then Vg AA-B becomes:
VA
E KG
(17)
|((B)]
JE((B)
Proof of stochasticity. Next, we show that é - is stochastic matrix where т (Ie(R))lvd) and @ (0._] is the HSA matrix for the nested signal æ in
This is equivalent to showing that @ is negative matrix whose rows sum - (V We prove the latter by induction on the depth of the signal hierarchy lxg.
The base case: Using the (17 for signal hierarchy h₂ of depth 1 (i.e simple signal), g reduces to:
exp (- (-b1,,)1,))
Σ
V q. ((Rx)
-kg
O.K (18)
val(R)))l
L)EBib(LL) El.emb(L..) exp v/1n))1))))
where
K [k1, Me((
- -exp (-,1,-L))
[e(Raa]
+. (a(e(RR)) ELNESibLL,) exp Lh -b,,-,,)
is the ith row of θ and L;, L, and La are the leaf nodes corresponding qis 9j and gk, respectively. From (18). it is clear that the elements of O are all negative and each row sums to - (I((R))))
The induction step: Now assume that the above statement holds for any θ matrix derived from signal hierarchy up to depth T 1, we show that also holds for the signal hierarchy h₂ of depth T. To this end, (9) can be written as:
IE(B') (B(B)) V d(B)) ECESMBB [(C))((()O) VBBBCC
V v(R
[E(R,))
a(B²) EcESMBB [(C)]B(BI,C)
I((B¹))
[u(B')V, ((*)
Σ ((C))
1((Re))
CEsib(Bi)
where,
a<B²)
M(B') a(Bi) + Σ Desit(Bi) [((D) D)
(19)
B(Bi,C C)
8(B² C) a aBBi) + Σ ED ((DD))((B,,)
(20)
and we have ((BB) + EcEnBBy [C)))((B' C) = 1, for all {(R,). On the other hand, since Bi is child of the root node Rg the depth of its corresponding sub-signal hierarchy inevitably less than T, and therefore its corresponding energy gradient Vo(Bi) induces an attention matrix
e that is negative with rows that sum to - (I((B)))v))= according to the induction hypothesis. With that in mind, we can write
1((B)) V (RR)) = [(()))
8(BAC)
(B O KB
Σ kj = O.K K
CE) Vd.I((B))
jEE(C)
34

---

where , is the ith row of O" .K
()) and we have:
[((B)]
concat 8(B',
j. concat
T A CESB((B.)
[E(R)))
vd-1e(R)>)
Then the sum of the elements of the row vector eg given by:
[(O 8(BB C)
Σ
Ois
(((B))((B"))
Σ
Σ
[e(Bx)]
Vd
JE((R.)
jE((B*
CEa
[-
Vd vdu(B")
1
[((C)
8(B',C) cC)
(e(Rx))
Vd
CES
vdle(R~)) ((B)) +
Σ
((C)))(('' o) (d
CEsib(B9)
where the second equality comes from the induction hypothesis that EjenBB 0a. ((( In other words, fh negative with rows that sum to -((((R))lV)) which in turn, implies that é - (I((R22) dVd @ is stochastie matrix.
O
Before proving the optimality HSA, we need to show that the KL-divergence admits optimal sub-structure in setting. To this end, let ş [pr] be categorical distribution over N items Furthermore, let
{ (R,.... Rx} be K-partition on the index set I == fl..... N) such that U R3 = I and R, n R3 0, Vi,j 1,. + j. We say categorical [w.]/_ admits the tie constraint w.r.t. R. iff we have w; wj if 3Rs E R s.t. i E R2. We refer to set of all such distributions as WR.
such that Σ Pi
distribution
WR and the sub sub-partition R' R, the projection of w R' is defined as WIR' = [ux/h] ( where I(R') U R, and h = EiE(r)) w, is the re-normalization constant. From this definition, is clear WIR is a categorical distribution restricted to the items in the partition R'.
Given distribution
Lemma H.1 (The optimal sub-structure of the Kl-divergence), Let R. and WR. be defined as above; furthermoe,, let wn E WR be the closest categorical distribution Wr to g in terms of the KL-diiergence; that is,
w = arg min (KI(wllo) wEWR
Then, for any R' C R. we have:
wiR = arg min DKI(WIlIRR) wEW
Proof. Let assume the closest distribution in WR to PLR is w/ that not equal wf Then we have,
DxI(w"llp)
Σ
1 log(w[/pp) +
Σ
wi log(w]/pp)
SEI(R))
VEI(R(RR)
In log(l1 (g/(1)//22)) + h1 Σ
wiRw log(wi
Σ
ue log(u] /p:)
iET(R!)
VEI(RiR))
hi log(hi/hz)) + h Σ
w/logww//pIRR +
wi log(w? (pg(
iET(R()
(ET(RAR))
Σ
logh +
w/ log(w; /ps) EEI(RAR)
SEI(R')
(KI(w")lp)
where
h1 w/w, wf,
for i E I(R') for i I(R R')
"n = [w such that w/""
35

---

and h1 = =EEIIRY ii and h2 = E Pt are the re-normalization coefficients. The inequality in the above derivation is the direct result of the fact that w' is the closest distribution pLR/ in WR. This further implies that we just found another distribution w" E WR that is closer to ş than w" is. And this contradicts assumption regarding the optimality of w". Therefore, WIR' must be the closest distribution to f R. in WR.
Intuitively speaking, Lemma H. states that any sub-structure of an optimal solution for the KL- divergence to a target distribution also optimal. With that, we now ready to show the optimality of HSA.
Proof of optimality. We would like show that the HSA formulation in results in self attention matrix @ that minimizes the total KL-divergence in (12). In order to do so, we derive the optimal solution for the total KL-divergence and show that it obeys the recurrence in (
For signal hierarchy ht rooted Rz, let è denote the closest HSA matrix in (the space of all matrices that admit the block constraint according to the signal hierarchy h2) to the flattened self-attention matrix O described by 13 That is,
= arg min BEB
Σ (E((R=)
(K2(0,,11(()) = arg min D MEB
(21)
Since each row of ② R straightforward to see that the diagonal blocks è n corresponding to the children of R2 are also (up re-normalization factor) the closest HSA matrices the restriction of the flattened self-attention matrix @ to the corresponding sub-hierane hieherrcrrhhs.. For the child node A chd(R (R,), the renormalized restriction of @ to A is denoted by A.A, The elements of O
is categorical distribution, by applying Lemma H.1 to the rows of O R
it is
are then can be written
A' = A A + A
Vi,j E ((Ro), 00
(22)
where A denotes that child of R2 that contains the ith leaf node, bis exp(- 23 E Eje((ay b.g. and 2, Ejee( bi.j. Similarly, we denote the renormalized restriction of Ô to A by ④ R.A the elements of of θ are then can be written as:
((AA))RAA if = A ((A'.A)),
Vi.j E ((R)). OR
(23)
Af A
where ıl( (*) and 8(A*, A²) are unknown coefficients. Note that unlike (22), for the case of A + A only have one number representing the attention weight between sub-trees A and A i.e.
8(A1, A This is due to the block constraint being enforced on O Similarly, the block constraint requires the renormalization coefficient for every child A', i.e. μ((A',, to be the same for all the rows E f(At). we assume we already know the optimal restricted HSA matrices e R.A . chd R.)), goal reduces computing the values of μ(A) and 8(A, B) for all A, B E chd(RR (Rx) such that the
36

---

total KL-divergence in 21
is minimized. By plugging Eqs. 22 and 23
into
21 we get:
(((A))
R..'
(2+ + 2,)
OKL (0))
Σ
Σ
Σ
m(A)0
". log
bg
AEchdd R₂) (EF(A)
-jE((A)
(2: + 2,)(A,
E
E
8(A, B), log
big
BEsib(A) jEa(B)
Σ
(M(A) 1K (10 (e(A)) (A) +
Σ
log
2i (i))))
24
(Echd(Ra.)
iE((A)
Σ
(I(A)II(()))CAAA))) B) ((A, B)
(insib(A)
D]
5(A, B3
(((B)] Σ log(z; +22 E Σ
log bi.)
(24)
iEÉ(A)
iE((A) JE((B)
Where DKL(G R.A 110rA) is the optimal value of the total KI.-divergence for the sub-problem induced by the child node A of Rg. Since é* ft
is the minimizer of (24). the values of μ((A) and
5(A, B), AA, B E chd( R.)) must be chosen such that they minimize 24 Furthermore, each of
the matrix ④ -f
must sum to 1, which results in the following set of constraints the values of p(A)
and 8(A, B):
Vi E {(Rx),
Σ
OR
=1
E
OR
+
E
OR =1
JE((R=)
JE((A))
JEC(R.)EECA)
+ m(AA)
Σ
@R.AA Ad
Σ
((B))((A'
jE((RZ)
BESib(A9)
M(A)
Σ
((B)))(A,
= AA E Rhd(
(25)
(E
where the second line is obtained by incorporating 23 and the last line the fact that the rows
of the restricted matrix O R.A are already normalized. To optimize 24) w.r.t. p(A) and 6(A,, B). VA, B chd(R,, while enforcing the constraints in (25). we form the Lagrangian as follows:
((((A), 8(A, B), VA, B E chd( (Rx))
= [x " 1
Σ
da M(A++
Σ ((B))8(A, -1] -1]
(26)
AEchd(R,))
BESb((A)
where AA, AA chil( R2x) are the Lagrange multipliers. By taking the partial derivatives of the Lagrangian w.r.t. (A(A) and 8(A, B) and solving for them, we get:
((A) exp
Aa DA-DKI( R.A
I +
Σ log
Z1
))
-1].
[((A)]
+
%
(EF(A)
-1]
8(A, B) exp
AA
Σ Σ
log bist
Σ
log(zi +x+)
(27)
[[(A)]
[(B)]
(E((A))EF(R)
iEB(A)
Now if we plug (27) into the constraints in (25). we can solve for AA'S. which can be further put back into derive the values of ((A) and 0AA, B) as:
(AA
((A,B)
(AA) 7(4) + Σ (ESiM(A) ((BB)((A,C)
5(A, B)= r(A) + E ECENi(A) (((B)C(AAC)
(28)
where
7(A) exp
Σ (E((A)
logz, DKL ( R.A (..AAIA1 ( |
(29)
[(((A)]
<(4, B) exp
Σ
log bigj =exp(-vAA-
(30)
[((A)][((B)
(EE(A) JE((B)
37

---

where the last equality directly results from the definition of buy and the definition of the interaction energy between A and B in 77. In case R₂ is of depth (that is, A is leaf node), ((A) simply defined to be O. By plugging these values into (24) and doing some algebra, we derive the optimal value of the total Kl.-divergence as follows:
@ (10)== = Σ
Σ (E((A)
log(z; + +2) -((A)) log ((A++ E (((B)K(A, B)
AEchd( R.y)
BESI((A)
(31)
On the other hand, using 29). we can derive ((Rz)) as:
((R,) exp
Σ
log 21
10)))))))))
(32)
(RR)))
iE((R,
Now by plugging 31 into 32
applying
30). and taking the logarithm of both sides, we arrive at:
|e(AA)
(RR,)
log exp (log (()) +
Σ
IC(B) exp(-
(33)
[((Rx))
AEchd(R..) R.)
(ESib(A)
By comparing 33 to the definition of the energy of the signal hierarchy in (8. it is clear that our proposed energy function o6.) and lod-logyy 7-) follow the exact same recurrence dynamic. Furthermore, since the initial values of these two functions at the leaf nodes are both equal to co, we can conclude > = (pp(()(())))))))) for all nodes A in the signal hierarchy hgg. In other words, 7-) and ((::)
are respectively the exact same functions as a(-) and (--,-) in (TO). This further means that the optimal coefficients u(-) and o(.-) in (28) to update the optimal self-attention matrix recurrence in (23) are the exact same coefficients in our proposed recurrence in ( to compute hierarchical self-attention. Since both methods result in the same attention matrix for the base case of one-level hierarchy (i.e. the standard Softmax attention), and also follow the exact same recurrence dynamic, we can conclude that they are equivalent. This means that our proposed HSA formulation is also optimal in the sense of the total KL-divergence, which concludes the proof.
H.3 Theorem 2: The Correctness and The Complexity of Algorithms 13
Proof. Before proving the correctness and the complexity of our proposed algorithm, we show the complexity of directly calculating (9) In order to compute V ((R)), we would need to first calculate the node energy function ((-) at every node in the signal hierarchy using the recursive formula in (8. For signal hierarchy with M internal nodes and the maximum b branching factor, we would have (M b) nodes in the hierarchy, at each one of them, we would need to compute the sum in in (8) over their O(b) siblings. This would make the total complexity of calculating s(-) O(M.b)). This is essentially the complexity of the recursive function Algorithm
Next, to compute ((.) from (9. we need to traverse the path from the root node to the leaf node corresponding qi which has Ollog, M)) nodes. In each node, we also need calculate sum over the (6) siblings of that node, which makes the cost of calculating V q. Ø(-) ( (bblog, M). However, since we would need to repeat this caleulation for all O(M b) leaf nodes q, 's, the total cost of computing HSA would become O M log, M).
Moving on with the proof, note that the recurrence relation in 9 can be written as:
V 20 RZ)) exp (J(B') V. ((B) + BB')
(34)
where
f(B') = log (B') g(B))
Σ |E(C))8(B', C) V UG
(35)
CEeib(B9)
and μ(-) and 6(..)) are given in 19 and 20) Furthermore, (34) first-order, non-homogeneous
recurrence relations with variable coefficients for which we can derive the following closed-form solution:
V q. ((x)) = Σ
exp)xxp
Σ
f(0))] = Σ
exp (u(Pa(B)))
BER
PP(B)
(36)
38

---

where
(AA) = E ((A)=((A)+u(Pp))
(37)
CER.-A
Rg rn A denotes the set of all nodes the path from the root node A excluding the root itself, L is the leaf node corresponding to and (a(B) denotes the parent of node B. Furthermore, define:
9(B) exp (u(Pa(B)) ((A) exp (u(Pa(4)) + 0 ((P(AA))
(AA) =
Σ BER
(38)
Then it is straightforward to see:
- Vi ((R,)
(39)
On the other hand, given that each Q and K variables normalized via LayerNorm layer,
plug 17 into (35). to get:
8(A,C) Re(A)lVa
g(A)= =
Σ
Σ Ky JEA(C)
CE))
1
Σ
((A.C) E
Σ{} 13) SEA())
[(())v/(a(A) + E EDESi((A) ((D))
CENiI(A)
Σ
(IC))((ACC)P(C)))
[((A)VV exp (- -((A) + +exp ( [-][A])]
CEaib(A)
Σ
[(A))VV][xxp) (-
((A) + exp ( (nn(A))]
CER(((A)
Ecesib(A) |((C)) exp En(A))en((C - Vd + (aPa Pu(C)) (a(CC)
(40)
le(A)lvd exp (()) + exp nn())
where the last equality is derived from (6).. and we have:
((A)= = - logg
Σ
((D)) D)]
DE
[((B)] exp [=(A)*E(B)) 1
- log Σ
Vā /d (a(A))*(((BB) - va))
(41)
(E
and
Pg (A) = [((A)]
Σ
4 (k(p)
Σ kj
(42)
[((A)]
JE((A)
SEE(A)
Furthermore, we can rewrite f(A) as:
a(A)
f(A) log ((A) log
(aAA) Σ DEMiB(A) ((D))B(A,D)
exp - -((A))
log
exp
x(A)) + -exp ( (-n((A))
log
1 + exp (o(A) - -)(A))
log Sigmoid (n4) [(A]]
(43)
Finally, by plugging 43) and (40) into the recurrence relations in and 338 we arrive at Lines 3-4
of Algorithm 2 This means that after completion of Algorithm 2 we can read off V, of R.,) = ((L))
at the leaf nodes of the hierarchy. This proves the correctness of proposed algorithm.
As for the complexity, since Algorithm visits each O(M b) nodes of the hierarchy once and performs the summation in Lines 16-18 over the O(b) siblings of each node, the complexity of Algorithm 3 is O(M b²), which means the total complexity of computing HSA using our dynamic programming approach is O(M 62). And this concludes the proof.
D
39

---

I Experimental Settings
In this appendix, we detail the experimental settings used for the reported experiments in the main paper, which are all completed on 2 Nvidia Titan V GPUs with 24GB GPU memory on local Lambda box.
L1 Datasets
For this experiment, we have chosen the text classification problem for the sentiment analysis task on two datasets: IMDB (57 I. and Elec [59 2 f for sentiment classification in movie reviews and Amazon electronics product reviews, respectively. The reason behind choosing these datasets lies in their inclusion of lengthy texts, which means they can benefit from hierarchical representation. Both datasets have 2 classes. Table summarizes some basic statistics for these datasets. For the validation set, we have used 10% of the training set.
Hierarchical Language:
# Classes
Train Size
Test Size
Avg. Word/Doc.
IMDB Elec
2
25K
25k
235
2
25K
25k
108
Table 5: The statisties for the IMDB and Elec datasets used for the sentiment classification task.
Multi-modal News Classification: For this task, we have performed experiments for the news classification task N24News dataset [93], where for each news article not only have language and image modalities present, but the text itself consists of multiple sub-modalities, i.e. headline, abstract, image caption and main body. N24News dataset consists of total of 61, 218 news stories and 24 total number of classes. The source of the news articles is the New York Times from 2010 2020. For training/vall vaiiaiii/viidaaiioneessnngggg splitting, we use random splitting of ratio 8:1:1 used by the original paper.
1.2 Model Architectures
All the competitor models in our experiments follow the same general architectural pattern: an attention layer, followed by global pooling layer, followed by multi-layer Perceptron (MLP).
The attention layer: For our main model, the attention layer is a single-layer, HSA-based transformer as depicted in Fig. 4 For brevity, we refer to this architecture simply as HSA. For the flattened self-attention (FSA) baseline,
same attention layer as Fig. 4 is applied, except that the input signal hierarchy to the layer is flattened into one level (simple) signal (For experiments using the standard transformer layers, see Appendix[] J. As shown by Proposition 1, a single level signal hicrarchy is mathematically equivalent to the standard Softmax attention mechanism, which means that view FSA representing the standard Softmax attention. For the DeepSet baseline in the multi-modal experiment, we apply the same architecture Fig. for the attention layer, except the attention operation itself is disabled. That is, all the other neural operations in the HTE layer applied except for the attention. This effectively means that individually transform each token in the signal hierarchy without letting them interact with each other through the attention mechanism. This operation followed by pooling and MLP layers effectively implements DeepSet architecture [T02] for combining the token representations in the input signal into a single, fixed sized vector. Note that in all of our experiments across differrnt models, the attention layer is simply the HTE layer Fig. or a variant of it, and as such we can specify the architectural details for each experiment/model as detailed in Table To ensure a fair comparison, we maintain an
using the same syper-parametee equal number of parameters across all models within each experiment.
The global pooling layer: The purpose of global pooling layer is aggregate the leaf representation across the hierarchy into a single, fixed-size vector. We have multiple options for this layer; in our experiments, have chosen the global mean pooling.
The MIP: After pooling the representation into a single vector, we apply a 1-hidden layer MI.P on the resulted vector, the dimensions of which are summarized in Table 6
40

---

Experiment
Hierarchical Language
Multi-modal News Classification
Model
FSA
HSA
Deep Set
FSA
HSA
# of Parameters
1.2M
1.2M
13.4M
11.8M
11.8M
Heads
3
3
3
HTE Layer Output dim
128
128
512
512
512
Position Embedding dim
768
768
768
768
768
Attention dim
128
128
768
256
256
MLP dim
128
128
512
512
512
Table 6: Configuration of model architectures employed all experiments/moddls
1.3 Training Hyper-parameters
Table summarizes the training hyper-parameters used for each experiment. We use the same
hyper parameters across different baselines for each experiment.
nyper-parameters
Experiment
Hicrarchical Language
Multi-modal News Classification
Loss Function
Standard Cross-Entrops Loss
Standard Cross Entsopy Loss
Train Batch Size
64
512
Test Batch Size
64
512
Optimizer
Adam W
Adam W
Max Tokens Training
512
512
Learning Rafe
2 10
1 x 10-4
Learning Rate Scheduler
LinearLR
LinearLR
Train Epochs
30
Table 7: The training hyperparameters used for each experiment.
Comparison to The Classical Transformer Architecture
The experimental results reported in Sections 5.1 and 5.2 aimed at comparing the performance of our HSA framework vs. that of the flat attention, where the rest of the architecture aside from the attention mechanism were the same one proposed in Appendix F However, more practical comparison would the one between the performance of these two mechanisms within classical transformer architecture proposed by [90]. To this end, we have conducted experiments where we train from scratch and compare a standard RoBERTa model and BaERTaa model (as proposed in Section 5.3 on two GLUE benchmarks. For HSA-Ros RoBERTa, operation with HSA, while the hierarchy is imposed a fixed four-level hierarchy where the branching factors from bottom to are 16, 8, 4, and 2.
simply replace the standard flat self-attention
Dataset
Model RoBERTa
Accuracy
Precision
Recall
Fl Score
0.8608
0.8872
0.9058
0.8964
MRPC
HSA-RoBERTa
0.8846
0.9165
0.9093
0.9129
RoBERTa
0.8158
0.7985
0.8167
0.8075
RTE
HSA-RoBERTa
0.8158
0,8076
0.8015
0.8045
RoBERTa
0.3681
0.3681
0.5381
0.4371
QQP (after epochs)
HSA-RoBERTa
0.9185
0,8764
0,9065
0.8911
Table 8: The comparison of training RoBERTa vs. HSA- from scratch three GLUE datasets.
Table shows the results on the evaluation set of each dataset after training. As these results show, the incorporation of HSA within standard transformer architecture not only can improve the computational complexity of self-attention computation, but it can also improve the evaluation metrics due to the regularization effects of our hierarchical framework. This result consistent with the ones in Sections 5.1 and 5.2 Furthermore, for the QQP dataset, we have shown the results just after 4 epochs; interestingly, these results show that HSA-RoBERTa converges much faster than the standard RoBERTa model.
41

---

K
Going Beyond Softmax Attention
One of the primary contributions of our work is generalizing Softmax attention from flat signals to the hierarchical structure of nested signals. This generalization is further confirmed by the theoretical result of Theorem [.2] However, there has been a significant effort in the literature to explore other forms of attention mechanisms than Softmax attention [20] [2 108
78 37]. One of the main motivations of departing from the Softmax attention lies in the fact that Softmax attention induces dense probability distribution over all tokens. Sparse attention (20 (2). on the other hand, organically induces sparse probability distributions over tokens which can greatly improve the interpretability and computational efficiency of transformer models. A natural question is then whether our hierarchical derivation can be applied to other forms of attention, in particular the sparse attention. In other words, can our formalism also generalize sparse attention from flat signals to the hicrarchical structure of nested signals?
K.1 Sparse Attention as Energy Minimization
The first step toward generalizing Sparse attention to the hierarchical setting is to formulate the flat case as an energy minimization problem, much like what we did in Proposition 3.1 for the Softmax attention. To this end, we would need to define an appropriate energy function for the sparse attention. But before that let us define generic form of energy function that can encompass various forms probability- based attentions.
Let Q and K be sets of query and key vectors with bounded norms (e.g. induced by LayerNorm) respectively; we define the generic energy function as:
Φ (Q, k)
Ei1, Z2,.... AN), where 2ij = =//kg
(44)
N
Then the gradient of 6= (Q, K) w.r.t the query token yE is:
N
Oo!
q: 60
k- = (V= K
(45)
Ozy
86 T is the attention weight where K is the key matrix (as defined in (1T) and V = Tazn JznN
009
vector. In (5). we defined of s to be the log-sum-exp function and that led the attention weight vector V.0f to be the Softmax function. Now in the general case, if 6e's are continuous and strictly convex, we can write (See [1] Proposition 1.3):
arg max ( JRRN
[p" 0g (p)]
(46)
where dfm(f) is the domain of function f(-), and of (p) = supsedom(d!) [p"2 - (z)] is the convex conjugate of of (z). For the log-sum-exp function (2) = log | j=1 exp(21))] }], the convex
conjugate is the negative Shannon Entropy of (p) = Σy= Pij log pij-
On the other hand, in sparse attention (22), the attention weight vector v0 is set to be the a-entmax function which has the exact same form as (46) with 6i" (p) = - H.T (p), where
a[a-I] プ (Ps p()) Ej=1Ph Pis log
+ 1 a a=
1t -
(47)
is the Tsallis continuous family of entropies [89]. It is straightforward to show that for a = 1 (i.e. the Shannon Entropy), the a-entmax function reduces to the Softmax function. However, as we saw before, we can alternatively derive the Softmax function by first deriving the energy component of (p) as the log-sum-exp function and then computing its gradient. Now by following the same process for the general Tsallis entropy, we can derive the equivalent energy component whose gradient would be the a-entmax function. In particular, by setting of" (p) = -HLT (p) (as done in the formulation of Sparse attention (22). we will have the energy component df( = [- H (p)]*, which can be further derived in closed form as:
N
(2 ala 1)
+
Uig
1/(a-1)
Zig
a(a 1) yio
1
(48)
j=1
42

---

where
Vij ReLUl(a 1)21j Tpl
(49)
and 7i is the Lagrange multiplier corresponding to the Σ 2== Pij = constraint. Note that, in general, is function of all 2ij 's; that is, T; (2i0...... 2,N). By plugging (48) into (44) we arrive at the equivalent energy function for the general a-entmax attention (i.e. the sparse attention):
8 (Q.
N
1/(0-1) Vij
2ij
a(a -1)
yig + C. where 2iz = o kg (50)
K.2 The Hierarchical Generalization
Now that we have the energy function for sparse attention in the flat case (550) ). we can generalize the hierarchical structure of nested signal by following similar recipe as 8 In particular, for node
A in the signal hierarchy ltz, the hierarchical sparse energy is recursively defined as:
[((B) (((A)) OB
(A)
Σ
log (c(c)))
(51)
(Echd(A)
where CI....,Cy Cs are the sibling nodes of B, BBBCC, is the interaction energy function defined in (7) and the multi-variate function 0% has the same functional form as (48). Then ( can be seen as special case of where = 1 and o3 reduces to the log-sum-exp function. Given the hierarchical sparse energy, we can derive the hierarchical sparse attention by taking the gradient of da (Rz) w.r.t to
each query vector qi, similar the derivation in (9) for the Softmax case. We leave further derivation of an efficient algorithm and theoretical optimality for sparse attention to future work.
I Lastly, it should be noted that similar to flat sparse attention, one can also learn the sparsity factor a via back propagation in the hicrarchical case. This can further extended learning different sparsity patterns for different levels of hierarchy, which can be useful depending the application.
L Zero-shot Approximation of Self-attention: Ablation Study
In this appendix, we further expand on the experimental results for the zero-shot HSA approximation of RoBERTa presented in the main paper. In particular, we study the effects of approximating different combination layers as well as different hierarchical structures for the datasets reported in the main paper.
L.1
Experimental Setup
Datasets:
We have experiments 5 GLUE datasets (SST-2, CoLA, MRPC, RTE and QNLI) as well as the AGNEWS and IMDB datasets.
Models: For each dataset, have used the appropriate pre-trained RoBERTa checkpoint and configuration that has been fine-tuned on the corresponding task. Table 9 lists the checkpoints and configuration used for each dataset. All of our experiments involve only evaluation of pre-trained RobERTa without any training of fine-tuning it.
Metries: We have computed Accuracy, Precision, Recall and F1 Score to measure the accuracy drop of pre-trained RoBERTa as various layers are approximated by HSA.
Impacted Layers: As mentioned the main paper, approximating all self-artention layers of RoBERTa typically leads to significant zero-shot accuracy drop across all tasks. However, approxi- mating subset of layers can introduce more reasonable gap while still benefiting from HSA speed up in terms of the number of FLOPs. Nevertheless, finding the best layer combination is combinatorial problem. To alleviate this issue, instead of examining all different combinations, we only look at certain combinations based on two empirical observations. In particular, we observed that earlier layers in the network are typically more sensitive to approximation, whereas the latter ones are more amenable to it. This observation intuitively makes sense because the sooner approximation takes place in the network, the higher approximation error accumulates along the network. Moreover, having consecutive layers approximated typically increases the accuracy gap whereas interleaving them with regular self-attention layers decreases the gap.
43

---

Dataset
RoBERTa Configuration
Checkpoint
IMDB
Large
AGNEWS
Base
SST-2
Base
testattckklobcrrr-bee--SS22
CoLA
Base
(extattaek/roberta--bes-----
MRPC
Base
Lextattek//oobetta-bsse--RPP
QNLI
Base
textattack/robertt--bsss---
RTE
Base
textattack/roberta--bsse
Table 9: Checkpoints and RoBERTa configurations used for evaluating each task.
hhuss//hugginggfcce
"htps://hugginggffce
M MMRP
https://hugginggfcc colteetaattaccc/oooottttaaaaaaaa base- QNLI
Based on these two observations, our experiments, we only examine combinations where a start layer (denoted by SL.) and every other layer after that are approximated by HSA. The X-axis for the bar plots in this section is associated with SL. Also, the right bar each plot represents the metrics for the original model without HSA approximation.
For these experiments we chose to use fixed hierarchies based on non-overlapping hopping windows rather than semantic hierarchies based on the text structure. The reason behind this choice that semantic hierarchies (such as sentences, paragraphs, etc.) are example dependant which means they would incur different number of FLOPs for different examples. But since our ultimate goal from this experiment to reduce the number of flops consistently across the data, we opted to use fixed hierarchies.
Hierarchy:
The fixed hierarchies here are characterized by having fixed branching factor for all the nodes belonging to the same level of the hierarchy. We then denote such hierarchy by the tuple (A, B, C,... where is the branching factor the lowest level of the hierarchy, B is the branching factor for the next level and so on. Having this notation in place, we have experimented with the following hierarchies:
1. (2, 2,2 A hierarchy with low branching factor at all levels.
2. (2, 4, 8, 16): A hierarchy with low branching factor on the bottom and high branching factor the top.
3. (7, 7,7,7), 7): A hierarchy with high branching factors at all levels.
4. (8, 4, 2) A hierarchy with high branching factor the bottom and low branching factor on the top.
L.2 Results
SST-2 Task: As Figure shows, the SST-2 task is relatively robust the choice of SL. (start layer for HSA approximation), where the accuracy gap widens if SL. falls below Layer 5. Also, the choice of hierarchy is relatively inconsequential except for the narrow hierarchy with low-branching factor across all its levels, which demonstrates slightly poorer results compared to the rest.
RTE Task:
As Figure shows, the RTE task exhibits the same behavior as the SST-2 task with major accuracy drop takes place when SL. falls below Layer 7. Different hierarchy structures seem to have similar behavior though.
the MRPC task, layers with even index seem to be way more sensitive to HSA approximation than the odd-index layers. Among odd index layers, the accuracy gap starts to widen for SI. below Layer 7. As for hierarchy structures, structures with low branching factor the bottom levels seem to do better than the other two candidates.
MRPC Task:
As FigureBssowss
44

---

2.2.22/9512
Auorasy 2.6.B.10rwt2
2.-
Viano
7,7.2.3/es12
14422
3.0
Acoraay
16
Figure 7: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the STT-2 task.
QNLI Task: As Figure 10 shows, for the QNLI task, there is sharp drop of accuracy when SL falls below Layer 8, whereas for the last four layers the accuracy drop is practically insignificant. This shows that in this case, the last 5 layers are quite amenable to approximation. As for the hierarchy structures, they do not exhibit any significant difference for this task.
CoLA Task: As Figure I1 shows, similar to the MRPC task, in CoLA task, the layers with even index seem to be way more sensitive to HSA approximation than the odd-index layers. However, unlike the MRPC task, the hierarchy structures with high branching factor the bottom seem to significantly perform better than the ones with low branching factor on the bottom.
AGNEWS Task:
As Figure 12 shows, for AGNEWS task, we can pretty much start SL at Layer 2 and as long as we approximate every other layer, the accuracy drop in insignificant. for hierarchy structures, we have tested only 2 of our structures with this datasets, but did not observe any significant difference.
IMDB Task: Unlike the previous tasks, for IMDB task, we use RoBERTa-large with 24 layers. As Figure shows, long as SL stays above Layer 15, the accuracy drop is insignificant. Also some layers like Layers 10 and 15 seem to be moresensitive we start the HSA approximation from them. As for hierarchical structure, among the two candidate we used for this task, the one with high branching factor on the bottom seems to do much better.
M Limitations
In this section, state some of the main limitations of the current work both in terms of the proposed framework itself as well as our process of evaluating it in this paper.
As for the HSA framework itself, we would like to emphasize that it only operates tree-based, compositional information hierarchies that are already given or extractable using preprocessing procedure. In other words, HSA does not automatically learn the hierarchical structure embedded in
45

---

2.2.2.2/te
2.4.&.
1.-
Acorasy
Viano
9h
7,2.7.70tt
8,4.2/mt
3.0
Acoraay
0h
16
Figure 8: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the RTE task.
the data, Furthermore, our proposed framework does not introduce any additional learnable parameter across the hierarchy on top of the standard self-attention parameters. While this is useful feature in certain usecases such as zero-shot HSA replacement post-training as described in Section 5.3 in some other scenarios, this would introduce a limitation in terms of the learning capacity of our framework. This also makes quite challenging to perform fair comparison of our framework against other hierarchical attention mechanisms as in pretty much all of those frameworks, extra hierarchy-realted, learnable parameters have been incorporated as a fundamental component of the framework.
As for our evaluation process in this paper, we understand that the scope of our empirical study does include mainstream applications. especially in the auto-regressive generation domain. However, we would like to note that this work primarily focuses on the theoretical foundations of the proposed nested signal data structure and its hierarchical attention mechanism. We hope that our work can be used as the foundation for many follow-up efforts focusing on various applications with potential technical extensions. Nevertheless, we have still provided the initial theoretical extensions for some of the major potential follow-ups to the present work, such hierarchical auto-regressive generation (Appendix G and hierarchical sparse attention (Appendix K)
Finally, we would like to note that the goal of training foundational models using HSA to incorporate the inherent hierarchical inductive biases of our human knowledge base is indeed inspirational, as we do not have the computational resources required for that scale of training. However, from theoretical point of view, our HSA method provides plausible framework to achieve this goal in practice.
46

---

2,2. 2.2.2.2mp 2hrge
24...16mppccc
10-
Accursc
Scom
Pradson
1.-
Acouaty
Precatn
0.3
1A
0.2
0.0
08
5.0
7...220mrpc
3,4.2mmpcc
1.0
Accursc
Nousenn
Rocal
2.0
ACOPaCC
Score
PIECODO
NOC
0.0
2A
0R
0.0
Figure 9: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the MRPC task.
2.2.2.2/eni Zienil
2.4.8.16Ggnl
Accuracy
Scare
Tradiiin
2.0
Acouracy
Precaio
local
1.5
0.5
0,2
22
0%
0.0
20
7.7.7.7oonl
8.4.2/qrii
Accupry
Scor
Predicon
Rocal
2.0
AcTuviTy
Secal
1.5
0
0.
CA
0.2
12
0.0
Figure 10: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the QNLI task.
47

---

2.2.2.000
08
7.7.7.7cc0
Figure 11: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the CoLA task.
7.777700gows
Snore
Recal
Fl Scate
Perzin
Secal
Figure 12: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the AGNEWS task.
48

---

2.6...6fmmbb
7.7.3.7mmmb
Pradicen
Scone
Precebn
Secal
1415
Figure 13: The accuracy metrics of HSA approximation of self-attention layers in RoBERTa for the IMDB task.
49


Reviews:
[{"review_id": "vSs7Mlq8pr", "reviewer": "Reviewer_edS5", "summary": "The paper introduces hierarchical self attention (HSA) to reason over data with hierarchies or multiple modalities. The paper introduces nested signal model for such data, and then perform entropy minimization to obtain HSA. They demonstrate that such a representation is KL-closest block-constrained approximation of standard soft max attention. The authors show that HSA increases accuracy on long-text sentiment tasks (ex: IMDB) and outperforms baselines such as DeepSet and vanilla transformer on multimodal news classification benchmark.", "questions": "Please see weakness section", "limitations": "None", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n1. Main strength of the paper it's clean derivation from first principles with strong grounding provided by KL-optimality proof. \n2. The dynamic programming implementation wit black box sub attention for plugging in fast approximation helps bring down the computation complexity significantly. \n3. It can be used as dtop in replacement for softmax attention without increasing memory footprint.\n\nWeaknesses\n1. Main supervised tasks are on a limited scale (ex: classification), and there is no token level evaluation of generation limiting the ability to understand the qualitative improvements obtained for generalized inference.\n2. The baselines such as DeepSet and flat attention for multimodal tasks are sparse, how does it perform on more complex tasks such as VQA, Image Description etc. where the model inference and training becomes complex. \n3. Main limiting factor of the proposed approach seem to be assumption of a fixed hierarchy. Is it possible to learn the hierarchy or how can this method be scaled where no existing parser is available.", "was_revised": true, "final_justification": "Although most of my asks were in the negative in the rebuttal, the paper assessment still stands. I would keep my initial rating. Hopefully the authors have ideas for improving and extending the work.", "rebuttal": "We appreciate the reviewer's insightful feedback and recognition of our fundamental derivation of hierarchical attention.\n\n1. **Re: the scope of our empirical study,** we agree with the reviewer that our empirical study is somewhat limited to classification tasks and has not covered generative tasks which are key to modern LLMs. However, as we also stated in the main paper, addressing token generation is beyond the scope of the current work which is already somewhat dense presenting the theoretical foundations of the HSA framework. Nevertheless, in Appendix G, we have laid the theoretical foundations for hierarchical generation by proposing techniques to perform efficient causal masking and key-value caching under the hierarchical setting (which are both crucial to any realistic transformer-based LLM). Further empirical exploration of hierarchical generation has been deferred to future work.\n\n2. **As for VQA and Image Description tasks**, we agree that both of these tasks involve more complex modality cross-interactions; however, those tasks are generative and, as mentioned above, outside the scope of our empirical study for this work, which is mainly focused on introducing the theoretical foundations for HSA and nested signals. Our goal is to address generative tasks in a follow up work dedicated to hierarchical generation.   \n\n\n3. **Re: learning the hierarchy,** as we have stated in our Limitations section, our HSA framework does not provide any mechanism to *learn* the hierarchal structure from data. In that sense, hierarchy is seen as a form of *domain knowledge* which injects strong inductive biases into the model and potentially reduces the statistical complexity of learning on top of improving the computational complexity. Nevertheless, depending on the scenario, one may be able to *impose* an artificial hierarchy in the absence of such domain knowledge or a parser. In fact, in Section 4.3 and Appendices J and L, we have *not* used any parser or domain knowledge for any of the experiments and all the reported results are based on purely imposed hierarchies (using various (fixed) branching factors at each level). As the results show, even with blindly imposed hierarchies for these problems, the HSA framework outperforms the baselines in terms of both accuracy and speed, highlighting the role of HSA as a form of regularization."}, {"review_id": "8LmSl7hvBh", "reviewer": "Reviewer_NPM6", "summary": "This paper introduces a novel mathematical concept termed nested signal to formally represent hierarchical and multi-modal information within the Transformer framework. Based on this, this paper derives Hierarchical Self-Attention (HSA), a generalization of the standard Softmax attention mechanism that incorporates hierarchical structure. The paper further presents an efficient dynamic programming algorithm to compute HSA and demonstrates its effectiveness on both unimodal and multimodal tasks. Empirical results on IMDB, Elec, and N24News datasets show performance improvements over baselines.", "questions": "Is it possible to dynamically infer or adapt the hierarchy depth rather than requiring it to be pre-specified? How would the method generalize to scenarios with unknown or evolving hierarchical structures?\n\nCan the authors provide further empirical evidence for the multi-modal applicability of the proposed method, ideally on more diverse and competitive benchmarks?\n\nHow does HSA interact with modern large-scale pre-trained vision-language models or other Transformer variants? Could HSA be integrated into such models, and if so, what would be the expected benefit?", "limitations": "The requirement to pre-define the hierarchy limits the method's flexibility and may restrict its applicability to tasks with fixed or clearly known hierarchical structures.\n\nThe experimental evaluation for multi-modal scenarios is limited, relying on a single dataset with relatively weak baselines, making it hard to draw strong conclusions about multi-modal performance.\n\nThe computational overhead of the proposed HSA, despite claimed efficiency improvements, is not thoroughly analyzed or compared to standard attention in practical large-scale settings.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "# Strengths:\n\nThe nested signal concept is new and interesting.\n\nThe design of Hierarchical Self-Attention appears principled and technically sound.\n\nAddressing multi-level, multi-modal learning within a unified Transformer-based framework is an important direction, especially given the increasing complexity of modern AI tasks.\n\nThe proposed method shows empirical improvements on several benchmarks (IMDB, Elec, N24News), indicating its potential practical utility.\n\n# Weaknesses:\n\nBased on my understanding, the proposed method requires the hierarchy depth to be explicitly pre-defined. This rigid design may limit its flexibility and scalability, especially when dealing with tasks or datasets where the hierarchical structure varies or is not clearly defined. This contrasts with traditional Transformers, which often benefit from scaling data or model size without making strong structural assumptions.\n\nAlthough the paper claims to address multi-modal scenarios, only the N24News dataset is used to demonstrate this aspect. Moreover, the reported improvements over baselines like FSA are relatively modest, and the chosen baselines are not particularly strong or recent.\n\nWhile I agree with the authors' experimental design principle of isolating contributing factors, limiting multimodal evaluation to a single dataset with weak baselines undermines the strength of the multi-modal claim. Additional, more convincing evidence is needed to support this aspect.", "was_revised": true, "final_justification": "I have read the rebuttal and the comments from other reviewers. While the rebuttal partially addresses my concerns, I believe there is still room for improvement—particularly in the experimental evaluation on multi-modal tasks. Based on the overall merits and the clarifications provided, I am updating my score to Weak Accept.", "rebuttal": "We are sincerely grateful to the reviewer for their insightful feedback and bringing up some important points that'd require our further explanation. \n\n1. **Re: \"predefined depths\",** we would like to emphasize that our proposed HSA framework does NOT require the input hierarchies to have any pre-defined depth, size or shape. In fact, within a minibatch, each nested signal example can have its own different depth, size and shape, and this is already supported via an efficient implementation using the *breath-wise tree concatenation* technique in our codebase, as explained in Appendix E.4. In other words, the same way that the classical transformer architecture is invariant to the length of the input (flat) sequence, the HSA-based transformer architecture is invariant to the depth, size and shape of the input nested signal. In that sense, HSA-based transformer can very well handle \"evolving\" hierarchies. Now, if one wants to define a separate $Q$, $K$ and $V$ projection for each different modality in each layer of transformer (as proposed in Appendix F), these modalities need to be predefined, so one can know how many of such projections are needed in each layer. However, once these modalities (e.g. language and vision) are known, one may compose arbitrarily-shaped hierarchies using these predefined modalities and feed them to the same HSA transformer model. \nThat said, as we explicitly stated in the paper, the HSA framework cannot \"infer\" the hierarchy; instead, the hierarchy and its implicit geometry are seen as external domain knowledge acting as a form of inductive bias. In that sense, hierarchy inference is completely tangential to the HSA framework, enabling us to use HSA in combination with any hierarchy extraction method. It should also be noted that, depending on the scenario, even when the hierarchical structure of the input isn't given or learned, we may still benefit from HSA by imposing \"artificial\" hierarchies on the input signal, and effectively using it as a form of regularization. In fact, all the results reported in Section 4.3 and Appendices J and L are obtained by incorporating imposed hierarchies.\n\n2. **Re: the multi-modal empirical study,** we understand that the reported results are somewhat limited to one task. However, we would like to note that most interesting multi-modal tasks such as VQA, visual-language reasoning, etc. almost always require auto-regressive generation (ARG) as a key component of the model. But in this work, we have deliberately tried to stay away from ARG because incorporating HSA with ARG capabilities is not trivial and merits its own dedicated, follow-up work. Nonetheless, we have laid the foundation for hierarchical ARG in Appendix G to be elaborated upon in future work. In the meantime, for the purpose of this work, we will make this point clear in the final version of the draft to avoid any empirically unsubstantiated claim or confusion.\n\n3. **Re: integrating with large-scale vision-language transformers:** The HSA framework can be in principle integrated with any multi-modal transformer architecture including large-scale vision-language transformers. Currently, many of such models naively put signals from various modalities on a simple, flat sequential geometry and instead rely on massive-scale pretraining to learn the underlying geometry of cross-modal interactions from data. However, if such geometrical information is partially available as domain knowledge, HSA can incorporate it within the proposed nested signal data structure to potentially reduce the statistical complexity of pretraining and finetuning large-scale transformers. For instance, assume that the input data consist of images with their corresponding scene graphs (where the scene graph's nodes are detected object bounding boxes with some textual descriptions while its edges model various relationships between the objects). Feeding such rich data structure to classical vision-language transforms often involves losing most of such rich geometrical and semantical information in the process of \"flattening\" the signal, and instead relying upon massive-scale pretraining to recover part of that knowledge. Whereas, it's quite seamless to encode such prior knowledge using the nested signal construct and perform self-attention operation on it using the HSA formalism.\n\n4. **Re: the computational overhead of HSA,** first please note that the main overhead of HSA is building the hierarchy for the input signal. But that can be seen as part of the tokenization process, and similar to tokenization, it can be done only once on the CPU before training. As for the rest of computational analysis, we have provided GFLOP metric on various tasks with different context lengths in Table 3. As shown in the table, HSA has significantly decreased the required GFLOPs compared to the baseline. Furthermore, the reduction is way more dramatic for larger context lengths, as shown in Table 3. As for the exact time comparisons of HSA and the baselines, we'd like to note that such comparisons are heavily implementation dependent and rely on various CUDA optimization tricks. However, efficient CUDA implementation of HSA is beyond the scope of our current paper, which is primarily meant to lay the theoretical foundations of the proposed framework."}, {"review_id": "OOWHTMbN9l", "reviewer": "Reviewer_HMnn", "summary": "The paper introduced hierarchical structure into the soft-max attention mechanism in a mathematically interpretable manner. The authors first derived the softmax attention layer as similar to one step gradient update of an entropy minimization problem. Then generalize the attention layer to nested signals by introduce hierarchical structure into the variables/nodes", "questions": "NA", "limitations": "+ Have the authors try to do more than 1 gradient step update ?\n+ How the performance vary according to different signal hierarchy ? Is the a measurement to quantify the robustness of one hierarchy to another aside from doing the training process.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n+ The network architecture are derived from unrolling, which leads to an interpretable network.\n+ Introduce the hierarchical structure into the attention map.\n+ Clear performance gain and also computations cost reduced due to hierarchical structure\n\nWeaknesses:\n+ Lack experiments to demonstrate the interpretability of the model, Ex: do more than 1 gradient steps ?, plot reduced entropy objective for various number step/layers ?\n+ The signal hierarchy are not always obvious for various datasets/application: graphs signals. The paper don't introduce any mathematical methods to create hierarchy from data and mainly use intuition.", "was_revised": true, "final_justification": "I appreciate the response/rebuttal from the authors. I will maintain my evaluation.", "rebuttal": "We truly appreciate the reviewer's constructive feedback and comments.\n\n1. **Regarding lack of experimentation for multi-step gradient-descent interpretation,** we'd like to note that the main angle and focus of this work is the theoretical derivation of hierarchical attention from the entropy minimization interpretation of self-attention. In that sense, the gradient-descent interpretation is a secondary byproduct of our theoretical framework which is, of course, worth exploring on its own in a separate work. At the same time, we'd like to mention that performing $n$-step gradient descent on conditional entropy using the attention operator is mathematically similar to having $n$ sequential self-attention layers that share weights. But weight sharing in transformers has been already studied in the literature and shown to improve the performance (e.g. see [A, B, C]).\n\n[A] Xiao, Tong, et al. \"Sharing attention weights for fast transformer.\" arXiv preprint arXiv:1906.11024 (2019).\n\n[B] Reid, Machel, Edison Marrese-Taylor, and Yutaka Matsuo. \"Subformer: Exploring weight sharing for parameter efficiency in generative transformers.\" arXiv preprint arXiv:2101.00234 (2021).\n\n[C] Kowsher, Md, et al. \"Does Self-Attention Need Separate Weights in Transformers?.\" arXiv preprint arXiv:2412.00359 (2024).\n\n2. **Re: \"a mathematical approach for creating the hierarchy\",** as we stated in the paper, the HSA framework by itself doesn't provide any systematic approach to infer the hierarchy from data. Instead, the hierarchical structure is seen as domain knowledge which can act as inductive bias and subsequently encoded via the HSA framework. This leaves the door open for HSA to integrate with any manual or machine-learned hierarchy extraction method. Nevertheless, as shown by the results in Section 4.3 and Appendices J and L, in the absence of the hierarchical knowledge, one may still impose an \"artificial\" hierarchy on the input signal and still benefit from the regularization and speedup characteristics of HSA. Also, to further understand the sensitivity of HSA to the choice of the hierarchical structure (as the reviewer suggested), we have conducted ablation studies in Appendix L using four different imposed hierarchies across multiple datasets (please see Figures 7-13 in the appendix). As these experiments show, certain tasks can be more sensitive to the choice of the imposed hierarchy, while others are more robust. This further shows the strong correlation between the data geometry and the nature of the classification task and highlights the importance of cross-validation for picking a suitable hierarchical structure when prior domain knowledge is limited."}, {"review_id": "WzPs54QtXd", "reviewer": "Reviewer_tJ56", "summary": "The paper presents a new mechanism attention that inherently supports both multi-modal and multi-scale problems. It starts by introducing a formalism of multi-scale problems as trees of nested signals. It then uses this formalism and the following logic to obtain a new formulation of the attention mechanism. First, given key K and queries Q, one can express the conditional entropy of Q given K, and try to minimize it to increase the information. Second, because this entropy is intractable, one expresses an upper bound using variational distribution. By considering a Boltzman distribution and computing the gradient descent for the query, we can obtain new formulations of the attention mechanism which depend on the energy function of the Boltzman prediction. With a particular energy function we can even retrieve an attention quite similar to the original softmax attention. Using another energy function one can rewrite the attention matrix as a block matrix. It means, still considering the hierarchical structure of the problems as a tree, attention value is uniform across leaf elements that belong to different branches. This effectively reduces the number of values to learn.\n\nThe first experiment evaluates the performance of this new hierarchical attention on sentiment analysis. Texts are split into paragraphs, sentences and tokens. The hierarchical attention outperforms the traditional one. The second experiment evaluate news classification in a multi-modal context. The proposed attention yield to slightly worse results but for a much cheaper computation.", "questions": "- Could comparison with models dedicated to hierarchical settings like Swin Transformer be added?\n- Would it be possible to reduce the section about finding back a similar attention as the traditional softmax one, to bring forward the per-block attention matrix? It would be also the opportunity to bring forward the dynamic programming method to compute the hierarchical attention, which is pointed as a major contribution but relegated to the appendix.", "limitations": "More experimental evaluation of the hierarchical self-attention would help assessing its limitations.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\n- The paper presents an original per-block attention that natively  support for multi-modal and hierarchical settings\n- Experimental results show the new attention mechanism performs on par with the traditional one for a cheaper cost.\n    \n## Weaknesses\n\n- The paper is dense, which makes it difficult to follow despite of the clarity of the different exposition steps. The paper may gain in readability with some reorganisation. Some parts of the mathematical justification behind the new attention mechanism may be superfluous. For instance, is it necessary to emphasize so much how to find back with the new formulation an attention that is close to the traditional one? The per-block attention matrix that supports hierarchical settings should be brought forward sooner in the paper.\n    \n- This new attention formalism has the potential to impact many problems due to its support for hierarchical and multi-modal settings. The experimental settings could be larger to assess this impact. For instance, it could include comparison with other methods that were design with hierarchy in mind like Swin Transformer.", "was_revised": true, "final_justification": "The paper presents an interesting hierarchical attention mechanism, which promises to leverage relational organization of the data at a lower computational cost than traditional self attention. The paper already good in its current form, would greatly benefit from some reorganization to convey a clearer message. For instance, it might be more important to expose the dynamic programming approach to reduce complexity than the ability of the proposed formalism to recover the traditional softmax attention. Additionally, the experimental settings could be strengthened to illustrate further the benefits of the proposed method. This could include use of the new attention mechanism in traditional model architectures on diverse tasks.\nFinally, I believe the idea is interesting enough for the paper to be accepted. Other reviewers made significant remarks that could be integrated in the final version.", "rebuttal": "We sincerely thank the reviewer for their insightful comments on our submission.\n\n1. **As for the organization of the paper,** we agree with the reviewer that our manuscript is somewhat dense as we have tried to cover many different bases and angles, leading to an additional 25 pages of appendix materials. In order to make the final manuscript more clear, we will move the notion of block attention higher to the beginning of the paper as the reviewer suggested. If possible, we will also try to move some of the background materials to the appendix to find more room for the Dynamic Programming section in the main text. As for the section on the KL-optimality of HSA (i.e. Theorem 3.2), we would like to note that this result is a crucial theoretical contribution of our work as it establishes the optimality of HSA approximation among other block approximations of the attention matrix. But why is this result important? Because it opens the door for another equally important application of HSA in practice: the zero-shot plug-n-play capability of HSA in integrating into the pretrained models, as it was empirically validated in Section 4.3 and Appendices J and L. Such plug-n-play capability effectively enables us to avoid building large-scale HSA models from scratch but from already pre-trained transformers. Our theoretical result guarantees the optimality of such zero-shot procedure.\n\n\n2. **As for comparison to some other hierarchy-based models such as Swin-Transformer,** we are planning to add such baselines to some of our experiments. However, making a fair comparison of HSA to these frameworks is somewhat tricky. In particular, most of such models including Swin-Transformer encode a fixed hierarchical structure within their architecture by assigning different sets of $Q$, $K$ and $V$ projection parameters to different levels of the hierarchy. While this restricts these models' capability in handling arbitrarily-shaped hierarchies (unlike HSA), it also introduces a completely different parameter structure than that of the HSA. In other words, any potential gain or loss cannot be properly attributed to the hierarchical attention mechanism only, which is the main focus of the current work. Please note that in all of our current experiments, the baselines as well as their rival HSA model have exactly the same parameter structure and count but different attention calculation procedures, which makes the comparison fair and the results meaningful. Nonetheless, we are currently trying to come up with a variation of some of these hierarchical baselines that have more similar parameter structure to HSA, but that would require training such baselines from scratch."}]

Author rebuttal:

Meta-review: Four reviewers recommend acceptance ratings. The main reasons for acceptance are the principled and original extension of attention to hierarchical and multi-modal data, the strong theoretical grounding, and the promising empirical results. The main weaknesses are the limited empirical scope, reliance on fixed hierarchies, and the need for broader evaluation and clearer presentation. The reviewers agree that the paper is technically solid and makes a valuable contribution, though there is room for improvement in future work.
On balance, the AC sees no basis to overturn the reviewer suggestions and recommends acceptance. The paper presents a theoretically grounded and practically promising approach to hierarchical self-attention, with clear potential for impact in multi-modal and hierarchical data scenarios. The AC highly recommends the authors to address the concerns of the reviewers and take into account their suggestions of improvement when preparing a revised version.
Decision: Accept (poster)