HexFormer: Hyperbolic Vision Transformer with Exponential Map Aggregation

TL;DR

This work tackles the challenge of representing hierarchical structure in visual data by adopting hyperbolic geometry, specifically the Lorentz model with negative curvature $K<0$. It introduces HexFormer, a hyperbolic Vision Transformer that uses exponential-map aggregation in hyperbolic attention, and HexFormer-Hybrid, which pairs a hyperbolic encoder with an Euclidean classification head. Across CIFAR-10, CIFAR-100, and Tiny-ImageNet, both variants outperform Euclidean ViTs and prior hyperbolic ViTs, with the hybrid variant achieving the strongest results. The study also shows that hyperbolic models exhibit more stable gradients and reduced warmup sensitivity, and that ExpAgg provides better numerical stability than centroid aggregation, making hyperbolic transformers a robust and practical option for vision tasks.

Abstract

Data across modalities such as images, text, and graphs often contains hierarchical and relational structures, which are challenging to model within Euclidean geometry. Hyperbolic geometry provides a natural framework for representing such structures. Building on this property, this work introduces HexFormer, a hyperbolic vision transformer for image classification that incorporates exponential map aggregation within its attention mechanism. Two designs are explored: a hyperbolic ViT (HexFormer) and a hybrid variant (HexFormer-Hybrid) that combines a hyperbolic encoder with an Euclidean linear classification head. HexFormer incorporates a novel attention mechanism based on exponential map aggregation, which yields more accurate and stable aggregated representations than standard centroid based averaging, showing that simpler approaches retain competitive merit. Experiments across multiple datasets demonstrate consistent performance improvements over Euclidean baselines and prior hyperbolic ViTs, with the hybrid variant achieving the strongest overall results. Additionally, this study provides an analysis of gradient stability in hyperbolic transformers. The results reveal that hyperbolic models exhibit more stable gradients and reduced sensitivity to warmup strategies compared to Euclidean architectures, highlighting their robustness and efficiency in training. Overall, these findings indicate that hyperbolic geometry can enhance vision transformer architectures by improving gradient stability and accuracy. In addition, relatively simple mechanisms such as exponential map aggregation can provide strong practical benefits.

Figures (10)

• Figure 1: Overview of HexFormer architecture. Hyperbolic embedding, transformer encoder, and classification head operate on the Lorentz manifold. Residual connections are applied to space-like components only.
• Figure 2: HexFormer multi-head attention. Queries, keys, and values are projected via Lorentz fully connected layers. Scores are computed using Lorentzian distances, aggregated in tangent space via exponential map, and recombined after scaling the space components (to normalize magnitudes and maintain stability in the Lorentz model), then passed as output through a Lorentz FC layer.
• Figure 3: The histogram of the absolute value of gradients on a log scale during the training of ViT-Small on CIFAR-100 comparing HexFormer vs. Euclidean ViT with and without warmup (offset mode). Top row shows close-up for first 200 iterations.
• Figure 4: Comparison of accuracy over epochs on CIFAR-100. The numbers 0 and 10 in the legend refer to the warmup.
• Figure 5: The histogram of the absolute value of gradients on a log scale during the training of ViT-Small on CIFAR-100 comparing HexFormer vs. Euclidean ViT with and without warmup (overlay mode). Top row shows close-up for first 200 iterations.