[논문 리뷰] Trackastra: Transformer-based cell tracking for live-cell microscopy

Benjamin Gallusser, Martin Weigert|arXiv (Cornell University)|2024. 05. 24.

Cell Image Analysis Techniques인용 수 4

한 줄 요약

Trackastra는 단일 시간 창 내에서 세포 검출 간의 쌍별 관계를 직접 학습하는 트랜스포머 기반 세포 추적 방법으로, 오직 객체의 위치와 기본 형태 특징만을 사용한다. 부모 소프트맥스를 도입하여 분열하는 세포의 추적 성능을 향상시켜, 단순한 게으른 연결 방식을 사용함에도 불구하고 복잡한 최적화 파이프라인에 의존하지 않도록 하여 다양한 생물학적 데이터셋(박테리아, 세포 배양, 형광 입자 포함)에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Cell tracking is a ubiquitous image analysis task in live-cell microscopy. Unlike multiple object tracking (MOT) for natural images, cell tracking typically involves hundreds of similar-looking objects that can divide in each frame, making it a particularly challenging problem. Current state-of-the-art approaches follow the tracking-by-detection paradigm, i.e. first all cells are detected per frame and successively linked in a second step to form biologically consistent cell tracks. Linking is commonly solved via discrete optimization methods, which require manual tuning of hyperparameters for each dataset and are therefore cumbersome to use in practice. Here we propose Trackastra, a general purpose cell tracking approach that uses a simple transformer architecture to directly learn pairwise associations of cells within a temporal window from annotated data. Importantly, unlike existing transformer-based MOT pipelines, our learning architecture also accounts for dividing objects such as cells and allows for accurate tracking even with simple greedy linking, thus making strides towards removing the requirement for a complex linking step. The proposed architecture operates on the full spatio-temporal context of detections within a time window by avoiding the computational burden of processing dense images. We show that our tracking approach performs on par with or better than highly tuned state-of-the-art cell tracking algorithms for various biological datasets, such as bacteria, cell cultures and fluorescent particles. We provide code at https://github.com/weigertlab/trackastra.

연구 동기 및 목표

• 세포가 시각적으로 유사하고 자주 분열하며 복잡한 운동 패턴을 보이는 생체세포 현미경에서 정확한 세포 추적의 과제를 해결하기 위해.
• 연결 단계에서 복잡한 이산 최적화 방법(예: 정수선형계획법, ILP)에 의존하는 것을 줄이기 위해 신경망을 통해 직접 연관 비용을 학습하기 위해.
• 데이터셋에 특화된 튜닝 없이도 다양한 생물학적 데이터셋에서 작동하는 일반 목적의 추적 프레임워크를 개발하기 위해.
• 새로운 부모 소프트맥스 정규화를 통해 어머니-자식 관계를 명시적으로 모델링하여 분열하는 세포의 추적을 향상시키기 위해.
• 데이터에서 생물학적으로 타당한 관계를 직접 학습함으로써 단순한 게으른 연결 방식으로도 견고한 추적을 가능하게 하기 위해.

제안 방법

• Trackastra는 짧은 시간 창(예: 3–6 프레임) 동안의 객체 검출을 트랜스포머 인코더-디코더 아키텍처의 토큰으로 처리한다.
• 각 검출 결과는 중심점, 면적, 타원도 등의 위치 및 기본 형태 특징을 사용해 임베딩되어 입력 토큰을 형성한다.
• 모델은 자기 및 크로스 어텐션 메커니즘을 사용하여 시간 간의 모든 검출 간 쌍방향 관계를 계산한다.
• 새로운 부모 소프트맥스 정규화가 관계 행렬에 적용되어 특히 세포 분열 기간 동안 생물학적으로 타당한 어머니-자식 연결을 우선순위로 지정한다.
• 예측된 관계 행렬은 후보 그래프를 구성하는 데 사용되며, 최종 트랙을 생성하기 위해 게으른 연결 또는 정수선형계획법(ILP)을 통해 정렬된다.
• 모델는 연관성에 대한 교차 엔트로피 손실과 부모 관계에 특화된 손실 성분을 조합하여 엔드 투 엔드로 훈련된다.

실험 결과

연구 질문

• RQ1 복잡한 이산 최적화에 의존하지 않고 순수한 트랜스포머 아키텍처가 생체세포 현미경에서 생물학적으로 타당한 세포 간 관계를 학습할 수 있는가?
• RQ2 제안된 부모 소프트맥스 정규화가 특히 분열하는 세포의 추적 정확도를 향상시키는가?
• RQ3 트랜스포머 기반 접근법에서 효과적인 추적을 위해 필요한 시간적 맥락(창 크기)은 어느 정도인가?
• RQ4 동일한 Trackastra 모델이 다양한 생물학적 데이터셋(다른 이미징 모odalities 및 세포 유형 포함)에서 얼마나 잘 일반화되는가?
• RQ5 복잡한 계산 비용이 큰 ILP 솔버가 필요 없이 단순한 게으른 연결 방식으로도 Trackastra가 경쟁 가능한 성능을 달성할 수 있는가?

주요 결과

• Trackastra는 박테리아, 세포 배양, 형광 입자 등 여러 데이터셋에서 최신 기술 수준의 성능을 달성하며, 정교하게 튜닝된 기존 기준 모델을 초월하거나 동등하게 성능을 내고 있다.
• 부모 소프트맥스는 게으른 연결 및 ILP 연결 방식을 모두 사용할 때 Bacteria 데이터셋에서 AOGM(실제 매칭의 평균 오버랩)을 약 20% 감소시켰다.
• 3–6 프레임의 창 크기를 사용할 경우 뛰어난 성능를 보이며, 창 크기를 s=2로 줄이면 정확도가 크게 떨어져 짧은 시간 맥락이 충분함을 시사한다.
• 위치 인코딩 또는 형태 기반 기술을 제거할 경우 오류가 크게 증가하여, 공간적 및 형태적 추론에서 이들이 중요한 역할을 한다는 것을 입증한다.
• 모델은 일반화 능력이 뛰어나, 한 번 학습된 Trackastra 모델이 피팅 튜닝 없이도 Bacteria 및 HeLa 세포 데이터셋 모두에서 잘 작동한다.
• 게으른 연결 방식을 사용하더라도 Trackastra는 Bacteria에서 AOGM 23.0을 달성하여 ILP를 사용할 경우의 18.8보다 낮지 않게 성능을 유지함으로써, 복잡한 최적화 없이도 높은 정확도를 달성할 수 있음을 보여준다.