[Paper Review] Trackastra: Transformer-based cell tracking for live-cell microscopy

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

Cell Image Analysis Techniques4 citations

TL;DR

Trackastra is a transformer-based cell tracking method that directly learns pairwise associations between cell detections within a short temporal window, using only object positions and basic shape features. It achieves state-of-the-art performance on diverse biological datasets—including bacteria, cell cultures, and fluorescent particles—by employing a parental softmax to improve tracking of dividing cells, enabling accurate results even with simple greedy linking, thus reducing reliance on complex optimization pipelines.

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.

Motivation & Objective

• To address the challenge of accurate cell tracking in live-cell microscopy, where cells are visually similar, divide frequently, and exhibit complex motion patterns.
• To reduce dependence on computationally expensive discrete optimization methods (e.g., ILP) in the linking step by learning association costs directly via a neural network.
• To develop a general-purpose tracking framework that works across diverse biological datasets without requiring dataset-specific tuning.
• To improve tracking of dividing cells by explicitly modeling mother-daughter relationships through a novel parental softmax normalization.
• To enable robust tracking with simple greedy linking by learning biologically plausible associations directly from data.

Proposed method

• Trackastra processes object detections from a short temporal window (e.g., 3–6 frames) as tokens in a transformer encoder-decoder architecture.
• Each detection is embedded using position and basic shape features (e.g., centroid, area, eccentricity), forming the input tokens.
• The model uses self- and cross-attention mechanisms to compute all-to-all pairwise associations between detections across time.
• A novel parental softmax normalization is applied to the association matrix to prioritize biologically plausible mother-daughter links, especially during cell division.
• The predicted association matrix is used to construct a candidate graph, which is pruned via either greedy linking or integer linear programming (ILP) to produce final tracks.
• The model is trained end-to-end using a combination of cross-entropy loss for associations and a dedicated loss component for parental relationships.

Experimental results

Research questions

• RQ1 Can a pure transformer architecture learn biologically plausible cell associations in live-cell microscopy without relying on complex discrete optimization?
• RQ2 Does the proposed parental softmax normalization improve tracking accuracy, particularly for dividing cells?
• RQ3 How much temporal context (window size) is necessary for effective tracking with the transformer-based approach?
• RQ4 To what extent can a single Trackastra model generalize across diverse biological datasets with different imaging modalities and cell types?
• RQ5 Can Trackastra achieve competitive performance with simple greedy linking, reducing the need for computationally heavy ILP solvers?

Key findings

• Trackastra achieves state-of-the-art performance on multiple datasets, including bacteria, cell cultures, and fluorescent particles, outperforming or matching highly tuned baselines.
• The parental softmax reduces AOGM (Average Overlap of Ground Truth Matches) by approximately 20% on the Bacteria dataset, with both greedy and ILP linkers.
• Using a window size of 3–6 frames yields strong performance, while reducing to s=2 leads to a significant drop in accuracy, indicating that short temporal context is sufficient.
• Removing positional encodings or shape descriptors increases errors substantially, demonstrating their importance for spatial and morphological reasoning.
• The model generalizes well across datasets: a single trained Trackastra model performs well on both Bacteria and HeLa cell datasets without fine-tuning.
• Even with greedy linking, Trackastra achieves AOGM of 23.0 on Bacteria (vs. 18.8 with ILP), showing that high accuracy can be achieved without complex optimization.