[Paper Review] TractOracle: towards an anatomically-informed reward function for RL-based tractography

Antoine Théberge, Maxime Descoteaux+1 more|arXiv (Cornell University)|Mar 26, 2024

Advanced Neuroimaging Techniques and ApplicationsMedicine3 citations

TL;DR

TractOracle introduces a novel reinforcement learning (RL)-based tractography framework that integrates an anatomically-informed reward function via a transformer-based streamline classification network (TractOracle-Net). By evaluating streamline plausibility in real time during tracking, the system simultaneously guides accurate tractography and prunes false positives, achieving up to 20% higher true positive rates and 3x fewer false positives compared to state-of-the-art methods on benchmark datasets.

ABSTRACT

Reinforcement learning (RL)-based tractography is a competitive alternative to machine learning and classical tractography algorithms due to its high anatomical accuracy obtained without the need for any annotated data. However, the reward functions so far used to train RL agents do not encapsulate anatomical knowledge which causes agents to generate spurious false positives tracts. In this paper, we propose a new RL tractography system, TractOracle, which relies on a reward network trained for streamline classification. This network is used both as a reward function during training as well as a mean for stopping the tracking process early and thus reduce the number of false positive streamlines. This makes our system a unique method that evaluates and reconstructs WM streamlines at the same time. We report an improvement of true positive ratios by almost 20\% and a reduction of 3x of false positive ratios on one dataset and an increase between 2x and 7x in the number true positive streamlines on another dataset.

Motivation & Objective

• To address the limitation of existing RL-based tractography methods that lack anatomical priors, leading to high false positive rates.
• To develop a unified framework that evaluates and reconstructs white matter streamlines simultaneously using real-time anatomical plausibility scoring.
• To reduce reliance on post-hoc filtering by embedding anatomical knowledge directly into the RL training process.
• To improve both tractogram generation and segmentation accuracy by leveraging a pre-trained streamline classification network as a dynamic reward and stopping signal.
• To achieve superior performance on in-silico and in-vivo datasets compared to classical and learning-based tractography baselines.

Proposed method

• A transformer-based neural network, TractOracle-Net, is trained to classify streamlines as anatomically plausible or implausible based on their shape and trajectory.
• During RL training, TractOracle-Net provides a continuous anatomical plausibility score for each streamline as it is being tracked, serving as a dense, real-time reward signal.
• The agent uses this reward to optimize for anatomically accurate streamline paths, minimizing 'reward hacking' and false positive tract generation.
• The system employs early stopping when the plausibility score drops below a threshold, preventing the tracking of anatomically implausible streamlines.
• The same TractOracle-Net model is used post-training for tractogram segmentation and filtering, enabling end-to-end anatomical evaluation.
• The framework is trained and evaluated on both in-silico (ISMRM2015) and in-vivo (TractoInferno) datasets using standard tractography metrics and tools like Recobundles and extractor_flow.

(a)

Experimental results

Research questions

• RQ1 Can a learned anatomical oracle improve the accuracy and reduce false positives in RL-based tractography without requiring annotated training data?
• RQ2 To what extent can real-time streamline plausibility scoring during tracking enhance the quality of reconstructed tractograms?
• RQ3 How does integrating a classification network as both a reward function and early stopping criterion affect tractography performance compared to standard RL or filtering pipelines?
• RQ4 Can the proposed method outperform state-of-the-art classical and learning-based tractography methods in terms of true positive rate, false positive rate, and anatomical fidelity?
• RQ5 Does the system preserve or improve the recovery of missing or subtle white matter pathways that are often missed by conventional methods?

Key findings

• On the ISMRM2015 in-silico dataset, TractOracle achieved a 22% valid connection rate, representing a 20% improvement over the second-best method.
• TractOracle reduced invalid connection rates by 3x and no-connection rates by 3x compared to the second-best method on the same dataset.
• On the TractoInferno in-vivo dataset, TractOracle-RL produced 184,355,412 streamlines classified as plausible by TractOracle-Net—7x more than sd_stream and 2x more than Track-to-Learn.
• The method reconstructed anatomically plausible bundles with superior fanning and coverage, including previously missed regions like the 'tail' of the cingulum and parts of the uncinate fasciculus.
• TractOracle-Net achieved 89.21% accuracy on the TractoInferno test split, outperforming state-of-the-art streamline classification models.
• Despite higher streamline counts, TractOracle maintained a low overlap-to-overreach ratio, indicating high anatomical specificity and reduced redundancy.

(b)