Curriculum Hierarchical Knowledge Distillation for Bias-Free Survival Prediction

1 Introduction

Survival prediction, the estimation of mortality risk within a given timeframe, represents a cornerstone of clinical oncology [Shedden et al., 2008]. This endeavor predominantly entails the analysis of whole slide images (WSIs), which encapsulate intricate spatial patterns and the complexities of the tumor microenvironment [Pantanowitz et al., 2011]. Deep learning has emerged as a transformative paradigm, automating WSI analysis and offering unprecedented prospects for alleviating pathologists' burdens and empowering physicians in critical decision-making [Hanna et al., 2020].

Existing WSI-based survival prediction models [Yao et al., 2020a; Zhu et al., 2017; Liu et al., 2023; Shao et al., 2023a; Tang et al., 2019] follow a two-stage paradigm: WSI preparation, which encompasses collecting WSIs, segmenting them into patches, and encoding these patches to address gigapixellevel resolution; and patient-level prediction, wherein patient outcomes are forecasted by leveraging both the encoded patches and their inter-patch relationships. Most works [Shao et al., 2023a; Liu et al., 2023] focus on developing sophisticated prediction models for the patient-level prediction stage, such as Transformers [Chen et al., 2022a], Graph Neural Networks [Gadiya et al., 2020] and Capsule Networks [Tang et al., 2019]. However, the substantial bias introduced during the WSI preparation stage is largely overlooked, which may limit the effectiveness of current approaches and their potential to serve as a panacea [Jaume et al., 2021].

∗Corresponding Author: Xi Zhang.

The efficacy of existing survival prediction models depends on the assumption that the data are independent and identically distributed (IID), implying that the WSIs of different patients are drawn from the same distribution. However, the inherent limitations in the WSI preparation process, coupled with the uncertain status of infected tissues and variability in pathologist expertise, often undermine the validity of the IID assumption in practical scenarios. For instance, biopsy strategies [Strassburg and Manns, 2006] and tissue sectioning practices [Fischer et al., 2008] are influenced by the heuristic expertise of clinicians, leading to variation in the numbers of WSIs across patients. In the TCGA-LUAD dataset [Tomczak et al., 2015], 20.1% of patients have more than 6 WSIs, while 30.3% have fewer than 3. This non-IID nature of WSI numbers may result in imbalanced training data and skewed performance. Additionally, the variability in lesion morphology and slice positioning leads to significant differences in resolution across WSIs. As shown in Fig. 1(a), WSIs with higher resolution tend to be more informative (e.g., containing more tissue patches) than those with lower resolution.

To highlight the impact of the non-IID phenomenon on the survival prediction task, a preliminary study is conducted using the popular baseline DeepAttnMISL [Yao et al., 2020a] on the NLST dataset [Team, 2011]. Patients within the testing set are divided into different groups based on two dimensions: (1) the average numbers of WSIs, as indicated by the vertical axis in Fig. 1(b), and (2) the average informative resolution of WSIs, as indicated by the horizontal axis. The average informative resolution is implemented as the average number of non-blank patches (500×500 pixels at a scaling factor of 20) for each patient's corresponding WSIs, calculated after applying the OTSU algorithm for filtering [Otsu, 1979]. Groups with larger labels correspond to a greater number of WSIs or patches. Fig. 1(b) demonstrates the performance of a well-trained model on different groups. One can clearly see that: (1) patients with more WSIs consistently outperform those with fewer WSIs; and (2) given the same number of WSIs, patients with WSIs of larger resolution (i.e., those with a larger number of patches) achieve better performance.

In this paper, we focus on the novel problem of mitigating bias within both the numbers and resolutions of WSIs across patients. A straightforward solution involves resampling techniques [Good, 2006], while undersampling may result in information loss, while oversampling can contribute to overfitting and noise [Zhang et al., 2024b], particularly in the context of small survival prediction datasets [Zhang et al., 2024a]. Inspired by the idea of transferring knowledge from head patients, who possess more informative WSIs, to tail patients with less informative WSIs, we adopt knowledge distillation [Hinton, 2015] as a core methodology. Specifically, a teacher model is trained on the head patients, and closeness minimization losses are employed to facilitate the transfer of knowledge [Zhao et al., 2023].

However, directly applying vanilla knowledge distillation methods to survival prediction presents two challenges. (1) Hierarchical Structure. Survival prediction operates at the patient level [Shao et al., 2023a], while biases in number and resolution are hierarchically structured at both the inter- and intra-WSI levels. Vanilla distillation models [Hinton, 2015] overlook the interdependencies across these levels, thus failing to capture the comprehensive hierarchical information. (2) Data Scarcity. Given the typically limited number of available patient cases, vanilla knowledge distillation methods often split datasets to train the teacher and student models, further exacerbating the scarcity challenge.

To address the aforementioned challenges, we propose a novel model, PathoKD, designed to enhance the robustness of survival prediction through hierarchical knowledge distillation within a triple curriculum framework. The learning process is organized into three curricula, each targeting progressively difficult stages determined by the number and resolution of each patient's WSIs. Specifically, patients are categorized into three groups: head patients, who exhibit a large number of extensive WSIs; virtual tail patients, which are synthesized by removing patches from the head patients to simulate fewer and smaller WSIs, bridging the gap between head and tail patients; and tail patients, who naturally possess smaller and fewer WSIs. Curriculum I initiates the process by training the teacher model to extract knowledge from head patients, employing a label-assistance strategy to mitigate data scarcity. Building upon this, Curriculum II facilitates hierarchical knowledge distillation across different levels, which transfers knowledge from the teacher model to the student models via virtual tail patients. To further alleviate the data scarcity and to leverage the full spectrum of available data, Curriculum III refines both the teacher and student models, incorporating retrieval and generative enhancement to integrate knowledge from tail patients. Extensive experiments across multiple datasets validate its superiority. The key contributions of this work are as follows:

• To the best of our knowledge, we are the first to examine the novel issue of bias stemming from the numbers and resolutions of WSIs across patients, and to investigate its impact on the task of survival prediction.
• A novel curriculum learning framework is proposed based on hierarchical knowledge distillation, which effectively bridges the gap between head and tail patients.
• The efficacy of our approach is demonstrated by showcasing its superior performance over existing methods across multiple real-world cancer datasets.

2 ProblemDefinition

Survival data refers to a set of patients' data, denoted as P = {P1, . . . , P|P|}, where each patient's data Pi consists of a set of WSIs and a survival outcome label, Pi = {Wi , Yi}. Each WSI W i j represents a tissue sample from patient Pi, and the survival outcome label Yi = (ti, ci) includes the observation time ti and survival status ci. The binary status ci ∈ {0, 1} indicates whether ti is a survival time (ci = 1) or a rightcensored time (ci = 0). Censoring occurs when a patient is still alive at the end of the study or is lost to follow-up. For computational efficiency, each WSI is partitioned into fixedsize patches, W i j = {Ei j,1 , . . . , Ei j,|W i j | }, which represent the smallest unit in the hierarchical survival data structure.

WSI-based survival prediction forecasts a patient's survival outcome given a set of WSIs Wi . For a patient Pi with a survival label comprising observed time ti and status ci, the deep survival model f(Wi ) predicts a survival time ˆ ti ∈ T :

3 Methodology

Our method is organized into four parts. We begin by quantifying biases in survival-prediction performance across patients with differing WSI coverage and quality. The next three sections describe the PathoKD framework (Fig. 2), structured as an easy-to-hard curriculum: (1) teacher training on head patients (green solid box), (2) student training on virtual tail patients (blue dotted box), and (3) progressive integration of real tail patients into both models (red dotted box).

3.1 Analysis of the Biases across Patients

In an ideal scenario, all patients would possess an equal number of uniformly-sized WSIs. However, real-world variability-arising from biopsy technique, scanner settings, and QC-introduces two noise terms: resolution noise ϵi,j r (resolution fluctuations for the j-th WSI of patient i) and perpatient noise ϵin (variation in total WSI count for patient i). Modeling these terms allows us to quantify biases in WSI coverage and quality across the cohort.

To mitigate the influence of ϵi,j r and ϵin , we assume that patients with large and more WSIs are less susceptible to noise, treating them as an ideal reference set. We aim to construct artificial noise ˆ ϵin and ˆ ϵi,j r , enabling a model to learn to predict outcomes ˆ ti consistently, with or without noise:

This objective can be expressed as the minimization of the prediction discrepancy between the clean and noisy sets:

By enforcing alignment between predictions in both settings, this framework results in more reliable survival predictions that remain invariant to noise artifacts.

3.2 Curriculum I: Teacher Model Training on Head Patient

Given a training set of patient survival data Ptrain, head and tail patients, Phead and Ptail, are identified based on average resolution and WSI counts via pre-defined thresholds. For each head patient P head i , each WSI W i j is segmented into patches W i j = {Ei j,1 , . . . , Ei j,|W i j | }, which are mapped to a representation space using a pre-trained histopathological encoder (e.g., HIPT [Chen et al., 2022a]). The representation of the k-th patch of the j-th WSI for the i-th head patient is denoted Xhead i,j,k . WSI-level features are derived by modeling patch relationships within each WSI, which are then aggregated into patient-level features. These features are passed through a Multilayer Perceptron (MLP) to predict the survival score. The process is formulated as:

The feature Xhead i,j,1 is input into an inter-WSI aggregator Ae, yielding Shead i,j , the representation of the j-th WSI for the ith head patient. Dhead i denotes the representation of the i-th head patient. The predicted survival score Ohead i is obtained by passing Dhead i through an MLP, fMLP.

Teacher Training with Noisy Survival Labels. The limited number of head patients may hinder teacher model training. Therefore, the survival label is incorporated to provide additional knowledge. Unlike traditional training, where test set labels are unavailable, our teacher model uses survival labels to guide the student model's learning rather than predicting them. The survival label is encoded into an embedding to incorporate it as input for training the teacher model. Given that the survival label Y head i = (thead i , chead i ) consists of continuous (thead i ) and discrete (chead i ) components, which differ in dimensionality from the pathological data, two distinct embedding functions are applied. The continuous observed time thead i is normalized and embedded into a learnable vector, while the discrete binary censored status chead i ∈ {0, 1} is embedded using a learned binary vector:

where Ehead Y represents the encoded survival label. Directly incorporating the label with the WSI may lead to overreliance, diminishing generalizability. To mitigate this, Gaussian noise is introduced into the label embedding, resulting in a noisy embedding: ˆ Ehead Y = Ehead Y +ϵy, where ϵy ∼N(0, σ2 ). The noisy embedding is added to the patient feature representation: ˆ Dhead i =Dhead i + ˆ Ehead Y .

Contrastive learning is employed to align augmented versions of the same patient's data, thereby reducing the impact of noise. The alignment loss Lalign is defined as:

where ˆ Dhead i,1 and ˆ Dhead i,2 are two versions of the patient representation, each incorporating different noisy label embeddings. ˆ Ohead i,1 and ˆ Ohead i,2 are the predicted survival outcomes, and fdist measures the dissimilarity between them.

The teacher model loss combines survival prediction losses and the alignment loss:

where Lsur is the survival loss and λ balances alignment. The Cox loss [Zhu et al., 2016] is used to optimize hazard risk predictions, assigning higher risks to patients with shorter times:

where Oi denotes the hazard risk for patient pi. This function encourages higher hazard risk predictions for patients with shorter survival times, improving overall accuracy.

3.3 Curriculum II: Hierarchical Knowledge Distillation On Virtual Tail Patient

Head patients with a greater number and larger size of WSIs exhibit increased robustness to noise (e.g., tissue information loss) and thus can be utilized as a reference group with relatively complete data [Hellmann et al., 2014]. To improve predictions for tail patients, we introduce noise into the WSIs of head patients, thereby simulating conditions characteristic of patients with limited data. The teacher-student distillation framework is formalized based on Eq. (7) as follows:

where fstu and ftch are the student and teacher models, and Wi and ˆWi represent the clean and noisy WSIs of the head patient. ˆWi is the virtual tail patient, aligning the head and tail patient distributions.

Virtual Tail Patient Construction. As shown in Eq. (3), the virtual tail patient is generated by introducing artificial noise into the data of the head patient. Given the WSIs of a head patient, W head i , noise is applied both to the number and the resolution of the WSIs. The parameter ϵn modifies the number of WSIs, yielding a modified set |W head i + ϵn|, while ϵr alters the resolution of each individual WSI, represented as |W i,head j + ϵr|. The noise addition follows a "simple-todifficult" strategy, wherein patches with low attention values are initially masked, and progressively, patches with higher attention values are also masked. Upon the application of this strategy to the head patient's WSIs W head i , the resulting virtual tail patient ˆ W head i is generated.

Hierarchical Knowledge Distillation. The distributional differences between head and tail patients span WSI representations, aggregated patient embeddings, and prediction distributions [Guan and Liu, 2021; Zhang et al., 2024c], highlighting variations in data quality, feature representation, and predictive outcomes, which challenge knowledge transfer and model generalization. Thus, hierarchical knowledge distillation is applied at each level. A student model is constructed to mirror the teacher model's architecture, excluding the label embeddings. The architecture is shown in Eq. (4).

WSI-level Distillation. For each head and virtual tail patient, WSIs are passed through both teacher and student models, producing feature embeddings S h i and S v i , respectively. A feature-based distillation loss is applied to the WSI pairs:

Here, |W head i | denotes the number of WSIs for the head patient. Given the data heterogeneity, directly minimizing the pairwise WSI distance is inefficient. Instead, we minimize the distance between individual WSIs and the domain-level discrepancy, focusing on mean and variance alignment.

To align the mean of the distributions, the MaximumMean Discrepancy (MMD) loss [Gretton et al., 2006] is used:

which measures the distance between the means of these two distributions in the reproducing kernel Hilbert space (RKHS) associated with the kernel function k.

To align the variance, the similarity matrices for WSIs within each model are computed as

The difference between these matrices quantifies the gap:

The total WSI-level loss is:

In our implementation, each loss function is assigned a corresponding weight to balance its contribution to the overall loss. However, to simplify the notation and avoid redundancy, we omit the explicit representation of these weights.

Patient-level Distillation: At the patient level, alignment of features, distribution mean, and dispersion is required. The patient feature alignment mirrors the WSI-level losses:

Additionally, the predicted risk depends not only on the score but also on the ranking within the batch. Thus, we minimize the distance between predicted scores and ranks:

where Lot represents the distance between predicted risk scores, and Lrank is the distance between their rankings.

The total loss for the student model in Curriculum II is:

3.4 Curriculum III: Student-centered Progressive Training On Tail Patient

Curriculum I extracts knowledge from head patients for the teacher model, while Curriculum II transfers this knowledge to the student model via virtual tail patients, reducing the performance gap. However, both curricula depend solely on head patient data, limiting tail patient usage. The challenge lies in the lack of identity-matched knowledge for direct transfer to tail patients, complicating learning. Incorporating tail patients with insufficient data risks performance degradation, especially in limited survival datasets. To address this, we propose constructing virtual head patients by augmenting tail patient data with potential WSIs. The teacher model learns from these virtual head patients, improving knowledge transfer to the student model. This student-centered knowledge distillation approach enhances tail patient representation using augmented data before guiding the student model.

Virtual Head Patient Construction. To transform a tail patient into a virtual head patient, we propose two methods: Retrieval-based Expansion: A retrieval-based approach is proposed to enhance the number and resolution of WSIs. For resolution expansion, representative patches from the tail patient's WSI are selected based on attention scores. The top-M patches Etail,k i,j , which capture critical information, are identified. Similar patches are then retrieved from the training set using a similarity function sim(·, ·), and the top-N1 patches with the highest average similarity are selected for data augmentation. The expandedWSI ˆW tail i,j is defined as:

For number expansion, the aggregatedWSI representations are compared to identify the top-N2 similar WSIs, which are incorporated into the tail patient's data:

Generative-based Expansion: To further expand WSI resolution and number, a generative approach is employed leveraging the Denoising Diffusion Probabilistic Model (DDPM). For resolution expansion, the DDPM is trained by masking representative patches, identified using attention values, and using the remaining patches as conditions to reconstruct the masked regions. This enables the DDPM to learn the latent feature space of critical regions, allowing it to generate synthetic patch features. For each tail WSI ˆW tail i,j , the DDPM D1 generates synthetic patch features { ˆEtail,k i,j }:

where θ represents the trained parameters of the DDPM.

For number expansion, another DDPM is trained to generate synthetic WSIs by masking complete WSIs and using the remaining WSIs as conditions to reconstruct the masked WSIs. The synthetic WSIs are obtained as:

Progressive Training. The expanded WSIs of tail patients ˆW tail i are first used to train the teacher model, and a progressive curriculum is employed to incorporate orignal WSIs of tail patients W tail i into the student model. The proportion and difficulty of the tail patients introduced increase over time. The difficulty of a tail patient is quantified by the change in survival risk before and after training the teacher model:

where ∆Rtail i reflects the improvement in the model's understanding of the patient. Tail patients are ranked by difficulty, and the batch inclusion function γ(t) determines the proportion of tail patients included at time t:

where Qγ(t) represents the difficulty threshold based on the quantile of the sorted tail patients. The proportion γ(t) increases over time following a sigmoid schedule:

This curriculum learning strategy enables a gradual introduction of simpler tail patients, followed by more challenging ones, facilitating effective learning progression from simple to complex cases.

4 Experiment

4.1 Experimental Settings

Datasets We evaluate our proposal on five real-world cancer datasets. The first is the National Lung Screening Trial (NLST) dataset [Team, 2011], which includes cases of adenocarcinoma (ADC) and squamous cell carcinoma (SCC). The remaining four datasets are from The Cancer Genome Atlas (TCGA) [Tomczak et al., 2015], encompassing lung cancer (LUSC and LUAD), breast cancer (BRCA), and bladder cancer (BLCA). Each dataset contains WSIs stained with hematoxylin and eosin (H&E) and corresponding survival labels.

Baselines Three types of SOTA baselines are considered: (1) Traditional CNN-based methods that using partial WSI data: WSISA [Zhu et al., 2017] and DeepAttnMISL [Yao et al., 2020b]; (2) Transformer-based methods utilizing full WSI data: HIPT [Chen et al., 2022b], SeTranSurv [Huang et al., 2021], ESAT [Shen et al., 2022], LongViT [Wang et al., 2023a], and Prov-GigaPath [Xu et al., 2024]; and (3) Methods addressing bias: AugDiff [Shao et al., 2023b], Resampling [Good, 2006], MoE [Masoudnia and Ebrahimpour, 2014], and Retrieval [Wang et al., 2023b].

Implementation Details All baseline survival models are re-implemented faithfully from their original publications and available open-source repositories. For a fair comparison, we adopt a 5-fold cross-validation scheme: in each fold, 10% of the training split is held out as a validation set for early stopping [Bai et al., 2021]. Models are trained for up to 200 epochs (with patience of 10 epochs on validation loss) using the AdamW optimizer (weight decay 1e-4) and a fixed batch size of 64. All experiments are carried out in PyTorch 2.2.0 on NVIDIA V100 GPUs with 32 GB of memory. Performance metrics are averaged over the five folds.

4.2 Quantitative Evaluation

Table 2 presents survival prediction results using the C-index and STAGE-5 metrics. (1) PathoKD outperforms all stateof-the-art baselines, with improvements of over 4.28% in Cindex and 4.58% in STAGE-5. (2) Transformer-basedmodels outperform earlier CNN-based methods, which rely on partial WSI data, by processing the full WSI. (3) Our method further improves performance by addressing the challenges of limited WSI number and resolution through knowledge distillation. (4) Compared to other models that tackle number and resolution bias without distillation, our approach offers superior performance by transferring high-level, task-relevant knowledge rather than manipulating low-level model inputs.

4.3 Performance onMitigating Bias

A controlled experiment is conducted to validate our approach in addressing resolution and number biases in WSIs and its impact on survival prediction. Patients are categorized into three groups based on WSI resolution and number: low (fewer than 3 WSIs and the bottom 20% in size), medium (exactly 3 WSIs and the middle 60%), and high (more than 3 WSIs and the top 20%). As shown in Table 3, applying knowledge distillation significantly reduces the performance gap between these groups, demonstrating that PathoKD effectively mitigates this challenge. Notably, even the high group benefits from knowledge distillation, underscoring the effectiveness of our hierarchical alignment loss.

4.4 Ablation Study

We conduct ablation studies on the NLST and LUAD datasets to evaluate the contribution of each module in PathoKD, with results presented in Table 4. The first row shows that using label embedding alone, without noise, leads to over-reliance on labels and a significant performance drop, emphasizing the need for noise incorporation. The second row reveals that excluding Curriculum I severely degrades performance, highlighting the importance of labels as external knowledge. Removing instance distillation LInst results in a notable performance decline, underlining its role in mitigating bias and facilitating knowledge transfer. Omitting MMD distillation Lmmd also decreases performance, stressing the importance of aligning distributions between head and tail patients. Excluding similarity distillation LSim further reduces performance, demonstrating its critical role in maintaining variance in the representation distribution. The removal of prediction score alignment-one of the highest-level features-has the most significant impact, emphasizing the importance of ranking in measuring prediction score distribution. Lastly, excluding Curriculum III and neglecting tail patients results in a performance decline, reinforcing the effectiveness of our progressive learning approach. Between the two methods in Curriculum III, the generative approach outperforms retrieval, showcasing the DDPM's ability to capture latent features.

4.5 Efficiency Study

The efficiency of the proposed model is evaluated on the NLST dataset [Team, 2011] and compared with SOTA models. All configurations follow the original papers, using official implementations where available. For WSISA, the implementation was recreated due to the absence of official code. This study is the first to evaluate bias mitigation strategies (Resampling, MoE, and Retrieval) on survival prediction. All models were run on two Xeon E5-2690 v4 processors (2.60 GHz) and four NVIDIA V100 GPUs. Results show that our approach outperforms others in both accuracy and inference speed. Specifically, (1) Our method achieves an 11.4% higher C-index and runs 1.40 times faster than clusteringbased models like WSISA and DeepAttnMISL, due to the transformer architecture's direct processing of patches without computationally expensive clustering. (2) Our approach also outperforms other bias mitigation methods, achieving 34.6% faster inference by directly using original data without the additional preprocessing required by other models.

4.6 Parameter Sensitivity Analysis

Influence of noise strength γ1 and γ2 of discarding WSIs and patches: The γ1 and γ2 specify the number and percentage of discarding WSIs and patches. As illustrated in the left side of Fig. 3, PathoKD performs best with a mask WSI number of 3 and a mask patch ratio of 50%. Lower mask ratios might necessitate masking a greater proportion of WSIs to align virtual tail patients with tail patients, while higher ratios may leave too few visible histopathological features to effectively capture prognosis.

Influence of noise strength γ3 and γ4 of expanding WSIs and patches: The γ3 and γ4 specify the number and perof-the-art models trained on the NLST dataset.

centage of expanding WSIs and patches. As illustrated in the right side of Fig. 3, PathoKD performs best with an expanding WSI number of 3 and a mask patch ratio of 30%. Lower expansion ratios may require masking a larger proportion of WSIs to align virtual head patients with actual head patients, while higher ratios could introduce excessive noise due to the reduced proportion of real WSIs.

5 Conclusion

In this paper, we propose PathoKD, a novel model for WSI-based survival prediction, designed to address the challenges of data variability and scarcity, particularly for tail patients. By integrating hierarchical knowledge distillation with curriculum learning, PathoKD effectively mitigates the performance gap between patients with limited and abundant data. Extensive experiments on multiple widely used datasets demonstrate its superiority in improving prediction accuracy across diverse patient groups, providing a promising approach for developing more robust and equitable survival prediction models.

Acknowledgments

Zhihao Tang is the first student author. This work was supported by the Natural Science Foundation of China (No. 62372057). This work was partially sponsored by the CAAIHuawei MindSpore Open Fund. Compute services from Hebei Artificial Intelligence Computing Center.