LaB-GATr: geometric algebra transformers for large biomedical surface and volume meshes (2403.07536v2)

Abstract: Many anatomical structures can be described by surface or volume meshes. Machine learning is a promising tool to extract information from these 3D models. However, high-fidelity meshes often contain hundreds of thousands of vertices, which creates unique challenges in building deep neural network architectures. Furthermore, patient-specific meshes may not be canonically aligned which limits the generalisation of machine learning algorithms. We propose LaB-GATr, a transfomer neural network with geometric tokenisation that can effectively learn with large-scale (bio-)medical surface and volume meshes through sequence compression and interpolation. Our method extends the recently proposed geometric algebra transformer (GATr) and thus respects all Euclidean symmetries, i.e. rotation, translation and reflection, effectively mitigating the problem of canonical alignment between patients. LaB-GATr achieves state-of-the-art results on three tasks in cardiovascular hemodynamics modelling and neurodevelopmental phenotype prediction, featuring meshes of up to 200,000 vertices. Our results demonstrate that LaB-GATr is a powerful architecture for learning with high-fidelity meshes which has the potential to enable interesting downstream applications. Our implementation is publicly available.

• The paper presents a scalable transformer architecture that uses geometric tokenization to efficiently process large biomedical surface and volume meshes.
• The model achieves a 3.5% error in coronary artery blood velocity prediction and a mean absolute error of 0.54 weeks in postmenstrual age estimation.
• The methodology leverages geometric algebra to maintain Euclidean symmetries while employing sequence compression and vertex clustering for enhanced performance.

LaB-GATr: Geometric Algebra Transformers for Large Biomedical Meshes

The paper presents LaB-GATr, a novel transformer architecture designed to address the challenges posed by large-scale biomedical surface and volume meshes. The authors extend the recently proposed geometric algebra transformer (GATr) by incorporating geometric tokenization strategies that enable the effective handling of high-fidelity meshes with significantly large vertex counts, such as those ranging up to 200,000 vertices. This advancement allows LaB-GATr to achieve state-of-the-art performance on tasks within biomedical domains, including cardiovascular hemodynamics modeling and neurodevelopmental phenotype prediction.

Key Contributions

The primary contribution of this paper is the development of a scalable architecture that respects Euclidean symmetries, thereby mitigating issues related to the non-canonical alignment of patient-specific meshes. The architecture integrates sequence compression and interpolation mechanisms, which enable an efficient and effective learning process even in scenarios where GPU memory limitations present a bottleneck, as is the case with GATr when applied to large meshes.

Numerical Results

The LaB-GATr model attains unprecedented accuracy in multiple biomedical tasks. On the task of predicting blood velocity in synthetic coronary artery meshes, it achieves an approximation error of 3.5%, a significant improvement over the previous state-of-the-art figure of 7.4%. Additionally, without the need for cortical surface morphing to an icosphere, LaB-GATr outperforms previous approaches on phenotype prediction tasks, specifically setting a new benchmark for postmenstrual age (PMA) estimation with a mean absolute error of 0.54 weeks.

Methodology

The geometric tokenization component of LaB-GATr is noteworthy. It employs a novel vertex clustering approach to reduce the complexity of transformer self-attention, achieved by pooling input features to a coarse subset of mesh vertices. This strategy allows LaB-GATr to maintain the original mesh resolution while reducing the number of active tokens needed in self-attention layers.

The authors leverage geometric algebra in the transformer design, taking advantage of its characterized symmetries. By embedding and processing input features as multivectors in the geometric space $G(3, 0, 1)$, LaB-GATr effectively models the global interactions necessary for learning tasks involving large and complex 3D meshes.

Implications and Future Research

The architectural innovations within LaB-GATr exhibit significant potential for various downstream biomedical applications. Its capabilities suggest notable improvements in areas such as high-dimensional mesh processing for precise medical diagnostics and enhanced biomechanical modeling. Future research may explore the adaptation of this framework to incorporate further efficiencies, such as patch merging or advanced forms of token sparsification, which could contribute to even greater scalability and performance.

Overall, the proposed LaB-GATr structure represents a substantial step forward in applying machine learning to complex biomedical datasets. Its theoretical underpinnings in geometric algebra, coupled with the efficient handling of large-scale meshes, underscore its practical insignificance in the increasingly data-heavy landscape of medical image and mesh analysis. The availability of the LaB-GATr implementation offers researchers the tools to further test and expand upon these findings, fostering continued exploration into the fusion of deep learning and biomedical engineering.