Tract-RLFormer: A Tract-Specific RL policy based Decoder-only Transformer Network

Ankita Joshi,Ashutosh Sharma,Anoushkrit Goel,Ranjeet Ranjan Jha,Chirag Ahuja,Arnav Bhavsar,Aditya Nigam

2024-11-09

Abstract:Fiber tractography is a cornerstone of neuroimaging, enabling the detailed mapping of the brain's white matter pathways through diffusion MRI. This is crucial for understanding brain connectivity and function, making it a valuable tool in neurological applications. Despite its importance, tractography faces challenges due to its complexity and susceptibility to false positives, misrepresenting vital pathways. To address these issues, recent strategies have shifted towards deep learning, utilizing supervised learning, which depends on precise ground truth, or reinforcement learning, which operates without it. In this work, we propose Tract-RLFormer, a network utilizing both supervised and reinforcement learning, in a two-stage policy refinement process that markedly improves the accuracy and generalizability across various data-sets. By employing a tract-specific approach, our network directly delineates the tracts of interest, bypassing the traditional segmentation process. Through rigorous validation on datasets such as TractoInferno, HCP, and ISMRM-2015, our methodology demonstrates a leap forward in tractography, showcasing its ability to accurately map the brain's white matter tracts.

Machine Learning

What problem does this paper attempt to address?

### What problems does this paper attempt to solve?

This paper aims to solve the complexity and error - proneness in tractography. Specifically, tractography uses diffusion magnetic resonance imaging (diffusion MRI, dMRI) to draw detailed maps of brain white matter pathways, which is crucial for understanding brain connectivity and function. However, existing tractography methods face the following challenges:

1. **Cross - fiber problem**: Since the algorithms rely on local diffusion information to reconstruct the complete brain fiber network, it sometimes leads to incorrect fiber path projections or false positive connections.
2. **Lack of accurate ground - truth data**: Supervised learning methods require accurate ground - truth data for training and validation, and such data are difficult to obtain.
3. **Insufficient generalization ability**: Existing methods have limited generalization ability on different datasets.

To solve these problems, the authors propose **Tract - RLFormer**, a decoder - only Transformer network with a two - stage strategy improvement process that combines supervised learning and reinforcement learning. The main goals of this model are:

- **Improve accuracy**: Significantly improve the accuracy of fiber bundle generation by combining supervised learning and reinforcement learning.
- **Enhance generalization ability**: Ensure that the model has good generalization performance on different datasets.
- **Simplify the process**: Directly generate specific fiber bundles, skipping the traditional segmentation step, thereby simplifying the tractography process.

### Main contributions

1. **Data - driven strategy learning**: Propose a Tract - RLFormer based on a hybrid framework, using reinforcement learning (RL) and supervised learning (SL) paradigms to approximate and improve fiber generation strategies.
2. **Innovative specific - fiber - bundle generation**: Train Tract - RLFormer to generate fiber bundles of interest and use the developed Mask Refinement Module (MRM) to generate tracking masks, avoiding segmentation overhead.
3. **Extensive generalization ability**: Through extensive testing on multiple datasets (such as TractoInferno, HCP, and ISMRM - 2015), demonstrate the superior performance and generalization ability of the model in different neuroimaging environments.

### Technical details of the solution

- **Mask Refinement Module (MRM)**: Used to generate tracking masks for specific fiber bundles, based on the fiber orientation information of each subject.
- **Reinforcement learning strategy learning**: Train RL agents (TD3 algorithm) to explore and generate fiber bundles of interest.
- **T - RLF model**: A Transformer network based on the GPT architecture, which generates optimized strategies through two - stage training (general pre - training and specific - fiber - bundle fine - tuning).

Through these methods, Tract - RLFormer significantly improves the accuracy and generalization ability of tractography, providing a powerful tool for neuroscience research and clinical applications.