POMP: Pathology-omics Multimodal Pre-training Framework for Cancer Survival Prediction

1 Introduction

Cancer is widely recognized as a major global public health concern. For instance, there were 19.3 million newly diagnosed cases of cancer in 2020, and approximately 10.0 million deaths attributed to cancer [Sung et al., 2021]. Cancer survival prediction is an important direction in precision medicine, which aims to predict the time from cancer diagnosis to death events in cancer patients, thereby helping clinicians to tailor optimal treatment regimens [Vale-Silva and Rohr, 2021]. Traditional survival prediction methods mainly relied on clinical features like symptoms, signs, tumor biomarkers, and biochemical tests, as well as follow-up notes recording recurrence, metastasis, and response to treatment [Gensheimer et al., 2019; Chicco and Jurman, 2020]. Unfortunately, these methods are laborious and unfeasible when applied in clinical practice [Wang et al., 2019]. With the swift advancements in whole slide imaging and high-throughput sequencing technologies, survival prediction has transitioned from clinical notes to multi-omics data and pathology images, which are also known as whole slide images (WSIs) [Herrmann et al., 2021; Srinidhi et al., 2021]. Computational pathology methods can extract gold-standard information on survival prediction from pathology images regarding tumor cells and their microenvironment [Herrmann et al., 2021], and integrate multi-omics data such as mRNA, miRNA, and DNA methylation to enhance prediction from a molecular perspective [Srinidhi et al., 2021].

∗Corresponding author: Qingchen Zhang.

In recent years, multimodal learning has made promising progress in survival prediction by capturing complementary information from various modal perspectives [Cui et al., 2023]. However, most works simply concatenate representations extracted from multimodal data, and employ crossmodal attention mechanism [Chen et al., 2021; Wang et al., 2023; Li et al., 2022] or tensor fusion method [Chen et al., 2020; Wang et al., 2021a] to learn potential interaction information. Due to the lack of multimodal alignment, these methods prevent the models from fully utilizing the information between different modalities. Specifically, the survival prediction model is trained to learn pathology representation and omics representation from their own spaces, which has two main limitations: (i) it poses a challenge for the multimodal fusion module to fuse different modalities from different semantic spaces; (ii) this would lose the pathology-histology alignment information that could potentially be used as a signal for survival prediction.

In this work, we break with traditional approaches to study how to align multimodal information from the same patient before fusion while serving as survival predictions. To do so, we propose a pathology-omics multimodal pre-training (POMP) framework for cancer survival prediction, leveraging the principle of pre-trained models and exploring the benefit of aligning multimodal information from the same patient. In practice, we first encode the pathology image and multiomics data independently with a pathology encoder and an omics encoder, and then use a multimodal encoder to fuse representations of image and omics data. In particular, we introduce a pathology-omics contrastive learning method based on cosine similarity to align multimodal information, making the two modal data from the same patient closer in latent space. To fully learn the potential interaction information between different modalities, POMP is jointly pre-trained with three tasks based on the self-supervised learning paradigm. During fine-tuning, we leverage the cosine similarity, calculated in the same way as the pathology-omics contrastive learning, as the survival risk score for cancer survival prediction.

We conduct extensive experiments on six cancer datasets from the Cancer Genome Atlas (TCGA) [Tomczak et al., 2015]. The experimental results show that POMP consistently outperforms existing state-of-the-art methods across six datasets.

Our primary contributions are summarized as follows:

1. We propose a novel multimodal pre-training framework jointly learned with three training tasks to integrate pathological images and multi-omics data for cancer survival prediction.
2. We introduce a pathology-omics contrastive learning method to align the multi-modalities before fusion. This enables POMP to pull pathology images and omics data from the same patients closer together in the latent space, thereby better performing cross-modal learning.
3. POMP achieves the best results on six cancer datasets; Extensive experiments show that leveraging the cosine similarity, calculated in the same way as the pathologyomics contrast learning, as the survival risk score achieves better performance compared to other commonly used survival risk computational methods.

2 RelatedWork

2.1 Multimodal Pre-training

Multimodal pre-training aims to improve the performance of downstream tasks by jointly pre-training the model with multiple pre-training objectives, garnering considerable attention in AI research [Xu et al., 2023a; Pei et al., 2024]. A substantial amount of work has been presented across various domains, for example, vision-language pre-training (VLP) [Radford et al., 2021; Kim et al., 2021; Li et al., 2021; Yin et al., 2024], vision-audio pre-training (VAP) [Deshmukh et al., 2023; Xu et al., 2023b], audio-language pre-training (ALP) [Elizalde et al., 2023; Wu et al., 2023b], audio-visual-text Pre-training (AVTP) [Guzhov et al., 2022; Wu et al., 2022; Xu et al., 2024]. Among them, VLP is a major research problem in this field. The pre-training tasks commonly used in VLP include masked language modeling (MLM) [Li et al., 2021], masked image modeling (MIM) [Kim et al., 2021], image-text matching (ITM) [Li et al., 2021], image-text contrastive learning (ITC) [Bao et al., 2022; Cheng et al., 2021; Jin et al., 2023], and word patch alignment (WPA) [Kim et al., 2021]. While these multimodal pre-training methods have achieved great success in natural images (with low pixels and uniform size) and texts, it is not feasible to directly copy them to the scenario of cancer survival prediction since the pathological image has gigapixels, non-uniform sizes, and large surfaces of worthless background regions.

2.2 Survival Prediction usingMulti-modality

Recently, survival prediction using multi-modality has achieved notable advancement, attributed to the success of deep learning-based multimodal fusion technologies [Cui et al., 2023; Wang et al., 2024]. For instance, DPDBN [Wang et al., 2021b] and CAMR [Wu et al., 2023a] concatenate representations learned from multi-modality, where the input features of pathological images, such as sizes, shapes, intensity distributions, textures, and brightness levels, are extracted through the hand-crafted CellProfiler tool [McQuin et al., 2018]. DeepCorrSurv [Yao et al., 2017] maximizes the correlation to learn the shared representation from two modalities, while PathOmics [Ding et al., 2023] minimizes the mean square error to strengthen the interaction between different modalities. Pathomic Fusion [Chen et al., 2020] introduces a tensor fusion method, Kronecker product, to fuse multi-modality. MCAT, CMTA, and HC-MAE [Chen et al., 2021; Zhou and Chen, 2023; Wang et al., 2023] employ a cross-modal attention mechanism to integrate multi-modality. However, most methods lack proper alignment of multimodal data before fusion, which makes it difficult to effectively learn pathology-omics interactions. We aim to overcome this by leveraging self-supervised learning and contrastive learning, where our framework POMP is jointly pre-trained with three tasks to align multi-modal information.

3 Method

We propose POMP, a pathology-omics multimodal pretraining framework for cancer survival prediction, as illustrated in Figure 1. Given a patient sample xi from the data set x = {x1, x2, . . . , xN}, containing a gigapixel pathological image Pi, multi-omics data Oi, survival status ei, and survival time ti, we aim to train POMP to predict survival risk scores. This section first introduces data preprocessing, and then details POMP's key components, followed by its pretraining and fine-tuning.

3.1 Data Preprocessing

Pathological Images. Building upon the previous works [Zhou and Chen, 2023; Wang et al., 2023], we start by employing CLAM [Lu et al., 2021] to automatically identify the high-value region and eliminate the background area within the pathological image. In computational pathology, gigapixel WSI can reach dimensions of 150,000 × 150,000 pixels at 20× magnification. A straightforward way is to crop WSIs into 256 × 256 pixel patches, but this will generate a substantial number of segments, greatly increasing the computational load for self-attention in Pathology-Transformer. We divide the high-value region into M non-overlapping subregions with 4096×4096 pixels, which are then downsampled by a factor of 16 to the image patches {ik}Mk=1 with 256×256

pixels. Although this might lead to the loss of fine-grained information, Pathology-Transformer can efficiently capture inter-patch relationships.

Multi-omics data. For multi-omics data including RNASeq, miRNA, and DNA methylation, we first employ the Knearest neighbor interpolation method [Troyanskaya et al., 2001] to fill in missing values. Then, we calculate the variance of each gene across all patient samples and remove the genes with zero variance. To perform differential gene expression analysis, we use the pydeseq2 package [Muzellec et al., 2023] to select the genes that exhibit significant changes. Finally, for each omics type in each cancer dataset, we use the random survival forest (RSF) [P¨olsterl, 2020] to calculate the feature importance of each gene across all patient samples and retain the top 300 genes for survival analysis.

3.2 POMP Architecture

Given a patient sample containing pathological images and multi-omics data, we encode them separately into vector representations using two encoders, which are then fed into a multimodal encoder to fuse them.

Pathology Encoder. Pathology encoder consists of M Vision Transformer (ViT) [Dosovitskiy et al., 2020] and a Pat-Transformer. We start by embedding the image patches {ik}Mk=1 with 256 × 256 pixels by M ViTs with shared weights, which can be written as:

where AvgPool is average pooling used to generate the image patch embeddings Ipat ∈ RM×d . The lower level patch size in ViT is 16 × 16.

Afterward, we randomly initialize a [CLS-P] token Icls ∈ R1×d , followed by concatenating Ipat and Icls, and then feed them to the Pathology-Transformer for fully learning the correlation between image patches within WSI, thereby obtaining the final representation of WSI:

where Epos denotes the relative position embedding of image patches. I ′final consists of a sequence of vectors: {hv1 , . . . , hvM , v}, where v is the final representation of the [CLS-P] token. We define Ifinal = {hv1 , . . . , hvM }, which corresponds exclusively to image patch tokens.

Omics Encoder. Let Opre ∈ R3×300 denotes preprocessed multi-omics data, including mRNA, miRNA, and DNA methylation. We first feed the omics data to a shared linear network to learn the intrinsic information of the omics. Subsequently, the outputs of the shared linear network, together with a randomly initialized vector of [CLS-O] token, are sent into a 2-layer Transformer to learn the interactive information between multi-omics data, thereby generating the final representations of multi-omics data:

Where Ofinal consists of four vectors: {w, hw1 , hw2 , hw3 }, which are the final representations of [CLS-O] token, mRNA, miRNA, and DNA methylation.

Multimodal Encoder. We propose a multimodal encoder, which is implemented by an omics-guided cross-attention module, a 2-layer Transformer, and a residual module, to fuse the representations of two modalities, thereby learning the final multimodal representations:

where Wq,Wk,Wv ∈ Rd×d denote learnable weights multiplied by the queries Ofinal, keys Ifinal, and values Ifinal, respectively. LN represents layer normalization. H are the final multimodal representations. Notably, there are two special representations in H corresponding to the [CLS-O] token and masked omics token, which are denoted as hcls and hmask .

3.3 Pre-training Tasks

We pre-train POMP using three tasks: pathology-omics contrastive learning (POC), pathology-omics match (POM), and masked omics modeling (MOM).

Pathology-Omics Contrastive Learning. Given a training batch of N pathology-omics pairs, there are N positive and N 2 - N negative pathology-omics pairs. Here, we introduce a pathology-omics contrastive learning method that aims to learn effective representation by pulling pathologyomics pairs from the same patients closer and pushing negative pairs farther away. Specifically, we use the representation of [CLS-P], v, and the representation of [CLS-O], w, as the aggregated representation of the pathology image and multiomics data, respectively. Then, the pathology-to-omics and omics-to-pathology similarities are computed between v and w within a training batch. Formally, the softmax-normalized pathology-to-omics and omics-to-pathology similarities can be written as:

where σ is the temperature parameter to scale the similarities. Let yp2o and yo2p represent the ground-truth one-hot similarity, where the positive pair has a probability of 1 and the negative pair has a probability of 0. The pathology-omics contrastive loss is calculated using cross-entropy (CE) between s and y:

Pathology-Omics Matching. Inspired by ALBEF [Li et al., 2021], we sample two hard negative pathologyomics pairs based on the max softmax-normalized pathologyto-omics and omics-to-pathology similarities for each pathology-omics pair in the training batch. As a result, we obtain N positive pairs and 2N hard negative samples. Pathology-omics matching aims to predict whether a pair of pathology and omics is positive or negative. Specifically, we use the final multimodal representation of the [CLS-O] token, hcls, as the integrated representation of the pathology-omics pair, and feed it into a linear classifier g with cross-entropy loss for binary classification. Let ypom denote the groundtruth label, the POM loss is written as:

Masked Omics Modeling. We randomly mask one of three omics data and replace it with a zero vector. Masked omics modeling aims to reconstruct omics from a corrupted version based on other unmasked omics and pathology cues. Specifically, we feed the final representation of masked omics, hmask, into a classifier f over the dimension of omics data for predicting the probability of masked omics token f (hmask). Let ymask denote the ground-truth masked omics data, the MOM loss is written as:

3.4 Fine-tuning POMP on Survival Prediction

After pre-training POMP, we fine-tune it for the downstream survival prediction. Traditionally, the Cox proportional hazards model is expressed as:

where t represents the survival time, h0(t) is the baseline hazard, h(t|X) is determined by a linear combination of covariates (x1, ..., xp) and their corresponding coefficients (β1, ..., βp).

In this work, we employ a neural network to fit the Cox proportional hazards function, where we consider the final pathology representation v as covariates and the final omics representation w as coefficients. We also standardize the linear relationship between covariates and coefficients, making the model less sensitive to the scale of the variables. Formally the Cox proportional hazards function can be written as:

Parameter estimates in the Cox proportional hazards model are often obtained by maximizing the partial likelihood:

Finally, the survival loss function is set to be the negative log partial likelihood:

Unlike vanilla Cox models that only deal with linear conditions in the hazard function, our model can better fit multimodal data and learn complex interactions. Particularly, since h(T|v,w) is an increasing function of c (v,w), we can take the cosine similarity c (v,w) as the survival risk score when leveraging the concordance index as the evaluation metric.

4 Experiments and Results Analysis

We implement our framework POMP using PyTorch and train it on 3 NVIDIA A100 GPUs. The Pathology-Transformer, Omics-Transformer, Multimodal-Transformer, and ViT in POMP all adopt the vanilla Transformer, each module has a 2-layer Transformer block with 384 hidden dimensions and 6 attention heads. POMP is trained for 500 epochs during pretraining and 80 epochs during fine-tuning. For both training phases, we use Adam optimization with a weight decay of 1e2 and a learning rate of 5e-4. Since the pathological images have different sizes and are cropped into various sub-region numbers, we use a batch size of 1 with 50 forward accumulation steps (i.e., the actual batch size is equivalent to 50), then calculate the loss function and update the weights once.

4.1 Datasets and EvaluationMetrics

Datasets. We perform experiments using six cancer datasets obtained from the Cancer Genome Atlas (TCGA) data portal. These datasets cover various cancer types, namely colon adenocarcinoma (COAD), hepatocellular carcinoma (LIHC), stomach adenocarcinoma (STAD), breast invasive carcinoma (BRCA), lung adenocarcinoma (LUAD), and lower grade glioma (LGG). Each patient sample consists of a complete set of data types, including pathological images, as well as multi-omics data composed of RNA-Seq, miRNA, and DNA methylation. To better train the multimodal model, the patient samples missing any data type are removed. Additionally, we exclude patients with a survival time of fewer than 30 days or lacking follow-up records. As a result, the sample counts for COAD, LIHC, STAD, BRCA, LUAD, and LGG are 250, 323, 298, 724, 374, and 451, respectively. We assess all investigated methods using the same 5-fold cross-validation splits on each cancer dataset.

Evaluation Metrics. Following previous works [Yao et al., 2017; Wu et al., 2023a; Chen et al., 2021], we leverage the concordance index (C-index) as the evaluation metric of survival prediction. The C-index is commonly used to calculate the concordance pairs between actual survival times and predicted survival risk scores, which is computed as follows:

where I [·] stands for the indicator function, n signifies the number of comparable pairs in which the survival status e of patients is dead, and t denotes the actual survival time. A higher C-index suggests superior predictive performance, whereas a C-index of 0.5 denotes that the model's predictions are similar to random chance outcomes.

4.2 ComparisonWith State-of-the-art Methods

We compare POMP with several state-of-the-art integration methods, as shown in Table 1. The first observation is that our proposed framework POMP consistently outperforms all existing methods across all cancer datasets, both with and without multi-modal pre-training. Specifically, POMP achieves C-index values of 71.5% on COAD, 75.5% on LIHC, 65.1% on STAD, 70.0% on BRCA, 69.5% on LUAD, and 85.6% on LGG, outperforming the previous best-performing methods by 1.2%, 1.8%, 2.0%, 0.5%, 0.4%, and 0.5%, respectively.

4.3 Survival Analysis

To further evaluate the effectiveness of POMP for survival prediction, we plot the Kaplan-Meier curves and calculate the log-rank test p-values across all six cancer datasets, which are shown in Figure 2. We utilize the median of the predictive survival risk scores as a risk indicator to stratify patients into high-risk (blue) and low-risk (orange) groups, and then perform Kaplan-Meier analysis to visualize the survival probability over time for different groups. Additionally, the logrank test p-value is used to measure the statistical significance and compare the survival distributions between the high-risk group (blue) and the low-risk group (orange). In Figure 2, it can be observed that the p-values for all datasets are below the commonly used significance level threshold of 0.05. This suggests a statistically significant difference between the survival curves of all compared groups.

4.4 Ablation Study

To examine the contribution of each component of POMP, we conduct a set of ablation studies, including input modality, survival risk computing method, and pre-training tasks.

Unimodal and Multimodal. After pre-training, we add a linear layer to the pathology encoder and omics encoder, respectively, computing survival risk under single-modality. The survival prediction results of unimodal and multimodal are shown in Figure 3a. Multimodal modeling surpasses any single modality modeling in all six cancer datasets, demonstrating the effectiveness of our proposed framework in fusing multi-modality. Interestingly, we observe that the performance of the omics modality is better than the pathology modality in five of six datasets, indicating that omics can provide more information than pathology in our framework.

Survival Risk. To investigate the impacts of various survival risk computing methods in our framework, we conduct experiments on three kinds of survival risk computing methods, as illustrated in Figure 4. In this work we propose to compute the cosine similarity between two modalities as a risk score. Here we compare our method with two other methods, including (i) computing the fusion risk by feeding the final multimodal representation into a Cox layer implemented by a fully connected network; and (ii) computing the joint risk by averaging the scores of two modalities. Figure 3b shows the experimental results of the three survival risk computing methods. The cosine similarity method achieves superior survival prediction compared to other methods, indicating that the survival risk calculation method consistent with the pathology-omics contrastive learning in pre-training is more suitable for our proposed POMP framework.

Pre-training Tasks. We analyze the performance under individual pre-training tasks and different combinations to evaluate their impact on survival prediction performance. The

(a) Unimodal vs. multimodal

(b) Risk computing methods

Figure 3: Ablation studies of modality and survival risk.

experimental results are presented in Table 2, where POMPB0 denotes that POMP has not pre-training, +L∗ represents the specific task is used to pre-train POMP. When looking at the individual pre-training tasks, we find that POMP-B1 with POC task obtains the superior C-index values on four out of six cancer datasets and achieves the best result in terms of macro average, demonstrating that the POC task contributes more than the POM and MOM tasks. Additionally, POMPB3 with the MOM task obtains the worst performance among all individual pre-training tasks, even worse than POMP-B0 trained from scratch. However, in the scenario where two pretraining tasks are combined including POMP-B4, POMP-B5, and POMP-B6, we observe that POMP-B5 (which merges POC and MOM) achieves superior performance in terms of the macro average. Specifically, POMP-B5 surpasses POMPB4 and POMP-B6 in four out of six datasets, suggesting that the MOM task together with the POM task can facilitate POMP learning to better semantic representation and fusion capabilities. Notably, the POMP (Ltotal) framework, which joints all three pre-training tasks, achieves the best C-index value in terms of the macro average and outperforms other settings on five out of six datasets, demonstrating the effectiveness of our proposed POMP framework to joint three various pre-training tasks.

Pre-training Loss. To evaluate the impact of pre-training tasks on the convergence of POMP training, we analyze the pre-training loss curves corresponding to different pretraining task settings. As shown in Figure 5a, we can see that (i) the total loss and POM loss decrease smoothly as the number of training papers increases; (ii) the POC loss decreases rapidly in the first 20 epochs, and then there are several fluctuations, but the fluctuations disappear after 300 epochs; (iii) it is interesting to see that MOM loss does not converge, which is consistent with POMP-B3 (Table 2) obtaining the worst performance in three individual pre-training tasks. We further verify whether the MOM loss will converge when the MOM task is combined with other tasks, as shown in Figure 5b. The MOM losses in POMP-B5 (which combines POC and MOM tasks), POMP-B6 (which combines POM and MOM tasks), and POMP (which combines three tasks) all decrease and converge normally, suggesting that the MOM task combined with other tasks contributes to the training of POMP.

5 Conclusion

We proposed a novel pathology-omics multimodal pretraining (POMP) framework for cancer survival prediction, by leveraging the principle of pre-trained models and exploring the benefit of aligning multimodal information of the same patient. POMP is jointly pre-trained on three tasks based on self-supervised learning and contrastive learning to align pathology image and omics data before fusion, allowing us to use their similarity as survival predictions. Experimental results show that POMP outperforms existing state-of-the-art methods on six cancer datasets.

(b) MOM loss under different tasks

Figure 5: Loss curves for different pre-training tasks.

While POMP shows promising results, there is still substantial room for improvement.

1. We will extend POMP to integrate other types of omics data associated with cancer survival, such as somatic mutations, copy number alterations, and clinical data.
2. We intend to pre-train the model based on a pancreatic cancer dataset to learn more prior knowledge about the potential relationship between multi-modalities.
3. We are going to explore alternative strategies for selecting the important genes in original multi-omics data, which will provide more valuable information in the omics modality.

Acknowledgments

This study was supported by the grants with No. 62162023 and No. KYQD(ZR)-21079.