A teacher-teacher framework for clinical language representation learning

Data Requirements and Availability: A notable limitation of the proposed LINE framework is its dependency on well-aligned and specific types of data sources, which may not be readily available or commonly found in practical settings. For example, integrating data from disparate modalities like CT and MRI requires the availability of cases that include both types of data, which may not always be feasible. This requirement could limit the framework's applicability across different clinical or real-world scenarios where such aligned data sets are scarce. See weakness, under such situation is it possible to apply your method?

This is an extremely insightful question, and we appreciate you bringing this up! Yes, our method can handle such situations by introducing an intermediary to bridge the gap between different data sources.

For instance, if we want to infer $\mathbf{x}_1$ and $\mathbf{x}_3$ about each other but lack paired data $(\mathbf{x}_1, \mathbf{x}_3)$, the LINE framework can use an intermediary $\mathbf{x}_2$. This forms pairs $(\mathbf{x}_1, \mathbf{x}_2)$ and $(\mathbf{x}_2, \mathbf{x}_3)$, enabling the use of three teachers to jointly learn the alignment between all data pairs.

In a concrete example involving CT and MRI scans, it is often rare for patients to have both scans simultaneously ($\mathbf{x}_1$ for CT and $\mathbf{x}_3$ for MRI). However, each type of scan typically comes with corresponding clinical notes ($\mathbf{x}_2$). These clinical notes can serve as intermediaries. By breaking down the clinical notes into concepts corresponding to specific scan regions, we can leverage Teacher 2, a pre-trained LLM, to incorporate factual relationships about these concepts. Teachers 1 and 3, pre-trained large vision transformers, can then embed the CT and MRI scans, respectively. The three teachers can train jointly, allowing the framework to effectively handle the alignment of disparate data modalities.

The authors discuss the potential improvement instead of the limitation of the current work. Bring more information and try other situations cannot be counted as an adequate discussion of limitations.

Thank you for your constructive comment. We recognize two main limitations of our framework:

First, our approach requires data to come in pairs. While we can introduce intermediaries to create artificial pairs, this remains a limitation. The reliance on paired data may restrict the framework's applicability in scenarios where such pairs are not readily available.

Second, effective training with a list of concepts necessitates a pre-existing factual knowledge graph. Access to such a knowledge graph is not always guaranteed, which can hinder the framework's performance and generalizability.

We appreciate your feedback and will incorporate a more detailed discussion of these limitations in the paper.

Thank you for your comment.We will address your comments point-by-point.

The problem's importance is not significant.

Thank you for raising this important question. While our motivating example originated from the medical domain, our LINE framework is broadly applicable to a wide range of scenarios. The core problem we address involves pairs of unlabeled data points $(\mathbf{x}_1, \mathbf{x}_2)$ from different domains, where $\mathbf{x}_2$ consists of items with known relationships derived from pre-existing factual knowledge graphs. Our objective is to learn an embedding $\mathbf{z}_2$ for $\mathbf{x}_2$ that closely matches the embedding $\mathbf{z}_1$ for $\mathbf{x}_1$. This general framework can be applied to various critical applications, including the alignment of text from structured EHR codes paired with clinical notes, images/audio paired with text. Below, we detail a few specific applications:

1. Text Summarization: In healthcare systems, clinical notes ($\mathbf{x}_1$) often need to be inferred from corresponding concepts ($\mathbf{x}_2$) due to regulatory constraints. Since raw clinical notes are frequently inaccessible, using a list of concepts to summarize and index notes enables de-identification and facilitates sharing among researchers. The concept list can be utilized instead of full clinical notes for downstream analysis.
2. Text Summarization for Image and Audio: In a cross-modality setting, the LINE framework can align audio/image data ($\mathbf{x}_1$) with corresponding textual information ($\mathbf{x}_2$). For instance, medical images are more challenging to anonymize than text due to embedded metadata and visual information that might indirectly reveal patient identity. In this context, LINE can employ a pre-trained vision transformer to embed images and a pre-trained language model to embed conceptual entities from clinical text. Consequently, LINE can provide embedding surrogates for images, such as CT scans, by leveraging only the concepts extracted from corresponding clinical notes. The text summaries can be used to both index the images and for direct downstream analyses. Additionally, our LINE algorithm has the potential to identify "residual" information in the images that is not captured by the paired data, which can then be projected into the text domain to provide meaningful insights.

We will incorporate this into the motivation and discussion sections of the paper.

There lacks the inclusion of SOTA models like llama, gpt, etc.

Thanks for your suggestion! During the rebuttal phase, we are able to use GPT4 as the strong teacher 1 to come up with LINE (GPT4+CODER). And then we further compared the new LINE results with the GPT4 baseline. The results are shown in the tables in the rebuttal pdf. We also provide a summary of the key results below for your easy reference. The following tables are additional results to be added to Tables 2, 3 and 5 of the paper. The better results are highlighted in bold.

From the tables above, it can be observed that in most of the setting, LINE model performs better than directly using GPT4.

The results improvement is limited as shown in Tab. 4,5.

As discussed in our response to your first question, our evaluation of the proposed framework is not solely focused on outperforming the benchmark model. It also emphasizes the potential of using a list of concepts to summarize notes, which enables de-identification and facilitates sharing among researchers. In Table 2, we evaluate whether the summarization from the list of concepts can be effectively aligned with the clinical text through a rank-based retrieval task. In Table 5, we assess whether the performance using only concepts can be comparable to the performance using raw text. The results from both the tables in the paper and the additional results presented above consistently demonstrate that the LINE model is more effective in summarizing text.

Figure 1 lacks details of the proposed method.

Thank you for your comment. We have revised Figure 1 to include more details, which can be seen in Figure 1(b) of the rebuttal PDF. Specifically, Teacher 2 utilizes the trainable LINE module to first summarize the set $\mathbf{x}_2$, guided by pre-existing knowledge (i.e., the graph attention module and relational contrastive loss). Then, Teacher 2 generates an embedding $\mathbf{z}_2$ for $\mathbf{\tilde{x}}_2$ that closely resembles the embedding $\mathbf{z}_1$ for $\mathbf{x}_1$ through the alignment loss. Simultaneously, Teacher 1, a stronger LLM, uses a fully connected layer to learn the relational knowledge provided by Teacher 2 via the alignment loss.

Dear Reviewer Hfn8, I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors already left rebuttals for your review.

We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC

Thank you for your detailed comments. In the following, we will address your questions point-by-point.

The project's scope is incredibly narrow.

Thank you for raising this important question. While our motivating example comes from the medical domain, our LINE framework is broadly applicable to various scenarios. The core problem we address involves pairs of unlabeled data points $(\mathbf{x}_1, \mathbf{x}_2)$ from different domains, where $\mathbf{x}_2$ consists of items with known relationships derived from factual knowledge graphs. Our goal is to learn an embedding $\mathbf{z}_2$ for $\mathbf{x}_2$ that closely matches the embedding $\mathbf{z}_1$ for $\mathbf{x}_1$. This framework can be applied to multiple critical applications, such as:

1. Text Summarization: In healthcare, clinical notes ($\mathbf{x}_1$) often need to be inferred from concepts ($\mathbf{x}_2$) due to regulatory constraints. Using a list of concepts to summarize notes enables de-identification and facilitates sharing among researchers. These concept lists can replace full clinical notes for downstream analysis.
2. Text Summarization for Image and Audio: The LINE framework can align audio or image data ($\mathbf{x}_1$) with textual information ($\mathbf{x}_2$). Medical images are harder to anonymize than text due to metadata and visual information. LINE can use a pre-trained vision transformer to embed images and a language model to embed conceptual entities from clinical text, providing embedding surrogates for images like CT scans. These text summaries can index the images and be used for downstream analyses. Additionally, LINE can identify "residual" information in the images not captured by paired data and project it into the text domain for insights.

This discussion will be included into the paper. While including proof-of-concept on other settings would be beneficial, we are unable to do so due to time constraints.

I'd strongly consider citing and discussing DA literature for this paper.

Thank you for your suggestion. We agree that our work is related to domain adaptation (DA) and will cite relevant literature, though our approach differs in key ways:

• Label Requirements: Unlike DA, our method does not require labels during training. DA typically needs task-related labels for the source domain and may use "pseudo" labels for the target domain.
• Information Retention: Our goal is to align latent embeddings to preserve both overlapping and complementary information, using a residual alignment training step to refine the alignment with additional concepts. DA focuses only on retaining overlapping information.
• Embedding Richness: We aim to create semantic-rich embeddings by leveraging a graph-attention module, preserving relational information and mitigating rank degeneracy. DA often produces domain-agnostic embeddings, whose information richness depends on the specific task.

We will further include this discussion in the paper.

Would it be possible to further test the LINE model on other, more varied, benchmarks?

Thank you for your suggestion! In response, we are currently running an additional benchmark to evaluate the performance of the LINE models. We anticipate being able to present these results in the early stages of the discussion phase.

It could also be exciting to explore this with generative models.

The LINE framework generates in the embedding space, not the original data space. Given an unlabeled data pair $(\mathbf{x}_1, \mathbf{x}_2)$ from spaces $\mathcal{X}_1$ and $\mathcal{X}_2$, respectively, LINE aligns their latent embeddings $\mathbf{z}_1 \approx \mathbf{z}_2$. For a new observation $\mathbf{x}_2^{\text{new}}$, we can map it to the latent space to generate $\mathbf{z}_1^{\text{new}} \approx \mathbf{z}_2^{\text{new}} = \text{LINE}(f({\text{teacher 2}}(\mathbf{x}_2))$ for the missing $\mathbf{x}_1^{\text{new}}$. To extend this, we could add a decoder to map embeddings back to the clinical note space using a simple reconstructive loss. We will discuss this potential extension in future work.

Why not align directly in embedding space without the grounding concepts?

Thank you for your thoughtful question. Grounding concepts in our framework both serve to integrate external factual knowledge to enhance model performance and safety, and also provide an independent data source for generating clinical note summaries, making them integral to our approach.

Why not also compare the BGE-->CODER projection?

Thank you for the suggestion! Following your advice, we have included the results for the BGE-->CODER projection across our tasks in the rebuttal PDF. Overall, the projection between BGE and CODER yields worse results compared to the proposed LINE model.

Conversely, could you just fine-tune the generalist model on the extracted concepts as a means of medically aligning it?

During the rebuttal phase, we attempted to fine-tune generalist models directly and with LoRA, but found it infeasible due to extremely long training times and high computational resource requirements.

This highlights the computational efficiency of our proposed method. As shown in Figure 1(a) of the rebuttal PDF, even with LoRA, training the generalist BGE model takes at least 11 days per epoch, with rapidly increasing computational overhead. In contrast, our proposed model requires only about 2 hours per epoch using the same resources.

Without additional experiments at the minimum a stated limitation should be the highly restricted domain of application

We acknowledge this concern and will include a discussion of the limitation that our current framework has only been tested on medical topics. We recognize the importance of evaluating our framework in other domains and will address this as a key area for future work.

Dear Reviewer TNnd, I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors already left rebuttals for your review.

We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC

Thank you for your insightful comments! In the following, we will address them point-by-point.

Figure 1 is somewhat confusing. From my understanding, Teacher 1 should be a strong LLM, while Teacher 2 should be an LLM with existing domain-specific knowledge. However, Figure 1 gives the impression that Teacher 2 serves merely as a database.

Thank you so much for pointing this out! We have modified Figure 1 into Figure 1(b) in the rebuttal PDF for better clarity. In summary, both Teacher 1 and Teacher 2 can generate embeddings when feeding in new data which are then aligned by LINE. When the data are paired, e.g. they come from the same hospital visit of a patient, their LINE-aligned embeddings will be close to each other.

From Figure 1, if Teacher 2 only serves to provide domain-specific knowledge, why not implement a RAG framework, which is training-free and potentially more reliable?

The main role of Teacher 2 is to learn an embedding using a list of concepts to summarize the clinical notes, effectively matching the embedding generated directly from the raw notes. Since raw clinical notes are frequently inaccessible, summarizing and indexing notes using a list of concepts enables de-identification and facilitates sharing among researchers. Additionally, by allowing two teachers to align knowledge with each other via LINE, we hope to also improve the quality of the embeddings .

While the RAG framework also leverages pre-existing knowledge networks to enhance learning, it is not a generative model and cannot be used to generate new embeddings. Therefore, it does not fulfill the same role as our proposed method, which not only aligns domain-specific knowledge but also generates new, robust embeddings.

We will include RAG in our discussion of related work in the paper.

Although the paper compares the LINE model against several strong baseline models, it lacks a detailed comparison with the latest strong general LLMs, such as GPT4, which should be considered a strong baseline. Consider adding a small comparative analysis or stating the advantages of the framework over simply using GPT4.

Thank you for your advice! During the rebuttal phase, we incorporated GPT4 as a strong Teacher 1 in our LINE framework, resulting in LINE (GPT4+CODER). We then compared the performance of this new model with the GPT4 baseline. The results, shown in the tables in the rebuttal PDF, indicate that LINE (GPT4+CODER) generally achieves better performance than GPT4 alone.

Have you addressed potential hallucination issues? Could one teacher potentially mislead the other during the knowledge exchange process?

Thank you for the insightful question! To mitigate potential hallucination issues, we integrate factual knowledge from external reputable sources (e.g., UMLS) into our framework. Specifically, in our framework, external factual knowledge is incorporated into the training of Teacher 2 via the multihead graph attention module and the relational contrastive loss. This knowledge is then further propagated to Teacher 1 through the alignment loss.

In this regard, Teacher 2 acts as a factually rigorous component that regularizes potentially misleading information from Teacher 1. We conducted an experiment to assess the quality of the LINE-aligned concept embeddings. The results, shown in Table 3 of the paper and the second table in rebuttal PDF, indicate that the generated concept embeddings faithfully preserve factual relationships and even slightly improve upon capturing these relationships compared to the baseline.

What are the potential computational and scalability challenges of implementing the teacher-teacher framework in real-world clinical settings? How do you propose to mitigate these challenges?

Thank you for the question! One of the notable advantages of our proposed framework is its computational and memory efficiency, making it feasible to deploy with limited computational resources. To illustrate this, please refer to Figure1(a) in the rebuttal PDF, where we compare the training time of our model to both direct fine-tuning of LLMs and low-rank approximated fine-tuning of LLMs. The comparison shows that our model can be trained within two day on a RTX8000 48GB card.

How can regulatory mechanisms be incorporated into the framework for safety? Consider to add the assessment of how the framework could impact patient care, data security, and trust in AI systems in healthcare.

Thank you for this important question! Below, we assess how the framework impacts patient care, data security, and trust in AI systems, and discuss the incorporation of regulatory mechanisms for safety. This discussion will be included in the paper.

• Data Security: We utilize the publicly accessible MIMIC-IV dataset, which aligns with real-world clinical notes without privacy concerns. During inference, real medical notes aren't required; instead, a concept extractor tool like NILE can be used to extract key concepts for Teacher 2, preventing the use of sensitive patient information. Regulatory mechanisms can monitor and update the medical concepts in the knowledge graph, ensuring they reflect the most up-to-date factual knowledge. For instance, newly defined concepts like COVID-19 can be added, outdated concepts depreciated, and misleading concepts corrected.
• Patient Care: The framework's computational efficiency allows for the effective use of state-of-the-art LLMs, improving downstream tasks such as disease diagnosis and lab results analysis, thereby enhancing patient care.
• Trust in AI Systems: LINE can be retrained on local servers with various combinations of up-to-date LLMs, necessitating continual quality control. For example, our clinical concept similarity task (Section 3.3.1) can initially assess the fidelity of learned concept embeddings. During deployment, regular user feedback should be collected to further improve the system.

Dear Reviewer Yopj, I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors left rebuttals for your review. We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC

Thank you for your insightful comments! Below are our responses to your questions, addressed point-by-point:

It is unclear if LINE will work on combinations of other LLMs.

Thank you for raising this important question. To address your concern, we have extended our experiments to include additional combinations of LLMs. Specifically, we tested the alignment of LINE with the combination of GPT4 and CODER. The results, presented in the tables in the rebuttal PDF, indicate that LINE (GPT4+CODER) generally achieves better performance than GPT4 alone.

The concepts themselves can be negated, historical and hypothetical in context, but the proposed method does not seem to consider this.

Thank you for your insightful comment! We agree that concepts can vary significantly depending on their context, and our approach does take some parts of this into account. Specifically, we have implemented mechanisms to handle negated concepts and will further include a discussion on the possibility of modeling for concepts in historical or hypothetical contexts.

Handling Negated Concepts:

In our approach, we account for negated concepts by creating a separate dictionary of negative concepts derived from the original positive concepts. For instance, the positive concept "pneumothorax" is associated with a negated concept represented as "concept pneumothorax unobserved".

To ensure that negated concepts are accurately represented, we use a multihead graph attention module (as detailed in Section 2.2) to update the embeddings of positive concepts, which are then adapted for their corresponding negative concepts through a projection layer. We introduced a loss function specifically designed to maintain a distinct representation for negated concepts. This loss function ensures that the cosine similarity between the updated negative $\mathbf{c}_p$ concept and its positive counterpart $\mathbf{c}_p$ remains below a predefined threshold, $\delta$, calculated as:

$$-\log\frac{e^{\delta - \text{cos}(\mathbf{c}_p, \mathbf{c}_n)}}{1 + e^{\delta - \text{cos}(\mathbf{c}_p, \mathbf{c}_n)}},$$

where $0 < \delta \leq 0.5$. This mechanism is detailed in the Appendix and has been implemented in our two-step training process.

Historical and hypothetical contexts:

Additionally, we recognize the importance of considering whether concepts appear in historical or hypothetical contexts. Using NILE, we can effectively identify and extract contextually relevant keywords such as "previous", "since", or "in the near future". This capability enables us to incorporate context-specific modeling based on the temporal or hypothetical relevance of concepts. However, due to the complexity of real-world contexts, designing appropriate loss functions to account for these nuances requires further exploration and discussion.

Why was NILE selected? Have any other extractors been compared? Does the selection of extractors have a significant impact on results?

We chose NILE primarily for its speed and convenience. Compared to cTAKES and MedTagger—both of which are popular extractors—NILE is about 2000 times faster than cTAKES and 400 times faster than MedTagger, while delivering comparable performance.

Additionally, our framework includes a residual refinement step to recover important clinical concepts that might be missed due to issues like misspellings. This means that while the extractors need to perform reasonably well, they do not need to achieve extremely high accuracy. Consequently, the choice of extractor has a minimal impact on the overall results, as long as it effectively extracts crucial concepts.

Line 222, which contrastive loss function was used eventually?

Thanks for allowing us to clarify! We use the triplet loss.

Dear Reviewer Yayo, I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors left rebuttals for your review. We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC