SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Albert Tseng
albert@cs.cornell.edu
Tao Yu
tyu@cs.cornell.edu
Toni J.B. Liu
jl3499@cornell.edu
Christopher De Sa
cdesa@cs.cornell.edu
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 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.
In recent years, attention networks have achieved highly competitive performance in a variety of settings, often outperforming highly-engineered deep neural networks [5, 8, 31]. The majority of these networks use dot product attention, which defines the similarity between two points  $u, v \in \mathbb{R}^d$  by their inner product  $u^\top v$  [31]. Although dot product attention empirically performs well, it also suffers from drawbacks that limit its ability to scale to 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, 6].
Many real world datasets exhibit complex structural patterns and relationships which may not be well captured by an inner product [30, 16]. For example, NLP tasks often contain hierarchies over tokens, and images may contain clusters over pixels [34, 14]. Motivated by this, we propose a new framework based on hyperbolic entailment cones to compute attention between sets of points [10, 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 (LCA) in the cone partial ordering, which is analogous to finding their LCA in a 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-
Figure 1: Overview of cone attention vs. dot product attention. In dot product attention (left), similarity scores are calculated with  $K = qk^{\top}$ . In cone attention (right),  $q$  and  $k$  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 are not possible with Euclidean embeddings [25]. Entailment cones, which rely on geometric cones to define partial orders between points, allow us to calculate explicit relationships between points, such as their LCA [10, 37]. To the best of our knowledge, we are the first to define a hierarchy-aware attention operator with hyperbolic entailment cones.
Functionally, cone attention is a drop-in replacement for dot product attention. We test cone attention in both "classical" attention networks and transformers, and empirically show that cone attention consistently improves end task performance across a 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, our contributions are:
- We propose cone attention, a 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_de_en and DeiT-Ti models, respectively.
- We test cone attention at low embedding dimensions and show that we can significantly reduce model size while maintaining performance relative to dot product attention. With cone attention, we can use  $21\%$  fewer parameters for the IWSLT'14 De-En NMT task.
In this section, we provide background on attention mechanisms, motivate the use of hyperbolic space to embed hierarchies, and describe our choice of entailment cones to encode partial orderings.
2.1 Attention
The attention operator has gained significant popularity as a way to model interactions between sets of tokens [31]. At its core, the attention operator  $A$  performs a "lookup" between a single query  $q_{i} \in \mathbb{R}^{d}$  and a set of keys  $k \in \mathbb{R}^{n \times d}$ , and aggregates values  $v \in \mathbb{R}^{n \times d}$  associated with the keys to "read out" a single value for  $q_{i}$ . Mathematically, this can be represented as
$$
A \left(q _ {i}, k\right) = C \sum_ {j} \left(K \left(q _ {i}, k _ {j}\right) v _ {j}\right) \tag {1}
$$
where  $C \sum_{j} K(q_{i}, k_{j}) = 1$ . In "traditional" dot product attention, which has generally superseded "older" attention methods such as Additive Attention and Multiplicative Attention [4, 15],
$$
K \left(q _ {i}, k _ {j}\right) = \exp \left(\frac {q _ {i} k _ {j}}{\sqrt {d}}\right) \quad C = \frac {1}{\sum_ {j} K \left(q _ {i} , k _ {j}\right)} \quad A \left(q _ {i}, k\right) = \operatorname {s o f t m a x} \left(\frac {q _ {i} k ^ {\top}}{\sqrt {d}}\right) v \tag {2}
$$
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 [38]. Some recent works, such as EVA attention, parameterize these approximations and effectively get a larger class of dot-product-esque attention methods [38]. Beyond this, Tsai et al. [29] replace  $K$  with compositions of classical kernels. Others extend dot product attention, such as by controlling the "width" of attention with a learned Gaussian distribution or by using an exponential moving average on inputs, although extensions do not usually depend on  $K$  [12, 16]. Closer to our work, Gulcehre et al. [11] introduced hyperbolic distance attention, which defines  $K(q_{i},k_{i}) = \exp (-\beta d_{\mathbb{H}}(q_{i},k_{i}) - c)$  where  $d_{\mathbb{H}}$  is the hyperbolic distance,  $\beta \in \mathbb{R}^{+}$ , and  $c \in \mathbb{R}$ .  $K$  can be interpreted as an analog of the distance from  $q_{i}$  to  $k_{i}$  on a 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 are actually needed is outside the scope of this work, and we compare cone attention accordingly.
2.2 Hyperbolic Space
$d$ -dimensional Hyperbolic space, denoted  $\mathbb{H}_d$ , is a simply connected Riemannian manifold with constant negative sectional curvature [2]. This negative curvature results in geometric properties that makes hyperbolic space well-suited for embedding tree-like structures [25, 35]. For example, the volume of a hyperbolic ball grows exponentially with respect to its radius; in a tree, the number of leaves grows exponentially with respect to depth. Furthermore,  $d_{\mathbb{H}}(u,v)\approx d_{\mathbb{H}}(u,O) + d_{\mathbb{H}}(O,v)$ , where  $O$  is the origin, which again mirrors a tree where  $d_{T}(u,v) = d_{T}(u,\mathrm{LCA}(u,v)) + d_{T}(\mathrm{LCA}(u,v),v)$ . Since hyperbolic space cannot be isometrically embedded into Euclidean space, it is usually represented on a subset of Euclidean space by a "model" of  $\mathbb{H}_d$  [11]. 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 is the manifold  $\mathsf{H}^d = (\mathcal{U}^d,g_u)$  where  $\mathcal{U}^d = \{x\in \mathbb{R}^d:x_d > 0\}$ ,  $g_{u}(x) = g_{e} / x_{d}^{2}$ , and  $g_{e}$  is the Euclidean metric [2]. In the Poincaré half-space, "special" points and curves have particularly nice Euclidean forms. Ideal points, or points at infinity, are the points where  $x_{d} = 0$  (the "x-axis") and the single point  $x_{d} = \infty$  at which all lines orthogonal to the  $x$ -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  $x$ -axis or a horizontal hyperplane when the ideal point is  $x_{d} = \infty$ .
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 \in$  the cone of  $v$ , then  $v \prec u$ . Ganea's cones (figure 2 right) are defined on the Poincaré ball by a radial angle function  $\psi(r)$ , with an  $\epsilon$ -ball around the origin where cones are undefined [2, 10]. This makes learning complicated models with Ganea's cones difficult, as optimization on the Poincaré ball is nontrivial and the  $\epsilon$ -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 2 left quadrant). In the penumbral setting,  $S$  is a ball of fixed radius and points are points. The shadow and cone of  $u$  are both the region enclosed by geodesics through  $u$  tangent to  $S$ . In the umbral setting,  $S$  is instead a 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 a subset of the shadow of  $u$  (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-setting Shadow Cones. For penumbral cones, when  $S$  is a ball with center at  $x_{d} = \infty$ ,  $S$ 's boundary is a horosphere of user-defined height  $h$ . Here, all shadows are defined by intersections of
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  $\mathsf{H}^2$ . Shadows are represented by shaded regions, and cones are enclosed in green. In all figures,  $u \prec v$  since  $v$  is in the cone of  $u$ . (Right) Ganea's entailment cones, as defined on the Poincaré ball. Note the  $\epsilon$ -ball where cones are not defined, which makes optimization nontrivial. Figure from [10].
Euclidean semicircles 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  $S$  maps to the  $\epsilon$ -ball [37]. For umbral cones, when  $S$  is  $x_{d} = \infty$ , shadows are regions bounded by Euclidean lines perpendicular to the  $x$ -axis.
Unlike Ganea'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 a 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 of two nodes  $u, v$  in a directed acyclic graph is the lowest (deepest in hierarchy) node that is an ancestor of both  $u$  and  $v$ . 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  $x, y$  in a partial ordering, denoted  $\sup(x, y)$ , is the point  $p$  such that  $p \preceq x, y$  and  $\forall q$  where  $q \preceq x, y, q \preceq p$ . The key difference is that all other upper bounds must precede the least upper bound, while not all ancestors of  $u$  and  $v$  must also be ancestors of  $\operatorname{LCA}(u, v)$ . Furthermore,  $p$  may not actually exist.
3 Hierarchy Aware Attention
Here, we describe cone attention using the shadow cone construction, and discuss projection functions onto  $\mathbb{H}^d$  that allow us to use cone attention 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 cones as "Ganea's cones" and the general set of cones that captures entailment relations (e.g. Ganea's cones and shadow cones) as "entailment cones" [10, 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, v \in \mathsf{H}^d$  by their LCA in some latent tree  $T$ , which is analogous to finding their LCA, denoted  $\sup_2(u,v)$ , in the partial ordering defined by entailment cones. Formally,
Figure 3: In this example using penumbral cones in  $\mathsf{H}^d$ ,  $\sup_2(u,v)$  is the lowest common ancestor of  $u$  and  $v$ . The red region is the set of points  $P$  s.t.  $p \in P \preceq u, v$ . Of these,  $\sup_2(u,v)$  is the lowest point whose cone (light gray) contains both  $u$  and  $v$ . Here,  $\mathcal{S}$  is the root of all hierarchies, and points closer to  $x_d = 0$  are closer to the "leaves" of hierarchies. If  $d = 2$ , then  $\sup_2(u,v)$  is also the least upper bound of  $u$  and  $v$  in the partial ordering defined by entailment cones.
$$
\sup  _ {2} (u, v) = r \left(\underset {C: u, v \in C} {\arg \max } d _ {H} (\mathcal {S}, r (C))\right) \tag {3}
$$
where  $r(C)$  denotes the root of a cone  $C$ . This corresponds to finding the cone that is farthest away from  $S$ , which is the root of all hierarchies in the shadow cones construction. When  $d = 2$ ,  $\sup_2(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
$$
K (u, v) = f \left(d _ {\mathbb {H}} \left(\sup  _ {2} (u, v), \mathcal {S}\right)\right) \tag {4}
$$
where  $f$  is a user-defined monotonically increasing function. If  $K(u,v) > K(u,w)$ , then  $d_{\mathbb{H}}(\sup_2(u,v),\mathcal{S}) > d_{\mathbb{H}}(\sup_2(u,w),\mathcal{S})$ , which implies that  $u$  and  $v$  have a more recent "ancestor" than  $u$  and  $w$ . Thus,  $K$  gives a higher similarity score to points who have a more recent "ancestor" in  $T$ . In the infinite-setting shadow cone construction,  $\sup_2(u,v)$  is root of the minimum height (literal lowest) cone that contains both  $u$  and  $v$ , or  $\sup_2(u,v) = r$  ( $\arg \min_{C:u,v\in C}r(C)_d$ ).
Using this, we provide definitions for  $K(u,v)$  in the infinite-setting shadow cone construction. Both the umbral and penumbral definitions correspond to the Euclidean height of  $\sup_2(u,v)$ . In the penumbral setting, when  $\sup_2(u,v)$  does not exist, we return the Euclidean height of lowest light source where  $\sup_2(u,v)$  exists.  $x_d$  denotes the last dimension of  $x$ , and  $x_{: -1}$  denotes the first  $d - 1$  dimensions of  $x$ .  $\gamma \in \mathbb{R}^+$  corresponds to the softmax "temperature" in attention. In the penumbral setting,  $r \in \mathbb{R}^+$  is the user-defined height of the horosphere light source. In the umbral setting,  $r \in \mathbb{R}^+$  is the user-defined radius of the ball centered at each point.
Definition 1. Penumbral Attention:
$$
K (u, v) = \exp \left(- \gamma \max  \left(u _ {d}, v _ {d}, \sqrt {r ^ {2} - \left(\frac {\sqrt {r ^ {2} - u _ {d} ^ {2}} + \sqrt {r ^ {2} - v _ {d} ^ {2}} - \| u _ {: - 1} - v _ {: - 1} \|}{2}\right) ^ {2}}\right)\right) \tag {5}
$$
when there exists a cone that contains  $u$  and  $v$ , and
$$
K (u, v) = \exp \left(- \gamma \sqrt {\left(\frac {\| u _ {: - 1} - v _ {: - 1} \| ^ {2} + u _ {d} ^ {2} - v _ {d} ^ {2}}{2 \| u _ {: - 1} - v _ {: - 1} \|}\right) ^ {2} + v _ {d} ^ {2}}\right) \tag {6}
$$
otherwise. There exists a cone that contains  $u$  and  $v$  when
$$
\left(\left\| u _ {: - 1} - v _ {: - 1} \right\| - \sqrt {r ^ {2} - u _ {d} ^ {2}}\right) ^ {2} + v _ {d} ^ {2} <   r ^ {2} \tag {7}
$$
Definition 2. Umbral Attention:
$$
K (u, v) = \exp \left(- \gamma \max  \left(u _ {d}, v _ {d}, \frac {\| u _ {: - 1} - v _ {: - 1} \|}{2 \sinh (r)} + \frac {u _ {d} + v _ {d}}{2}\right)\right) \tag {8}
$$
These definitions possess a rather interesting property - when  $u_{d} = v_{d}$  and  $K$  is normalized across a set of  $v_{s}$ , cone attention reduces to the Laplacian kernel  $K(u_{: - 1},v_{: - 1}) = \exp (-\gamma \| u_{: - 1} - v_{: - 1}\|)$ . Since Euclidean space is isomorphic to a horosphere, and  $u$  and  $v$  are on the same horosphere if  $u_{d} = v_{d}$ , cone attention can also be seen as an extension of the Laplacian kernel [19].
Cone attention and dot product attention both take  $O(n^{2}d)$  time to compute pairwise attention between two sets of  $n$ $d$ -dimensional tokens  $q, k \in \mathbb{R}^{n \times d}$  [33]. However, cone attention takes more operations than dot product attention, which computes  $qk^{\top}$ . 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 is not optimal, and a 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  $\mathrm{Exp}_x(v)$ , which maps a vector  $v$  from the tangent space of a manifold  $\mathcal{M}$  at  $x$  onto  $\mathcal{M}$  by following the geodesic corresponding to  $v$  [24, 2]. In the Poincaré half-space [36],
$$
\operatorname {E x p} _ {x} (v) _ {: - 1} = x _ {: - 1} + \frac {x _ {d}}{\| v \| / \tanh  (\| v \|) - v _ {d}} v _ {: - 1} \quad \operatorname {E x p} _ {x} (v) _ {d} = \frac {x _ {d}}{\cosh (\| v \|) - v _ {d} \sinh (\| v \|) / \| v \|} \tag {9}
$$
While  $\mathrm{Exp}_x(v)$  is geometrically well motivated, it suffers from numerical instabilities when  $\| v\|$  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/underflows, but fp64 significantly reduces performance on GPUs, which is highly undesirable [1].
Instead, we use maps of the form  $(x_{: - 1},x_d)\to (x_{: - 1}f(x_d),f(x_d))$  that preserve the exponential volume of hyperbolic space while offering better numerics for large-scale optimization. Our use of alternatives to  $\mathrm{Exp}_x(v)$  follows Gulcehre et al. [11]'s psuedopolar map onto the Hyperboloid. Geometrically, since  $n$ -dimensional Euclidean space is isomorphic to a horosphere in  $\mathbb{H}^{n + 1}$ , these maps correspond to selecting a horosphere with  $f(x_{d})$  and then projecting  $x_{: - 1}$  onto that horosphere [19]. To achieve exponential space as  $x_{d}\rightarrow -\infty$ , we use functions  $f$  of the form  $\exp (\cdot)$ .
For the infinite-setting umbral construction, since there is no restriction on where points can go in the Poincaré half-space, we map  $x \in \mathbb{R}^d$  to  $\mathsf{H}^d$  with  $\psi : \mathbb{R}^d \to \mathsf{H}^d$ :
$$
\psi (x) _ {: - 1} = x _ {: - 1} \exp (x _ {d}) \quad \psi (x) _ {d} = \exp (x _ {d}) \tag {10}
$$
For the infinite-setting penumbral construction, since the mapped point cannot enter the light source at height  $h$ , we map  $x \in \mathbb{R}^d$  to area below the light source with  $\xi : \mathbb{R}^d \to \mathsf{H}^d$
$$
\xi (x) _ {: - 1} = x _ {: - 1} \frac {h}{1 + \exp (- x)} \quad \xi (x) _ {d} = \frac {h}{1 + \exp (- x)} \tag {11}
$$
While  $\xi$  is structurally the sigmoid operation, note that  $\mathrm{sigmoid}(x) = \exp (-\mathrm{softplus}(-x))$ . Since  $\mathrm{softplus}(x)\approx x$  for large values of  $x$ ,  $\xi$  preserves the exponential volume properties we seek.
Here, we present an 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 a 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 state-of-the-art (SOTA) when new.
4.1Baselines
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. [11]'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. [11]'s original hyperbolic distance attention formulation not only computed the similarity matrix  $K$  in the Hyperboloid, but also aggregated values  $\mathcal{V}$  in hyperbolic space by taking the Einstein midpoint with respect to weights  $\alpha$  in the Klein model (see appendix for definitions) [24]. However, the Einstein midpoint, defined as
$$
m (\alpha , v) = \sum_ {i} \left(\frac {\alpha_ {i} \gamma (v _ {i})}{\sum_ {j} \alpha_ {j} \gamma (v _ {j})}\right) \tag {12}
$$
where  $\gamma (v) = 1 / \sqrt{1 - \|v\|^2}$ , does not actually depend on how  $\alpha$  is computed. That is,  $\alpha$  could be computed with Euclidean dot product attention or cone attention and  $m(\alpha ,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. [11]'s original pseudopolar projection onto the Hyperboloid model and with our  $\psi$  map onto the Poincaré half-space.
4.2 Models