CAT: Curvature-Adaptive Transformers for Geometry-Aware Learning
R. Y. Lin, S. Ojha, Nicholas Bai
NeurIPS 2025 Workshop on Non-Euclidean Foundation Models and Geometric Learning
Workshop paper
Transformers that route tokens across three geometric attention branches with a mixture-of-experts gate.
Transformers achieve strong performance across diverse domains but implicitly assume Euclidean geometry in their attention mechanisms, limiting their effectiveness on data with non-Euclidean structure.
We introduce the Curvature-Adaptive Transformer (CAT), a novel architecture that dynamically learns per-token routing across three geometric attention branches through a lightweight, differentiable gating mechanism. Unlike fixed-geometry approaches, CAT enables adaptive geometric specialization, routing tokens to the appropriate curvature based on their local relational structure.
On knowledge graph completion benchmarks (FB15k-237, WN18RR), CAT achieves approximately 10% improvements in MRR and Hits@10 over fixed-geometry baselines with minimal overhead.