Title:	Rotational Equivariance for Reinforcement Learning in Tractography
DNr:	Berzelius-2024-159
Project Type:	LiU Berzelius
Principal Investigator:	Fabian Sinzinger <fabiansi@kth.se>
Affiliation:	Kungliga Tekniska högskolan
Duration:	2024-04-17 – 2024-11-01
Classification:	20603
Keywords:

Abstract

This work extends recent advances in reinforcement learning-based tractography. The focus here lies on encoding rotational equivariance into the model. To the best of our knowledge, this is the first attempt at making RL rotationally equivariant in general and for tractography in particular. Integral to our methodology is combining novel tractography with methods from geometrical deep learning and introducing custom components. Brain tractography involves mapping diffusion-weighted images (DWI) onto streamlines representing neural fibre bundles. Recent research avenues suggest tractography with actor-critic models in a reinforcement learning (RL) framework. However, learning-based methods may compromise geometrical relations between input (DWI) and output (tractogram). We note that transformations like 3D rotations applied to the input of learning-based tractography are not adequately reflected in the output, indicating a lack of SO(3) equivariance. This study aims to restore the equivariance present in previous non-learning-based methods like IFOD2 to RL-based tractography. To achieve this, we introduce SO(3) equivariant and invariant components for the actors (direction prediction model) and critics (q-value prediction model), respectively. These models employ tensor-product convolutions in the spherical harmonics domain, followed by gated nonlinearities. The fact that both the input DWI and the output directional update can be represented as spherical tensors and transform under representations of SO(3) makes this formulation a natural fit for the present problem. Another benefit of RL-based tractography is that incorporating local neighbourhoods can help mitigate well-known tractography problems (e.g. kissing, crossing, fanning, etc.). Our proposed algorithm extracts neighbourhood information as a k-nearest neighbour graph with the spherical signals on the nodes, facilitating continuous streamline updates from spatially discrete DWI. The contribution of this work is threefold; firstly, we exhibit the presence of the equivariance breaking problem in learning-based tractography; second, we provide a method how for connecting RL with equivariant graph neural networks; third, we demonstrate the performance of the novel method of phantom data as well as real brain dwi."