MLIP: Enhancing Medical Visual Representation with Divergence Encoder and Knowledge-guided Contrastive Learning

Medical reports are pivotal in unsupervised medical visual representation learning, with two primary methods dominating the field. The first method involves extracting disease labels from radiology reports using manually designed rules [34, 33], followed by pre-training image models for downstream tasks. However, defining the rules requires considerable human effort and domain expertise. On the other hand, the second method adopts image-text contrastive learning methods to integrate text and vision in an unsupervised manner [19, 61, 31, 32, 53]. These methods have been shown remarkable performance in diverse downstream tasks, including medical object detection [4], image classification [32, 61], and semantic segmentation [61]. However, they have not effectively explored visual representations at different granularities and rely on partial semantic information.

To address these limitations, MGCA [53] proposes to leverage multiple visual features at different granularities during the pre-training phase, enhancing the performance of models in downstream tasks. However, it overlooks the challenging sample issue in medical radiology. In this work, we propose a divergence encoder that manually updates its parameters based on the similarity between the output features and those of a common encoder. By increasing divergence between the two encoders, we enhance feature diversity and train the model to discriminate among similar samples effectively.

To enhance the model’s knowledge and understanding ability by leveraging a broader background, numerous vision-and-language pre-training methods have been devised to incorporate domain-specific knowledge. These methods can be categorized into four distinct knowledge-guided schemes: embedding combination [63], data structure compatibility [25, 39], knowledge supervision [55], and neural-symbolic methods [2]. For instance, ERNIE-ViL [59] introduces a vision and language alignment technique by utilizing a scene graph extracted from the input text. Similarly, KB-VLP [11] incorporates object tags from images and knowledge graph embeddings from texts to enhance the acquisition of knowledge-aware representations. ARL [15] utilizes expert knowledge as an intermediate medium to align images and reports. Additionally, a recent study [42] proposes the automatic generation of visual and textual prompts, injecting expert medical knowledge into the prompt for pre-training.

In contrast to existing works, we propose an alignment method that leverages domain-specific knowledge as an intermediate mediator for aligning texts and images, along with a knowledge-guided prototype clustering contrastive learning. This approach integrates expert domain knowledge derived from the Unified Medical Language System (UMLS) [6]. By incorporating UMLS knowledge into both vision and language modalities, our approach leverages knowledge as a medium to achieve improved alignment between images and text, facilitating more effective clustering of image-text pairs. Importantly, our method effectively mitigates the influence of disease-level false negatives without relying on object detectors or scene graph parsers.

Recently, it has been demonstrated in [53, 32] that learning medical visual representation learning without labels can achieve competitive performance. In this study, we follow the setting in [53], given a training set of $N$ medical image-report pairs $\mathcal{D}=\{(x_{i},y_{i})\}_{i=1,...,N}$, we use an image encoder $f_{v}$ and a text encoder $f_{t}$ encode $\mathcal{D}$ to a global feature set $\mathcal{E}_{il}\!=\!\{(v_{i},t_{i})\!\!\!\mid\!\!\!v_{i}=f_{v}(x_{i}),t_{i}=f_{t}(y_{i})\}_{i=1,...,N}$, and a local feature set $\mathcal{E}_{tl}\!=\!\{(\mathcal{P}_{i},\mathcal{S}_{i})\}_{i=1,...,N}$, where $\mathcal{S}_{i}=\{s^{1}_{i},s^{2}_{i},...,s^{V}_{i}\}\in\mathbb{R}^{V\times d}$ and $\mathcal{P}_{i}=\{p^{1}_{i},p^{2}_{i},...,p^{M^{2}}_{i}\}\in\mathbb{R}^{M^{2}\times d}$. $V$ denotes the length of the sentence and $M^{2}$ denotes the number of image patches.

Furthermore, we incorporate expert knowledge into our model by constructing an extracted knowledge graph, as described in [15]. This knowledge graph is denoted as $\mathcal{G}=\{(he_{i},re_{i},ta_{i})\}_{i=1}^{N_{\mathcal{G}}}$, where $N_{\mathcal{G}}$ represents the number of graph triples, and $he_{i}$, $re_{i}$, and $ta_{i}$ correspond to the head entity, relation, and tail entity, respectively. The inclusion of this expert knowledge enhances the model’s understanding and reasoning capabilities, enabling more informed alignment and representation learning.

To pull correct samples closer and push random samples apart in the latent space, we follow [49, 30], present a comprehensive discussion on global image-text contrastive learning by maximizing mutual information $\mathcal{I}(X,Y)$ between the vision element $X$ and the language component $Y$:

$${\mathcal{I}(X,Y)\!=\!{\sum\limits}_{y\in Y}{\sum\limits}_{x\in X}P(x,y)\log\frac{P(x|y)}{P(x)}}.$$

(1)

Eq.1 suggests that the fraction $\frac{P(x|y)}{P(x)}$ collapses to zero when $x$ and $y$ are incompatible with each other. Therefore, we hypothesize that $\frac{P(x|y)}{P(x)}$ is proportional to the similarity between $x$ and $y$. Further, the maximization of mutual information corresponds to the maximization of the similarity $\textrm{sim}(x,y)$ between $x$ and $y$, which can be represented as:

$$\mathcal{I}(v,t)\propto\mathcal{I}(X,Y)\propto\textrm{sim}(x,y)\propto\textrm{sim}(v,t).$$

(2)

Specifically, inspired by [12], we firstly utilize two projection layers $h_{v}$ and $h_{t}$ to map $v_{i}$ and $t_{i}$ into a normalized shared feature space, yielding $v_{i}^{\ast}\in\mathbb{R}^{d}$ and $t_{i}^{\ast}\in\mathbb{R}^{d}$, respectively. Then, we apply the dot product to model the similarity between $v_{i}^{\ast}$ and $t_{i}^{\ast}$. To obtain more effective features, we perform Self-Attention [51] and LayerNorm [3] on features:

	$\displaystyle v_{i}^{\ast}=\text{LN}(\text{SA}\{h_{v}(v_{i})\});$		(3a)
	$\displaystyle t_{i}^{\ast}=\text{LN}(\text{SA}\{h_{t}(t_{i})\}),$		(3b)
	$\displaystyle\textrm{sim}\{v_{i}^{\ast},t_{i}^{\ast}\}=v_{i}^{\ast}{t_{i}}^{T},$		(3c)

where SA denotes Self-Attention module and LN denotes LayerNorm module.

We optimize this process via image-text contrastive loss based on InfoNCE loss [50], which are designed to maximize the mutual information between the correct image-text pairs in the latent space:

	$\displaystyle\mathcal{L}^{il}_{v2t}(v_{i},t_{i})=-\log(\frac{\phi_{il}(v_{i},t_{i})}{{\sum\limits}_{k=1}^{B}\phi_{il}(v_{i},t_{k})}),$		(4a)
	$\displaystyle\mathcal{L}^{il}_{t2v}(v_{i},t_{i})=-\log(\frac{\phi_{il}(v_{i},t_{i})}{{\sum\limits}_{k=1}^{B}\phi_{il}(v_{k},t_{i})}),$		(4b)

where $\phi_{il}(v_{i},t_{i})=\exp(\frac{\textrm{sim}(v_{i}^{\ast},t_{i}^{\ast})}{\tau_{1}})$, $B$ is the batch size and $\tau_{1}$ is the global temperature hyper-parameter.

Directly optimizing $\mathcal{I}(v,t)$ is a challenging task. As an alternative, [50] has proposed an alternative method to optimize the lower bound of mutual information:

$${\mathcal{I}(v,t)}\geq\log N^{{}^{\prime}}-\mathcal{L}^{\text{NCE}}(v,t),$$

(5)

where $N^{{}^{\prime}}$ is the number of negative samples. In Eq.5, minimizing $\mathcal{L}^{\text{NCE}}(v,t)$ is equivalent to maximizing the lower bound of the mutual information between the medical image and the corresponding report.

To increase the number of samples and enhance the feature diversity, we perform a divergence encoder to achieve data augmentation and extend the gap between samples. We define image divergence encoder $o_{v}$ and text divergence encoder $o_{t}$, initialized by $f_{v}$ and $f_{t}$, respectively. Then we obtain features incrementally differentiated from $v_{i}$ and $t_{i}$:

$$v_{i}^{aug}=o_{v}(x^{rt}_{i});t_{i}^{aug}=o_{t}(y_{i}),$$

(6)

where $x^{rt}_{i}$ denotes randomly transformed images. We manually update divergence encoders’ parameters instead of relying on backpropagation:

	$\displaystyle\theta_{o_{t}}\!\!=\!\!s_{t}\!*\!\theta_{f_{t}}+(1-s_{t})\!*\!\theta_{o_{t}},$		(7a)
	$\displaystyle\theta_{o_{v}}\!\!=\!\!s_{v}\!*\!\theta_{f_{v}}+(1-s_{v})\!*\!\theta_{o_{v}},$		(7b)

where $s_{t}=\text{cosine}(t_{i},t_{i}^{aug})$ and $s_{v}=\text{cosine}(v_{i},v_{i}^{aug})$, and $\theta_{o_{t}},\theta_{o_{v}},\theta_{f_{t}},\theta_{f_{v}}$ are the parameters of $o_{t},o_{v},f_{t},f_{v}$, respectively. In this way, as the $s_{v}(s_{t})$ increases, we aim to retain fewer parameters from $o_{v}(o_{t})$ and incorporate more parameters from $f_{v}(f_{t})$, in order to generate more diverse features. Then we use Eq.4a, 4b to compute $\mathcal{L}_{v2a}^{il}$ and $\mathcal{L}_{avt}^{il}$.

We compute the objective $\mathcal{L}_{ita}$ as the average of the four loss values:

$$\begin{split}\mathcal{L}_{ita}=\frac{1}{2N}{\sum\limits}^{N}_{i=1}(\mathcal{L}_{v2t}^{il}(v_{i},t_{i})+\mathcal{L}_{t2v}^{il}(v_{i},t_{i}))\\ +\frac{\lambda_{0}}{2N}{\sum\limits}^{N}_{i=1}(\mathcal{L}_{v2a}^{il}(v_{i},t_{i}^{aug})+\mathcal{L}_{avt}^{il}(v_{i}^{aug},t_{i})),\end{split}$$

(8)

where $N$ is the total number of samples and $\lambda_{0}$ denotes the weight for augmented image-text contrastive learning.

In medical images, pathologies are often visually subtle and occupy a small fraction of the overall image, while only a few disease-related tags in the associated report accurately depict the critical medical condition. Given this observation, we employ a local image-text contrastive learning method to maximize the mutual information between local features and achieve cross-modal alignment between images and texts, inspired by [53, 17].

However, traditional token-patch alignment contrastive learning is utilizing the local features of the image and text to compute the attention matrix, and then perform contrastive learning after aligning the images and texts. Since medical radiology is highly professional and there is a certain bias between different datasets, we regard professional knowledge from the UMLS [6] as a medium between vision and language. To achieve more accurate token-patch alignment, we align the knowledge with radiographs and reports.

Similar to global feature, we apply Self-Attention and LayerNorm module on every features:

$$p_{i}=\text{LN}(\text{SA}\{h_{v}(p_{i})\});s_{i}=\text{LN}(\text{SA}\{h_{t}(s_{i})\}).$$

(9)

We apply the knowledge representation learning algorithm TransE [7] to the knowledge graph $\mathcal{G}$ to obtain entity embeddings. Subsequently, we utilize the Graph Attention Network [52] to capture local information in the graph neighborhood for each node. This allows us to obtain knowledge representations, denoted as $\{e_{i}\}_{i=1}^{N_{e}}\in\mathbb{R}^{N_{e}\times d_{e}}$, where $d_{e}$ represents the feature dimension and $N_{e}$ denotes the number of entity.

We adopt cross-modal attention mechanism [14, 41] to explore the matching between knowledge and image:

	$\displaystyle attn_{j,k}$	${}^{vk}=\text{softmax}(\frac{(Qp_{i}^{j})^{T}(Ke_{i}^{k})}{\sqrt{d}}),$		(10a)
	$\displaystyle{zv}_{i}^{j}$	$\displaystyle={\sum\limits}_{k=1}^{N}attn_{j,k}^{vk}(Ve_{i}^{k}),$		(10b)

where $Q,K,V\in\mathbb{R}^{d\times d}$ are trainable matrices. $e_{i}$ is mapped to $\mathbb{R}^{M^{2}\times d}$. ${zv}_{i}^{j}$ is cross-modal knowledge embedding corresponding to $p_{i}^{j}$.

Lying in the purpose of maximizing the lower bound of mutual information, we leverage InfoNCE loss [50] to pull $p_{i}^{j}$ and ${zv}_{i}^{j}$ closer and push $p_{i}^{j}$ and other cross-modal knowledge embeddings apart. However, given that irrelevant information only occupies a vast majority of medical images, we employ $w_{i}^{j}$ to balance the weights of different patches. The loss $\mathcal{L}_{v2t}^{tl}$ is designed symmetrically as:

$$\begin{split}\mathcal{L}_{v2t}^{tl}=-\frac{1}{2NM^{2}}&\sum\limits_{i=1}^{N}\sum\limits_{j=1}^{M^{2}}w_{i}^{j}(\log\frac{\phi_{tl}({p_{i}^{j}},{{zv}_{i}^{j}})}{{\sum\limits}_{k=1}^{M^{2}}\phi_{tl}({p_{i}^{j}},{{zv}_{i}^{k}})}\\ &+\log\frac{\phi_{tl}({{zv}_{i}^{j}},{p_{i}^{j}})}{{\sum\limits}_{k=1}^{M^{2}}\phi_{tl}({{zv}_{i}^{k}},{p_{i}^{j}})}),\end{split}$$

(11)

where $\phi_{tl}(p_{i}^{j},{zv}_{i}^{j})=\exp(\frac{\textrm{sim}(p_{i}^{j},{zv}_{i}^{j})}{\tau_{2}})$, $\tau_{2}$ is the local temperature hyper-parameter. To establish the correlation between the $j$-th visual patch and the [CLS] token, we assign the weight $w_{i}^{j}$ using the last-layer attention mechanism averaged across multiple heads.

Similarly, for the $j$-th text token, we calculate corresponding cross-modal knowledge embedding ${zt}_{i}^{j}$ and construct local contrastive loss $\mathcal{L}_{t2v}^{tl}$ to maximize the lower bound of mutual information between $s_{i}^{j}$ and ${zt}_{i}^{j}$. The objective $\mathcal{L}_{tl}$ can be defined as the average of these two losses:

$$\mathcal{L}_{tl}=\frac{1}{2}(\mathcal{L}_{v2t}^{tl}+\mathcal{L}_{t2v}^{tl}).$$

(12)

For a given radiograph-report pair, traditional contrastive learning approaches treat other radiograph-report pairs within the same batch as negative samples. However, in the context of category-level analysis, samples that belong to different batches but exhibit highly similar semantics should be considered positive samples. In our approach, we aim to select representative samples in each iteration, emphasizing their ability to capture meaningful disease-related information. In the medical domain, expert knowledge plays a crucial role in representation learning. We purpose to bridge the gap between the vast knowledge learned from general visual and textual data and its effective application in the intricate realm of medical radiology. Therefore, we incorporate expert knowledge from UMLS [6] as an auxiliary signal. Drawing inspiration from [8, 42], we propose a knowledge-guided clustering-based approach to improve the efficacy of learned representations. We bring together highly similar samples with high-level semantics, even when originating from different batches, and ensure their proximity in the feature space, rather than increasing their distance from one another.

Motivated by [38], we realize to filter out irrelevant information and explore more fine-grained relations between images and text. To achieve this, we employ a mechanism that identifies the most relevant topic in a given context. Specifically, we utilize $v_{i}^{\ast}$ to find the most relevant topic in $t_{i}^{\ast}$, resulting in $\dot{t_{i}}$. Then, we use $\dot{t_{i}}$ to find the relevant topic in $v_{i}^{\ast}$, leading to $\dot{v_{i}}$. The process is mathematically defined as follows:

$$\dot{t_{i}}\!=\!\text{LN}(\text{softmax}(\frac{{v_{i}^{\ast}}^{T}t_{i}^{\ast}}{\sqrt{d}})t_{i}^{\ast});\dot{v_{i}}\!=\!\text{LN}(\text{softmax}(\frac{{v_{i}^{\ast}}^{T}v_{i}^{\ast}}{\sqrt{d}})\dot{t_{i}}),$$

(13)

then we utilize tucker fusion [5] to seamlessly integrate visual and textual features, further fuse with knowledge representations:

$$\mathcal{Q}=(({\mathcal{T}_{c}}\times_{1}{\dot{v_{i}}})\times_{2}{\dot{t_{i}}})\times_{3}\mathcal{W}_{o},$$

(14)

where $\mathcal{W}_{o}$ represents a mapping matrix which is trainable and maps fused features to a certain dimensional space, and $\mathcal{T}_{c}$ denotes the core tensor.

To further integrate knowledge with modality-specific features, we employ a linear mapping layer to project the knowledge representation $e_{i}$ into a $d$-dimensional space and incorporate it with fused features using cross-modal attention, thereby facilitating the fusion of information across modalities:

$$vkt_{i}=\text{SA}(\text{softmax}(\frac{{\mathcal{Q}}^{T}e_{i}}{\tau_{3}})\cdot e_{i}),$$

(15)

where $\tau_{3}$ is the temperature hyper-parameter we set to scale the attention.

For image-text features pair $(\dot{v_{i}},\dot{t_{i}})$ and knowledge-fused features, we apply the iterative Sinkhorn-Knopp clustering algorithm [18] to generate a cluster assignment code $u^{vkt,i}\in\mathbb{R}^{C}$, by assigning $vkt_{i}$ to $C$ clusters separately. To facilitate this, we introduce a set $\mathcal{J}={j_{1},...,j_{C}}$ that contains $C$ trainable cross-modal prototypes, where each prototype $j_{c}\in\mathbb{R}^{d}$. We calculate the visual softmax probability $p^{v,i}$ by computing the cosine similarity between the visual feature vector $\dot{v_{i}}$ and all cross-modal prototypes in $\mathcal{J}$. Similarly, the textual softmax probability $p^{t,i}$ is obtained by measuring the cosine similarity between the textual feature vector $\dot{t_{i}}$ and all cross-modal prototypes in $\mathcal{J}$:

$$p^{v,i}_{c}=\frac{\exp({\dot{v_{i}}}^{T}j_{c}/\tau_{4})}{{\sum\limits}_{l}\exp({\dot{v_{i}}}^{T}j_{l}/\tau_{4})};p^{t,i}_{c}=\frac{\exp({\dot{t_{i}}}^{T}j_{c}/\tau_{4})}{{\sum\limits}_{l}\exp({\dot{t_{i}}}^{T}j_{l}/\tau_{4})},$$

(16)

where $\tau_{4}$ is a category-level temperature hyper-parameter and $c$ denotes the $c$-th element of the vector.

To enable knowledge-guided category-level contrastive learning, we employ $u^{vkt,i}$ as the pseudo-label for training $\dot{t_{i}}$ and $\dot{v_{i}}$. This allows the three features to interact in the latent space and guide the shifting of positive and negative samples with the assistance of domain-specific knowledge. The objective loss $\mathcal{L}_{cl}$ is formulated as follows:

$$\mathcal{L}_{cl}\!=\!\frac{1}{2N}\sum\limits_{i=1}^{N}\sum\limits^{C}_{c=1}u^{vkt,i}_{c}\log p^{v,i}_{c}\!+u^{vkt,i}_{c}\log p^{t,i}_{c}.$$

(17)

In order to identify the alignment between radiographs and their corresponding reports, we propose two pretext tasks aimed at bridging the semantic divide between visual and linguistic information within the feature space: 1) computing relevance scores between image patch and contextualized sentence to evaluate the degree of correlation between the image and text elements; 2) randomly substituting medical reports corresponding to the image with a predetermined probability, improving the discriminative ability on mismatched samples of the model.

We assume that the text features $t$ and image features $v$ have been normalized. Therefore, we construct the similarity between the two modalities as a relevance score:

$$r(v,t)=v^{T}\cdot t,$$

(18)

subsequently, we randomly select another image $v^{{}^{\prime}}$ and obtain its corresponding relevance score $r(v^{{}^{\prime}},t)$. To ensure that the difference between $r(v,t)$ and $r(v^{{}^{\prime}},t)$ is greater than a pre-specified margin $\mathcal{G}$, we utilize the hinge loss function to compute image-text match loss:

$$\mathcal{L}_{itm}=\max(0,\mathcal{G}-r(v,t)+r(v^{{}^{\prime}},t)).$$

(19)

Similarly, we propose a text swapping task, which involves randomly replacing text with a predefined probability $\gamma$. We employ a bidirectional similarity Hinge loss to penalize the model for insufficient discriminative ability. This task aims to enhance the model’s ability to distinguish between different reports. We employ a cross-modal attention mechanism to fuse the text and image modalities, then compute the relevance score by performing a weighted summation of the similarity between the fused representation and the original text-image pair. Our objective is to ensure that this score exceeds the score obtained after replacing the text by a margin $\mathcal{G}^{{}^{\prime}}$:

	$\displaystyle r$	${}_{ts}(v,t)=v^{T}\cdot t+\alpha\cdot\text{CA}(v,t)^{T}\cdot\text{CA}(t,v),$		(20a)
	$\displaystyle r_{ts}$	$\displaystyle(v,t^{{}^{\prime}})=v^{T}\cdot t+\alpha\cdot\text{CA}(v,t^{{}^{\prime}})^{T}\cdot\text{CA}(t^{{}^{\prime}},v),$		(20b)
		$\displaystyle\mathcal{L}_{ts}=\max(0,\mathcal{G}{{}^{\prime}}-r_{ts}(v,t)+r_{ts}(v^{{}^{\prime}},t)),$		(20c)

where $\text{CA(x,y)}=\text{softmax}(\frac{x^{T}\cdot y}{\sqrt{d}})\cdot y$. Through these two designed proxy tasks, we compute the image-text matching loss $\mathcal{L}_{itm}$ and the text swapping loss $\mathcal{L}_{ts}$. These losses quantify the model’s ability to accurately match radiographs to their appropriate reports, thereby providing a measurable objective for the optimization process.

Our training approach involves joint optimization of the five losses, aiming to promote the acquisition of effective and generalizable medical image representations by the network. The overall training objective can be expressed as follows:

$$\mathcal{L}\!=\!\lambda_{1}\mathcal{L}_{il}+\lambda_{2}\mathcal{L}_{tl}+\lambda_{3}\mathcal{L}_{cl}+\lambda_{4}\mathcal{L}_{itm}+\lambda_{5}\mathcal{L}_{ts},$$

(21)

where $\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, $\lambda_{4}$ and $\lambda_{5}$ are hyper-parameters employed to balance the weights associated with each respective loss.

Our MLIP framework is initially pre-trained on the MIMIC-CXR 2.0.0 dataset [34], with data consistency ensured through preprocessing methods from [61]. Lateral views are excluded from the dataset as downstream datasets only include frontal-view chest images. Inspired by [53], we extract impression and finding sections from free-text reports, providing comprehensive descriptions of medical diseases. We filter out empty or short reports, resulting in approximately 217,000 image-text pairs. Details about our implementation can be found in the supplement LABEL:sec:rationale.

Table 4 presents ablation results on semantic segmentation for both RSNA and SIIM datasets. We observe that leveraging knowledge as an intermediate medium for aligning image-text pairs in contrastive learning substantially enhances the model’s performance. Moreover, category-level contrastive learning aids in mitigating false negatives, thereby improving the model’s generalization. Global contrastive learning acts as a performance lower bound, complementing local and category-level approaches and yielding promising outcomes. Table 5 presents the ablation results of our main contributions on the object detection task. Our proposed divergence encoder enhances feature diversity and enables the model to better adapt to challenging samples. With the assistance of expert knowledge, the alignment between medical images and medical reports becomes more efficient. Lastly, the proxy tasks designed in our approach strengthen the model’s ability to discriminate negative samples.

To further understand the inner workings of MLIP, we present learned local correspondences between radiographs and medical reports in the form of heatmaps and showcase the performance of MLIP on downstream tasks (semantic segmentation and object detection) in the supplementary LABEL:sec:Visualization. The visual evidence supports that MLIP excels in fine-grained feature extraction, boosting accuracy.