Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Rotation-invariant graph message passing enables acquisition protocol generalisation in learning-based brain microstructure estimation

Leevi Kerkelä, Hui Zhang

TL;DR

This work presents a graph neural network that represents input data as a point cloud in the 3D space where diffusion-weighted measurements are made and performs rotation-invariant message passing with permutation-invariant pooling, producing fixed-size embeddings that encode microstructure.

Abstract

Estimating brain microstructure has important applications in medicine and neuroscience. Diffusion-weighted magnetic resonance imaging enables measuring microstructure \textit{in vivo}. Conventional biophysical model fitting can be accurate but is slow and impractical for time-critical clinical use, where machine learning can offer a potential route to rapid estimation. We address the problem of microstructure estimation under arbitrary acquisition protocols where most existing learning-based methods fail due to protocol assumptions, requiring retraining when the protocol changes. We present a graph neural network that represents input data as a point cloud in the 3D space where diffusion-weighted measurements are made and performs rotation-invariant message passing with permutation-invariant pooling, producing fixed-size embeddings that encode microstructure. The inductive biases of our relatively small model were guided by the underlying physics and symmetries of the problem rather than by generic model architectures. Trained on randomised simulated data, our model demonstrates domain generalisation, accurately estimating microstructure from data with unseen real-world protocols without retraining. This approach represents a step towards a "train once, deploy anywhere" architecture, bringing rapid learning-based microstructure mapping closer to clinical deployment.

Rotation-invariant graph message passing enables acquisition protocol generalisation in learning-based brain microstructure estimation

TL;DR

This work presents a graph neural network that represents input data as a point cloud in the 3D space where diffusion-weighted measurements are made and performs rotation-invariant message passing with permutation-invariant pooling, producing fixed-size embeddings that encode microstructure.

Abstract

Estimating brain microstructure has important applications in medicine and neuroscience. Diffusion-weighted magnetic resonance imaging enables measuring microstructure \textit{in vivo}. Conventional biophysical model fitting can be accurate but is slow and impractical for time-critical clinical use, where machine learning can offer a potential route to rapid estimation. We address the problem of microstructure estimation under arbitrary acquisition protocols where most existing learning-based methods fail due to protocol assumptions, requiring retraining when the protocol changes. We present a graph neural network that represents input data as a point cloud in the 3D space where diffusion-weighted measurements are made and performs rotation-invariant message passing with permutation-invariant pooling, producing fixed-size embeddings that encode microstructure. The inductive biases of our relatively small model were guided by the underlying physics and symmetries of the problem rather than by generic model architectures. Trained on randomised simulated data, our model demonstrates domain generalisation, accurately estimating microstructure from data with unseen real-world protocols without retraining. This approach represents a step towards a "train once, deploy anywhere" architecture, bringing rapid learning-based microstructure mapping closer to clinical deployment.
Paper Structure (23 sections, 14 equations, 4 figures, 1 table)

This paper contains 23 sections, 14 equations, 4 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related work
  3. Geometric deep learning
  4. Learning-based microstructure estimation
  5. Preliminaries
  6. Diffusion MRI data
  7. Problem statement
  8. Method
  9. Graph construction
  10. Message passing
  11. Pooling and parameter prediction
  12. Experiments
  13. Network architecture
  14. Training data generation
  15. Training
  16. ...and 8 more sections

Figures (4)

  • Figure 1: Illustration of the model input data. (a) Axial slices from five dMRI data volumes, each acquired with specific acquisition settings. (b) The measurements from the voxel highlighted in (a) across all volumes in the 3D space where the dMRI measurements are made, revealing the geometric structure of the two-shell protocol ($b$ = 1 and 2 ms/µ m$^2$, 50 directions per shell); colour encodes relative signal intensity. (c) The graph constructed from the voxel-level data in (b) used as the GNN input: each measurement is mirrored across the origin to enforce antipodal symmetry, and a $k$-nearest neighbour graph ($k = 8$) is constructed; colour encodes edge feature norm magnitude.
  • Figure 2: NODDI parameter estimation errors for DSI (a), HCP (b), and UKBB (c) protocols on the test dataset.
  • Figure 3: t-SNE on the pooled embeddings produced by the GNN from test data with DSI, HCP, and UKBB protocols, revealing microstructure-aligned manifolds. (a)--(c) Embeddings coloured by ground-truth parameter values. (d) Embeddings coloured by protocol.
  • Figure 4: Axial slices of NODDI parameter maps produced by the GNN on DSI (a), HCP (b), and UKBB (c) protocols.