VBA: Vector Bundle Attention for Intrinsically Geometric Representation Learning
Abstract
Learning from geometrically structured data is central to applications in biology, physics, and computer vision. In many tasks, meaningful comparisons depend on how features are aligned in space. Graph Neural Networks capture local structure but are constrained by message passing. Transformers model long-range dependencies but largely ignore geometry. We introduce the Vector Bundle Attention Transformer (VBA-Transformer), a framework that redefines attention as an intrinsic geometric operator. Each token couples a base manifold coordinate with a fiber feature vector, following vector bundle theory. A principled parallel transport mechanism aligns fiber features across local coordinate systems before similarity is computed. This embeds geometry directly into the attention operator. Unlike prior methods that inject geometry as an external bias or positional encoding, VBA integrates geometry natively inside attention. On challenging single-cell RNA sequencing benchmarks, VBA achieves state-of-the-art accuracy, outperforming Transformer baselines by over 3--5\%. On spatial transcriptomics, it demonstrates superior clustering performance. On 3D point clouds, it achieves competitive accuracy, validating broad generalization across domains. Beyond empirical gains, we provide theoretical analysis of invariance and perturbation stability. We also demonstrate robust transport behavior empirically. Together, these results establish intrinsic geometric alignment as a powerful principle for scalable representation learning.
Lay Summary
Many real-world datasets have hidden structure. For example, cells in a tissue are not just isolated measurements: their meaning depends on where they are located, how they relate to nearby cells, and how their features should be compared. Standard AI models are powerful, but they often compare data points directly without fully considering this kind of structure. In this paper, we introduce a new AI model called Vector Bundle Attention. The main idea is simple: before comparing two pieces of information, the model first “aligns” them into a shared frame of reference, so the comparison is more meaningful. This is similar to comparing directions on a curved surface only after making sure they are expressed in the same local viewpoint. We test the model on several types of structured data, including single-cell gene expression, spatial tissue measurements, and 3D object shapes. The model improves cell type prediction, performs strongly in identifying tissue regions, and remains competitive on 3D shape recognition. These results suggest that building geometric awareness directly into AI models can make them more reliable for complex scientific data, especially in biology and medicine.