Coneheads: Hierarchy Aware Attention
Authors: Albert Tseng
Tao Yu
Toni J.B. Liu
Christopher De Sa
Venue: NeurIPS 2023 poster (2023)
Keywords: Hyperbolic Entailment Cones
Hyperbolic Space
Entailment Cones
Attention
Dot Product
Hierarchy
Transformers

Abstract: Attention networks such as transformers have achieved state-of-the-art performance in many domains. 
These networks rely heavily on the dot product attention operator, which computes the similarity between two points by taking their inner product.
However, the inner product does not explicitly model the complex structural properties of real world datasets, such as hierarchies between data points.
To remedy this, we introduce cone attention, a drop-in replacement for dot product attention based on hyperbolic entailment cones.
Cone attention associates two points by the depth of their lowest common ancestor in a hierarchy defined by hyperbolic cones, which intuitively measures the divergence of two points and gives a $\textit{hierarchy aware}$ similarity score.
We test cone attention on a wide variety of models and tasks and show that it improves task-level performance over dot product attention and other baselines, and is able to match dot-product attention with significantly fewer parameters.
Our results suggest that cone attention is an effective way to capture hierarchical relationships when calculating attention.

Coneheads: Hierarchy Aware Attention
Albert Tseng Cornell University al albert@cs. cornell.ede
Tao Yu Cornell University tyuêcs. cornell.... edu
Toni J.B. Liu
Christopher De Sa
Cornell University j13499@corne11.co edu
Cornell University cdesa@cs. cornell edu
Abstract
Attention networks such transformers have achieved state-of-the-ii performance in many domains. These networks rely heavily on the dot product attention opera- tor, which computes the similarity between two points by taking their inner product. However, the inner product does not explicitly model the complex structural prop- erties of real world datasets, such hierarchies between data points. To remedy this, we introduce cone attention, drop-in replacement for dot product attention based hyperbolic entailment cones. Cone attention associates two points by the depth of their lowest common ancestor in a hierarchy defined by hyperbolic cones, which intuitively measures the divergence of two points and gives hierarchy aware similarity score. We test cone attention on a wide variety of models and tasks and show that it improves task-leve performance over dot product attention and other baselines, and is able to match dot- not-product attention with significantly fewer parameters. Our results suggest that cone attention an effective way capture hierarchical relationships when calculating attention.
Introduction
In recent years, attention networks have achieved highly competitive performance in a variety of settings, often outperforming highly-engineered deep neural networks [5] 8 [I]. The majority of these networks use product attention, which defines the similarity between two points u, v Bd by their inner product u [31] Although dot product attention empirically performs well, it also suffers from drawbacks that limit its ability scale and capture complex relationships in large datasets [33 [27] The most well known of these issues is the quadratic time and memory cost of computing pairwise attention. While many works on attention mechanisms have focused on reducing the computational cost of dot product attention, few have considered the properties of the dot product operator itself [33
Many real world datasets exhibit complex structural patterns and relationships which may not be well captured by inner product (30 T6]. For example, NLP tasks often contain hierarchies over tokens, and images may contain clusters over pixels [34 [4]. Motivated by this, propose new framework based on hyperbolic entailment cones to compute attention between sets of points [IO] 37]. Our attention mechanism, which we dub "cone attention", utilizes partial orderings defined by hyperbolic cones to better model hierarchical relationships between data points. More specifically, we associate two points by the depth of their lowest common ancestor (L.CA) in the cone partial ordering, which is analogous finding their LCA in latent tree and captures how divergent two points are.
Cone attention effectively relies on two components: hyperbolic embeddings and entailment cones. Hyperbolic embeddings, which use the underlying geometric properties of hyperbolic space, give low-
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

kt
Kig, upper bound)
fgm
Hyperbolic
softmax
qk?
x
V
softmax
Space
X
Figure 1: Overview of cone attention vs. dot product attention. In dot product attention (left), similarity scores are calculated with = gg. In cone attention (right), q and i are first projected onto hyperbolic space. Then, pairwise similarity is calculated from the lowest common ancestor of points in the partial ordering defined by entailment cones. Cone attention allows us to explicitly encode notions of hierarchy in attention, and empirically gives better performance than dot product attention.
distortion embeddings of hierarchies that cones, which rely on geometric cones to define partial orders between points, allow us to calculate explicit relationships between points, such as their LCA TO 37]. To the best of our knowledge, we are the first to define hiirarchy-awarr attention operator with hyperbolic entailment cones.
possible with Euclidean embeddings 25
Entailment
Functionally, attention is drop-in replacement for dot product attention. We test in both "classical" attention networks and transformers, and empirically show that cone attention consistently improves end task performance across variety of NLP, vision, and graph prediction tasks. Furthermore, we are able to match dot product attention with significantly fewer embedding dimensions, resulting in smaller models. To summarize, contributions
attention
We propose cone attention, hierarchy aware attention operator that uses the lowest common ancestor of points in the partial ordering defined by hyperbolic entailment cones.
We evaluate cone attention on NLP, vision, and graph prediction tasks, and show that it consistently outperforms dot product attention and other baselines. For example, we achieve +1 BLEU and +1% % ImageNet Top-1 Accuracy on the transformer, iwslt_ddd DeiT-Ti models, respectively.
and
We test cone attention at low embedding dimensions and show that we can significantly reduce model size while maintaining performance relative dot product attention. With cone attention, we can use 21 % fewer parameters for the IWSLT 14 De-En NMT task.
2
Background
In this section, we provide background on attention mechanisms, motivate the use of hyperbolic space embed hierarchies, and describe choice of entailment cones to encode partial orderings.
2.1 Attention
The attention operator has gained significant popularity way to model interactions between sets of tokens B1.. At its core, the attention operator performs a "lookup" between single query qi E d and set of keys k gaxd and aggregates values E gnxd associated with the keys to "read out" a single value for q; Mathematically, this can be represented as
A(q, k) C. > (K()))))us)
where CE ((q, = 1. In "traditional" dot product attention, which has generally superseded "older" attention methods such as Additive Attention and Multiplicative Attention 15
giki Vd
Aik
K K(g,,y) ky) = exp
C=
A(g. k) softmax
(2)
E, k(g,k)
Vd
similarity is scored with a combination of cosine similarity and embedding magnitudes.

---

Existing works have proposed replacing dot product attention to various degrees. A large body of works focus on efficiently computing dot product attention, such as with Random Fourier Features and low-rank methods 23 33 These methods generally perform worse than dot product attention, as they are approximations XX Some recent works, such as EVA attention, parameterize these approximations and effectively get larger class of dot-product-esque attention methods [3]]. Beyond this, Tsai et al. [29] replace K with compositions of classical kernels. Others extend dot product attention, such by controlling the "width" of attention with learned Gaussian distribution or by using an exponential moving average on inputs, although extensions do not usually depend on K 112 16]. Closer our work, Gulcehre et al. [I] introduced hyperbolic distance attention, which defines K(g;, k2) = exp(- (9. k;) - c) where dg is the hyperbolic distance, / E R+, and CE R. K can be interpreted as analog of the distance from q: to K latent tree. Finally, in an orthogonal direction, Tay et al. [26] ignore token-token interactions and synthesize attention maps directly from random alignment matrices. Whether token-token interactions actually needed is outside the scope of this work, and we compare cone attention accordingly.
2.2 Hyperbolic Space
d-dimensional Hyperbolic space, denoted HH,, is simply connected Riemannian manifold with con- stant negative sectional curvature [2]. This negative curvature results in geometric properties that makes hyperbolic space well-suited for embedding reee structures 22 35). For example, the vol- ume of a hyperbolic ball grows exponentially with respect to its radius; in tree, the number of leaves grows exponentially with respect to depth. Furthermore, dar du, v) R du, 0) + dn(O,v), where O is the origin, which again mirrors a tree where du, v) = dr(u,, ICA, u)) + dI.A,,, v), Since hyperbolic space cannot be isometrically embedded into Euclidean space, it is usually repre- sented on subset of Euclidean space by "model" of H. [T]. These models are isometric to each other, and the key differences between them lie in their different parameterizations, which allow for cleaner computations and visualizations for certain tasks.
In this work, we primarily use the Poincaré half-space model, which the manifold H (21 9w) where ج = {ar Rd : xa > 0}, (1) = ge/ ge/ and ge is the Fuclidean metric 2 In the Poincaré half-space, "special" points and curves have particularly nice Fuclidean forms. Ideal points, or points infinity, are the points where a 0 (the "z-axis") and the single point xa oo at which all lines orthogonal to the z-axis converge. Geodesics, the shortest path between two points, are Euclidean semicircles with the origin on the x-axis or vertical rays orthogonal to the x-axis. Horospheres, curves where all normal curves converge at the ideal point, are represented by either an Euclidean ball tangent to the z-axis or horizontal hyperplane when the ideal point is xd = oc.
2.3 Entailment Cones
Entailment cones in hyperbolic space were first introduced by Ganea et al. 10 to embed partial orders. The general concept of Ganea's entailment cones is to capture partial orders between points with membership relations between points and geodesically convex cones rooted at said points. That is, if u E the cone of v, then u < u. Ganea's 's cones (figure right) are defined on the Poincaré ball by radial angle function (r)),, with c-ball around the origin where cones undefined 2 10). This makes learning complicated models with Ganea's 's cones difficult, as optimization on the Poincaré ball is nontrivial and the ball negatively impacts embedding initializations [37 10].
In this work, we instead use the shadow cone construction introduced in [37 and operate on the Poincaré half-space, which makes computing our desired attention function numerically simpler. Shadow cones are defined by shadows cast by points and a single light source S. and consist of the penumbral and umbral settings (figure left quadrant). In the penumbral setting, S is ball of fixed radius and points are points. The shadow and cone of u are both the region enclosed by geodesics through tangent to S. In the umbral setting, S is instead point, and points are centers of balls of fixed radius. Here, the shadow of u is the region enclosed by geodesics tangent to the ball around u that intersect at S. However, to preserve transitivity, the cone of u is subset of the shadow of (see figure(2, The shadow cone formulation can also be achieved with subset relations between shadows instead of membership relations between points and cones, which may be conceptually clearer.
Infinite-seting Shadow Cones. For penumbral cones, when S a ball with center at Za = 20, S's boundary is a horosphere of user-defined height h. Here, all shadows are defined by intersections of
3

---

X = 00
X =
8
P
X
X 00
8(xa
Kg
X =
Figure 2: (Left Quadrant) Clockwise from top left: finite setting penumbral cone, infinite setting penumbral cone, infinite setting umbral cone, and finite setting umbral cone. All figures are in H Shadows are represented by shaded regions, and cones are enclosed in green. In all figures, since the cone of u. (Right) Ganea's entailment cones, as defined on the Poincaré ball. Note the e-ball where cones are not defined, which makes optimization nontrivial. Figure from [0]
Euclidean semicireles of Euclidean radius h. Notably, the infinite setting penumbral cone construction is similar to Ganea's cones under an isometry from the Poincaré half-space to the Poincaré ball where maps the e-ball (37] For umbral cones, when xa = oo, shadows are regions bounded by Euclidean lines perpendicular to the r-axis.
Unlike Ganea's 's cones and penumbral cones, umbral cones are not geodesically convex 37 That is, the shortest path between two points in an umbral cone may not necessarily lie in the cone. In tree, this corresponds to the shortest path between two nodes not being in the subtree of their LCA, which is not possible. Empirically, while still better than dot product attention, umbral attention usually performs worse than penumbral attention.
2.4 Lowest Common Ancestor (LCA) and Least Upper Bound
The LCA two nodes u, v directed acyclic graph the lowest (deepest in hierarchy) node that is an ancestor of both and Although the two terms are similar, the LCA is not the same as the least upper bound of a partial ordering. The least upper bound of two points æ, y in partial ordering, denoted sup(æ, y), the point p such that p < æ, y and Vq where difference is that all other upper bounds must precede the least upper bound, while not all ancestors of and v must also be ancestors of LCA(uu, v). Furthermore, p may not actually exist.
x. y. q p. The key
Hierarchy Aware Attention
Here, we describe cone attention using the shadow cone construction, and discuss projection functions onto Hd that allow to use cone atlention within attention networks. All definitions use the infinite- setting shadow cones, and proofs and derivations are in the appendix. For clarity, we refer to Ganea's entailment
"Ganea's cones" and the general of cones that captures entailment relations (e.g. Ganea's 's cones and shadow cones) as "entailment cones" [IO 37]. Our methods are agnostic entailment cone choice, and can also be used with Ganea's cones.
3.1 Cone Attention
We wish to associate u, Hd by their LCA in some latent tree T, which analogous to finding their LCA, denoted sup v), in the partial ordering defined by entailment cones. Formally,

---

X 00
sups (u,
Xg
Figure 3: In this example using penumbral cones in Hd. sup2 v) is the lowest common ancestor of u and v. The red region is the set of points P s.t. P 3 u, v. these, supy (u, v) lowest point whose cone (light gray) contains both and Here, S is the root of all hierarchies, and points closer to Zd = 0 are closer to the "leaves" of hierarchies. d'== 2, then sup (u, v) is also the least upper bound of u and in the partial ordering defined by entailment cones.
(u, =r arg max du (S,((())) C.C,eEC
where r(C) denotes the root of cone C. This corresponds to finding the cone that is farthest away from S, which the root of all hierarchies the shadow cones construction. When d = 2, sup2 (u, v) also has the nice property of being the least upper bound sup(u, v) in the partial ordering defined by shadow cones. Then, we define the similarity between u and v as
(uu, 8) = f (da o),
(4)
where f is user-defined monotonically increasing function. If K(u,v) du (s(uupy v), > dn(sup, (u, w), S), which implies that u and v have a more recent "ancestor" than and w. Thus, K gives higher similarity score points who have more recent "ancestor" in T. In the dlimite-setting shadow cone construction, sup2 u, v) is root of the minimum height (literal lowest) cone that contains both u and v, or sup2 (u, v) = r (arg minc...vec r(C)a).
K(u, w). then
Using this, we provide definitions for K(u, 8) in the udfnite-settinn shadow cone construction. Both the umbral and penumbral definitions correspond to the Euclidean height of sup2 v). In the penumbral setting, when supy(u, v) does not exist, we return the Euclidean height of lowest light source where sup2 (u, exists. Fa denotes the last dimension of x, and x; denotes the first d - dimensions æ. 7 E IR corresponds to the softmax "temperature" in attention. the penumbral setting, E Bt is the user-defined height of the horosphere light source. In the umbral setting, r E RT is the user-defined radius the ball centered at each point.
Definition 1. Penumbral Attention:
V12 uz
22
U
U.
K(u, u) = exp
-7 max
Udd Udy
x2-
2
(5)
when there exists a cone that contains u and v, and
(-vV
||u U:
22
2
2
K K(u,v) exp
+
(6)
0--1
otherwise. There exists cone that contains and when
2
-Iu-1- Vr2 1
+ x²
(7)
Definition 2. Umbral Attention:
K(u, v) = -exp
(-y -ax
Udd Vdl
|0. 2 1/ 2sinh(r)
+
Ud + Ud
(8)
5

---

These definitions possess a rather interesting property when ud = and K is normalized across set of us, cone attention reduces the Laplacian kernel K (u, U;-1) exp(->llu-- u - I)- Since Euclidean space is isomorphic to a horosphere, and u and v: are on the same horosphere if Ud = Vid, cone attention can also seen as an extension of the Laplacian kernel [19].
Cone attention and dot product attention both take Onndd) time to compute pairwise attention between two sets of n d-dimensional tokens q, k E IR "xd [33] However, cone attention takes more operations than dot product attention, which computes gk In transformers, our PyTorch cone attention implementations with torch. compile were empirically 10-20% slower than dot product attention with torch. bmm (cuBLAS))21 (see section ??). torch. compile l is not optimal, and raw CUDA implementation of cone attention would likely be faster and narrow the speed gap between the two methods 21
3.2 Mapping Functions
Here, we discuss mappings from Euclidean space to the Poincaré half-space. These mappings allow us to use cone attention within larger models. The canonical map in manifold learning is the exponential map Exp,(v), which maps a vector y from the tangent space of manifold M at æ onto In the Poincaré half-space (36.
M by following the geodesic corresponding to [24 2]
Exp,()):: 2:-1+
V Expa(v)a=
||||// tanh([]u[)) vd
cosh([|v|]) - U. sinh (I
While Exp_(") (v) is geometrically well motivated, it suffers from numerical instabilities when |ea| is very large or small. These instabilities make using the exponential map in complicated models such as transformers rather difficult. Using fp64 reduces the risk of numerical over/underfloos, but f664 significantly reduces performance on GPUs, which is highly undesirable [..
Instead, we use maps of the form (z.-1, za) - (x:=1f(zz), f(z)) that preserve the exponential volume of hyperbolic space while offering better numerics for large-scale optimization. Our use of alternatives to (v) follows Gulcehre et al. (IT]s 's psuedopolar map onto the Hyperboloid. Geometrically, since n-diiensionna Euclidean space is isomorphic
horosphere maps corresond to selecting a horosphere with f*a) and then projecting z:-1 onto that horosphere 19 To achieve exponential space as Fd -00, we use functions of the form exp(.).
H these
For the infinite-setting umbral construction, since there is no restriction on where points can go in the Poincaré half-space, we map x E IR to Hd with sb IRd Hd:
(x)).-1 = x.- exp(x.)
(x)x)) exp(xa)
(10)
For the snfinite-setting penumbral construction, since the mapped point cannot enter the light source at height h, we map x E Id to area below the light source with & Rdd + Hd
h
{(x)-1 -.- 1+exp(-xxp
{(x)a
(11)
1 1+ exp(-x)
While & is structurally the sigmoid operation, note that sigmoid (z) = exp(-softplus(-a)) Since softplus(x)) ss for large values of z, & preserves the exponential volume properties we seek.
4 Experiments
Here, present empirical evaluation of cone attention in various attention networks. For each model we test, our experimental procedure consists of changing K in attention and training new model from scratch. Unless otherwise noted in the appendix, we use the code and training scripts that the authors of each original model released. We assume released hyperparameters are tuned for dot product attention, as these models were sfate-of-the-an art (SOTA) when new.
4.1 Baselines
Our main baseline is dot product attention, as it is the most commonly used form of attention in modern attention networks. Additionally, we compare cone attention against Gulcehre et al. [''s

---

hyperbolic distance attention and the Laplacian kernel. To the best of our knowledge, few other works have studied direct replacements of the dot product for attention.
Gulcehre et al. [1'ss 's original hyperbolic distance attention formulation not only computed the similarity matrix K in the Hyperboloid, but also aggregated values in hyperbolic space by taking the Einstein midpoint with respect to weights a in the Klein model (see appendix for definitions) 24]. However, the Einstein midpoint, defined
an(()) ΣS (n)(()
m(a, (00,))=) =
(12)
where 7(v) -1/VI-lui| does not actually depend on how a is computed. That is, could be computed with Euclidean dot product attention or cone attention and m m(a, v) would still be valid. The focus of our work is on computing similarity scores, so we do not use hyperbolic aggregation in our experiments. We test hyperbolic distance attention using both Gulcehre et al. [I]]'s original pseudopolar projection onto the Hyperboloid model and with & map onto the Poincaré half-space,
4.2
Models
We use the following models to test cone attention and the various baselines. These models span graph prediction, NLP, and vision tasks, and range size from a few thousand almost 250 million parameters. While we would have liked to test larger models, our compute infrastructure limited us from feasibly training billion-pant partmmeer models from scratch.
Graph Attention Networks. Graph attention networks (GATs) were first introduced by Veličković al. 332 for graph prediction tasks. GAT's use self-attention layers to compute node-level attention maps over neighboring nodes. The original GAT used a concatenation-base attention mechanism and achieved SOTA performance on multiple transductive and inductive graph prediction tasks [32 We test GATs on the transductive Cora and inductive multi-graph PPI datasets [17 [3
Neural Machine Translation (NMT) Transformers. Transformers were first applied to NMT in Vaswani et al. 31'ss 's seminal transformer paper. We use the fairseq transformer_isee iwslt_de____.....il de en architecture to train German to English translation model on the IWSLT' 14 De-En dataset 20 9. This architecture contains 39.5 million parameters and achieves near-SOTA performance on the IWSLT' 14 De-En task for vanilla transformers (20) As this model is the fastest transformer train out of the tested models, we use it for ablations.
Vision Transformers. Vision transformers (ViT) transformer-lik architectures to perform image classification [8] In a VIT, image patches are used as tokens in a transformer encoder to classify the image. We use the Data Efficient Vision Transformer (DeiT) model proposed by FAIR, which uses a student-teacher setup to improve the data efficiency of ViTs [28]. DeiT's and ViTs share the same architecture, and the only differences train DeiT-Ti Ti models with 5 million parameters on the ImageNet- KK dataset for 300 epochs (28] 77]. We also train cone and dot product attention for 500 epochs, as observed that training for more iterations improves performance.
how they are trained and the distillation token. We
daptive Input Representations for Transformers. Adaptive input representations were introduced by Baevski and Auli for transformers in language modeling tasks [3]. In adaptive inputs, rarer tokens use lower dimensional embeddings, which serves as a form of regularization. We use the fairseq transformer ltt___m_wiiii language modeling dataset with block size of 512 tokens 220 [8]. We also test the same architecture without adaptive inputs, which has 520M parameters. This version converges in significantly fewer iterations, allowing to train such a large model.
wiki103 103 architecture (246.9M parameters) and train models on the WikiText- 103
Diffusion Transformers. Diffusion transformers (DiT) are diffusion models that replace the U-Net backbone with transformer 22 DiTs operate on latent space patches, so we expect there to be less hierarchical information vs. taking image patches. We use DiT's to test cone attention when the data is less hicrarchical. We train DiT-B/4 models with 130M parameters on ImageNet-IK -K [22 7
7

---

Table 1:
Performance of various attention methods across models and tasks. iidicatess the model failed to converge or ran into NaN errors. 1 indicates higher is better, and 4 indicates lower is better, Cone attention methods (i) generally outperform dot product attention and other baselines. Model default refers to a model's default attention method, which is dot product attention except for GATs.
DeiT-Ti Imagenet Top-1 5 (Acc.
DeïT-Ti Imagenet
NMT IWLST
GAT Cora/ / PPI (Acc. 1)
Method
Top-1 / (Acc. f) 500 Epochs
(BLEU f)
300 Epochs
Model Default
34.56
72.05 /91.17
73.65 / 91.99
0.831 / 0.977
Dot Product
34.56
72.05 191.17
73.65 / 99.99
0.834 0.985
Penumbral
35.56
72.67 191.12
74.34 92.38
0.835 0.990
Umbralt
35.07
73.14/ /9182
74.46 74.56/92.54
0,836 0.989
dg H g du Hyperboloid
32.54
49.19* 74.68
0.834 0.987
33.80
0.10 0.45*
0.13* 0.989
Laplacian Kernel
34.68
71.25 / 90.84
0.823 0.986
Adaptive Inputs WikiText- 103
Without Adaptive Inputs
DīT-B/4 @ 400K Steps (FID-50K !)
Method
0/ 480 Context Window (Ppl. L)
Dot Product
20.86 / 19.22
26.62 24.73
68.9
Penumbral
20.72/ / 19.01
26.44 / 24.31
67.7
Umbral
21.16 / 19.59
27.82 / 26.73
67.6
Results
Table summarizes the performance of cone attention across various models and tasks. Both penumbral and umbral attention significantly outperform baselines on the NMT IWSLT and DeïT-Ti Imagenet tasks. For DeiT-Ti, umbral attention achieves 73. 14% top-1 and 91.82% top-5 accuracy
at 300 epochs and 74.46% top-1 and 92.54% top-5 accuracy at 500 epochs, which matches distilled DeiT-Ti with more parameters (74.5% 91.9%) (28]. On the GAT tasks, cone attention again outperforms baselines and achieves Graph Convolutional Network-level performance on the PPI dataset. Interestingly, almost all baselines, including dot product attention, outperform the imncatenation-based attention in the original GAT paper. The two GAT tasks also reveal some interesting differences between penumbral and umbral attention. The Cora citation dataset is more tree-like with average of edges per node and chronological dependencies between nodes (a paper cannot cite a newer paper), while the PPI dataset is closer to clustered graph with 14.3 edges per node [32 13 17). As noted in (37], umbral cones appear to be better for strict hierarchies, while penumbral cones capture more complicated relationships such as those in the PPI dataset.
The adaptive inputs method regularizes models by reducing d for words. We expect rarer words to be closer to the leaves of word hierarchy, so the adaptive inputs method also acts as a hierarchical prior on the data [4.. We use adaptive inputs to test cone attention when combined with other hierarchical priors. Penumbral attention outperforms dot product attention with or without adaptive inputs, but the gap between the two methods is larger without adaptive inputs. On the other hand, umbral attention performs worse than dot product attention and penumbral attention on the WikiText 103 task, with without adaptive inputs. Since umbral cones not geodesically convex, which means that the shortest path between two points in a cone may not necessarily lie entirely in that cone, we suspect that convexity is important for the WikiText- 103 language modeling task. In the DiT model, patches are taken at the latent space-level, which we suspect gives less hierarchical information than when patches are taken from an input image (22). Here, cone attention still outperforms dot product attention, but to lesser degree. This mirrors our expectation that the DiT model does not benefit as much from hierarchical attention, and verifies that in such cases, using cone attention does not hurt performance. Interestingly, umbral cones slightly outperform penumbral cones in the DiT- B/4 task, which may be result of patches being over the latent space, and not the input image.
5.1 Effect of Mapping Functions
Table 2 compares how de and & compare to the exponential map and psueodpolar map (section 3.2 on the NMT IWSLT task. Here, we use the exponential map at the origin of Hd, O (0. 0, ..., To compare Expo (v) zi, we first take (xpo and then project the point onto the boundary of
8

---

Table 2: Comparison of various mappings from Euclidean space to hyperbolic models for the NMT IWSLT task (BLEU scores, higher better). The exponential map generally performs worse than yb and the pseudopolar map. While z also performs worse than the psuedopolar map, table indicates that the pseudopolar map is less numerically stable than
Method
Expo (v) - Hd
- Hd
V + H
Pseudopolar + Hyperboloid
Penumbral
35.12
35.56
Umbral
34.63
35.07
da
30.59
32.54
33.80
Performance at Low q, Dimensions for NMT IWSLT
36
34
BLE BLEU
30
(30.2M params)
4 (30.3M)
(30.5M)
(30.6M)
12 (30.9M)
16 (31.2M)
Dimension
Dot Product
Umbral
Penumbral
Dot Product d=128 (39. SM) Umbral d= 128 (39.5M)
Penumbral d=128 (39.5M)
Figure 4: Performance of cone and dot product attention at low dimensions on NMT IWSLT. The base architecture uses 128 dimensions, and cone attention is able to match dot product attention with only dimensions. For this model, this allows us to use 21% % fewer parameters to reach performance parity, indicating that cone attention more efficiently captures hierarchical information.
if Expo (v) is in L. In the infinite penumbral setting, this corresponds to taking Ex)o = min (v)a, h). and & generally perform much better than taking the exponential map the origin O, which suggests that t and have better optimization properties. For du attention, Gulcehre al.'s pseudopolar map slightly outperforms Ç. However, table I indicates that outside of this specific task, using the psuedopolar map and Hyperboloid is less numerically stable.
5.2 Attention Efficiency and Model Size
Figure 4 shows the performance of cone attention vs. dot product attention at low token (q, embedding dimensions (d from before) on the NMT IWSLT task. Both umbral and penumbral attention are able to achieve significantly better performance than dot product attention in this regime, matching dot product attention at d 128 with only d 16. For this model, using 16 dimensions reduces the number of parameters from 39.5M to 31.2M 1..M. Table shows the performance of DeiT-Ti at d 16. Here, penumbral attention d 16 is able outperform dot product attention at d 64, again indicating that cone attention able to more efficiently capture hierarchies.
5.3 Sensitivity to Initialization
To test cone attention's sensitivity to initialization, we trained seeds for the IWSLT De2En task for cone attention and dot
Hyperbolic embeddings known to be sensitive initialization 10 37
Table 3: Performance of DeiT-Ti at 64 (default) and 16 dimensions, 300 epochs.
Method
d= 64
d 16
Dot Product 72.05 /91.17
71.29 90.54
Penumbral
72.67 /91.12
72.25 91.20
Umbral
73.14/91.82
71.67 990711

---

Penumbral
Umioral
Dot Product
Kh
sauo
do
Se
letee
Ieten
oederken
cedanten
gedanien
ventegi
venragg
varate
100
100
vige
de
mem
meri
Bonge
Figure 5: Attention heatmaps from an attention head a trained IWSLT De2En translation model for the tokenized validation sequence "Und ich meine, "Zieh sie aus! Mach dir keine gedanken, verstehst du.', which translates to "I'm like, "Get it off! Don't 't worry about it. you know." The cone attention heatmaps (left and center) have clear distinction between the two parts of the sequence separated by "!", whereas the dot product heatmap (right) does not have a clear separation.
product attention. Dot product had an average BLEU score of 34.59 and standard deviation 0. 0.12, penumbral achieved 35.41 0.14, and umbral achieved 35.03 = 0.30. There was one outlier in the unbral attention trials, and with the outlier removed umbral achieved 1 The cone attention methods appear have slightly higher variance than dot product attention, but significantly so.
5.4
Attention Heatmaps
natural question arises about what cone attention methods actually learn. Figure showws heatmaps from attention head in trained IWSLT De2En translation model. The heatmaps for penumbral and umbral attention show clearer separation than the dot product attention heatmap. Furthermore, the attention methods happens the exclamation mark, natural segmentation of the sequence into two parts. Intuitively, cone attention can be seen as an attention method that satisfies certain "logical constraints," such "if y and 4 x, then relations between attention scores. For example. if K (æx, y) and K Ky, 2) are both high, then K K(x, =) should also be relatively high in cone attention. In dot product attention, this is not guaranteed. If (1, 1, O, 0), y (10, 10, 10, 10), and z (0, 0, 1, 1). then (x, y) (v, 2) = 20, but (x, 2) 0. We suspect this a reason why cone attention methods show better separation than dot product attention, which can aid performance.
separation in the
x., which leads to
6 Conclusion
We introduce cone attention, a hierarchy-aware method for calculating attention. Cone attention relies on entailment cones and the geometric properties of hyperbolic space to capture complex structural patterns that dot product attention does not explicitly model. We test cone attention in variety of attention networks ranging from few thousand to few hundred million parameters, and achieve consistent performance improvements in NLP, vision, and graph prediction tasks over dot product attention and baselines. Cone attention also matches dot product attention with significantly fewer embedding dimensions, opening the potential for smaller models. These results suggest that cone attention is effective way encode hierarchical relationships in attention, and can potentially improve task-level performance a wide variety of models and task domains.
Future Work. It remains to be seen how cone attention scales to very large models. Beyond this, 10 and 337 suggest that hyperbolic embeddings are sensitive to initialization, which implies that different transformer weight initializations may affect cone attention.
Acknowledgements
This work was supported by NSF Award IIS-2008 102. Compute resources were provided by the Comell G2 Cluster.
10

---

References
[I] NVIDIA A100 Tensor Core GPU Architecture. Technical report, NVIDIA.
[2] James W. Anderson. Hyperbolic geometry. Springer, 2007.
[3] Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint arXiv: 1809. 1etp::089..0000 10853, 2018.
[4] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly 2014.
learning to align and translate. arXiv preprint
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, 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.
[6] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
[7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hicrarchical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pages 248-255, 2009. doi: 10.1 10.1109/CVPR22209550066
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint M:Xiv:2010.11990 2020.
[9] Marcello Federico, Sebastian Stüker, and François Yvon, editors. Proceedings of the 11th Inter- national Workshop on Spoken Language Translation: Evaluation Campaign, Lake Tahoe, Cali- fornia, December 4-5 2014. httpttts///caaantooooll
org/2014, wwwsstteelllttete
0
[10] Octavian Ganea, Gary Bécigneul, and Thomas Hofmann. Hyperbolic entailment cones for learning hierarchical embeddings. In International Conference on Machine Learning, pages 1656-1655. PMLR, 2018.
[11] Caglar Gulcehre, Misha Denil, Mateusz Malinowski, Ali Razavi, Razvan Pascanu, Karl Moritz Hermann, Peter Battaglia, Victor Bapst, David Raposo, Adam Santoro, et al. Hyperbolic attention networks. arXiv preprint 1eXivv18055 1805.09786, 2018.
[12] Maosheng Guo, Yu Zhang, and Ting Liu. Gaussian transformer: A lightweight approach for natural language inference. Proceedings of the AAAI Conference Artificial Intelligence, 33 (01):6489-64)) 6496. Jul. 2019. doi: w0..660/a.ii.v w33i:aaa.iii33333....3333113333. 333000.333011449, URL httph: aaai org/inder: php/ TAAT/arttile/viee///55
[13] Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
[14] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pages 2961-2969, 2017.
[15] Minh-Thang Luong, Hicu Pham, and Christopher Manning. Effective approaches to attention- based neural machine translation. arXiv preprint 15X8::150..0422 2015.
[16] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209.1065 10655, 2022.
[17] Andrew Kachites McCallum, Kamal Nigam, Jason Rennie, and Kristie Seymore. Automating the construction of internet portals with machine learning. Information Retrieval, 3(2) Jul 2000. ISSN 1573-7659, doi: 10.3//// 10099/// (00995381990)) 19988. URL httpt:://ot.ii.... //tp::///ooooo///0///0/////////0 1023 A 1009953844888

---

[18] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, arXiv preprint arXiv: 07843, 2016.
[19] Hee Oh. Euclidean traveller in hyperbolic worlds. arXiv preprint arXiv:2209.01.000 2022.
[20] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint WrXiv:190..01036 2019.
[21] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- dighpperformanc deep learning library. Advances in neural information processing systems, 32, 2019.
[22] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09778, 2022.
[23] Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.021433 2021.
[24] Wei Peng. Tuomas Varanka, Abdelrahman Mostafa, Henglin Shi, and Guoying Zhao. Hyper- bolic deep neural networks: A survey, IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(12)::0023310044 2021.
[25] Frederic Sala, Chris De Sa, Albert Gu, and Christopher Ré. Representation tradeoff's for hyperbolic embeddings. In International conference on machine learning, pages 4460-4469. PMLR, 2018.
[26] Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng, Synthesizer: Rethinking self-attention for transformer models. In International conference machine learning, pages 10183-10192. PMLR, 2021.
[27] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.
[28] 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.
[29] Yao-Hung Hubert Tsai, Shaojie Bai. Makoto Yamada, Louis-Philippe Morency, and Rus- lan Salakhutdinov. Transformer dissection: An unified understanding for transformer's at- tention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Meth- ods Natural Language Processing and the 9th International Joint Conference on Natu- ral Language Processing (EMNLP-IJCNLL)), pages 4344 4353, Hong Kong, China, Novem- ber 2019. Association for Computational Linguistics. doi: 10.18653/1/DD1911443 1443. URL org/ D
[30] Albert Tseng, Jennifer J. Sun, and Yisong Yue. Automatic synthesis of diverse weak supervision sources for behavior analysis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 2211-2220, June 2022.
[31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[32] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 17 2017.
[33] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 2704422006.04768, 2020.
[34] Hanqi Yan, Lin Gui, and Yulan He. Hierarchical interpretation of neural text classification. Computational Linguistics, 1804),987-1020, 2022.
12

---

[35] Tao Yu and Christopher M De Sa. Numerically accurate hyperbolic embeddings using tiling- based models. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Ban Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL /ttp (4411118822244466cccccC040111111 pdf
cw/paper.p........ fatrrrffleeeeeeeee fftess/ppere////////m//paa /aper/ cttpppperfalleee/e//eeeeeeeeesee
[36] Tao Yu and Christopher M De Sa. Representing hyperbolic space accurately using multi- component floats. Advances in Neural Information Processing Systems, 34:15570-15581, 2021.
[37] Tao Yu, Toni J.B. Liu, Albert Tseng, and Christopher De Sa. Shadow cones: Unveiling partial orders hyperbolic space. arXiv preprint, 2023.
[38] Lin Zheng, Jianbo Yuan, Chong Wang, and Lingpeng Kong. Efficient attention via control variates. arXiv preprint arXiv:2302.0455: 2023.
13


Reviews:
[{"review_id": "eUXCmiywvF", "reviewer": "Reviewer_Rx82", "summary": "This paper proposes replacing the traditional softmax attention kernel with one that tries to captures the hierarchical relationship between two vectors by mapping them to hyperbolic space. ", "questions": "There are a few details of the presentation that I couldn't quite understand.\n- On line 191, you introduce $\\xi$ as a way of satisfying the requirements of the penumbral setting, but in Table 2 $\\xi$ only appears for the *Umbral* method. Likewise on line 190 $\\psi$ is associated with the umbral construction, but in Table 2 $\\psi$ is paired with *Penumbral*. Could you explain?\n- Figure 2 showcases both finite setting and infinite setting entailment cones. What about the actual attention method and Definitions 1 & 2? Am I correct in understanding that only the infinite setting is being considered from that point of the paper onward?\n- On lines 158/159, a user-defined quantity $r$ is introduced, which appears to be a hyperparameter. Is that right? I found no mention in the paper of how this hyperparameter is set or tuned. Could you say more on this, and whether choice of hyperparameter might have any impact on the relative performance of the umbral and penumbral methods?\n\nThe best-performing penumbral attention method appears to have a piece-wise definition. One regime is when an entailment cone exists that contains both u and v with respect to a fixed light source, and the other regime is when the cone does not exist. Have you checked which regime is actually used by your trained models? I could see several possibilities here:\n- Perhaps after training attention mostly falls into equation (5), and the piecewise definition is only needed to avoid issues during optimization\n- Or maybe a trained model mostly falls under equation (6)? In that case, wouldn't it be possible to dispense with the piecewise definition and use equation (6) as the entirety of the attention mechanism? Have you tried this?\n- Or perhaps the boundary criterion (7) is being used as a critical source of nonlinearity when comparing attention queries and keys?\nKnowing this would help me make sense of whether the actual operation of the attention method is in line with the proposed intuition of respecting hierarchy.\n\nAre there any plans to release code? The paper is unfortunately not very accessible to neural network practitioners unfamiliar with hyperbolic embeddings and related literature, and having drop-in code would be helpful to allow the broader community to try out the proposed approach.\n", "limitations": "No concerns regarding limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "The idea of using hyperbolic attention cones to define hierarchy-aware attention is (to my knowledge) novel in the literature. The specific choices of cone construction and mapping into hyperbolic space is also a contribution of the paper.\n\nThe numerical results appear quite strong, where the proposed attention method outperforms dot product attention in all settings, and retains better performance at lower dimensions.", "weaknesses": "While the paper validates the proposed attention method in terms of numerical results, one thing that's not studied is whether the behavior of the trained models actually matches the intuition that gave rise to the method, namely the use of hierarchy in the attention process. It would help complete the story if there were some results of this nature, such as attention heatmaps that exhibit hierarchical behavior, or a demonstration of how learned vectors induced by the attention mechanism correspond to meaningful tree structures. (It is unfortunately the case with neural networks that even when an idea works well, the *reason* it works may have little to nothing to do with the original motivation)\n\nThere are also certain aspects of the presentation that could be made more clear to the reader (or at least to me, see my questions below).\n\nIt might also be helpful to have a comparison with hyperbolic distance attention (Gulcehre et al) in a results table, to clarify whether the benefits come just from the use of hyperbolic geometry, or whether the proposed contributions related to entailment cones are in fact critical to achieving the reported results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- Hierarchical behavior of trained models\n    - Please see the general response.\n- Comparison with hyperbolic distance attention, hyperbolic geometry, and entailment cones\n    - We suspect that most of the performance difference between the two methods comes from our use of entailment cones vs. Gulcehre et al.’s use of hyperbolic distance. Both hyperbolic distance attention with the hyperboloid map and the $\\xi$ map failed to converge on multiple tasks, which suggests that hyperbolic distance attention is less robust during optimization. The hyperbolic distance between two points grows exponentially with respect to the distance between the points and the hyperbolic origin, which may be the reason for these instabilities. Table 2 in the paper shows a comparison between the two attention methods under various mappings into hyperbolic space. The mapping and hyperbolic model used can also influence the end result, so this may be another factor in the performance difference as well.\n- $\\xi$ and $\\psi$\n    - Table 2 in the paper should read $\\xi$ for penumbral and $\\psi$ for umbral. This is a typo and we will fix it in future versions.\n- Infinite setting cones in the paper?\n    - Yes, from that point on everything uses infinite setting cones.\n- How to tune or set hyperparameters?\n    - We used h=1 for penumbral attention and r=0.1 for umbral attention. For some of the experiments, we experimented with different values of h and r but did not see any significant differences between runs. Using a higher h for penumbral attention reduces the chance that a random initialization falls into equation (6), but this may not actually matter much for a converged model. Using a lower h for a fixed $\\gamma$ (temperature) increases the lowest similarity score achievable by penumbral attention with equation (5).\n- Piecewise definition of penumbral attention\n    - With an appropriately set combination of h and $\\gamma$ (i.e. h$\\gamma$ is not too low), points fall into both equations (5) and (6). We have not tried using only equation (6) as the attention kernel, which is an interesting idea. The boundary condition in equation (7) is to ensure piecewise definition is continuous. Using a different definition of equation (6) would also be interesting. We leave these directions for future work.\n"}, {"review_id": "E8zOv6hyRe", "reviewer": "Reviewer_AaRb", "summary": "This paper introduces hyperbolic shadow cone similarity and induced Hierarchy Aware Attention. The main contribution lies in their proposal to utilize hyperbolic shadow cones (specifically, umbral cone and penumbral cone) as an entailment setting to induce pairwise similarity scores denoted as $K(\\cdot,\\cdot)$ in hyperbolic space $\\mathbb{H}^d$. By replacing the euclidean inner-product in Transformers with this hyperbolic similarity score based on shadow cones, the model achieves comparable performance across various tasks such as graph prediction, natural language processing (NLP), and vision tasks, while significantly reducing the number of parameters required.", "questions": "As the word ``Hierarchical`` is in the title, can the authors clarify on why with this similarity mechanism, the model becomes ``Hierarchical`` some qualitative results on how the hierarchy is built?\n\nWhat is the cost of using hyperbolic cones? More specifically, How much slower cone attention is compared with euclidean dot-product in different tasks over graph prediction, NLP, and vision domain?", "limitations": "The research is about fundamental mechanisms of machine learning models, so it does not have potential negative societal impact.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The concept of utilizing shadow-cone based similarity and extending it to attention is innovative.\n\nApplying this shadow-cone based hyperbolic similarity score to Transformers results in a notable reduction in the number of parameters required, while maintaining performance comparable to the Euclidean baselines.\n\nThe experiments conducted in this study are extensive, covering graph prediction, natural language processing (NLP), and vision tasks.", "weaknesses": "While entailment cones do enforce the encoding of hierarchical information, the learning scheme employed in this paper differs fundamentally from that of [12]. In [12], a contrastive loss embeds the hierarchy of data points, while an entailment loss ensures the preservation of partial order relationships. The extent to which this method is \"hierarchical\" is neither explicitly explored within the method nor through experiments.\n\nThe use of \"Attention\" in the title may be a bit over-claiming, as this paper primarily contributes to defining pairwise similarity within hyperbolic cones, rather than addressing other aspects of the attention mechanism. It is worth noting that in the closely related work [13] proposed an attention framework that covers various aspects, including the construction of K, Q, and V in hyperbolic space, as well as their operations.\n\nThe speed of calculating similarity is a crucial aspect of the attention mechanism, but it is only briefly mentioned in Section 3.1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- Why is the model hierarchical?\n    - We apologize if our use of “hierarchical” caused confusion and are open to using other terms in the final paper. Please see the general response for more information about hierarchies and the model.\n- The use of \"Attention\" in the title may be a bit over-claiming, as this paper primarily contributes to defining pairwise similarity within hyperbolic cones, rather than addressing other aspects of the attention mechanism. It is worth noting that in the closely related work [13] proposed an attention framework that covers various aspects, including the construction of K, Q, and V in hyperbolic space, as well as their operations.\n    - Thank you for your feedback, we are open to using other terms in the paper. However, we would like to point out that we do define multiple mappings ($\\xi, \\psi$) into hyperbolic space that are critical to our method. We did not explore hyperbolic aggregation as Gulcehre et al. did, since hyperbolic aggregation is independent of how similarity between data points is measured. However, our method is fully compatible with hyperbolic aggregation, which we leave for future work. \n"}, {"review_id": "fW9QWueygX", "reviewer": "Reviewer_hcYy", "summary": "The paper explores a variant of hyperbolic attention to induce a hierarchical inductive bias by exploiting the corollaries of points in a hyperbolic space and nodes in a tree. The paper integrates shadow cone construction for entailment cones to define partial order relations between points (the partial order relation can be interpreted as hierarchical relations in a tree structure) replacing a prior method from Ganea et a. 2018. The authors then model the similarity function of attention over any two vectors roughly based on the distance of the lowest common ancestor of them based on the partial ordering defined by shadow cone constructions. Intuitively this would lead to giving higher attention to representations that have close-by levels of hierarchy as per their locations in the hyperbolic space.", "questions": "Will the code be open-sourced? Also, how is the empirical time/memory trade-off?", "limitations": "Limitations seem fairly addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Reasonable contribution in the space of hyperbolic embeddings/hyperbolic attention. Can be of interest to the community in that space. \n* Fair improvement in different domains and over prior works like Gulcehre et al. 2018. \n* Hyperbolic embeddings tend to be more effective with less dimensions. Improving in this space can be meaningful for development of smaller more effecient models - although there is a question if there is a diminishing return as we increase the dimensions.", "weaknesses": "* The limitation says that the model can be sensitive to initialization. While I am not going to take that itself as a negative point but this fact may put some concern given that the models are only trained once per data. However, since the improvement is quite consistent across multiple tasks and datasets (even if run only once), I think the concern is mitigated a bit. One potential way to show some more insights related to sensitivity to initialization within resource constraints - could be to try running multiple seeds in only a subset of tasks (maybe even just one of them) and show the comparisons with error bars.  \n\n(I have read the rebuttal)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- “The limitation says that the model can be sensitive to initialization. While I am not going to take that itself as a negative point but this fact may put some concern given that the models are only trained once per data. However, since the improvement is quite consistent across multiple tasks and datasets (even if run only once), I think the concern is mitigated a bit. One potential way to show some more insights related to sensitivity to initialization within resource constraints - could be to try running multiple seeds in only a subset of tasks (maybe even just one of them) and show the comparisons with error bars.” \n    - Hyperbolic embeddings have been noted to be sensitive to initialization in prior works, such as Ganea et al.’s entailment cones. As you noted, the empirical results suggest that cone attention may be relatively robust to initialization. We suspect this is because in most of our tasks, cone attention is used within a much larger model (e.g. a transformer) as opposed to traditional graph embedding tasks, where embeddings are directly optimized with few additional parameters.\n    - We ran an experiment where we trained the IWSLT De2En model with 5 different seeds to measure sensitivity to initialization. Below are the mean and standard deviation of the 5 trials. There was one outlier (34.51 BLEU) for umbral attention. The cone attention methods have slightly higher variance between runs than dot product attention, but not significantly so.\n\nMultiseed experiments\n- Dot product: mean 34.59 std 0.12 \n- Penumbral: mean 35.41 std 0.14 \n- Umbral with outlier: mean 35.03 std 0.30\n- Umbral without outlier: mean 35.16 std 0.09\n"}, {"review_id": "PEwOFm9S23", "reviewer": "Reviewer_F7eB", "summary": "This paper introduces a simple drop-and-replace operation for dot-product attention, based on cones. The new operation is designed to model hierarchical information better than standard attention. The proposed method outperforms standard attention on various tasks. I will say upfront that I am not very familiar with the mathematical background used in this work, so my evaluation is mostly on the general framework and experiments.\n\n", "questions": "- What are the model parameters? It would be nice to see a comparison between the proposed method and standard attention.\n", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- An interesting idea to present alternatives to the dot-product attention. Most existing works in this field target efficiency and simplicity, while this work is actually slower to run, but seems to improve performance. It also allows using smaller models and reaching similar performance, which can be thought of as a more efficient variant after all.\n- Extensive experiments, showing that the proposed approach outperforms dot-product attention.\n\n", "weaknesses": "- The paper title is “Hierarchy Aware Attention”, and the mathematical model, to the best of my understanding, indeed promotes hierarchy. However, there is no explicit discussion about the desired hierarchy to be learned. Are the authors expecting the model, e.g., to learn the syntactic trees in sentences? The event graph in an image? Importantly, none of the experiments quantify this learned hierarchy either, which is a bit disappointing. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- Hierarchical nature of model.\n    - Please see the general response. Beyond that, discovering hierarchical structures in an unsupervised or semi-supervised manner is an active research area, and it remains to be seen if cone attention can be applied there. Using cone attention for structure discovery is an interesting idea, and we leave that for future work.\n- “What are the model parameters? It would be nice to see a comparison between the proposed method and standard attention\n    - Figure 4 in the attached PDF shows the cosine between parameter tensors for 15 weight layers between dot product, penumbral, and umbral attention for the same seed. The cosines are all somewhat low, which suggests the parameter tensors are somewhat orthogonal and different. Furthermore, taking a trained dot product attention model and swapping out the attention calculation with cone attention (and vice versa) results in poor performance, which suggests the models are learning different things. For example, using dot product attention on a penumbral model gives 4.75 BLEU, and the other way around gives 3.62 BLEU\n"}]

Author rebuttal:
Dear reviewers, thank you for your detailed and thorough reviews. As reviewers Rx82 and AaRb noted, our use of entailment cones to define an attention mapping is “novel” and “innovative.” All reviewers noted that throughout our “extensive” suite of experiments across graph prediction, NLP, and vision tasks, cone attention demonstrated empirical improvements over dot product and prior hyperbolic attention works. Furthermore, reviewers hcYy and F7eB noted that cone attention matches dot product attention with fewer parameters, allowing for “smaller more efficient models.” Below, we have written responses to common questions about our paper. Namely, we have included a clarification about how hierarchies are used by cone attention. For responses to individual questions, please see our individual responses.

Beyond this, after the submission deadline and during the review period, we discovered an interesting result with our DeiT-Ti baseline. Unlike previously reported in the original DeiT-Ti paper, we found that DeiT-Ti did not plateau, and instead kept getting better with more epochs. In the version of the paper submitted for review, we reported the DeiT-Ti numbers at 300 epochs from the original paper, and reported penumbral attention numbers at 500 epochs, as we found that penumbral attention got better with more epochs. For transparency, we have compiled numbers from our DeiT-Ti experiments below for dot product and cone attention methods at 300 and 500 epochs. We note that penumbral and umbral attention both still outperform dot product attention at both 300 and 500 epochs, so our final conclusions are still the same. We will update Table 1 with the numbers from both 300 and 500 epochs. We also re-ran the experiments in Table 3, and the updated results are below. All methods do better than previously reported, and penumbral attention at d=16 now actually outperforms dot product attention at d=64. At d=16, penumbral and umbral attention still outperform dot product attention, and our conclusions do not change.

We are letting you know this updated information now in case this affects your reviews. We stress that cone attention still outperforms dot product attention for the DeiT-Ti task, and that our conclusions are still the same.

DeiT-Ti Table 1 update
- 300 epochs (top 1 acc / top 5 acc)
    - Dot Product: 72.05 / 91.17
    - Penumbral: 72.67 / 91.12
    - Umbral: **73.14 / 91.82**
- 500 epochs (top 1 acc / top 5 acc)
    - Dot Product: 73.65 / 91.99
    - Penumbral: **74.34 / 92.38**
    - Umbral: Still running, will be updated in a few days.



DeiT-Ti Table 3 update
- d = 64 (top 1 acc / top 5 acc)
    - Dot Product: 72.05 / 91.17
    - Penumbral: 72.67 / 91.12
    - Umbral: **73.14 / 91.82**
- d = 16 (top 1 acc / top 5 acc)
    - Dot Product: 71.292 / 90.542  
    - Penumbral: **72.250 / 91.190**
    - Umbral: 71.668 / 90.714

## Responses to Common Questions

- Hierarchical nature of the model: 
    - Our method relies on the implicit hierarchy defined by hyperbolic entailment cones to measure the similarity between two points. It does not take into account explicit ground truth hierarchies between points, such as a manually labeled syntax tree. Furthermore, such ground truth information was not available for the tasks we tested on. Essentially, we are learning a latent hierarchy in which the observed data points are only a subset of the nodes. 
    - Figures 1-3 in the attached PDF show heatmaps from an attention head in a trained IWSLT De2En translation model for the tokenized sequence from the validation set ['und', 'ich', 'meine', ',', '&quot;', 'zie@@', 'h', 'sie', 'aus', '!', 'mach', 'dir', 'keine', 'gedanken', ',', 'verste@@', 'h@@', 'st', 'du', '.', '
'], which translates to ['i', '&apos;m', 'like', ',', '&quot;', 'get', 'it', 'off', '!', 'don', '&apos;t', 'worry', 'about', 'it', ',', 'you', 'know', '.', '
']. Figure 1 is from penumbral attention, figure 2 is from umbral attention, and figure 3 is from dot product attention. The cone attention heatmaps have a clear separation between the two sentences in the sequence (separated by “!”), whereas the dot product heatmap does not have a clear separation. This separation can be considered a hierarchy between the two parts of the sequence. 
    - Furthermore, cone attention can be seen as an attention formulation that satisfies certain “logical constraints,” such as “if $z \prec y$ and $y \prec z$, then $z \prec x$," which leads to relations between attention scores. For example, if $K(x, y)$ and $K(y, z)$ are both high, then $K(x, z)$ should also be relatively high in cone attention. In dot product attention, this is not guaranteed. For example, if a = [1, 1, 0, 0], b = [10, 10, 10, 10], and c = [0, 0, 1, 1], <a, b> = <b, c> = 20 but <a, c> = 0. We suspect this is a reason why cone attention methods show better separation than dot product attention, which can aid performance.
- Will the code be open sourced?
    - Yes
- What is the runtime of cone attention?
    - Table 5 in the supplementary section has a simple benchmark on most of the tasks we ran experiments on. In general, using torch compile / writing an optimized kernel significantly reduces the runtime gap between cone attention and dot product attention. 



Meta-review: Reviewers seem to be happy with this paper and there seems to be healthy discussions. Authors make an effort to run new experiments and improve the comprehensiveness of the work. This work looks kind of interesting though and perhaps a little novel. I recommend acceptance.
Decision: Accept (poster)