[논문 리뷰] TractOracle: towards an anatomically-informed reward function for RL-based tractography

Antoine Théberge, Maxime Descoteaux외 1명|arXiv (Cornell University)|2024. 03. 26.

Advanced Neuroimaging Techniques and ApplicationsMedicine인용 수 3

한 줄 요약

TractOracle는 트랜스포머 기반의 스트림라인 분류 네트워크(TractOracle-Net)를 통해 해부학적 정보를 반영한 보상 함수를 통합한 새로운 강화학습(RL)-기반 연장도 프레임워크를 제안한다. 추적 과정에서 스트림라인의 타당성을 실시간으로 평가함으로써 정확한 연장도 추적과 동시에 잘못된 양성 결과를 제거하며, 기준 데이터셋에서 최신 기술 대비 최대 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.

연구 동기 및 목표

• 기존의 RL 기반 연장도 방법이 해부학적 사전 지식을 갖추지 못해 높은 참 양성 결과 비율을 보이는 데 대한 한계를 해결하기 위해.
• 실시간 해부학적 타당성 평가를 통해 스트림라인을 동시에 평가하고 재구성하는 통합 프레임워크를 개발하기 위해.
• 해부학적 지식을 직접 RL 학습 과정에 통합함으로써 사후 필터링에 대한 의존도를 줄이기 위해.
• 사전 훈련된 스트림라인 분류 네트워크를 동적 보상 및 정지 신호로 활용해 연장도 생성 및 세그멘테이션 정확도를 향상시키기 위해.
• 기존의 고전적 및 학습 기반 연장도 기준 대비 시뮬레이션 및 생체 내 데이터셋에서 뛰어난 성능을 달성하기 위해.

제안 방법

• 트랜스포머 기반 신경망인 TractOracle-Net을 사용하여 스트림라인의 형태와 궤적을 기반으로 그 해부학적 타당성 여부를 분류하도록 훈련한다.
• RL 학습 과정 중 TractOracle-Net은 추적 중인 각 스트림라인에 대해 연속적인 해부학적 타당성 점수를 제공하며, 이는 조밀하고 실시간 보상 신호로 기능한다.
• 에이전트는 이 보상을 활용해 해부학적으로 정확한 스트림라인 경로를 최적화하며, '보상 유도 공격(reward hacking)'과 잘못된 양성 결과 생성을 최소화한다.
• 해부학적으로 타당성이 떨어지는 스트림라인이 추적되지 않도록, 타당성 점수가 임계값 이하로 떨어지면 조기에 정지한다.
• 훈련 후에도 동일한 TractOracle-Net 모델을 사용해 연장도 세그멘테이션 및 필터링을 수행함으로써 종단 간 해부학적 평가가 가능해진다.
• 표준 연장도 평가 지표와 Recobundles, extractor_flow 등의 도구를 사용하여 ISMRM2015(시뮬레이션) 및 TractoInferno(생체 내) 데이터셋에서 프레임워크를 훈련 및 평가한다.

(a)

실험 결과

연구 질문

• RQ1 학습용 레이블이 필요 없이도 학습된 해부학적 오라클이 RL 기반 연장도의 정확도를 향상시키고 참 양성 결과를 줄일 수 있는가?
• RQ2 추적 과정 중 실시간 스트림라인 타당성 평가가 재구성된 연장도의 품질을 얼마나 향상시킬 수 있는가?
• RQ3 분류 네트워크를 보상 함수이자 조기 정지 기준으로 통합함으로써 기존의 RL 또는 필터링 파이프라인 대비 연장도 성능에 어떤 영향을 미치는가?
• RQ4 제안된 방법이 기존의 최신 기술 수준의 고전적 및 학습 기반 연장도 방법보다 참 양성률, 참 양성 결과 비율, 해부학적 정확도 측면에서 뛰어난 성능을 보일 수 있는가?
• RQ5 기존 방법에서 자주 간과되는 부족하거나 미세한 백질 경로의 복원이 이 시스템을 통해 유지되거나 향상되는가?

주요 결과

• ISMRM2015 시뮬레이션 데이터셋에서 TractOracle는 22%의 유효 연결 비율을 기록했으며, 이는 두 번째로 우수한 방법 대비 20% 향상된 결과이다.
• 동일한 데이터셋에서 TractOracle는 유효하지 않은 연결 비율과 연결이 없는 비율을 각각 두 번째로 우수한 방법 대비 3배씩 감소시켰다.
• TractoInferno 생체 내 데이터셋에서 TractOracle-RL은 TractOracle-Net에 의해 타당성 있는 것으로 분류된 184,355,412개의 스트림라인을 생성했으며, 이는 sd_stream 대비 7배, Track-to-Learn 대비 2배 많다.
• 이 방법은 이전에 간과되었던 영역인 쿠니큘럼의 '꼬리' 부분과 언디네이트 피ascิ큘러의 일부 영역까지 포함해 해부학적으로 타당한 번들을 뛰어난 산란과 커버리지로 재구성했다.
• TractoInferno 테스트 세트에서 TractOracle-Net은 89.21%의 정확도를 기록했으며, 최신 기술 수준의 스트림라인 분류 모델들을 능가했다.
• 더 높은 스트림라인 수에도 불구하고 TractOracle는 낮은 오버랩 대 오버드라이브 비율을 유지하여 높은 해부학적 특이성과 낮은 중복을 보였다.

(b)