Paper: BioBridge: Bridging Biomedical Foundation Models via Knowledge Graphs (ICLR-2024)

Reviewer comment (Theory):
Looking at equation 1, the transformation seems dependent on the (1) target modality node type and (2) relation type and (3) individual node in the target modality. For example, given the same node drug, it could have one transformed representation for protein A, a different one for protein B, and a different one for disease, etc. Is that right? Could the authors describe the design choice of this? Why not map all of them into a single unified space? This brings a separate question: how do you conduct a similarity search if the embedding is dependent on individual node?

Author response:
**1. Transformation module**
> given the same node drug, it could have one transformed representation for protein A, a different one for protein B, and a different one for disease, etc. Is that right?

Thanks for pointing out the phrases that cause confusion. Refer to Eq. (1), this drug will be transformed to the protein space with $\hat{h} = z + \psi(z, c_i, c_j, r_{ij})$, where $z$ is the drug’s raw embedding, $c_i$ indicates the **source modality** is drug, c_j indicates the **target modality** is protein, and $r_{ij}$ indicates the relation is `associate with`. In this sense, this drug is transformed to the protein space, yielding $\hat{h}$, which is not entangled with any individual protein embeddings, such that we can use $\hat{h}$ to match all the candidate proteins.

**2. Transformation module**
> why not map all of them into a single unified space?

Our transformation module is relation-aware, i.e., it maps the same entity to different spaces considering the target relation. This allows the model to differentiate between drug-indicate-disease and drug-contraindicate-disease. A single unified space is not able to make such a distinction.

Reviewer comment (Experiments):
The transformation loss function (eq.2) makes intuitive sense, but there seem to be numerous other options to achieve similar goals such as self-supervised link prediction. Have the authors experimented with other transformation techniques? If so, could the authors provide any additional information? If not, could the authors describe the motivation for choosing this particular approach?

Author response:
**3. Choice of the bridge formula**
> could the authors describe the motivation on why choosing this particular approach?

In this paper, we aim to make a proof of concept about bridging multiple unimodal FMs, supervised with knowledge graphs. We chose InfoNCE loss since it was used by numerous multimodal learning algorithms, such as CLIP [1] and ImageBind [2]. In addition, we tried out some alternative transformation modules for Eq. (1). There are two variants:
- Variant 1: in Eq. (1), removing the residual connection, i.e., using $h = \psi(z,c_i,c_j,r)$.
- Variant 2: in Eq. (1), using RotatE transformation, i.e., using $h = z \circ \psi(z,c_i,c_j,r)$.

In our experiments, Variant 1 failed to converge; Variant 2 obtained a worse performance than the additive transformation in Eq. (1). 


[1] Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International conference on machine learning. PMLR, 2021: 8748-8763.

[2] Girdhar R, El-Nouby A, Liu Z, et al. Imagebind: One embedding space to bind them all[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 15180-15190.

Reviewer comment (Experiments):
The author uses a transformer model to achieve the transformation. It seems unnecessary since it is just encoding 4 embeddings. Have the authors experimented with other simpler approaches?

Author response:
**4. Choice of transformers in the bridge module**
> The author uses a transformer model to achieve the transformation. It seems unnecessary.

We build the bridge module with transformer layers for the following reasons: 
- (1) transformers ensure the input and output embeddings are with the same dimension, which provides a convenient way to achieve the transformation across modalities; 
- (2) transformer offers an effective way of feature interactions between the inputs; 
- (3) transformers were proven with high learning capacity, as they are used in many large language models, and we choose transformers by default to keep the architecture minimum viable, as a proof of the validity of the bridging.

Reviewer comment (Experiments):
For semantic similarity task, there is a concern about potential data leakage. Since the protein node is connected to these GO terms in the PrimeKG, the embedding already implicitly knows the labels during training. Have the authors addressed this concern by conducting some holdout protein-GO links?

Author response:
**6. Data leakage in semantic similarity task**
> For semantic similarity task, I worry that there is data leakage

Thank you for this question. In order to explicitly avoid data leakage, we have removed all proteins in the testing set from the training KG. In this case, all testing protein-go links were not visible to the BioBridge model. We reason that the improvement led by projecting protein embedding to the GO term space stems from the discriminative patterns of GO terms formed by PubMedBERT for associated proteins. This discriminative pattern in GO term space is more pronounced than the original protein embedding space, which only exploits proteins’ internal information.

Reviewer comment (Theory):
For PPI, since it is the same modality between the head and tail nodes, why is a bridge module needed?

Author response:
**7. Bridge Module for PPI**
> since it is the same modality between the head and tail nodes, why do we need bridge module?

There are 7 types of relations representing different types of protein-protein interactions (PPIs): reaction, binding, post-translational modifications (ptmod), activation, inhibition, catalysis, and expression. So, it is still useful if we can perform a bridge within the same modality considering different relations. We also proved bridging proteins with PPI relation leads to better performance in PPI prediction tasks than the original embedding by ESM2 models, as shown in Sec 4.3 Table 3 Page 7, where the transformed protein embeddings considering PPI relations induce better performance than the protein embedding baselines.

Reviewer comment (Experiments):
For the cross-species task, why is phenotype not available? It seems to be available in PrimeKG.

Author response:
**8. Phenotype in PrimeKG**
> For cross-species task, why is phenotype not available? it seems to be available in PrimeKG?

Thanks for pointing it out. This paper aims to make a proof of concept in bridging uni-modal FMs, so it simplifies the KG in the experiment.  Nonetheless, protein, disease, drug, and gene ontology terms are the main biomedical entities covering the majority of real-world biomedical tasks (e.g., drug discovery, repurposing, PPI, protein function prediction, drug target identification, drug-target interaction prediction). Please refer to Table 8 for the data statistics. We also added the explanation to Appendix B Page 15 in the new version.

Reviewer comment (Experiments):
The baselines in Section 4.1 are several years ago. No recent studies are included.

Author response:
In cross-modal prediction tasks, various methods are trained using the same Knowledge Graph (KG) incorporating a range of node types. This approach is evident in a series of multimodal graph learning research, including studies like IKRL [1], TBKGE [2], TransAE [3], and MMKRL [4]. However, these methods primarily focus on visual and textual modalities, which may not directly suit our specific application needs. A critical limitation of these methods is their confinement to predicting nodes within the training KG during the testing phase because these methods all assume the KG is fixed. In contrast, BioBridge stands out by leveraging base unimodal Feature Models (FMs) to process inputs from any modality present in the KG. It can then convert these inputs into any desired target modality. We have added this discussion to Appendix H Page 22, in the new version.

[1] Xie R, Heinrich S, Liu Z, et al. Integrating image-based and knowledge-based representation learning[J]. IEEE Transactions on Cognitive and Developmental Systems, 2019, 12(2): 169-178.

[2] Mousselly-Sergieh H, Botschen T, Gurevych I, et al. A multimodal translation-based approach for knowledge graph representation learning[C]//Proceedings of the Seventh Joint Conference on Lexical and Computational Semantics. 2018: 225-234.

[3] Wang Z, Li L, Li Q, et al. Multimodal data enhanced representation learning for knowledge graphs[C]//2019 International Joint Conference on Neural Networks (IJCNN). IEEE, 2019: 1-8.

[4] Lu X, Wang L, Jiang Z, et al. MMKRL: A robust embedding approach for multi-modal knowledge graph representation learning[J]. Applied Intelligence, 2022: 1-18.

Reviewer comment (Novelty):
What are the contributions of the contrastive learning?

Author response:
Contrastive learning has been widely used in multimodal representation learning, such as in CLIP and ImageBind. Here, we leverage contrastive learning to push the transformed head embeddings close to the target tail node embeddings by optimizing the parameters of the transformation model. We do not intend to propose a new contrastive learning algorithm but choose InfoNCE loss to validate the idea of bridging.

Reviewer comment (Experiments):
The term "multimodal" as mentioned in this paper is confined to different types of biomedical entities. While the authors compare their work with "ImageBind," the experimental section lacks tasks that bridge text and image modalities, which are more complex and crucial for multimodal foundation models.

Author response:
We agree that texts and images are two crucial modalities for multimodal learning, including in the biomedical domain. In this paper, we do include text and prove the feasibility of bridging text with other modalities, such as molecules and proteins, via bridging. The proposed method leverages modalities mentioned in the used KG. As such, it can be extended to image modality if we have a KG with image nodes. For instance, medical images can be connected to such KG via patient-image-disease relations. However, no multi-condition/disease medical imaging datasets are publicly available. Although there are a few single-disease medical image datasets across patients (e.g., MIMIC-CXR), it is not meaningful to connect all of the images to just one disease node in a KG.

Reviewer comment (Writing):
The learning process is guided by knowledge graphs, limiting the scope of "modalities" to those represented within biomedical knowledge graphs. Therefore, instead of the broad term "biomedical foundation model," it would be more accurate to describe it as a "biomedical knowledge graph foundation model."

Author response:
We thank the reviewer for the suggestion. In our paper, KG is leveraged as the supervision for the bridge module to apply transformation across biomedical uni-modal foundation models (FMs). As the KG is not involved in the testing phase since we encode inputs using the base FMs and apply transformation via the bridge module,  we believe that “bridging biomedical foundation models” is an accurate description of the proposed method. We have added the discussion in Sec 3.1, Page 3, to clarify further. Please check the edits in red in the new version.

Reviewer comment (Experiments):
The paper does not present ablation studies, such as evaluations of the contrastive learning objectives or hyper-parameter tuning.

Author response:
We agree that ablation is important. In this paper, we aim to prove that bridging uni-modal FMs is feasible. As such, we deliberately choose simple architectures, like transformers, to transform input embeddings and original InfoNCE loss to learn the transformation modules. All are kept to a minimum but necessary to make BioBridge work. 

We did experiments to check the sensitivity of BioBridge w.r.t. hyperparameters; we chose the batch size to be in $\{512, 1024, 4096\}$, and training epochs to be $\{10, 50, 100\}$. We found the performance was not significantly different for the tried batch size. The method turned out to be converging within 50 epochs, and training with more epochs does not lead to further improvement. As such, we keep the same set of hyperparameters for BioBridge across all experiments: batch size 4096, training epochs 50, and learning rate 1e-4.

In our experiments, we also did the ablations for the bridge module:
- Variant 1: in Eq. (1), removing the residual connection, i.e., using $h = \psi(z,c_i,c_j,r)$.
- Variant 2: in Eq. (1), using RotatE transformation, i.e., using $h = z \circ \psi(z,c_i,c_j,r)$.

In our experiments, Variant 1 failed to converge; Variant 2 obtained a worse performance than the additive transformation in Eq. (1). 

We have added these ablations and analyses to Appendix I on Page 22.

Reviewer comment (Evaluation):
The case study focusing on molecule generation is intriguing. Quantitative assessments of generation performance would be beneficial, for instance, by making direct comparisons with general-domain foundation models, or by offering more qualitative examples to demonstrate the efficacy of the proposed framework.

Author response:
We agree that a quantitative assessment of generation performance supported by BioBridge is interesting. In our testing case, we used biobridge as a retriever to enable conditional molecule generation guided by text prompts. To the best of our knowledge, we did not find testing sets suitable for our evaluation setups. Meanwhile, the generation performance is also influenced by the generator used; in our case, we use Galactica. This is not the core contribution of this paper, and we would like to leave it in for future research to build on with many applications. Below is the comparison of Galactica generating from scratch and augmented by BioBridge retrieval:

- Question: Lymphoma is a broad term for cancer that begins in cells of the lymph system. The drug treats malignant lymphoma by the inhibition of mitosis at metaphase of cancer cells, via polychemotherapy. Generate the most possible SMILES structure of this drug.
  - Galactica with BioBridge: C#CC(C)(C)NC(=O)C1=CC(Cl)=CC(Cl)=C1, most similar drug: Procarbazine DB01168, which is an antineoplastic agent indicated for the treatment of stage III and stage IV Hodgkin's disease in combination with other chemotherapeutic agents.
  - Galactica without retrieval: CCOC(=O)C1(C2=CC=CC=C2)CCN(CCC2=CC=C(O)C(OCC3=CC=CC=C3)=C2)CC1, most similar drug: Meperidine DB00454 is an opioid agonist with analgesic and sedative properties used to manage severe pain.

- Question: Depressive disorder is a common mental disorder. It involves a depressed mood or loss of pleasure or interest in activities for long periods of time. The drug treats depressive disorder by being a selective dopamine receptor antagonist, hence inhibiting dopaminergic hyperactivity. Generate the most possible SMILES structure of this drug that treats depressive disorder.
  - Galactica with BioBridge: CN1CC[C@H]2CN(C[C@@H]3Cc4ccc(F)cc4[C @@H]3[C@H]2C1, most similar drug: Levorphanol DBH00854 which is an opioid analgesic used to treat moderate to severe pain.
  - Galactica without retrieval: C1=CC=C(C2=C(C3=CC=NC=C3)C(C3=CC=NC=C3)=NO2)C=C1, most similar drug: Irospan 24/6 Kit DB00139 which is a water-soluble, colorless crystal with an acid taste that is used as a chemical intermediate.

It can be seen that without retrieval with BioBridge, Galactica tends to generate irrelevant molecules. We have added more examples to Table 21, Page 21, in the new version. We also added the explanation of the most similar drugs to Table 6, Page 9.

Reviewer comment (Presentation):
Presentation of Related work: Currently the submission only has one paragraph on knowledge graph learning and barely describes the embedding alignment literature e.g. in the context of cross modal retrieval (one of the application of proposed method), one can also mention deep CCA literature as Canonical correlation analysis (CCA) is the core of many cross modal retrieval methods.

Author response:
We acknowledge the reviewer’s notice of prior literature on cross-modality retrieval, such as canonical correlation analysis (CCA). We have added more discussion to the prior cross-modal retrieval literature to the new appendix due to the page limit. Since we are working on biomedical tasks, our previous discussion was mostly on multimodal learning for biomedical modalities, as described in the 2nd paragraph of Sec 2. Their main weaknesses are they only focus on two modalities and require large paired datasets.

Reviewer comment (Presentation):
Presentation of Methodology: the submission should motivate the solution somewhat intuitively. Section 3.2 on encoding and transformation is very to the point and concise. The proposed methodology could be motivated better e.g. by contextualizing with respect to some prior work on KGE literature. Although there is a paragraph in the related work on KGE completion, the proposed method is not contextualized.

Author response:
We have extended the discussion in related work about KGE embeddings to the method section to provide a sensible motivation for the proposed method. Although traditional KG embedding (KGE) methods also enable cross-modal prediction via link prediction, they do not generally extrapolate to nodes that are not in the training KG. Instead, BioBridge learns from the triplets to bridge the modalities by transforming the head-modality FM embeddings to the tail's space. In the testing phase, KG is no longer needed for the inputs. Also, we have made it explicit that the parameters of the transformation module and the projection head are optimized in training, with all unimodal FMs frozen. Please find the edits in red in the new version in Sec 3.1: “Although traditional KG embedding (KGE) methods also enable cross-modal prediction via link prediction, they do not extrapolate to nodes not in the training KG….”.

Reviewer comment (Reproducibility):
In the methodology, it is not very clear what parameters are optimized with contrastive loss since the submission keeps the pretrained model frozen.

Author response:
We have extended the discussion in related work about KGE embeddings to the method section to provide a sensible motivation for the proposed method. Although traditional KG embedding (KGE) methods also enable cross-modal prediction via link prediction, they do not generally extrapolate to nodes that are not in the training KG. Instead, BioBridge learns from the triplets to bridge the modalities by transforming the head-modality FM embeddings to the tail's space. In the testing phase, KG is no longer needed for the inputs. Also, we have made it explicit that the parameters of the transformation module and the projection head are optimized in training, with all unimodal FMs frozen. Please find the edits in red in the new version in Sec 3.1: “Although traditional KG embedding (KGE) methods also enable cross-modal prediction via link prediction, they do not extrapolate to nodes not in the training KG….”.