Multimodal Representation Learning via Maximization of Local Mutual Information

We seek such image and text encoders and learn their representations by maximizing MI between the image representation and the text representation:

We employ MI as a statistical measure that captures dependency between images and text in the joint representation space. Maximizing MI between image and text representations is equivalent to maximizing the difference of the entropy and the conditional entropy of image representation given text: $I(z^{\text{I}},z^{\text{R}})=H(z^{\text{I}})-H(z^{\text{I}}|z^{\text{R}})$. This criterion encourages the model to learn feature representations where the information from one modality reduces the entropy of the other data modality, which is a better choice than solely minimizing the conditional entropy, where the image encoder could generate identical features for all data to achieve the conditional entropy minimum.

Estimating mutual information between high-dimensional continuous variables from finite data samples is challenging. We leverage the recent advances that employ neural network discriminators for MI estimation and maximization [2, 23, 18, 26]. The key idea is to construct a discriminator $f(\text{z}_{i}^{\text{I}},\text{z}_{j}^{\text{R}};\theta_{\text{D}})$, parameterized by $\theta_{\text{D}}$, that estimates the likelihood (or the likelihood ratio) of whether a sample pair $(\text{z}_{i}^{\text{I}},\text{z}_{j}^{\text{R}})$ is sampled from the joint distribution $p(z^{\text{I}},z^{\text{R}})$ or from the product of marginals $p(z^{\text{I}})p(z^{\text{R}})$. The discriminator is commonly found by maximizing the lower bound of the MI approximated by the likelihood ratio in Eq. (1) [2, 23].

We train the discriminator $f(\text{z}_{i}^{\text{I}},\text{z}_{j}^{\text{R}};\theta_{\text{D}})$ jointly with image and text encoders $z^{\text{I}}(x^{\text{I}};\theta_{\text{E}}^{\text{I}})$ and $z^{\text{R}}(x^{\text{R}};\theta_{\text{E}}^{\text{R}})$ via MI maximization:

where $\hat{I}(z^{\text{I}},z^{\text{R}};\theta_{\text{D}})$ is a lower bound on $I(z^{\text{I}},z^{\text{R}})$. We consider two MI lower bounds: Mutual Information Neural Estimation (MINE) [2] and Contrastive Predictive Coding (CPC) [23]. In our experiments, we empirically show that our method is not sensitive to the choice of the lower bound. MINE estimates the MI lower bound by approximating the log likelihood ratio in Eq. (1), using the Donsker-Varadhan (DV) variational formula of the KL divergence between the joint distribution and the product of the marginals, which yields the lower bound

CPC computes the MI lower bound by approximating the likelihood of an image-text feature pair being sampled from the joint distribution over the product of marginals, which leads to the objective function

Both methods sample from the matched image-text pairs and from shuffled pairs (to approximate the product of marginals), and train the discriminator to differentiate between these two types of sample pairs.

We propose to maximize MI between local features of images and sentence-level features from text. Given a sentence-level feature in the text, we estimate the MI values between all local image features and this sentence, select the image feature with the highest MI, and maximize the MI between that image feature and the sentence feature (Fig. 2). We train the image and text encoders, as well as the MI discriminator based on all the image-text data:

where $\text{z}_{j,(n)}^{\text{I}}$ is the $n$-th local feature in image $\text{x}_{j}^{\text{I}}$, and $\text{z}_{j,(m)}^{\text{R}}$ is the $m$-th sentence feature in text $\text{x}_{j}^{\text{R}}$. We use this one-way maximum, because in image captioning, every sentence was written to describe some finding in the corresponding image. In contrast, not every region in the image has a related sentence in the text that describes it.

We demonstrate our approach on the MIMIC-CXR dataset v2.0 [13] that includes around 250K frontal-view chest radiographs with their associated radiology reports. We evaluate our representation learning methods on two downstream classification tasks:

• •

  Pathology9. Detecting 9 pathologies from the chest radiographs against the labels that were extracted from the corresponding radiology reports using a radiology report labeler CheXpert [12, 15, 14]. Note that there are 14 findings available in the repository [14]. We only train and evaluate 9 for which there are more than around 100 images available in the test set.

• •

  EdemaSeverity. Assessing pulmonary edema severity from chest radiographs against the labels that were annotated by radiologists on the images [11, 17, 19, 28]. The severity level ranges from 0 to 3 with a higher score indicating higher risk.

The test sets provided in MIMIC-CXR with CheXpert labels [14] and with edema severity labels [17] are used to evaluate our methods. The patients that are in either of the two test sets are excluded from the model training. Table 1 summarizes the size of the (labeled) training data and test data.

Our goal is to learn representations that are useful for downstream classification tasks. Therefore, we use a fully supervised image model trained on the chest radiographs with available training labels as our benchmark. We compare two ways to use our image representations when re-training the image classifier: 1) freezing the image encoder; 2) fine-tuning the image encoder. In either case, the image encoder followed by a classifier is trained on the same training set that is used to train the fully supervised image model.

We compare our MI maximization approach on local features with the global MI maximization. We test both MINE [2] and CPC [23] as MI estimators. To summarize, we evaluate the variants of our model and training regimes as follows:

• •

  image-only-supervised: An image-only model trained on the training data provided in [14, 17].

• •

  global-mi-mine, global-mi-cpc: Representation learning on the chest radiographs and the radiology reports using global MI maximization.

  • –

    encoder-frozen, encoder-tuned: Once representation learning is completed, the image encoder followed by a classifier is re-trained on the labeled training image data, with the encoder frozen or fine-tuned.

• •

  local-mi-mine, local-mi-cpc: Representation learning using local MI maximization in Eq. (5).

  • –

    encoder-frozen, encoder-tuned: The resulting image encoder followed by a classifier is re-trained, with the encoder frozen or fine-tuned.

At the image model training or re-training time, all variants are trained on the same training sets. Note that the local-mi approach makes use of lower level image features. To make the encoder-frozen experiments comparable between local-mi and global-mi, we only freeze the same lower level feature extractor in both encoders.

Chest radiographs are downsampled to 256$\times$256. We use a 5-block resnet [9] as the image encoder in the local MI approach and the image feature representation $z^{\text{I}}$ is 16$\times$512 (4$\times$4$\times$512) feature vectors. We use a 6-block resnet as the image encoder for the global MI maximization, where the image representation $z^{\text{I}}$ from this encoder is a 768-dimensional feature vector. We use the clinical BERT model [1] as the text encoder for both report-level and sentence-level feature extraction. The [CLS] token is used as the text feature $z^{\text{R}}$, which is a 768-dimensional vector. The MI discriminator for both MINE and CPC is a 1280$\rightarrow$1024$\rightarrow$512$\rightarrow$1 multilayer perceptron to estimate local MI and a 1536$\rightarrow$1024$\rightarrow$512$\rightarrow$1 multilayer perceptron to estimate global MI. The image feature and the text feature are concatenated to construct the input for the discriminator for MI estimation. The image models in all training variants at the image training or re-training time have the same architecture (6-block resnet followed by a fully connected layer).

The AdamW [30] optimizer is employed for the BERT encoder and the Adam [16] optimizer is used for the other parts of the model. The initial learning rate is $5{\cdot}10^{-4}$. The representation learning phase is trained for 5 epochs and the image model re-training phase is trained for 50 epochs. The fully supervised image model is trained for 100 epochs. Data augmentation including random rotation, translation, and cropping is performed on the images during training.

In Table 2 and Table 3, we present the area under the receiver operating characteristic curve (AUC) statistics for the variants of our algorithms on the EdemaSeverity classification task and the Pathology9 binary classification tasks. For most classification tasks, the local MI approach with encoder tuning performs the best and significantly improves the performance over solely supervised learning on labeled images. The local MI approach brings in noteworthy improvement compared to global MI. Both CPC and MINE perform similar in most tasks. Remarkably, the classification results from the frozen encoders approach the fully supervised learning results in many tasks, suggesting that the unsupervised learning captures useful features for image classification tasks even before supervision is provided.

The local MI offers substantial improvement in performance when the features are fine-tuned with the downstream model, while its performance is comparable with global MI if the features are frozen for the subsequent classification. In our experiments, training jointly with the downstream classifier (fine-tuning) typically improves performance of all tasks, with greater benefits for local MI. This suggests that local MI yields more flexible representations that adjust better for the downstream task. Our results are also supported by the analysis in Section 3 that shows certain structural assumptions lead to the local MI approach, which is easier to discover useful representations due to its lower dimension and more training samples.