Leveraging rotational equivariance for reinforcement learning in tractography

Area of Science:

• Neuroimaging
• Computational Neuroscience
• Machine Learning

Background:

• Brain tractography maps neural fibre bundles using diffusion-weighted images (DWI).
• Reinforcement learning (RL) frameworks, particularly actor-critic models, are emerging for tractography.
• Existing RL methods lack SO(3) equivariance, failing to accurately reflect 3D input rotations in output tractograms.

Purpose of the Study:

• To restore rotational equivariance to RL-based brain tractography.
• To integrate SO(3) equivariant and invariant components into RL actor-critic models for tractography.
• To theoretically discuss and practically implement rotational equivariance in streamline tractography.

Main Methods:

• Introduced SO(3) equivariant components for actor (direction prediction) and invariant components for critic (Q-value prediction).
• Utilized an SE(3)-equivariant transformer as the next direction prediction function.
• Formulated DWI and directional updates as spherical signals transforming under SO(3) representations.

Main Results:

• The proposed method successfully restores SO(3) equivariance to RL-based tractography.
• Evaluated equivariance locally and globally using phantom and in vivo data.
• Demonstrated restoration of equivariance in Track-to-Learn, a state-of-the-art RL tractography method.

Conclusions:

• The developed method integrates rotational equivariance into RL tractography, addressing limitations of previous approaches.
• This work provides a theoretical and practical framework for equivariant streamline tractography.
• The findings advance RL-based tractography by ensuring geometric consistency with input data transformations.