Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LaB-GATr: geometric algebra transformers for large biomedical surface and volume meshes

Julian Suk, Baris Imre, Jelmer M. Wolterink

TL;DR

LaB-GATr is 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 and achieves state-of-the-art results on three tasks in cardiovascular hemodynamics modelling and neurodevelopmental phenotype prediction.

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.

LaB-GATr: geometric algebra transformers for large biomedical surface and volume meshes

TL;DR

LaB-GATr is 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 and achieves state-of-the-art results on three tasks in cardiovascular hemodynamics modelling and neurodevelopmental phenotype prediction.

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.
Paper Structure (19 sections, 1 theorem, 10 equations, 3 figures, 2 tables)

This paper contains 19 sections, 1 theorem, 10 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Methods
  3. Geometric algebra
  4. Embedding meshes in $\mathbf{G}(3, 0, 1)$
  5. Geometric algebra transformers
  6. Learned, geometric tokenisation
  7. Pooling
  8. Interpolation
  9. Neural network architecture
  10. Experiments and results
  11. Coronary artery models
  12. Surface-based WSS estimation
  13. Volume-based velocity field estimation
  14. Postmenstrual age prediction from the cortical surface
  15. Discussion and conclusion
  16. ...and 4 more sections

Key Result

proposition thmcounterproposition

Consider a set of multivectors $x^i \in \mathbf{G}(3, 0, 1)$ and denote by the extraction of point coordinates from a multivector. Assume $x^i_{123} > 0$. Then the point extracted from convex combination of $x^i$ is an element of the convex hull of $\{t(x^i)\}_i$.

Figures (3)

  • Figure 1: LaB-GATr takes input features in the form of multivectors, which are constructed by embedding, e.g., mesh vertices as points and surface normal vectors as planes (see Section \ref{['sec:embed']}). In the tokenisation module, the features are pooled to a coarse subset of mesh vertices via message passing (see Section \ref{['sec:pool']}). The tokenisation allows control over the number of tokens which are processed by the GATr module. Downstream, the interpolation module lifts the tokenisation back to original mesh resolution (see Section \ref{['sec:interp']}). An optional ($^*$) class token is appended to the token sequence for mesh-level output. Subsequently, scalar or vector-valued output features are extracted.
  • Figure 2: Qualitative results of cardiovascular hemodynamics estimation on test-split arteries. Left: (top) surface and (bottom) volume mesh. Middle: (top) wall shear stress and (bottom) velocity field, via computational fluid dynamics (CFD). Right: LaB-GATr prediction.
  • Figure 3: Postmenstrual age (PMA) prediction (values in weeks) on test-split subjects. LaB-GATr predictions based on the newborn's cortical surface correlated well with the reference values.

Theorems & Definitions (2)

  • proposition thmcounterproposition
  • proof