[论文解读] 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 ApplicationsMedicine被引用 3

一句话总结

TractOracle 提出了一种基于强化学习（RL）的新型纤维束成像框架，通过基于 Transformer 的纤维束分类网络（TractOracle-Net）整合了解剖学先验的奖励函数。在追踪过程中实时评估纤维束的合理性，该系统同时引导准确的纤维束成像并修剪假阳性结果，在基准数据集上相比最先进方法实现了高达 20% 的真阳性率提升和 3 倍的假阳性减少。

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.

研究动机与目标

• 为了解决现有基于强化学习的纤维束成像方法缺乏解剖学先验，导致假阳性率较高的局限性。
• 开发一个统一框架，通过实时解剖学合理性评分，同时评估并重建白质纤维束。
• 通过将解剖学知识直接嵌入强化学习训练过程，减少对事后过滤的依赖。
• 通过利用预训练的纤维束分类网络作为动态奖励和停止信号，提升纤维束图生成与分割的准确性。
• 在仿真与在体数据集上，实现优于经典方法与学习型基线的卓越性能。

提出的方法

• 训练一个基于 Transformer 的神经网络 TractOracle-Net，根据纤维束的形态与轨迹，将其分类为解剖学上合理或不合理。
• 在强化学习训练过程中，TractOracle-Net 为每条正在追踪的纤维束提供连续的解剖学合理性评分，作为密集的、实时的奖励信号。
• 智能体利用该奖励优化解剖学上准确的纤维束路径，最大限度减少“奖励欺骗”与假阳性纤维束的生成。
• 当合理性评分低于阈值时，系统执行早期停止，防止追踪解剖学上不合理的纤维束。
• 训练后，同一 TractOracle-Net 模型被用于纤维束图的分割与过滤，实现端到端的解剖学评估。
• 该框架在仿真数据集（ISMRM2015）与在体数据集（TractoInferno）上进行训练与评估，使用标准纤维束成像指标与工具（如 Recobundles 和 extractor_flow）。

(a)

实验结果

研究问题

• RQ1 学习得到的解剖学‘圣杯’是否能在无需标注训练数据的情况下，提升基于强化学习的纤维束成像精度并减少假阳性？
• RQ2 在追踪过程中实时进行纤维束合理性评分，能在多大程度上提升重建纤维束图的质量？
• RQ3 将分类网络同时作为奖励函数与早期停止标准，与标准强化学习或过滤流程相比，对纤维束成像性能有何影响？
• RQ4 所提出的方法是否能在真阳性率、假阳性率与解剖学保真度方面超越最先进经典与学习型纤维束成像方法？
• RQ5 该系统是否能保留或改善对传统方法常遗漏的缺失或细微白质通路的恢复？

主要发现

• 在 ISMRM2015 仿真数据集上，TractOracle 实现了 22% 的有效连接率，较第二好的方法提升了 20%。
• 在相同数据集上，TractOracle 将无效连接率与无连接率均降低了 3 倍，相比第二好的方法。
• 在 TractoInferno 在体数据集上，TractOracle-RL 生成了 184,355,412 条由 TractOracle-Net 识别为合理的纤维束——是 sd_stream 的 7 倍，也是 Track-to-Learn 的 2 倍。
• 该方法重建了解剖学上合理的纤维束束，具有更优的发散性与覆盖范围，包括此前被遗漏的区域，如扣带的‘尾部’及钩束的部分区域。
• TractOracle-Net 在 TractoInferno 测试集上达到 89.21% 的准确率，优于当前最先进纤维束分类模型。
• 尽管纤维束数量更高，TractOracle 仍保持较低的重叠-过度延伸比，表明其具有高解剖学特异性与较低冗余。

(b)