Paper: Prototypical Information Bottlenecking and Disentangling for Multimodal Cancer Survival Prediction (ICLR-2024)

Reviewer comment (Writing):
The explanation around PIB and PID need to be much more straightforward for the work to be appreciated. In its current form, it is really hard to understand and Section 3 lacks clarity. The authors should understand that typical readership of this article will either be pathologist or ML practitioner who would be less-versed in probabilistic and variational language. This, in my opinion, significantly reduces the impact of the paper as any other researcher hoping to reuse the framework for their own experiments. Authors should try to inject more pathology-related intuition throughout the work.

Author response:
In response to your valuable feedback, we have revised **Section 3** of the paper. The modifications include removing the introduction to Markov chains, emphasizing the changes made by the new approach more straightforward for conciseness and clarity.

Reviewer comment (Theory):
I am not sure if "discriminative" approach (despite its good performance as shown in the paper) is the right approach for survival. Different from diagnosis problems where subtype/grades are indeed decided with discriminative instances (patches), survival prediction problems are not as clear-cut. Certain morphological characteristics will overlap between different risk groups (or bands), which the current approach, centered around sampling positive/negative prototypes for maximum discrimination, is not well suited for.

Author response:
1. We agree that survival prediction problems present greater complexity than diagnosis. In practical scenarios, there may exist deciding biomarkers within different risk bands, which leads to the “discriminative” method having demonstrated good performances in many cases [1,2].  Although some instances **related** to different risk groups may exhibit some overlap, these overlapping instances may not significantly contribute to **distinguishing** risk bands. Similarly, the same situation of overlapping can be observed in subtyping/grading as well. Therefore, our method, through the use of prototypes, excels in identifying discriminative features that are specifically associated with the positive risk level. Hence, discriminative assumption may not be a weakness but an advantage for biomarker discovery of different risk levels.

2. Indeed, the prototype learning in the current approach primarily relies on cross-entropy(CE)-like loss. This is because the prototypes are constructed based on information bottleneck theory, in which the first term of the derivation in Eq. (8) is formulated as a form of cross-entropy. Fortunately, NLL loss, a CE-like loss, makes it feasible to employ IB for survival prediction task. 
    
    Although the prototype component is based on CE-like loss, the task loss in our framework is still flexible and can be easily switched to Cox or rank loss. Furthermore, even in rank-based losses, which focus on relative comparisons between samples, discriminative information is still needed. Although not directly involving classification, the model needs to learn the differences between samples to correctly rank them. Hence, the assumption of PIB may not conflict in such scenarios.
    
    Despite the discriminative hypothesis, this work first attempts to directly model the bag-level representations for a bag structure through the probabilistic model in survival prediction task. There is undoubtedly significant value in expanding this approach to various hypotheses like rank-based hypothesis. Hence, further investigation is still warranted in the future.

**REFERENCES**

[1] Echle A, Rindtorff N T, Brinker T J, et al. Deep learning in cancer pathology: a new generation of clinical biomarkers[J]. British journal of cancer, 2021, 124(4): 686-696.

[2] Wiegrebe, Simon, et al. "Deep Learning for Survival Analysis: A Review." *arXiv preprint arXiv:2305.14961* (2023).

Reviewer comment (Theory):
This "discriminative" approach seems only applicable to NLL loss, which treats the problem as a classification problem and thus have "different" bins (i.e., discretized timeline). What if more common loss functions, such as Cox PH or rank losses, were to be used that do not rely on such discriminative hypothesis?

Author response:
1. We agree that survival prediction problems present greater complexity than diagnosis. In practical scenarios, there may exist deciding biomarkers within different risk bands, which leads to the “discriminative” method having demonstrated good performances in many cases [1,2].  Although some instances **related** to different risk groups may exhibit some overlap, these overlapping instances may not significantly contribute to **distinguishing** risk bands. Similarly, the same situation of overlapping can be observed in subtyping/grading as well. Therefore, our method, through the use of prototypes, excels in identifying discriminative features that are specifically associated with the positive risk level. Hence, discriminative assumption may not be a weakness but an advantage for biomarker discovery of different risk levels.

2. Indeed, the prototype learning in the current approach primarily relies on cross-entropy(CE)-like loss. This is because the prototypes are constructed based on information bottleneck theory, in which the first term of the derivation in Eq. (8) is formulated as a form of cross-entropy. Fortunately, NLL loss, a CE-like loss, makes it feasible to employ IB for survival prediction task. 
    
    Although the prototype component is based on CE-like loss, the task loss in our framework is still flexible and can be easily switched to Cox or rank loss. Furthermore, even in rank-based losses, which focus on relative comparisons between samples, discriminative information is still needed. Although not directly involving classification, the model needs to learn the differences between samples to correctly rank them. Hence, the assumption of PIB may not conflict in such scenarios.
    
    Despite the discriminative hypothesis, this work first attempts to directly model the bag-level representations for a bag structure through the probabilistic model in survival prediction task. There is undoubtedly significant value in expanding this approach to various hypotheses like rank-based hypothesis. Hence, further investigation is still warranted in the future.

**REFERENCES**

[1] Echle A, Rindtorff N T, Brinker T J, et al. Deep learning in cancer pathology: a new generation of clinical biomarkers[J]. British journal of cancer, 2021, 124(4): 686-696.

[2] Wiegrebe, Simon, et al. "Deep Learning for Survival Analysis: A Review." *arXiv preprint arXiv:2305.14961* (2023).

Reviewer comment (Evaluation):
In Kaplan-Meier analysis section, authors state that "our approach demonstrates significantly improved discrimination.... as indicated by lower p-values". While it is tempting to compare between p-values as a metric for separation of the risk groups, it is statistically not correct. Once these are below significance threshold, it shouldn't be over-interpreted. Maybe, use median survival days for each risk group?

Author response:
Thank you very much for highlighting this concern! Based on your suggestion, we have revised the relevant section on **Page 8** and incorporated median survival months as a new and more meaningful reference point.  Since we use the five-fold cross-validation, we get the median survival months for each risk group in each validation set. We report both the average median values and their standard deviation.  From the table below, it can be observed that PIBD, compared to Survpath, exhibits a more pronounced distinction in median survival months between the two risk groups across most datasets, more closely with the median values in actual grouping.

| High-Risk Group/Low-Risk Group | HNSC | STAD | BLCA | COADREAD | BRCA |
| --- | --- | --- | --- | --- | --- |
| Ground Truth | 12.21(0.68)/41.18(10.55) | 10.13(1.67)/26.37(6.53) | 10.76(1.13)/33.74(6.00) | 14.02(1.36)/37.03(4.17) | 15.77(1.64)/60.35(6.18) |
| SurvPath | 18.17(2.37)/23.05(1.78) | 15.02(1.35)/20.63(1.40) | 15.33(1.91)/22.82(4.72) | 22.77(2.95)/24.49(4.98) | 25.12(3.48)/33.84(8.39) |
| PIBD | 18.03(3.43)/24.55(4.28) | 13.08(3.34)/18.95(3.53) | 14.84(1.20)/23.10(5.25) | 20.15(3.26)/25.13(6.69) | 24.33(3.97)/33.37(6.88) |

Reviewer comment (Experiments):
Since both PID and PIB rely on sampling from distribution, it does seem that the performance will indeed be affected by which samples are chosen or how many of them are sampled. The discussion around this point needs to be made explicit.

Author response:
Thanks for this very helpful comment! The discussion on sampling is indeed crucial. In the revised version, we have incorporated experimental results and discussions on this aspect into **Appendix C.2 on Page 20**. 

In our work, we adopt the commonly used Monte Carlo sampling method following the information bottleneck-based approaches [1-3], leveraging the reparameterization trick and randomly sampling from a standard Gaussian distribution. Moreover, the choice of the number of samples affects the models’ performances. Therefore, we gradually increase the number of samples from a small to a larger number, as illustrated in the table below. The experiments demonstrate that as the number increases, the model's performance gradually improves. Increasing the number of samples effectively aids in learning prototypes. In the paper, we strive to balance performance and sampling complexity and thus set the sample number uniformly to 50. 

| Sample Number | BLCA | COADREAD | STAD |
| --- | --- | --- | --- |
| 1 | 0.614(0.032) | 0.763(0.129) | 0.634(0.059) |
| 10 | 0.659(0.074) | 0.804(0.093) | 0.673(0.082) |
| 50 | $\underline{0.667(0.061)}$ | $\underline{0.768(0.124)}$ | **0.684(0.035)** |
| 100 | **0.689(0.064)** | **0.811(0.128)** | $\underline{0.678(0.060)}$ |

Reviewer comment (Experiments):
While there are detailed information about training procedure, not much is written about the actual inference step. For instance, how many samples for each prototype are required for reliable performance?

Author response:
Thank you for pointing out this issue. In the new revised version, we have included more details about the inference process on **Page 7**. About the number of samples, as discussed in response to Q5), we set it to 50.

Reviewer comment (Theory):
Motivation leading up to Eq. (5) is confusing. Why is large bag size an issue, if you can learn parametric mapping (via neural network) between z and x, similar to Eq. (4)? Wouldn't this simply be a matter of backpropagation in the training, which is scalable to numerous instances?

Author response:
1. The $\bf{x}$ in Eq.(5) is organized as a “bag” structure containing numerous instances.  If we directly employ the variational approximation of Eq.(4) in VIB to learn a compact representation for each instance $x \in \bf{x}$ in the bag, every instance in a bag would have its own predicted distribution.  **The drawbacks of this solution are two-fold.** First, it is challenging to derive the overall distribution of the entire bag $p(z|\bf{x})$ for a bag $\bf{x}$ based on such a large number of individual instance distributions, leading to a high-dimensional computational challenge. That is, the posterior distribution $p(z|\bf{x})$ with respect to high-dimensional $\bf{x}$ of the second term in Eq.(3) is intractable. Second, since the distribution of each instance is individually learned, it is difficult to capture bag-level information for representing a compact bag. Therefore, we propose using prototypes to directly approximate the bag-level distribution. This simplifies the issue and allows us to capture the essence of the bag's distribution more effectively.

Reviewer comment (Reproducibility):
Variational approaches involving NNs are known to be prone to mode collapse - What are the tricks/methods used in the work to prevent it?

Author response:
We agree that mode collapse in variational approaches is a common issue, which is investigated in many previous researches. In this work, we mainly follow the tricks used in VIB[1] and CLUB [4]. 

- In PIB, some regularization techniques such as dropout are employed. Furthermore, the second term in Eq. (9) is essentially considered a form of regularization to the learnable distribution of prototypes by pushing it forward to a prior distribution r(z). Hence, carefully tuning the coefficient of this term is beneficial to preventing it from mode collapse.
- In PID, variational approaches primarily involve mutual information estimation in Eq. (10). Here we used vCLUB estimator which is illustrated in section B.3 of appendix, where training a good predictor for the conditional distribution $p(b|a)$ of Eq. (31) is critical. In this context, one trick to avoid mode collapse is predicting a lower dimensional distribution conditioned on a higher dimensional distribution [5]. Following this, our work predicted $q_\theta(c|s)$ instead of $q_\theta(s|c)$ for estimating mutual information in Eq. (10), where $s$ is modality-specific distribution integrating two modalities with higher dimension, while $c$ is modality-common distribution with lower dimension.

Regarding the details of training, we have elaborated on it in the appendix on **Page 18** to provide a reference for readers.

Reviewer comment (Presentation):
I think the introduction of Markov chains confuses the paper more - Perhaps better to remove it.

Author response:
In response to your valuable feedback, we have revised **Section 3** of the paper. The modifications include removing the introduction to Markov chains, emphasizing the changes made by the new approach more straightforward for conciseness and clarity.

Reviewer comment (Reproducibility):
Usually the use of "Top-k" refers to small number of instances being selected. However, from implementation section, it seems the authors are using 50~80% of the entire bag? Am I misunderstanding something?

Author response:
Thank you very much for pointing out the ambiguity in our expression. To avoid confusion, we have removed the mention of using "Top-k" and replaced it with the concept of information retention rate (Irr) to represent the proportion of selected instances in the revised manuscript.

Reviewer comment (Evaluation):
**Limited findings w.r.t. clinical application / multimodal interpretability**: PID/PIB is intuitive for solving multimodal integration problems of WSIs and gene sets, but the demonstrated findings of this work appears limited to only improvement of c-Index performance. Methodologies such as MOTCat and SurvPath, though not approaching survival analysis from an information bottleneck perspective, similarly resolve issues regarding redundancy of patch features in MIL and reach similar findings already established in prior works. The interpretability experiments, shown as attention heatmaps in the supplement, are difficult to interpret and draw conclusions from. As the improvement in c-Index is minor, it would be valuable to investigate the interpretability of PIBD, for example: (1) performing local and global interpretability to assess individual and cohort-level image-omic drivers of disease survival, or (2) experiment+showcase unique clinical applications that would arise from PIBD that cannot be performed in other methods, which would produce new scientific findings and expand this study's significance. In particular, I found the prototypical aspect of this work to be intriguing, with concurrent works also realizing (and finding new applications) of using prototypical patterns for pathology [5].

Author response:
Thank you for your feedback and suggestions! In order to present the concerns in a more organized manner, we will state our response point-by-point.

**Contribution**. We'd like to emphasize the contributions of our method, as highlighted by Reviewers bKh7, HNZs, and G86X. Our approach achieves innovation by reducing redundant information within and between modalities for survival prediction through the information bottleneck in a quite novel prototypical method. Methods like MOTCat and SurvPath, which utilize co-attention mechanisms, focus specifically on gene-related instances of WSIs, emphasizing modality-common information. They failed to address the issues that we’re concerned about in the following aspects. First, these approaches ignored the distinct value of modality-specific information. Second, helpful instances in a bag only occupy an extremely small portion, which means massive intra-modal redundancy. Without explicit redundancy removal, such as MOTCat and Survpath, useful information will be flooded. Lastly, they cannot directly model at the bag-level, as noted by Reviewer G86X, "probabilistic modeling on bags opens up new avenues for probabilistic approaches for CPath, which have been under-explored thus far.”

**Interpretability**. For local interpretability of individual patients, we have visualized similarity scores of instances to each prototype of risk level and presented top-3 relevant patches and biological pathways, as well as top-3 irrelevant biological pathways. Through analyzing these instances, we may discover biomarkers of various risk levels. This is because prototypes are designed to be discriminative, which enforces selecting instances that are most distinctive for each risk level, potentially corresponding to biomarkers for different risk levels. In addition to what we have done, another potential way is to analyze the modality-specific features extracted from PID using CAM-based approaches, which will facilitate the discovery of unique biomarkers within each modality.

Since prototypes represent the distribution for each risk band, they inherently possess a global nature. As a result, the global interpretability can be demonstrated in the following aspects. First, by collecting and organizing instances with high similarity to prototypes across a large patient cohort, we can assess cohort-level disease survival. Second, in PID, the joint distribution of histology and genomic data containing modality-common information, contributes to identifying paired instances related to both modalities and further establishing their relationship between patches and pathways. This also helps to understand their collaborative mechanisms.

**Clinical applications.** We truly appreciate your suggestions and believe that the aforementioned way to interpret the model has the potential to provide some clinical value. However, the investigation of clinical applications requires collaborative efforts from experts across disciplines, such as pathologists, statisticians, AI experts, etc. For example, interpretability in clinical applications demands extensive review of pathology slides by pathologists and continuous enhancements in algorithms by AI engineers[1-3]. Therefore, it definitely takes a lot of time and effort in development and elaboration, which has gone beyond the scope of the current work. This work still primarily focuses on technical innovation in methodology. We acknowledge the importance of clinical applications, and we will continue to explore this area and hope to report the findings soon.

**REFERENCES**

[1] Hosseini, Mahdi S., et al. "Computational Pathology: A Survey Review and The Way Forward." *arXiv preprint arXiv:2304.05482* (2023).

[2] Berbís, M. Alvaro, et al. "Computational pathology in 2030: a Delphi study forecasting the role of AI in pathology within the next decade." *EBioMedicine* 88 (2023).

[3] Lu, Ming Y., et al. "AI-based pathology predicts origins for cancers of unknown primary." *Nature* 594.7861 (2021): 106-110.

Reviewer comment (Experiments):
**Ablation Experiments for PIB/PID**: Some baselines do not seem fair in the PID/PIB ablation experiments.

Author response:
Thank you for your valuable feedback! First, we would like to clarify the rationale behind our ablation experiments in our original version. In our method, PID relies on the prototypes obtained from PIB to guide the disentangling of modality-common and modality-specific information by computing the joint prototypical distribution. Therefore, using PID without PIB is not feasible. To assess the effectiveness of each module, we started from a basic baseline without PIB and PID, gradually incorporating these proposed components to evaluate their individual contributions.

However, your suggestion is highly appreciated, prompting us to reconsider the comparison using average pooling may not be strong. Following your advice, we have employed a stronger alternative, the TransMIL-encoder:

1. In the PID ablation experiment, as you suggested, the baseline now uses a non-disentangled Transformer, TransMIL, on top of the prototypical features. The results are shown in variant (b) of the below table.
2. In the PIB ablation experiment, since using PID without PIB is not applicable, we opted to compare with a baseline using a TransMIL-like encoder on the **original features**. The performance can be found in variant (a) of the following table. This comparison allows us to evaluate the effectiveness of PIB by examining the impact of using prototypical features versus original features.

The table below illustrates the results of the updated ablation experiments on five datasets. It is evident that, with the adoption of a stronger baseline, the effectiveness of both PIB and PID remains validated. We have incorporated these findings into the revised version **on Page 9** and provided a more in-depth analysis for better clarity.

| Variants | PIB | PID | BRCA | BLCA | COADREAD | HNSC | STAD | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (a) TransMIL |  |  | 0.672(0.088) | 0.636(0.059) | 0.750(0.133) | 0.591(0.080) | $\underline{ 0.662(0.090)}$ | 0.662 |
| (b) PIBD (PIB+TransMIL) | √ |  |$\underline{ 0.696(0.069)}$ | $\underline{ 0.648(0.074) }$| $\underline{ 0.757(0.176) }$ | $\underline{ 0.615(0.062)}$ | 0.643(0.074) | $\underline{ 0.672}$ |
| (c) PIBD | √ | √ | **0.736(0.072)** | **0.667(0.061)** | **0.768(0.124)** | **0.640(0.039)** | **0.684(0.035)** | **0.699** |

Reviewer comment (Experiments):
In ablating PID, the baseline model considered was average pooling the prototypical features - a more appropriate baseline would be to use a non-disentangled Transformer like TransMIL on top of the prototypical features.

Author response:
Thank you for your valuable feedback! First, we would like to clarify the rationale behind our ablation experiments in our original version. In our method, PID relies on the prototypes obtained from PIB to guide the disentangling of modality-common and modality-specific information by computing the joint prototypical distribution. Therefore, using PID without PIB is not feasible. To assess the effectiveness of each module, we started from a basic baseline without PIB and PID, gradually incorporating these proposed components to evaluate their individual contributions.

However, your suggestion is highly appreciated, prompting us to reconsider the comparison using average pooling may not be strong. Following your advice, we have employed a stronger alternative, the TransMIL-encoder:

1. In the PID ablation experiment, as you suggested, the baseline now uses a non-disentangled Transformer, TransMIL, on top of the prototypical features. The results are shown in variant (b) of the below table.
2. In the PIB ablation experiment, since using PID without PIB is not applicable, we opted to compare with a baseline using a TransMIL-like encoder on the **original features**. The performance can be found in variant (a) of the following table. This comparison allows us to evaluate the effectiveness of PIB by examining the impact of using prototypical features versus original features.

The table below illustrates the results of the updated ablation experiments on five datasets. It is evident that, with the adoption of a stronger baseline, the effectiveness of both PIB and PID remains validated. We have incorporated these findings into the revised version **on Page 9** and provided a more in-depth analysis for better clarity.

| Variants | PIB | PID | BRCA | BLCA | COADREAD | HNSC | STAD | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (a) TransMIL |  |  | 0.672(0.088) | 0.636(0.059) | 0.750(0.133) | 0.591(0.080) | $\underline{ 0.662(0.090)}$ | 0.662 |
| (b) PIBD (PIB+TransMIL) | √ |  |$\underline{ 0.696(0.069)}$ | $\underline{ 0.648(0.074) }$| $\underline{ 0.757(0.176) }$ | $\underline{ 0.615(0.062)}$ | 0.643(0.074) | $\underline{ 0.672}$ |
| (c) PIBD | √ | √ | **0.736(0.072)** | **0.667(0.061)** | **0.768(0.124)** | **0.640(0.039)** | **0.684(0.035)** | **0.699** |

Reviewer comment (Experiments):
In ablating PIB, the reviewer is not certain what the baseline looks like (using PID without PIB), but would think that this comparison would be similar to having each modality being processed by a TransMIL-like encoder + PID (or SurvPath + PID).

Author response:
Thank you for your valuable feedback! First, we would like to clarify the rationale behind our ablation experiments in our original version. In our method, PID relies on the prototypes obtained from PIB to guide the disentangling of modality-common and modality-specific information by computing the joint prototypical distribution. Therefore, using PID without PIB is not feasible. To assess the effectiveness of each module, we started from a basic baseline without PIB and PID, gradually incorporating these proposed components to evaluate their individual contributions.

However, your suggestion is highly appreciated, prompting us to reconsider the comparison using average pooling may not be strong. Following your advice, we have employed a stronger alternative, the TransMIL-encoder:

1. In the PID ablation experiment, as you suggested, the baseline now uses a non-disentangled Transformer, TransMIL, on top of the prototypical features. The results are shown in variant (b) of the below table.
2. In the PIB ablation experiment, since using PID without PIB is not applicable, we opted to compare with a baseline using a TransMIL-like encoder on the **original features**. The performance can be found in variant (a) of the following table. This comparison allows us to evaluate the effectiveness of PIB by examining the impact of using prototypical features versus original features.

The table below illustrates the results of the updated ablation experiments on five datasets. It is evident that, with the adoption of a stronger baseline, the effectiveness of both PIB and PID remains validated. We have incorporated these findings into the revised version **on Page 9** and provided a more in-depth analysis for better clarity.

| Variants | PIB | PID | BRCA | BLCA | COADREAD | HNSC | STAD | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (a) TransMIL |  |  | 0.672(0.088) | 0.636(0.059) | 0.750(0.133) | 0.591(0.080) | $\underline{ 0.662(0.090)}$ | 0.662 |
| (b) PIBD (PIB+TransMIL) | √ |  |$\underline{ 0.696(0.069)}$ | $\underline{ 0.648(0.074) }$| $\underline{ 0.757(0.176) }$ | $\underline{ 0.615(0.062)}$ | 0.643(0.074) | $\underline{ 0.672}$ |
| (c) PIBD | √ | √ | **0.736(0.072)** | **0.667(0.061)** | **0.768(0.124)** | **0.640(0.039)** | **0.684(0.035)** | **0.699** |

Reviewer comment (Novelty):
Current application and findings of PIBD overlap with other multimodal survival analysis works wr.t. to addressing patch feature redundancy, and having similar interpretability experiments. How does the contributions and findings of this work advance the field further?

Author response:
Thank you for your feedback and suggestions! In order to present the concerns in a more organized manner, we will state our response point-by-point.

**Contribution**. We'd like to emphasize the contributions of our method, as highlighted by Reviewers bKh7, HNZs, and G86X. Our approach achieves innovation by reducing redundant information within and between modalities for survival prediction through the information bottleneck in a quite novel prototypical method. Methods like MOTCat and SurvPath, which utilize co-attention mechanisms, focus specifically on gene-related instances of WSIs, emphasizing modality-common information. They failed to address the issues that we’re concerned about in the following aspects. First, these approaches ignored the distinct value of modality-specific information. Second, helpful instances in a bag only occupy an extremely small portion, which means massive intra-modal redundancy. Without explicit redundancy removal, such as MOTCat and Survpath, useful information will be flooded. Lastly, they cannot directly model at the bag-level, as noted by Reviewer G86X, "probabilistic modeling on bags opens up new avenues for probabilistic approaches for CPath, which have been under-explored thus far.”

Reviewer comment (Experiments):
Ablation experiments for PID/PIB can have stronger baselines.

Author response:
Thank you for your valuable feedback! First, we would like to clarify the rationale behind our ablation experiments in our original version. In our method, PID relies on the prototypes obtained from PIB to guide the disentangling of modality-common and modality-specific information by computing the joint prototypical distribution. Therefore, using PID without PIB is not feasible. To assess the effectiveness of each module, we started from a basic baseline without PIB and PID, gradually incorporating these proposed components to evaluate their individual contributions.

However, your suggestion is highly appreciated, prompting us to reconsider the comparison using average pooling may not be strong. Following your advice, we have employed a stronger alternative, the TransMIL-encoder:

1. In the PID ablation experiment, as you suggested, the baseline now uses a non-disentangled Transformer, TransMIL, on top of the prototypical features. The results are shown in variant (b) of the below table.
2. In the PIB ablation experiment, since using PID without PIB is not applicable, we opted to compare with a baseline using a TransMIL-like encoder on the **original features**. The performance can be found in variant (a) of the following table. This comparison allows us to evaluate the effectiveness of PIB by examining the impact of using prototypical features versus original features.

The table below illustrates the results of the updated ablation experiments on five datasets. It is evident that, with the adoption of a stronger baseline, the effectiveness of both PIB and PID remains validated. We have incorporated these findings into the revised version **on Page 9** and provided a more in-depth analysis for better clarity.

| Variants | PIB | PID | BRCA | BLCA | COADREAD | HNSC | STAD | Overall |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (a) TransMIL |  |  | 0.672(0.088) | 0.636(0.059) | 0.750(0.133) | 0.591(0.080) | $\underline{ 0.662(0.090)}$ | 0.662 |
| (b) PIBD (PIB+TransMIL) | √ |  |$\underline{ 0.696(0.069)}$ | $\underline{ 0.648(0.074) }$| $\underline{ 0.757(0.176) }$ | $\underline{ 0.615(0.062)}$ | 0.643(0.074) | $\underline{ 0.672}$ |
| (c) PIBD | √ | √ | **0.736(0.072)** | **0.667(0.061)** | **0.768(0.124)** | **0.640(0.039)** | **0.684(0.035)** | **0.699** |

Reviewer comment (Experiments):
Evaluating PIBD with a ResNet-50 encoder (ImageNet transfer) or a non-TCGA-pretrained encoder would strength the study and its findings.

Author response:
Thank you for your valuable feedback! We appreciate your understanding of our decision to use a pretrained encoder for pathology. For your concern about using CTransPath, we have added additional experiments with a ResNet50 encoder pretrained only on ImageNet. Also since the limitation of computing resources and time, we temporarily evaluated these experiments on three datasets. We selected the well-performing TransMIL in the unimodal group and the SurvPath in the multimodal group as the comparison methods, as shown in the table below.  

| Method | BLCA | COADREAD | STAD | Overall |
| --- | --- | --- | --- | --- |
| TransMIL | $\underline{ 0.588(0.072) }$| 0.699(0.170) | 0.592(0.095) | 0.626 |
| SurvPath | 0.570(0.052) | $\underline{ 0.749(0.100)}$ | $\underline{ 0.624(0.057)}$ | $\underline{ 0.648 }$|
| PIBD | **0.595(0.061)** | **0.784(0.124)** | **0.644(0.063)** | **0.674** |

The results clearly demonstrate that even when employing encoders not pretrained on TCGA data, PIPD exhibits great improvement compared to the comparison methods, which further strengthens the superiority of our method. We will conduct more comprehensive comparative experiments in the future and incorporate them into subsequent versions.

Reviewer comment (Theory):
How are CLAM-SB / CLAM-MB adapted for survival analysis? To the reviewer's knowledge, the implementation of the CLAM framework is mostly situated for slide classification (not survival), as the clustering constraints in CLAM are most appropriate for subtyping problems.

Author response:
Yes, the CLAM framework was originally designed for slide classification. To adapt CLAM-SB and CLAM-MB for survival analysis, we followed the setting of comparative experiments conducted in the works [1,2]. Since NLL loss used for survival prediction is a cross-entropy-like loss, CLAM is easily adapted to this loss by modifying the classifier head.

**REFERENCES**

[1] Xu Y, Chen H. Multimodal Optimal Transport-based Co-Attention Transformer with Global Structure Consistency for Survival Prediction[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2023, pp. 21241-21251

[2] Zhou F, Chen H. Cross-Modal Translation and Alignment for Survival Analysis[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 21485-21494.

Reviewer comment (Writing):
The Prototypical Information Disentanglement (PID) module could benefit from improved clarity in its description.

Author response:
Regarding the clarity of the PID module, we have revised the corresponding **section 3.3 on Page 6** in the manuscript to enhance the understanding for readers.

Reviewer comment (Evaluation):
How does PIB save the computation? Please make quantitative analysis.

Author response:
Thank you for your feedback. For the computation efficiency of PIB, it can be known that directly employing VIB would lead to intractable bag-level distribution $p(z|\bf{x})$ due to the large number of individual instance distributions within a bag $\bf{x}$, based on the analysis in the second paragraph **on Page 5**.  The proposed PIB can simplify this issue by modeling bag-level distribution with prototypes.

We can coarsely quantify the reduction in computational complexity compared to the original VIB. Assuming the sampling frequency is $K$, the number of samples per iteration is $N$, and there are $M$ instances of $D$-dimension in a bag, the computational complexity is $O(K\cdot N\cdot M\cdot D)$. However, in PIB, we only need to sample from prototypes. Assume that the number of prototypes is $P$, the complexity becomes  $O(K\cdot N\cdot P\cdot D)$, where $P << M$. For example, P is 8 in our work, while M is usually larger than 10,000.