Whole Heart 3D+T Representation Learning Through Sparse 2D Cardiac MR Images

We propose a simple but effective solution to correlate the "interlaced" multi-view CMR sequences and learn a whole-heart representation from them. We assume that each scan includes a set of sparse 2D+T CMR sequences with $M$ SA planes $\{S_{1},...,S_{M}\}$ and $N$ LA planes $\{L_{1},...,L_{N}\}$. Directly stacking all 2D+T sequences into a 4D tensor and conducting a 4D operation is computationally prohibitive. Nevertheless, LA and SA planes are acquired from different views and do not share spatial feature similarity. Thus, instead of forming a 4D tensor we first decompose each plane into $P$ small 2D+T patches, denoted as $s^{m}_{p}$ for the $p$-th patch of $m$-th SA plane and $l^{n}_{p}$ for the $p$-th patch of $n$-th LA plane. Our method takes all these individual patches as an input vector, denoted as $X=[s^{1}_{1},...,s^{M}_{P},l^{1}_{1},...,l^{N}_{P}]$, containing in total $(M+N)\times P$ patches. To capitalize on the temporal redundancy inherent in CMR data and further reduce computational demands, we use a larger patch size along the temporal than the spatial dimension. To enhance localization, each patch is added with a 4D positional embedding [11], indicating the patch’s $x\textrm{--}y\textrm{--}t$ spatiotemporal index and the corresponding plane of origin. Then a random mask is applied on the input to mask away $q\%$ of the patches. The shortened input with remaining patches, $\hat{X}$ (refer to Fig 1), are then fed into an encoder $\mathcal{E}$ to learn a dense 3D+T cardiac representation. A high maskout ratio $q\%$ is chosen here to force the encoder to find the underlying spatiotemporal correlation across different cardiac planes, rather than easily extrapolate missing pixel intensity. Afterwards, a lightweight decoder $\mathcal{D}$ is applied to predict the masked-out patches and reconstruct $X$. We applied the mean squared error (MSE) for optimization with a loss function formulated as $\mathcal{L}=\left\|X-\mathcal{D}\left(\mathcal{E}\left(\hat{X}\right)\right)\right\|_{2}^{2}.$

The 3D+T whole heart representation can be used to solve different downstream tasks. We fine-tune our model on a small-scale dataset where corresponding labels are available. In phase II, we forward all available CMR planes to the model without maskout. The encoder remains unchanged and the reconstruction decoder is replaced with a different decoder according to the task we aim to solve. For cardiac phenotype (e.g. ejection fraction and volume) prediction, we use shallow linear layers since we assume the spatiotemporal representation learned in phase I is already comprehensive and can represent cardiac morphology to some extent. For all-planes CMR sequences segmentation, a U-Net-wise decoder with skip-connections [34] is adopted. It is noteworthy that, in contrast to conventional cardiac segmentation (refer to section  1.1) which can only conduct specific 2D plane segmentation at one inference time, our model can segment all CMR planes/frames in an end-to-end fashion at a single inference time. Furthermore, since the model can exploit the correlation between different CMR planes, the final segmentation is further enhanced and is expected to perform better than single plane segmentation [5] (more details refer to Table 2 and Fig. 3).

In clinic scenarios, variations in acquisition protocols or artifacts from patient motion can lead to missing/corrupted CMR planes. Our method, however, remains robust against these challenges due to its contrastive learning inherence. Contrastive learning is usually used in representation learning to pull positive (similar) data and push dissimilar data [9, 7]. Our method employs a form of implicit contrastive learning, aligning positive pairs through diverse mask patterns [32]. During different training epochs in phase I, our model is forced to generate the same representation/reconstruction for a subject, regardless of different random mask applications. This ensures that even in scenarios where certain CMR planes are missing, our model can still deliver the same representation as complete CMR plane inputs.

We first evaluate reconstruction using only SA sequences, only LA sequences, and all available views in pre-training. The quantitative results and reconstructed CMR images are shown in Appendix Table 3 and Fig. 4. The superior results of all-view reconstruction to that of using only limited spatial planes imply the benefit of involving all-view spatiotemporal information for whole-heart representation learning. Incorporating different views can help the model build spatial knowledge of the heart and therefore enhance the representation/reconstruction.

We further provide the t-SNE visualization of the learned representations in Fig. 2. We label the latent embeddings after pre-training with right ventricular ejection fraction (RVEF) and left ventricle mass (LVM). The generation of these two phenotypes relies on 3D spatial and temporal information of the heart. Notably, even without using any labels during pre-training, our model can already generate a well-clustered latent space reflective of spatiotemporal differences across the subjects. This visualization serves as strong evidence for the effectiveness of the learned representations.

We assess the efficacy of learned whole-heart representations by predicting critical phenotype values. Table 1 shows the mean absolute error of our predictions compared to ResNet50 and ViT for LVM, RVEF, right atrial ejection fraction (RAEF), right ventricular end-diastolic volume (RVEDV), and left atrial stroke volume (LASV). For accurate prediction of these values, both adequate spatial and temporal cardiac information is essential. The superior performance to the baseline methods underscores its capability to capture high-level spatiotemporal features of the entire heart.

The competence of our model to perform end-to-end segmentation across all planes is shown in Table 2 and Fig. 3. Our model not only exhibits comparable quantitative dice scores with nnUNet (that leverages exhaustive parameters tuning) but also shows superior performance over UNETR+ in all regions for both SA/LA planes. This benefit is gained from the learned whole-heart representation. Moreover, the superior performance against SA-only and LA-only segmentation also highlights the significance of integrating multi-view CMR information for more accurate segmentation outcomes.

To simulate scenarios with incomplete or defective CMR data, we randomly remove 1 or 2 CMR planes from each subject and generate a new latent representation. Calculating the cosine similarity between this representation and the one with all planes available, we observe an average value of 1.0 in both cases, indicating a total alignment. Further, we conduct downstream LVM and RVEF predictions with 2 slices masked out and compare them with the regular setting (no maskout). The predicted value difference is $\pm 0.3$g for LVM and $\pm 0.4\%$ for RVEF, emphasizing its robustness and significant potential in practical scenarios with missing planes.

While the presented evaluation is limited to prediction and segmentation tasks, future work will further explore tasks requiring whole 3D+T information, such as cardiac motion tracking and other modality mapping (e.g. genetics) to widen the scope of cardiac analysis. While our model demonstrates meaningful representations across diverse subjects, it has not been evaluated on longitudinal data. This presents an avenue for tracking potential cardiac disease progression by observing patient trajectories in the latent space. Additionally, the capability of our method to offer comprehensive cardiac representations with fewer CMR images opens opportunities to reduce CMR scan times and alleviate patient discomfort.