HEALNet: Multimodal Fusion for Heterogeneous Biomedical Data

Comments on weaknesses:

Usefulness of iterative fusion:

In short, having multiple fusion layers (i.e., iterations) helps to prevent overfitting. During development, we found that the modality-specific updates only learn how to operate on the latent embedding $S$ directly after the previous modality update if you only have one fusion layer. Increasing the number of fusion layers means that $S$ is updated with all modalities multiple times such that information from each modality becomes part of the context vector for every other modality. We share the weights between the fusion layers to prevent layer specificity such that the modality-specific updates become context-agnostic. This is why the model also delivers robust performance when skipping a modality at inference time (Table 2).

Regarding the number of fusion layers needed, we ran some initial ablation studies during development, where we did not find a clear pattern on the exact number of fusion layers needed, but found that $d>1$ was an effective regulariser.

Cross-modal interactions & other architectures:

Using the shared latent space projection in a “monolithic” architecture is something that we originally considered and decided against for three main reasons:

• Don’t align what shouldn’t be aligned: First, we would like to point out a crucial distinction of multimodal modelling in the Vision & Language (V&L) domain and on biomedical modalities beyond V&L.
  • Vision & Language models rely on the assumption that the same semantic concept is explicitly expressed across modalities (e.g., an image and its description). In other words, both modalities follow the same training distribution and V&L models consequently align the latent representations to reflect this distribution across both modalities, such as through the contrastive objectives in the linked Chameleon paper.
  • For Biomedical modalities beyond V&L, we don’t want to make this assumption. Often, the signal from different modalities (e.g., morphological features, genomics, and transcriptomics) follows different distributions and we want to maintain this separation to build a more predictive model. If we used projections into the same latent space combined with contrastive objectives, this would align the signal from the separate distributions. This defeats the purpose of cross-modal learning for these modalities as it becomes more difficult to learn from both distributions.
  • We have benchmarked several approaches which leverage projection and alignment in the same latent space. The MCAT baseline also aligns both modalities directly through the “genomics-guided cross-attention” where the dot product of the WSIs and the multiomic data is minimised. Another example is the Perceiver (Early Fusion) baseline which is first combining the input modalities to a common tensor representation and then passes them through multiple cross-and self-attention layers. HEALNet consistently outperforms these methods (Table 1).
• Scaling up to many modalities: The iterative attention mechanism alleviates the use of fusion operators that can result in high-dimensional latents (e.g., Kronecker product [1]). This allows for combining many modalities while preserving the structural information of each as well as learning cross-modal interactions.
• Improved handling of missing modalities: Another benefit of an iterative architecture instead of a monolithic one is that each layer pass does not require all modalities at the same time.

Larger datasets:

We would like to point out that for cancer pathology tasks The Cancer Genome Atlas (TCGA) contains by far the largest patient cohorts available in the public domain with each cancer cohort being over 500GB in size (L192-198).

To show HEALNet’s performance on larger sample sizes, we ran some additional experiments comparing HEALNet to a subset of baselines (due to resource constraints in the rebuttal period) on a disease classification (ICD9) and patient mortality (MORT) prediction tasks of MIMIC-III (n=32,616). The modalities in MIMIC are tabular EHR features and time series on vital signs and test results. The results are shown in the general comment of the rebuttal.

Responses to questions

• […]unclear how the features from different patches in the WSI aggregated? […]
  • The final embedding dimension is a 2D Vector of size num_patches x encoding_dim . The encoding_dim refers to the dimensions of the Kather100k pre-trained encoder, which results in a 2048-dimensional feature vector per patch (see L195).
• Do the early-fusion models have more parameters? […]
  • This highly depends on the early fusion method used. A commonly used early fusion methods is the Kronecker product [1] which is constructing a high-dimensional input tensor, meaning that models using Tensor fusion would require a higher number of parameters in the earlier layers than HEALNet. However, other early fusion models use aggregation methods (e.g., concatenation and averaging across excess channels, as done in the Perceiver). In these cases, the parameters required in the early layers would be lower than HEALNet’s, but notably this approach also collapses an entire dimension, which may remove salient signal.
• […] how much does the order of fusing modalities impact the result?
  • During development of HEALNet, we performed preliminary analyses on the influence of the order of modalities on the downstream performance: We did not find any significant difference. Since all modalities become the context of each other’s update as long as $d>1$, the order of the modalities does not matter.

References:

[1] Zadeh, A. et al (2017). Tensor fusion network for multimodal sentiment analysis.

Thank you very much for responding to our rebuttal! Per your suggestion, we will further highlight the role of having multiple fusion layers in the manuscript. We briefly wanted to clarify the open points below:

“If the modalities are complementary its unclear why a concatenation approach doesn't work as well, while it may not handle missing modalities.”

We agree that, in principle, the concatenation approach would work well if all modalities had the same tensor dimensions. However, if we are processing structurally different modalities (e.g., 1D vs 2D vs 3D tensors), one would need to flatten or aggregate the excessive dimensions to ensure shape consistency for the projection. This leads to two problems: 1) we may end up with excessively large 1D tensors which 2) remove spatial signal if we e.g., flatten an image representation. In practice, we cannot always assume shape consistency between our modalities and the presence of all modalities for every sample (especially in clinical practice) - both of these aspects informed our design choices. We touch on these aspects in L89 & L111-114.

“Are the tokens from different modalities concatenated and passed through iterative fusion? Wouldn't it have similar paramters to HEALNet then (the Q,K,V projection dims remain similar)”

Contrary to the Perceiver, in HEALNet each modality has its own Q,K,V projection (as indicated in Figure 1), which allows us to keep the input data intact irrespective of the tensor dimensions. That is, HEALNet does not require to aggregate excessive dimensions which helps to maintain the structural signal of each modality.

Comments on weaknesses

• “methodology is not clearly written and difficult to understand”: We would like to point to the step-by-step methodology, which is discussed in detail in Section 3, illustrated in Figure 1, and presented as pseudocode in Appendix A. We believe this to be a very detailed description of HEALNet. Nevertheless, we would be grateful if the reviewer could specify which parts of the methodology are unclear, so we can further improve the clarity of our manuscript.
• “[…] Figure 3 features a bar plot that is unreadable”: We would kindly like to point out that the aim of Figure 3 is to provide a high-level illustration of the general inspection capabilities of HEALNet that help users validate the model, as well as understand the results. We believe we mention this in the caption of Figure 3. In this work, Figure 3 is meant to showcase the capabilities of HEALNet, rather than to deep dive into the biology of UCEC. The size/formatting of the Figure 3 is due to page limitations. Still, Figure 3 is sufficiently high resolution that the individual feature names can be read on the PDF version of the paper. Nevertheless, we appreciate the feedback and will provide a larger version of the image in the Appendix.
• “iterative attention mechanims […] can be computationally expensive”: The computational cost of the iterative attention mechanism is alleviated by sharing weights across the fusion layers, which makes the number of parameters independent of the fusion layers. It is worth noting that the existing tasks of the paper are already dealing with very high-dimensional datasets (see L181 onwards). To show that HEALNet also works for datasets with a larger sample size, we ran some additional experiments comparing HEALNet to relevant baselines on a disease classification (ICD9) and patient mortality (MORT) prediction tasks of MIMIC-III (n=32,616). The table below shows the AUC and Macro AUC for the respective tasks. The modalities in MIMIC are tabular electronic health record features and time series on vital signs and test results.

• “[...] number of fusion layers is a critical parameter that can significantly affect the computational efficiency of the model”: As previously mentioned, the number of fusion layers does not affect the computational efficiency of the model due to the implementation of weight sharing.

Responses to questions

• Handling discrepancies in feature scales and semantic differences between modalities: We would like to point out a crucial distinction between multimodal modelling tasks in the “standard” Vision & Language (V&L) setting and ones over Biomedical modalities (beyond V&L):
  • Current Vision & Language approaches rely on the assumption that the same semantic concepts are expressed within the modalities (e.g., an image and its text description). Moreover, these semantic concepts are typically explicit (e.g., an image with a bird and a description that includes ‘bird’). Therefore, both modalities follow the same training distribution. In turn, V&L approaches align the latent representations to reflect that distribution across both modalities, and leverage these explicitly expressed semantic concepts through a contrastive objectives, such as the one in the linked Tangled paper [1].
  • For Biomedical modalities beyond V&L, we generally cannot make such assumptions. Often times, the signal from different modalities (e.g., morphological features, genomics, and transcriptomics) follows different distributions and we want to maintain that separation to build a more predictive model. If we used projections into the same latent space combined with contrastive objectives, this would likely mute the signal from the separate modalities. This defeats the purpose of cross-modal learning for these modalities as it becomes more difficult to learn from both distributions. It is worth noting that while we separately align the latent representation to each modality, we never align the modality representation with each other to prevent such a muting effect. Moreover, on a broader note, if we were to make a similar assumption as in [1], such an approach would be applicable and scale only to two modalities - which is the case in [1], where the authors learn from WSI and gene expression profiles only. HEALNet, on the other hand, is robust and can readily scale to many different modalities of any size, as shown in the MIMIC experiments above.
• While we appreciate and thank the reviewer for the reference [1], we would like to point out that paper [1] was published just three days before the NeurIPS deadline. Nevertheless, we have benchmarked some approaches that perform a projection and alignment in the same latent space. Specifically, the MCAT [2] baseline also aligns both modalities directly through the “genomics-guided cross-attention” where the dot product of the WSIs and the multiomic data is minimised. Another example is the Perceiver (Early Fusion) [3] baseline which is first combining the input modalities to a common tensor representation and then passes them through multiple cross-and self-attention layers. Note that HEALNet consistently outperforms both of these methods, as shown in Table 1 and Section 5.

[1] Jaume et al. (2024) Transcriptomics-guided Slide Representation Learning in Computational Pathology

[2] Chen, R.J. et al. (2021). Multimodal co-attention transformer for survival prediction in gigapixel whole slide images.

[3] Jaegle, A., et al. (2021) Perceiver: General perception with iterative attention.

Thank you for responding to our rebuttal. In our initial response we briefly discuss the computational complexity of HEALNet, and specifically the role of the fusion layers (as per the last weakness point). We also added new experiments to show that HEALNet is able to scale seamlessly to larger datasets, where the computational cost of the iterative attention mechanism is alleviated by sharing weights across the fusion layers. This makes the number of parameters independent of the number of fusion layers (i.e., iterations).

More formally, for any multimodal model, it is important to scale with both number of samples $n$ and the number of modalities $m$. For instance, MCAT (our most competitive baseline) is natively designed for two modalities. To scale this approach to $m>2$, one would need to calculate the modality-guided cross-attention for all unique pairwise combinations ${m \choose 2} = \frac{m(m-1)}{2}$ which is $\mathcal{O}(m^2)$.

A key advantage of HEALNet’s sequential setup is that it scales linearly. Since the cross-attention and SNN layers used within the fusion layers scale with $\mathcal{O}(n)$, each fusion layer has a complexity of $\mathcal{O}(mn)$ such that the runtime then mainly depends on the number of fusion layers $d$. Note that, HEALNet we consider $d$ to be a hyper-parameter with the values provided in Appendix E. We plan to post the number of parameters of HEALNet and the baselines on the KIRP dataset later today.

We acknowledge that the actual runtime will depend on several other factors such as the size of the modality-specific embeddings within each model. As we didn't measure FLOPs at this instance, we provide the BigO analyses of our method with respect to the other baselines. Assuming we use the same encoders for each model, the time complexity for each fusion operation is summarized below.

We hope that this clarified the remaining concerns and will be reflected in your updated score.

Responses to Weaknesses and Questions

Hyperparameter optimization: We would like to clarify that we ran the Bayesian Hyperparameter optimisation for all baselines. The shared parameters in Table 5 were ran equally for all baselines from Table 1. We acknowledge that Table 5 currently does not show the final hyperparameters of the baselines, but we will add these for the final manuscript. Please see the relevant hyperparameters that we tuned for the baselines below. Note that this choice was mostly informed by the original implementations of the respective papers.

• MCAT/MOTCAT/Porpoise:
  • model_size_omic : Embedding size of early layers for the omic data (pre-set choices between ‘large’ and ‘small’ from original papers)
  • model_size_wsi : Same as model_size_omic, but for the WSI embedding
  • dropout : Dropout applied after each layer of original implementation
  • fusion_method : choice between concatenation and bilinear fusion
• Perceiver (early):
  • layers : total number of iterations
  • latent_dims : dimensions of latent bottleneck
  • attention_heads : number of attention heads
  • attn_dropout : dropout after each self-attention layer
  • ff_dropout : dropout after each feedforward pass
• MultiModN:
  • embedding_dims : size of latent embedding

Latent Update in Eq. 2:

• Clarification on $\psi(\cdot)$: In the main body, the update function $\psi(\cdot)$ represents the learnable function (i.e., the cross-attention update) instead of a fixed closed-form expression. In the pseudocode in Algorithm 2, given the attention matrix $attn \in \mathbb{R}^{b \times i \times j}$ for batch size $b$, number of queries $i$, and keys/values $j$, for modality $m$ at time step $t$ and the value matrix $v \in \mathbb{R}^{b \times j \times d}$ for the dimensions of the modality value vector $d$, the update mechanism is $\psi(S_t, a^{(t)}_m) = \sum_j attn_{b,i,j} \cdot v_{b,j,d}$. We will make this part clearer in the updated version.
• Clarification on $\rho$: Thanks for catching this. It is a leftover notation (typo) that we missed in the last version of the manuscript. We will correct this accordingly.

Comments on Weaknesses

Clarification on Novelty: As you rightly point out, the general idea of iterative cross-attention with latent state passing has been previously suggested, which we acknowledge as a starting point for HEALNet’s architecture (L89, L113, L151). While other iterative modelling architectures (such as the Perceiver) handle a range of uni-modal tasks, HEALNet introduces a novel way of how iterative attention can be leveraged for learning effective multi-modal representations. Concretely, the Perceiver architecture requires concatenation of the input modalities. This becomes particularly problematic for modalities with misaligned dimensions and channels, as you either need to flatten all tensors or aggregate across channels, both of which trades off information for shape compatibility. The Perceiver paper demonstrates multimodal capabilities by concatenating audio and video embeddings, reducing the dimensions of the video and treats them as a single modality thereafter — this is identical to the Early Fusion baseline shown in Table~1, which HEALNet outperforms. HEALNet learns modality-specific latent updates to ensure that the modality shapes can be kept intact. Additionally, learning modality-specific updates becomes valuable in the handling of missing modalities. In such scenarios, an early fusion approach like the Perceiver would need to impute the missing values, which is noisy. In contrast, HEALNet’s setup allows to skip the missing updates at train and inference time.

Classification and Regression tasks: We agree that the method is general which was a key design objective of the architecture. The focus on survival analysis in this paper was chosen for two main reasons.

First, this paper focuses on multimodal data in the context of biomedical tasks. In these domains, most tasks that are a) clinically relevant and b) have sufficient data available to deal with time-to-event labels, where the outcomes of some study subjects is unobserved (right-censored data). Neither classification nor regression on these task labels is suited to handle time-to-event or the concept of censorship accurately.

Second, it is worth noting that the survival analysis task was chosen to ensure comparability with relevant baselines. Under the hood, the survival analysis implementation is using a classifier to predict the hazard bins of a proportional hazards survival model.

Architecturally, the only thing that would change to adapt to a classification task is the very final layer (”Head”) in Figure 1A.

Sample size: In the field of cancer pathology, The Cancer Genome Atlas (TCGA) contains by far the largest patient cohorts. It is worth noting that given the high dimensions of the whole slide images (see L185), even a single cancer cohort (e.g., TCGA-BLCA) contains over 500GB of data. This is also the case for the benchmarked previously published studies.

Confidence Intervals: We reported standard deviation instead of confidence intervals for two reasons. First, the target variable does not follow a t-distribution (see this plot), which is a crucial assumption for the valid interpretation of confidence intervals. As such, reporting standard deviations across the repeated random sub-sampling folds is a more valid estimate of the model’s confidence. Second, using standard deviations was in line with the benchmarked methods, which improves the comparability when directly contrasting the approaches.

Metrics beyond c-Index: We agree that Kaplan-Meier plots provide an in-depth representation of the survival probabilities on the test set, but the actual visual differences are often subtle and barely detectible by a naked eye. Therefore, the c-Index coupled with the standard deviation across folds provides a better way of comparing results, as it is also commonly found in the published baselines.

Responses to questions

Higher performance reported in MOTCat [1] paper:

To make the results across experimental setups comparable, we re-ran all experiments under the same experimental conditions and the tunable sets of hyperparameters from the original paper, which we show in Table~1. The different results can be explained by the fact that this paper’s experimental setup is quite different from the one reported in [1]. While [1] predicts disease-specific survival (DSS), HEALNet is trained to predict general survival (GS). General survival does not filter for case deaths unrelated to the disease and is therefore a noisier prediction task. This filtering is evident by the smaller sample sizes on the same cancer cohorts in [1]. Getting higher c-Index results is a general trend we observed in studies predicting DSS instead of GS[2].

HEALNet follows the same experimental setup, evaluation, and cohort selection as in several multimodal benchmarks (Porpoise, MCAT).

Rationale for weight sharing:

We would would like to clarify two points: First, the weights are only shared between the sets of modality-specific weights across the fusion layer. For example, the cross-attention update for modality 1 and 2 would be shared across all the fusion layers, but no weights are shared between the modality 1 and 2 within or across fusion layers. Second, the weight sharing helps to 1) improve the parameter efficiency of the model and prevent exploding parameter size with a high number of layers and modalities and 2) we found it to prevent overfitting.

[1] Xu, Y., Chen, H., 2023. Multimodal Optimal Transport-based Co-Attention Transformer with Global Structure Consistency for Survival Prediction, in: 2023 IEEE/CVF International Conference on Computer Vision (ICCV). https://doi.org/10.1109/ICCV51070.2023.01942

[2] Jaume, G., et al., 2023. Modelling dense multimodal interactions between biological pathways and histology for survival prediction. arXiv preprint arXiv:2304.06819.

Dear reviewer a9ao,

Thank you again for your thoughtful review and positive attitude toward our work! We have addressed all of your questions and concerns. As the discussion period is drawing to a close, we are keen on getting your feedback.

We are looking forward to your response. Thank you for your time.

Kind regards, The Authors

Dear AC,

We really appreciate your time and effort in managing our submission.

Even though we addressed all of reviewer a9ao questions and concerns, we've noted that they haven't responded nor acknowledged our rebuttal thus far, even though they showed positive outlook on our work. We believe further interaction can be beneficial, and we are keen to discuss any remaining clarifications in the final hours of the discussion period.

We would appreciate your assistance in reminding the reviewer to respond/acknowledge our rebuttal.

Kind regards, Authors