Fine-Grained_Features_Alignment_via_Constrastive_Learning

code slides paper

Chest Image

Generated by GPT-4

Abstract

Leveraging multimodal data, the Vision-Language Model(VLM) demonstrated the impressive capability of bridging the knowledge among multi-modalities. VLMs like CLIP, Flamingo, and DALL-E, which are trained based on the tremendous amount of data and computational resources show good performance in many different downstream tasks due to the good generalizability. However, like a double-edged sword, the generalizability of pre-trained VLMs limits their performance in the customized setting. In this project, I try to leverage the prior knowledge in pre-trained VLMs and customize the embedding generation for my general classification task. Using the simple contrastive learning method proposed in the report, a robust generalist classifier is available with the deficiency of training data, which is a ubiquitous context in the biomedical setting.

A feature representation task.
Currently using the Classification task as a utility task, aiming to build a general disease detecter via chest X-ray images.
This project tries to explore the data cooperation between multimodal data – text and image, and then improve AI downstream tasks.

Problem Description:

To alleviate the workload of radiologists, researchers develop algorithms that can automatically classify X-ray images into different classes (corresponding to the existence of different diseases discovered in X-ray images).
even though the current SOTA specialist model (customized model for one certain disease) gets very good performance, a generalist model (capable of handling multiple diseases simultaneously) is still weak because of the difficulty of this task and the deficiency of training data.
In this context, what I want to do is establish an algorithm which capable of detecting multiple diseases from X-ray images.

Challenges:

To solve this problem, I need to confront 2 main challenges

  1. data scarcity – clinical data is always scarce due to private code. The deficiency of training data limits the performance of AI models.
  2. generalizability – Training generalist models that are robust in the detection of multiple diseases is challenging.

Potential Solutions

  1. leveraging the knowledge in the powerful pre-trained LLMs to guide the feature extraction during my image classification task; using text data to alleviate the data scarcity issue.
  2. using a knowledge graph to inject the prior knowledge of the correlation among different diseases, facilitating the algorithm to grasp the latent relationship among diseases.

Model

methodology_blank Image

Using contrastive learning to align the diseases’ representation between text and image, leveraging the power of LLMs to guide the feature extraction in the image branch.

Pathology correlation:

In reality, the diseases diagnosed from X-ray images are supposed to be correlated with each other to some extent. Therefore, the totally orthogonalizing may not make sense.
To inject the prior knowledge of this correlation, I use the graph to represent the hierarchical relationship between my 14 labels and hope this prior knowledge can guide the model learning.

graph relationship

Hierarchical relation tree of 14 labels

AUC comparison among 14 labels

Chest Image

AUC comparison among 14 labels(config: BiomedCLIP+grpah+Orth+NoContrastive)

AUC improves by 3.38%

AUC

AUC comparision among 14 labels(config: BiomedCLIP+grpah+NoOrth+NoContrastive)

AUC improves by 6.29%. By injecting customized knowledge, the model performance in fracture improves obviously, which is the most challenging one for discrimination [3].

Reference

[1] Towards long-tailed, multi-label disease classification from chest X-ray: Overview of the CXR-LT challenge
[2] Xplainer: From X-Ray Observations to Explainable Zero-Shot Diagnosis
[3] CheXclusion: Fairness gaps in deep chest X-ray classifiers
[4] A Simple General Approach to Balance Task Difficulty in Multi-Task Learning
[5] Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics
[6] Attentional Mixtures of Soft Prompt Tuning for Parameter-efficient Multi-task Knowledge Sharing
[7] A Pilot Study of Query-Free Adversarial Attack against Stable Diffusion