[論文レビュー] Trackastra: Transformer-based cell tracking for live-cell microscopy

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

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 モデルが、異なる画像モalityおよび細胞タイプを有する多様な生物学的データセットにどの程度一般化できるか？
• RQ5 グリーディリンクを用いても、Trackastra は計算コストの高い ILP ソルバーの必要性を低減しつつ、競争力のある性能を達成できるか？

主な発見

• Trackastra は、細菌、細胞系、蛍光粒子を含む複数のデータセットで最先端の性能を達成し、高度にチューニングされたベースラインを上回るか同等の性能を示した。
• 親子関係用ソフトマックス正規化により、Bacteria データセットでは AOGM（正解マッチの平均オーバーラップ）が約 20% 減少した。グリーディリンクおよび ILP リンカーの両方で同様の効果が得られた。
• 窓サイズを 3〜6 フレームに設定すると高い性能が得られ、s=2 に短縮すると顕著な精度低下が生じた。これは、短い時間的文脈が十分であることを示している。
• 位置エンコーディングや形状記述子を削除すると誤差が著しく増加し、空間的および形態的推論においてそれらの重要性が示された。
• モデルは一般化性能に優れている：1つのトレーニング済み Trackastra モデルは、微調整なしに Bacteria および HeLa 細胞データセットの両方で良好な性能を発揮した。
• グリーディリンクでさえ、Bacteria で AOGM 23.0 を達成（ILP では 18.8）した。これは、複雑な最適化を伴わずに高精度な追跡が可能であることを示している。