Source-linked AI summary
A foundation for systematic analysis of transformers and RNNs for tractography
Emmanuelle Renauld, Philippe Poulin, Hugo Larochelle, Antoine Théberge, Maxime Descoteaux
TL;DR
Iterative tractography must reconcile local diffusion information with globally plausible pathways, while local training loss can diverge from generated streamline quality. The paper evaluates RNNs and Transformers and introduces generation-validation supervision; its models outperform prior methods on ISMRM data.
Problem
Iterative tractography must reconcile local diffusion information with global anatomical plausibility, while local training loss may not reflect whole-streamline quality.
Source: §1 · p. 1, Local loss vs. tractography quality · p. 5Method
The paper evaluates RNNs and Transformers for iterative tractography and uses generation-validation to assess complete streamlines during validation.
Source: §4 · p. 30, Conclusion · p. 35Results
The authors report that their models outperform previous state-of-the-art methods on the ISMRM dataset.
Source: §4 · p. 27Takeaways & Limitations
The models reproduce training-set structure, so carefully curated training sets matter, particularly for bundles absent from training.
Source: §4 · p. 31The connectivity FP:TP validation metric can score a streamline as acceptable based on its endpoint region despite an invalid overall pathway.
Source: §2.3 · p. 10
Abstract
from arXiv · showhide
Machine learning (ML) has emerged as a promising approach for improving diffusion MRI (dMRI) tractography, a task that remains limited by the intrinsic tension between local diffusion information and global anatomical plausibility. In this work, we systematically evaluate recurrent neural networks (RNNs) and Transformer models for iterative tractography, with particular attention to training strategies, input representations (including convolutional neural network (CNN)-based embeddings and end-of-sequence (EOS) tokens), and hyperparameter selection. We introduce a generation-validation phase enabling supervision at the streamline level during training, allowing supervision despite the mismatch between local loss functions and global streamline quality. Using the ISMRM2015 tractography challenge dataset, our models achieve the highest reported performance to date. Through controlled experiments, we quantify the impact of missing bundles, noisy or imperfect training streamlines, and invalid fibers in the training set. Finally, we demonstrate the applicability of our best-performing models for in vivo data from the Tractoinferno database. Overall, our results highlight both the potential and the limits of sequence-based deep learning models such as Transformers and RNNs for tractography, and emphasize the need for improved phantoms and evaluation methods for in vivo validation. We provide takeaways and recommendations for future researchers training and validating sequence-based supervised methods for tractography.
1. Introduction
dMRI tractography must reconcile locally supported directions with globally plausible pathways, but image limitations prevent fully satisfying both objectives. The paper motivates machine learning as a way to incorporate broader anatomical knowledge into iterative whole-brain tractography.
1. Introduction
- ML-based iterative whole-brain tractography is the paper’s target, building on machine learning’s use across dMRI tasks.The paper lists local reconstruction, white-matter and bundle segmentation, tractogram filtering, tractometry, and connectomics among those tasks.Source: §1 · pp. 1–2
- Tractography advances streamlines step by step from seed points, and choices such as direction, step size, and stopping criteria strongly influence the result.Tracking repeats from the seed in the reversed direction to create the streamline’s second half.Source: §1 · p. 2
- Local diffusion information and global pathway plausibility are competing objectives that cannot both be fully satisfied, even with noise-free imaging.Following the main anisotropy direction can reduce spurious streamlines but also produce fewer pathways and poor spatial coverage than probabilistic tractography.Source: §1 · p. 2
- Global tractography is difficult to formulate and optimize, while iterative methods can follow locally supported but anatomically implausible paths, especially in crossings and bottlenecks.The paper notes that incorporating expert anatomical knowledge rigorously into whole-brain tracking remains difficult; machine learning may offer a way to use broader bundle-shape and brain-organization information.Source: §1 · pp. 2–3
RNNs and Transformers: the natural choices for tractography?
Machine-learning tractography models differ in how much streamline history informs their next-direction prediction. The paper focuses on sequence-native RNNs and Transformers, which retain information across the streamline.
RNNs and Transformers: the natural choices for tractography?
- Type 1 models use local dMRI features to predict direction, so they lack additional global information for ambiguous or poor-quality local regions.The paper characterizes them as improved local reconstruction algorithms that cannot outperform classical tractography in such regions.Source: RNNs and Transformers: the natural… · p. 3
- Type 2 models concatenate information from a fixed number of previous points, adding history but not explicitly representing the sequence.The resulting input is large, and the number of previous directions must be fixed.Source: RNNs and Transformers: the natural… · p. 3
- Figure 1 distinguishes local-only, fixed-history, and sequence-native inputs for next-direction prediction.Type 1 uses features at P4, type 2 uses concatenated previous points, and type 3 processes sequential data.Source: Figure 1 · p. 4
- RNNs retain streamline information through hidden states, while Transformers use attention; this work tests both with careful training-set and hyperparameter selection.The authors provide updated Learn2track and a Transformer model through the dwi_ml library.Source: RNNs and Transformers: the natural… · p. 4
Local loss vs. tractography quality
Iterative sequence models are trained with local losses but judged on complete generated streamlines, creating a mismatch compounded by exposure bias. The paper addresses validation with a generation-validation phase and evaluates its models against prior methods on ISMRM data.
Local loss vs. tractography quality
- At training time, the model predicts from reference-sequence history, but at generation it uses its own outputs, allowing errors to accumulate as exposure bias.This mismatch separates pointwise local loss from the quality of the final generated streamline.Source: Local loss vs. tractography quality · p. 5
- The section frames whether local loss can train tractography models and how to validate and tune them when that loss may not reflect generation quality.Figure 2 contrasts recurrent one-input-at-a-time RNN training with whole-streamline Transformer training; repeated Transformer inputs make tracking more computationally expensive.Source: Local loss vs. tractography quality · p. 5, Figure 2 · p. 6
- The generation-validation phase generates complete streamlines during validation, enabling assessment of model quality beyond local loss.It also allows on-the-fly comparison of models using different local loss functions, while being quicker than tracking at every training epoch.Source: Local loss vs. tractography quality · p. 5
- The authors report that their models outperform state-of-the-art methods on the ISMRM 2015 Tractography Challenge dataset after an extensive hyperparameter search.They also validated the models in vivo on the Tractoinferno database.Source: Local loss vs. tractography quality · p. 6
2. Methods
The methods introduce generation-validation for assessing generated streamlines during training, evaluate RNN and Transformer variants and training choices, and test models on challenge and in vivo data. Experiments also examine attention weights and the effects of incomplete or imperfect training streamlines.
2.1 Preparation of the GV phase
- Generation-validation re-tracked validation streamlines from their first coordinate and scored endpoint connectivity against a validation-derived matrix, using FP:TP while excluding TN and FN rates.With N = 1, the procedure evaluated generated segments each epoch while controlling streamline count, rather than requiring whole-brain tractography.Source: §2.1 · pp. 6–7
2.2 Model architecture
- The models included Learn2track RNNs and causal full encoder–decoder or decoder-only Transformers, with direction inputs and optional SOS/EOS tokens.Future inputs were masked because their dMRI information depends on generated coordinates.Source: Figure 3 · p. 7, §2.2 · pp. 8–9
- The tested local objectives included direction regression, classification, smooth classification, and Gaussian or Fisher-von-Mises parameter regression.Inputs comprised MRI data sequences, including raw dMRI or fODF features and optional T1 or FA values; targets were sequences of directions.Source: Table 1 · p. 8, §2.2 · p. 8
2.3 Experiments
- The ISMRM hyperparameter search varied losses, sequence and embedding choices, input modalities, optimizers, model settings, and data augmentation, retaining configurations with consistent GV validation improvements.Training used noisy dMRI inputs and ground-truth streamline targets, with training and validation subsets downsampled to 70,000 streamlines each.Source: §2.3 · pp. 9–10
- Whole-brain test tracking measured invalid-streamline ratio and F1, while in vivo quality was assessed visually and described as exploratory rather than precisely quantifiable.The GV FP:TP rate could score an endpoint connection positively despite an invalid overall pathway.Source: §2.3 · pp. 10, 13
- Controlled ISMRM experiments tested recovery of three omitted bundles and training on deterministic, probabilistic, or combined classical-tractography streamlines, including bundle-specific tracking for difficult bundles.The omitted bundles were the corpus callosum, left superior cerebellar peduncle, and anterior commissure; the anterior commissure and corpus callosum were not recovered in initial classical tracking for the imperfect-streamline experiment.Source: §2.3 · p. 11
- A separate ISMRM experiment added back 5%, 10%, 20%, or 50% invalid streamlines; Tractoinferno training used 18 subjects, six validation subjects, and three test subjects.For Tractoinferno, the authors used T1, FA, and fODF inputs and included available bundles even when they had not passed the authors’ quality control.Source: §2.3 · p. 12
2.4 Analysis of the Transformer’s attention
- The attention analysis visualized TractoTransformer weights per streamline and as voxel-level averages in in vivo data from one subject.The authors also colored duplicated streamline segments according to attention weights used when predicting the next point.Source: §2.4 · p. 13
3. Results
Learn2track with regression produced the strongest results, while experiments showed how training data, missing bundles, and in-vivo evaluation affect tractography outcomes. In vivo results reproduced most bundles, but visual comparison suggested lower quality than on ISMRM data.
Experiment 1: Hyperparameter search
- Learn2track with regression achieved the best results, and all models outperformed previous work; its generated tractogram resembled ground truth but was less dense.The comparison particularly favored the models on F1 score.Source: §3 · p. 14, Figure 5 · p. 16, Figure 6 · p. 17
Experiment 2: Recovering missing bundles
- Bundles omitted from training were not recovered or were poorly reconstructed, even when tracking began extensively from their endpoints.This was observed for missing bundles including the CA and SCP.Source: §3 · p. 19, Figure 8 · p. 19
Experiment 3: Deterministic and probabilistic streamlines
- Probabilistic training produced the best results, but lower F1 than ground-truth training; deterministic training reduced overreach while improving alignment with local information.Probabilistic training also showed stronger overfitting of local loss, with high validation errors at points that varied from training streamlines.Source: §3 · p. 19, Table 5 · p. 20, Figure 9 · p. 21
Experiment 4: Adding invalid streamlines
- Adding invalid streamlines to training decreased valid-streamline ratios and F1, although F1 was less affected partly because overreach improved.Scores remained similar to those generally obtained with classical tractography.Source: §3 · p. 20, Table 6 · p. 22
Experiment 5: In-vivo results
- The models reproduced most bundles in vivo, including for a subject from a site excluded from training, but visual results appeared lower quality than ISMRM results.The valid-streamline ratio could not be evaluated for in-vivo data.Source: §3 · p. 22, Figure 12 · p. 24, Figure 13 · p. 25
Visualisation of Transformer attention
- Transformer attention spanned points across tracking time rather than focusing only on the current point, and average attention varied across voxels.The visualized attention averaged the two heads of the first layer; the voxel map averaged all three layers.Source: §3 · p. 26, Figure 14 · p. 26, Figure 15 · p. 27
4. Discussion
Sequence-based tractography models perform strongly on the ISMRM benchmark, but their validation, training-data coverage, and in-vivo evaluation remain important constraints.
4. Discussion
- RNN and Transformer models achieved the authors’ highest-reported ISMRM performance; Learn2track outperformed TractoTransformer and tracked faster.The authors attribute their results to the generation-validation phase and extensive hyperparameter search.Source: §4 · pp. 27–28
- Adding an EOS token had an important effect and removed the need for a tracking mask; hyperparameter choices substantially affected quality, and loss functions differed in speed and stability.Classification trained faster, regression gave better results, and Gaussian and Fisher-von-Mises models were less stable; the latter did not include later entropy regularization.Source: §4 · p. 28
- Increasing the tracking step size from 1 mm to 2 mm helped the models without reducing tractogram quality.Smaller steps create longer sequences, which can challenge RNN long-term dependencies and dilute Transformer attention across points.Source: §4 · pp. 28–29
- Input-feature choices had limited effects: adding previous directions helped little, T1 and FA helped moderately, raw dMRI was not better, and CNN embeddings did not improve results.The authors suggest neighborhood position may be useful for estimating direction, unlike translation-invariant CNN processing.Source: §4 · p. 29
- Local losses trained models effectively but correlated weakly with other quality metrics, motivating generation-validation against global measures.The GV phase completed validation streamlines through tractography and scored generated batches; connectivity false-positive rate correlated well with final tractometry scores in Experiment 1.Source: §4 · pp. 27, 29, 30, 31
- Reference-direction losses can overfit training directions, while peak-based losses may favor deterministic-like tracking; sequence-level learning remains a possible but unstable alternative.The authors note that sequence-level learning requires initialization with local-loss training and has not yet been tested for tractography.Source: §4 · p. 30
- Models largely reproduce training-set anatomy, while in-vivo experiments showed formed bundles but noisier results; those experiments are feasibility evidence, not quantitative generalization.The discussion calls for curated training coverage and more precise in-vivo quality-control metrics.Source: §4 · pp. 31, 34
Conclusion
The paper evaluates RNN- and Transformer-based iterative tractography models and reports successful models for ISMRM and in vivo data. It identifies the mismatch between local training objectives and global tractography quality, notably due to exposure bias, as a central limitation.
Conclusion
- The study presents successful tractography models for ISMRM data and in vivo data.It systematically evaluates RNN- and Transformer-based models for iterative tractography.Source: Conclusion · p. 35
- Sequence-based models better capture streamline dependencies than purely local approaches, but local training objectives remain mismatched with global tractography quality.The paper identifies exposure bias as a notable source of this limitation.Source: Conclusion · p. 35
- The findings provide guidance for future work toward models that better integrate global anatomical consistency during generation.Source: Conclusion · p. 35
Ethics
The study used only publicly available, de-identified datasets and collected no new human-participant data. Under institutional and national guidelines, this secondary analysis required no additional ethical approval.
Ethics
- The study used exclusively publicly available, de-identified datasets.Source: Ethics · p. 41
- The authors collected no new data involving human participants.Source: Ethics · p. 41
- The secondary analysis required no additional ethical approval under institutional and national guidelines.Source: Ethics · p. 41