VBA: Vector Bundle Attention for intrinsically geometric representation learning

Fang, Shenglei, Sun, Xianfang ORCID: https://orcid.org/0000-0002-6114-0766 and Zhou, You ORCID: https://orcid.org/0000-0002-1743-1291 2026. VBA: Vector Bundle Attention for intrinsically geometric representation learning. Presented at: ICML 2026, Seoul, South Korea, 6-11 July 2026. Proceedings of Machine Learning Research.

PDF - Published Version Available under License Creative Commons Attribution. Download (4MB)

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 (VBATransformer), 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. Our code is available at: https://github.com/ yzlab1/Vector-Bundle-Attention

Item Type:	Conference or Workshop Item - published (Paper)
Status:	In Press
Schools:	Schools > Medicine Schools > Computer Science & Informatics
ISSN:	1938-7228
Date of First Compliant Deposit:	21 July 2026
Last Modified:	21 Jul 2026 11:54
URI:	https://orca.cardiff.ac.uk/id/eprint/188360

Actions (repository staff only)

Edit Item