SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We use the following models to test cone attention and the various baselines. These models span graph prediction, NLP, and vision tasks, and range in size from a few thousand to almost 250 million parameters. While we would have liked to test larger models, our compute infrastructure limited us from feasibly training billion-parameter models from scratch.
Graph Attention Networks. Graph attention networks (GATs) were first introduced by Velicković et al. [32] for graph prediction tasks. GATs use self-attention layers to compute node-level attention maps over neighboring nodes. The original GAT used a concatenation-based 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, 13].
Neural Machine Translation (NMT) Transformers. Transformers were first applied to NMT in Vaswani et al. [31]'s seminal transformer paper. We use the faireq transformer_iwslt_de_en architecture to train a 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 to train out of the tested models, we use it for ablations.
Vision Transformers. Vision transformers (ViT) use transformer-like 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]. DeiTs and ViTs share the same architecture, and the only differences are how they are trained and the distillation token. We train DeiT-Ti models with 5 million parameters on the ImageNet-1K dataset for 300 epochs[28, 7]. We also train cone and dot product attention for 500 epochs, as we observed that training for more iterations improves performance.
Adaptive 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 faireq transformer_lm_wiki103 architecture (246.9M parameters) and train models on the WikiText-103 language modeling dataset with a block size of 512 tokens [20, 18]. We also test the same architecture without adaptive inputs, which has 520M parameters. This version converges in significantly fewer iterations, allowing us to train such a large model.
Diffusion Transformers. Diffusion transformers (DiT) are diffusion models that replace the U-Net backbone with a transformer [22]. DiTs operate on latent space patches, so we expect there to be less hierarchical information vs. taking image patches. We use DiTs to test cone attention when the data is less hierarchical. We train DiT-B/4 models with 130M parameters on ImageNet-1K [22, 7].
Table 1: Performance of various attention methods across models and tasks. * indicates the model failed to converge or ran into NaN errors. ↑ indicates higher is better, and ↓ indicates lower is better. Cone attention methods (†) 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.
Method	NMT IWLST (BLEU ↑)	DeiT-Ti Imagenet Top-1 / 5 (Acc. ↑) 300 Epochs	DeiT-Ti Imagenet Top-1 / 5 (Acc. ↑) 500 Epochs	GAT Cora / PPI (Acc. ↑)
Model Default	34.56	72.05 / 91.17	73.65 / 91.99	0.831 / 0.977
Dot Product	34.56	72.05 / 91.17	73.65 / 91.99	0.834 / 0.985
Penumbral†	35.56	72.67 / 91.12	74.34 / 92.38	0.835 / 0.990
Umbral†	35.07	73.14 / 91.82	74.46 / 92.54	0.836 / 0.989
dHnξ	32.54	49.19* / 74.68*	-	0.834 / 0.987
dHHyperboloid	33.80	0.10* / 0.45*	-	0.13* / 0.989
Laplacian Kernel	34.68	71.25 / 90.84	-	0.823 / 0.986
Method	Adaptive Inputs WikiText-103 0 / 480 Context Window (Ppl. ↓)	Without Adaptive Inputs	DiT-B/4 @ 400K Steps (FID-50K ↓)
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
Table 1 summarizes the performance of cone attention across various models and tasks. Both penumbral and umbral attention significantly outperform baselines on the NMT IWSLT and DeiT-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 a 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 concatenation-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 an average of 2 edges per node and chronological dependencies between nodes (a paper cannot cite a newer paper), while the PPI dataset is closer to a 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 rare words. We expect rarer words to be closer to the leaves of a word hierarchy, so the adaptive inputs method also acts as a hierarchical prior on the data [34]. 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 WikiText103 task, with or without adaptive inputs. Since umbral cones are 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 a 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 a result of patches being over the latent space, and not the input image.
5.1 Effect of Mapping Functions
Table 2 compares how  $\psi$  and  $\xi$  compare to the exponential map and psuedpolar map (section 3.2) on the NMT IWSLT task. Here, we use the exponential map at the origin of  $\mathsf{H}^d$ ,  $O = (0,0,\dots,1)$ . To compare  $\mathrm{Exp}_O(v)$  to  $\psi$ , we first take  $\mathrm{Exp}_O(v)$  and then project the point onto the boundary of
Table 2: Comparison of various mappings from Euclidean space to hyperbolic models for the NMT IWSLT task (BLEU scores, higher is better). The exponential map generally performs worse than  $\psi$  and the pseudopolar map. While  $\psi$  also performs worse than the psuedopolar map, table 1 indicates that the pseudopolar map is less numerically stable than  $\psi$ .
Method	ExpO(v) → Hd	ξ → Hd	ψ → Hd	Pseudopolar → Hyperboloid
Penumbral	35.12	35.56	-	-
Umbral	34.63	-	35.07	-
dH	30.59	-	32.54	33.80
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 16 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.
$L$  if  $\mathrm{Exp}_O(v)$  is in  $L$ . In the infinite penumbral setting, this corresponds to taking  $\mathrm{Exp}_O(v)_d = \min (\mathrm{Exp}_O(v)_d,h)$ .  $\psi$  and  $\xi$  generally perform much better than taking the exponential map at the origin  $O$ , which suggests that  $\psi$  and  $\xi$  have better optimization properties. For  $d_{\mathbb{H}}$  attention, Gulcehre et al.'s pseudopolar map slightly outperforms  $\xi$ . However, table 1 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, k)$  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. Table 3 shows the performance of DeiT-Ti at  $d = 16$ . Here, penumbral attention at  $d = 16$  is able to outperform dot product attention at  $d = 64$ , again indicating that cone attention is able to more efficiently capture hierarchies.
5.3 Sensitivity to Initialization
Hyperbolic embeddings are known to be sensitive to initialization [10, 37]. To test cone attention's sensitivity to initialization, we trained 5 seeds for the IWSLT De2En task for cone attention and dot
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 / 90.71
Figure 5: Attention heatmaps from an attention head in 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 worry about it, you know." The cone attention heatmaps (left and center) have a 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.12, penumbral achieved  $35.41 \pm 0.14$ , and umbral achieved  $35.03 \pm 0.30$ . There was one outlier in the umbral attention trials, and with the outlier removed umbral achieved  $35.16 \pm 0.09$ . The cone attention methods appear to have slightly higher variance than dot product attention, but not significantly so.
5.4 Attention Heatmaps
A natural question arises about what cone attention methods actually learn. Figure 5 shows heatmaps from an attention head in a trained IWSLT De2En translation model. The heatmaps for penumbral and umbral attention show clearer separation than the dot product attention heatmap. Furthermore, the separation in the two cone attention methods happens at the exclamation mark, a natural segmentation of the sequence into two parts. Intuitively, cone attention can be seen as an attention method that satisfies certain "logical constraints," such as "if  $z \prec y$  and  $y \prec x$ , 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. If  $x = (1,1,0,0), y = (10,10,10,10)$ , and  $z = (0,0,1,1)$ , then  $\langle x,y\rangle = \langle y,z\rangle = 20$ , but  $\langle x,z\rangle = 0$ . We suspect this is a reason why cone attention methods show better separation than dot product attention, which can aid performance.
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 a variety of attention networks ranging from a few thousand to a 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 an effective way to encode hierarchical relationships in attention, and can potentially improve task-level performance in 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 [37] suggest that hyperbolic embeddings are sensitive to initialization, which implies that different transformer weight initializations may affect cone attention.
This work was supported by NSF Award IIS-2008102. Compute resources were provided by the Cornell G2 Cluster.
[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é-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file82c2559140b95ccda9c6ca4a8b981f1e-Paper.pdf.
[36] Tao Yu and Christopher M De Sa. Representing hyperbolic space accurately using multicomponent 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 in hyperbolic space. arXiv preprint, 2023.
[38] Lin Zheng, Jianbo Yuan, Chong Wang, and Lingpeng Kong. Efficient attention via control variates. arXiv preprint arXiv:2302.04542, 2023.