Joint Modeling of Chest Radiographs and Radiology Reports for Pulmonary Edema Assessment

Abstract

We propose and demonstrate a novel machine learning algorithm that assesses pulmonary edema severity from chest radiographs. While large publicly available datasets of chest radiographs and free-text radiology reports exist, only limited numerical edema severity labels can be extracted from radiology reports. This is a significant challenge in learning such models for image classification. To take advantage of the rich information present in the radiology reports, we develop a neural network model that is trained on both images and free-text to assess pulmonary edema severity from chest radiographs at inference time. Our experimental results suggest that the joint image-text representation learning improves the performance of pulmonary edema assessment compared to a supervised model trained on images only. We also show the use of the text for explaining the image classification by the joint model. To the best of our knowledge, our approach is the first to leverage free-text radiology reports for improving the image model performance in this application. Our code is available at: https://github.com/RayRuizhiLiao/joint_chestxray.

Introduction

We present a novel approach to training machine learning models for assessing pulmonary edema severity from chest radiographs by jointly learning representations from the images (chest radiographs) and their associated radiology reports. Pulmonary edema is the most common reason patients with acute congestive heart failure (CHF) seek care in hospitals [1,9,15]. The treatment success in acute CHF cases depends crucially on effective management of patient fluid status, which in turn requires pulmonary edema quantification, rather than detecting its mere absence or presence.

Chest radiographs are commonly acquired to assess pulmonary edema in routine clinical practice. Radiology reports capture radiologists’ impressions of the edema severity in the form of unstructured text. While the chest radiographs possess ground-truth information about the disease, they are often time intensive (and therefore expensive) for manual labeling. Therefore, labels extracted from reports are used as a proxy for ground-truth image labels. Only limited numerical edema severity labels can be extracted from the reports, which limits the amount of labeled image data we can learn from. This presents a significant challenge for learning accurate image-based models for edema assessment. To improve the performance of the image-based model and allow leveraging larger amount of training data, we make use of free-text reports to include rich information about radiographic findings and reasoning of pathology assessment. We incorporate free-text information associated with the images by including them during our training process.

We propose a neural network model that jointly learns from images and free-text to quantify pulmonary edema severity from images (chest radiographs). At training time, the model learns from a large number of chest radiographs and their associated radiology reports, with a limited number of numerical edema severity labels. At inference time, the model computes edema severity given the input image. While the model can also make predictions from reports, our main interest is to leverage free-text information during training to improve the accuracy of image-based inference. Compared to prior work in the image-text domain that fuses image and text features [5], our goal is to decouple the two modalities during inference to construct an accurate image-based model.

Prior work in assessing pulmonary edema severity from chest radiographs has focused on using image data only [14,18]. To the best of our knowledge, ours is the first method to leverage the free-text radiology reports for improving the image model performance in this application. Our experimental results demonstrate that the joint representation learning framework improves the accuracy of edema severity estimates over a purely image-based model on a fully labeled subset of the data (supervised). The joint learning framework uses a ranking-based criterion [7,12], allowing for training the model on a larger dataset of unlabeled images and reports. This semi-supervised modification demonstrates a further improvement in accuracy. Additional advantages of our joint learning framework are 1) allowing for the image and text models to be decoupled at inference time, and 2) providing textual explanations for image classification in the form of saliency highlights in the radiology reports.

Related Work.

The ability of neural networks to learn effective feature representations from images and text has catalyzed the recent surge of interest in joint image-text modeling. In supervised learning, tasks such as image captioning have leveraged a recurrent visual attention mechanism using recurrent neural networks (RNNs) to improve captioning performance [28]. The TieNet used this attention-based text embedding framework for pathology detection from chest radiographs [26], which was further improved by introducing a global topic vector and transfer learning [29]. A similar image-text embedding setup has been employed for chest radiograph (image) annotations [20]. In unsupervised learning, training a joint global embedding space for visual object discovery has recently been shown to capture relevant structure [11]. All of these models used RNNs for encoding text features. More recently, transformers such as the BERT model [8] have shown the ability to capture richer contextualized word representations using self-attention and have advanced the state-of-the-art in nearly every language processing task compared to variants of RNNs. Our setup, while similar to [26] and [11], uses a series of residual blocks [13] to encode the image representation and uses the BERT model to encode the text representation. We use the radiology reports during training only, to improve the image-based model’s performance. This is in contrast to visual question answering [2, 3, 19], where inference is performed on an image-text pair, and image/video captioning [16, 22, 24, 28], where the model generates text from the input image.

2. Data

For training and evaluating our model, we use the MIMIC-CXR dataset v2.0 [17], consisting of 377,110 chest radiographs associated with 227,835 radiology reports. The data was collected in routine clinical practice, and each report is associated with one or more images. We limited our study to 247,425 frontal-view radiographs.

Regex Labeling.

We extracted pulmonary edema severity labels from the associated radiology reports using regular expressions (regex) with negation detection [6]. The keywords of each severity level (none=0, vascular congestion=1, interstitial edema=2, and alveolar edema=3) are summarized in the supplementary materials. In order to limit confounding keywords from other disease processes, we limited the label extraction to patients with congestive heart failure (CHF) based on their ED ICD-9 diagnosis code in the MIMIC dataset [10]. Cohort selection by diagnosis code for CHF was previously validated by manual chart review. This resulted in 16,108 radiology reports. Regex labeling yielded 6,710 labeled reports associated with 6,743 frontal-view images¹. Hence, our dataset includes 247,425 image-text pairs, 6,743 of which are of CHF patients with edema severity labels. Note that some reports are associated with more than one image, so one report may appear in more than one image-text pair.

3. Methods

Let x^I be a 2D chest radiograph, x^R be the free-text in a radiology report, and y ∈ {0, 1, 2, 3} be the corresponding edema severity label. Our dataset includes a set of N image-text pairs $\text{X}={\left\{{\text{x}}_j\right\}}_{j=1}^N$, where ${\text{x}}_j=({\text{x}}_j^{\text{I}},{\text{x}}_j^{\text{R}})$. The first N_l image-text pairs are annotated with severity labels $Y={\left\{y_j\right\}}_{j=1}^N$. Here we train a joint model that constructs an image-text embedding space, where an image encoder and a text encoder are used to extract image features and text features separately (Fig. 1). Two classifiers are trained to classify the severity labels independently from the image features and from the text features. This setup enables us to decouple the image classification and the text classification at inference time. Learning the two representations jointly at training time improves the performance of the image model.

Fig. 1:

The architecture of our joint model, along with an example chest radiograph x^I and its associated radiology report x^R. At training time, the model predicts the edema severity level from images and text through their respective encoders and classifiers, and compares the predictions with the labels. The joint embedding loss ${\mathcal{J}}_{\text{E}}$ associates image embeddings I with text embeddings R in the joint embedding space. At inference time, the image stream and the text stream are decoupled and only the image stream is used. Given a new chest radiograph (image), the image encoder and classifier compute its edema severity level.

Joint Representation Learning.

We apply a ranking-based criterion [7, 12] for training the image encoder and the text encoder parameterized by ${\theta }_{\text{E}}^{\text{I}}$ and ${\theta }_{\text{E}}^{\text{R}}$ respectively, to learn image and text feature representations $I\left(x^{\text{I}};{\theta }_{\text{E}}^{\text{I}}\right)$ and $R\left(x^{\text{R}};{\theta }_{\text{E}}^{\text{R}}\right)$. Specifically, given an image-text pair $({\text{x}}_j^{\text{I}},{\text{x}}_j^{\text{R}})$, we randomly select an impostor image ${\text{x}}_s^{\text{I}}{}_j$) and an impostor report ${\text{x}}_{sj}^{\text{R}}$₎ from X. This selection is generated at the beginning of each training epoch. Map s(j) produces a random permutation of {1, 2, …, N}.

We encourage the feature representations between a matched pair (I_j, R_j) to be “closer” than those between mismatched pairs (I_{s j)}, R_j) and (I_j, R_{s j)}) in the joint embedding space. Direct minimization of the distance between I and R could end up pushing the image and text features into a small cluster in the embedding space. Instead, we encourage matched image-text features to be close while spreading out all feature representations in the embedding space for downstream classification by constructing an appropriate loss function:

$$\begin{gathered} {\mathcal{J}}_{\text{E}}({\theta }_{\text{E}}^{\text{I}},{\theta }_{\text{E}}^{\text{R}};{\text{x}}_j,{\text{x}}_{sj)})=\max \left(0,\mathrm{Sim}({\text{I}}_j,{\text{R}}_{sj)}\right)\mathrm{Sim}\left({\text{I}}_j,{\text{R}}_j\right)+) \\ +\max \left(0,\mathrm{Sim}\left({\text{I}}_{sj)},{\text{R}}_j\right)\mathrm{Sim}\left({\text{I}}_j,{\text{R}}_j\right)+\right), \end{gathered}$$ (1)

where Sim(·, ·) is the similarity measurement of two feature representations in the joint 0embedding space and is a margin parameter that is set to |y_j y_{s j)}| when both j N_L and s(j) N_L; otherwise, = 0.5. The margin is determined by the difference due to the mismatch, if both labels are known: otherwise the margin is a constant.

Classi cation.

We employ two fully connected layers (with the same neural network architecture) on the joint embedding space to assess edema severity from the image and the report respectively. For simplicity, we treat the problem as multi-class classification, i.e. the classifiers’ outputs ${\widehat{y}}^{\text{I}}\left(\text{I};{\theta }_{\text{C}}^{\text{I}}\right)$ and ${\widehat{y}}^{\text{R}}\left(\mathrm{R};{\theta }_{\text{C}}^{\text{R}}\right)$ are encoded as one-hot 4-dimensional vectors. We use cross entropy as the loss function for training the classifiers and the encoders on the labeled data:

$$\begin{array}{l}\stackrel{3}{\underset{\begin{array}{l}i=0\\ 3\end{array}}{{\mathcal{J}}_{\text{C}}({\theta }_{\text{E}}^{\text{I}},{\theta }_{\text{E}}^{\text{R}},{\theta }_{\text{C}}^{\text{I}},{\theta }_{\text{C}}^{\text{R}};{\text{x}}_j,{\text{y}}_j)={\text{y}}_{ji}\log {\widehat{y}}_i^{\text{I}}({\text{I}}_j(x_j^{\text{I}};{\theta }_{\text{E}}^{\text{I}});{\theta }_{\text{C}}^{\text{I}})}}\\ {\text{y}}_{ji}\log {\widehat{y}}_i^{\text{R}}({\text{R}}_j(x_j^{\text{R}};{\theta }_{\text{E}}^{\text{R}});{\theta }_{\text{C}}^{\text{R}}),\\ i=0\end{array}$$ (2)

i.e., minimizing the cross entropy also affects the encoder parameters.

Loss Function.

Combining Eq. (1) and Eq. (2), we obtain the loss function for training the joint model:

$$\mathcal{J}({\theta }_{\text{E}}^{\text{I}},{\theta }_{\text{E}}^{\text{R}},{\theta }_{\text{C}}^{\text{I}},{\theta }_{\text{C}}^{\text{R}};\text{X},\text{Y})\underset{j=1}{\overset{N}{=}}{\mathcal{J}}_E({\theta }_{\text{E}}^{\text{I}},{\theta }_{\text{E}}^{\text{R}};{\text{x}}_j,{\text{x}}_{sj)})\underset{j=1}{\overset{N}{+}}{\mathcal{J}}_C({\theta }_{\text{E}}^{\text{I}},{\theta }_{\text{E}}^{\text{R}},{\theta }_{\text{C}}^{\text{I}},{\theta }_{\text{C}}^{\text{R}};{\text{x}}_j,{\text{y}}_j).$$ (3)

Implementation Details.

The image encoder is implemented as a series of residual blocks [13], the text encoder is a BERT model that uses the beginner [CLS] token’s hidden unit size of 768 and maximum sequence length of 320 [8]. The image encoder is trained from a random initialization, while the BERT model is fine-tuned during the training of the joint model. The BERT model parameters are initialized using pre-trained weights on scientific text [4]. The image features and the text features are represented as 768-dimensional vectors in the joint embedding space. The two classifiers are both 768-to-4 fully connected layers. The neural network architecture is provided in the supplementary materials.

We employ the stochastic gradient-based optimization procedure AdamW [27] to minimize the loss in Eq. (3) and use a warm-up linear scheduler [25] for the learning rate. The model is trained on all the image-text pairs by optimizing the first term in Eq. (3) for 10 epochs and then trained on the labeled image-text pairs by optimizing Eq. (3) for 50 epochs. The mini-batch size is 4. We use dot product as the similarity metric in Eq. (1). The dataset is split into training and test sets. All the hyper-parameters are selected based on the results from 5-fold cross validation within the training set.

4. Experiments

Data Preprocessing.

The size of the chest radiographs varies and is around 3000×3000 pixels. We randomly translate and rotate the images on the fiy during training and crop them to 2048×2048 pixels as part of data augmentation. We maintain the original image resolution to capture the subtle differences in the images between different levels of pulmonary edema severity. For the radiology reports, we extract the impressions, ndings, conclusion and recommendation sections. If none of these sections are present in the report, we use the nal report section. We perform tokenization of the text using ScispaCy [21] before providing it to the BERT tokenizer.

Expert Labeling.

For evaluating our model, we randomly selected 531 labeled image-text pairs (corresponding to 485 reports) for expert annotation. A board-certified radiologist and two domain experts reviewed and corrected the regex labels of the reports. We use the expert labels for model testing. The overall accuracy of the regex labels (positive predictive value compared against the expert labels) is 89%. The other 6,212 labeled image-text pairs and around 240K unlabeled image-text pairs were used for training. There is no patient overlap between the training set and the test set.

Model Evaluation.

We evaluated variants of our model and training regimes as follows:

• image only: An image-only model with the same architecture as the image stream in our joint model. We trained the image model in isolation on the 6,212 labeled images.

• A joint image-text model trained on the 6,212 labeled image-text pairs only. We compare two alternatives to the joint representation learning loss:

  • ranking dot, ranking 12, ranking cosine: the ranking based criterion in Eq. (1) with Sim(I, R) defined as one of the dot product I R, the reciprocal of euclidean distance ∥I R∥, and the cosine similarity $\frac{IR}{\Vert I\Vert \Vert R\Vert }$;
  • dot, l2, cosine: direct minimization on the similarity metrics without the ranking based criterion.

• ranking dot semi: A joint image-text model trained on the 6,212 labeled and the 240K unlabeled image-text pairs in a semi-supervised fashion, using the ranking based criterion with dot product in Eq. (1). Dot product is selected for the ranking-based loss based on cross-validation experiments on the supervised data comparing ranking-dot, ranking-l2, ranking-cosine, dot, l2, and cosine.

All reported results are compared against the expert labels in the test set. The image portion of the joint model is decoupled for testing, and the reported results are predicted from images only. To optimize the baseline performance, we performed a separate hyper-parameter search for the image only model using 5-fold cross validation (while holding out the test set).

We use the area under the ROC (AUC) and macro-averaged Fl-scores (macro-Fl) for our model evaluation. We dichotomize the severity levels and report 3 comparisons (0 vs l,2,3; 0,l vs 2,3; and 0,l,2 vs 3), since these 4 classes are ordinal (e.g., $(\text{severity}=0\text{or}1)={\widehat{y}}_0^{\text{I}}+{\widehat{y}}_1^{\text{I}}$, $(\text{severity}=2\text{or}3)={\widehat{y}}_2^{\text{I}}+{\widehat{y}}_3^{\text{I}})$.

Results.

Table 1 reports the performance statistics for all similarity measures. The findings are consistent with our cross-validation results: the ranking based criterion offers significant improvement when it is combined with the dot product as the similarity metric.

Table 1:

Performance statistics for all similarity measures.

Method	AUC (0 v 1,2,3)	AUC (0,1 v 2,3)	AUC (0,1,2 v 3)	macro-F1
12	0.78	0.76	0.83	0.42
ranking 12	0.77	0.75	0.80	0.43
cosine	0.77	0.75	0.81	0.44
ranking cosine	0.77	0.72	0.83	0.41
dot	0.65	0.63	0.61	0.15
ranking dot	0.80	0.78	0.87	0.45

Table 2 reports the performance of the optimized baseline model (image only) and two variants of the joint model (ranking dot and ranking dot semi). We observe that when the joint model learns from the large number of unlabeled image-text pairs, it achieves the best performance. The unsupervised learning minimizes the ranking-based loss in Eq. (1), which does not depend on availability of labels.

Table 2:

Performance statistics for the two variants of our joint model and the baseline image model.

Method	AUC (0 v 1,2,3)	AUC (0,1 v 2,3)	AUC (0,1,2 v 3)	macro-F1
image only	0.74	0.73	0.78	0.43
ranking dot	0.80	0.78	0.87	0.45
ranking dot semi	0.82	0.81	0.90	0.51

It is not surprising that the model is better at differentiating the severity level 3 than other severity categories, because level 3 has the most distinctive radiographic features in the images.

Joint Model Visualization.

As a by-product, our approach provides the possibility of interpreting model classification using text. While a method like Grad-CAM [23] can be used to localize regions in the image that are “important” to tire model prediction, it does not identify tire relevant characteristics of tire radiographs, sudi as texture. By leveraging tire image-text embedding association, we visualize tire heatmap of text attention corresponding to tire last layer of tire [CLS] token in tire BERT model. Tiris heatmap indicates report tokens that are important to our model prediction. As slrown in Fig. 2, we use Grad-CAM [23] to localize relevant image regions and tire lrighlighted words (radiograplric findings, anatomical structures, etc.) from the text embedding to explain the model’s decision making.

Fig. 2:

Joint model visualization. Top to bottoni: (Level 1) The lhghlight of the Grad-CAM image is centered around the right hilar region, which is consistent with findings in pulmonary vascular congestion as shown in the report. (Level 2) The lhghlight of the Grad-CAM image is centered around the left hilar region which shows radiating interstitial markings as confirmed by the report heatmap. (Level 3) Grad-CAM liighlights bilateral alveolar opacities radiating out from the hila and sparing the outer lungs. This pattern is classically described as “batwing” pulmonary edema mentioned in the report. The report text is presented in the form of sub-word tokenization performed by the BERT model, starting the report with a [CLS] token and ending with a [SEP].

5. Conclusion

In this paper, we presented a neural network model that jointly learns from images and text to assess pulmonary edema severity from chest radiographs. The joint image-text representation learning framework incorporates the rich information present in the free-text radiology reports and significantly improves the performance of edema assessment compared to learning from images alone. Moreover, our experimental results show that joint representation learning benefits from the large amount of unlabeled image-text data.

Expert labeling of the radiology reports enabled us to quickly obtain a reasonable amount of test data, but this is inferior to direct labeling of images. The joint model visualization suggests the possibility of using the text to semantically explain the image model, which represents a promising direction for future investigation.