VBA: Vector Bundle Attention for Intrinsically Geometry-Aware Learning

ICLR 2026 Conference Submission1109 Authors

Keywords: Vector Bundle Attention, Differential Geometric Transformer, Structured Attention, Single-cell RNA-seq, Spatial Transcriptomics, Representation Learning

TL;DR: We propose Vector Bundle Attention, a geometry-inspired Transformer that enforces isometric parallel transport and curvature correction, achieving robust and generalizable representation learning across biological and 3D domains.

Abstract: Learning from geometrically structured data is fundamental in biology, physics, and computer vision. Graph Neural Networks capture local structure but are limited by message passing, while Transformers model global dependencies yet often neglect 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 before similarity is computed, embedding geometry directly into the attention operation. Unlike prior models that inject geometry as an external bias, VBA integrates it natively. On challenging single-cell RNA sequencing benchmarks, our model achieves state-of-the-art accuracy, with significant gains over Transformer-based baselines (over 3%-5%). In spatial transcriptomics task, it demonstrates superior clustering performance. On 3D point clouds, VBA achieves competitive accuracy, validating its broad generalization capabilities. Beyond empirical gains, VBA provides theoretical analyses (invariance of the ideal operator and perturbation bounds), and empirical evidence of stability for the practical implementation, establishing geometric disentanglement as a powerful new principle for a versatile architecture poised for broad impact.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 1109