HexFormer: Hyperbolic Vision Transformer with Exponential Map Aggregation

TL;DR

HexFormer introduces a hyperbolic vision transformer utilizing exponential map aggregation, which improves accuracy and gradient stability in image classification tasks by effectively modeling hierarchical data structures.

Contribution

This work presents a novel hyperbolic vision transformer with exponential map aggregation, demonstrating improved performance and training stability over Euclidean models.

Findings

Consistent performance improvements over Euclidean baselines.

Hyperbolic models exhibit more stable gradients.

Hybrid variant achieves the strongest results.

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.