Semi-supervised Learning for Quantification of Pulmonary Edema in Chest X-Ray Images
Ruizhi Liao Jonathan Rubin Grace Lam Seth Berkowitz Sandeep Dalal William Wells Steven Horng Polina Golland




Semi-supervised Learning for Chest X-Ray Images



 

R. Liao et al.



Massachusetts Institute of Technology, Cambridge, MA, USA 

ruizhi@mit.edu Philips Research North America, Cambridge, MA, USA Beth Israel Deaconess Medical Center, Harvard Medical School, Boston, MA, USA Brigham and Women's Hospital, Harvard Medical School, Boston, MA, USA

Semi-supervised Learning for Quantification of Pulmonary Edema in Chest X-Ray Images
    Ruizhi Liao1 Jonathan Rubin2 Grace Lam1 Seth Berkowitz3 Sandeep Dalal2 William Wells1,4 Steven Horng3 Polina Golland1 
    December 30, 2023
==========================================================================================================================




We propose and demonstrate machine learning algorithms to assess the severity of pulmonary edema in chest x-ray images of congestive heart failure patients. Accurate assessment of pulmonary edema in heart failure is critical when making treatment and disposition decisions. Our work is grounded in a large-scale clinical dataset of over 300,000 x-ray images with associated radiology reports. While edema severity labels can be extracted unambiguously from a small fraction of the radiology reports, accurate annotation is challenging in most cases. To take advantage of the unlabeled images, we develop a Bayesian model that includes a variational auto-encoder for learning a latent representation from the entire image set trained jointly with a regressor that employs this representation for predicting pulmonary edema severity. Our experimental results suggest that modeling the distribution of images jointly with the limited labels improves the accuracy of pulmonary edema scoring compared to a strictly supervised approach. To the best of our knowledge, this is the first attempt to employ machine learning algorithms to automatically and quantitatively assess the severity of pulmonary edema in chest x-ray images.








§ INTRODUCTION


We propose and demonstrate a semi-supervised learning algorithm to support clinical decisions in congestive heart failure (CHF) by quantifying pulmonary edema. Limited ground truth labels are one of the most significant challenges in medical image analysis and many other machine learning applications in healthcare. It is of great practical interest to develop machine learning algorithms that take advantage of the entire data set to improve the performance of strictly supervised classification or regression methods. In this work, we develop a Bayesian model that learns probabilistic feature representations from the entire image set with limited labels for predicting edema severity. 

Chest x-ray images are commonly used in CHF patients to assess pulmonary edema, which is one of the most direct symptoms of CHF <cit.>. CHF causes pulmonary venous pressure to elevate, which in turn causes the fluid to leak from the blood vessels in the lungs into the lung tissue. The excess fluid in the lungs is called pulmonary edema. Heart failure patients have extremely heterogenous responses to treatment <cit.>. The assessment of pulmonary edema severity will enable clinicians to make better treatment plans based on prior patient responses and will facilitate clinical research studies that require quantitative phenotyping of the patient status <cit.>. While we focus on CHF, the quantification of pulmonary edema is also useful elsewhere in medicine. Quantifying pulmonary edema in a chest x-ray image could be used as a surrogate for patient intravascular volume status, which would rapidly advance research in sepsis and other disease processes where volume status is critical.

Quantifying pulmonary edema is more challenging than detection of pathologies in chest x-ray images <cit.> because grading of pulmonary edema severity relies on much more subtle image findings (features). Accurate grading of the pulmonary edema severity is challenging for medical experts as well <cit.>. Our work is grounded in a large-scale clinical dataset that includes approximately 330,000 frontal view chest x-ray images and associated radiology reports, which serve as the source of the severity labels. Of these, about 30,000 images are of CHF patients. Labels extracted from radiology reports via keyword matching are available for about 6,000 images. Thus our image set includes a large number of images, but only a small fraction of images is annotated with edema severity labels.

We use variational auto-encoder (VAE) to capture the image distribution from both unlabeled and labeled images to improve the accuracy of edema severity grading. Auto-encoder neural networks have shown promise for representational modeling <cit.>. Earlier work attempted to learn a separate VAE for each label category from unlabeled and labeled data <cit.>. We argue and demonstrate in our experiments that this structure does not fit our application well, because pulmonary edema severity score is based on subtle image features and should be represented as a continuous quantity. Instead, we learn one VAE from the entire image set. By training the VAE jointly with a regressor, we ensure it captures compact feature representations for scoring pulmonary edema severity. Similar setups have also been employed in computer vision <cit.>. The experimental results show that our method outperforms the multi-VAE approach <cit.>, the entropy minimization based self-learning approach <cit.>, and strictly supervised learning. To the best of our knowledge, this paper demonstrates the first attempt to employ machine learning algorithms to automatically and quantitatively assess the severity of pulmonary edema from chest x-ray images. 



§ METHODS





Let x∈ℝ^n×n be a 2D x-ray image and y∈{0,1,2,3} be the corresponding edema severity label. Our dataset includes a set of N images 𝐱={x_i}_i=1^N with the first N_L images annotated with severity labels 𝐲={y_i}_i=1^N_L. Here, we derive a learning algorithm that constructs a compact probabilistic feature representation z that is learned from all images and used to predict pulmonary edema severity. Fig. <ref> illustrates the Bayesian model and the inference algorithm.



  
Learning.
The learning algorithm maximizes the log probability of the data with respect to parameters θ:

    log p(𝐱, 𝐲; θ)=∑_i=1^N_Llog p(x_i, y_i; θ) + ∑_i=N_L+1^Nlog p(x_i; θ).

We model z as a continuous latent variable with a prior distribution p(z), which generates images and predicts pulmonary edema severity. Unlike <cit.> that constructs a separate encoder q(z|x,y) for each value of discrete label y, we use a single encoder q(z|x) to capture image structure relevant to labels. Distribution q(z|x) serves as a variational approximation for p(z|x, y) for the lower bound:

    ℒ_1(θ; x_i, y_i) =    log p(x_i, y_i; θ) - D_KL(q(z_i|x_i; θ)||p(z_i|x_i, y_i)), 
    
    =    _q(z_i|x_i; θ)[log p(x_i, y_i; θ) + log p(z_i|x_i, y_i) - log q(z_i|x_i; θ) ] 
    
    =    _q(z_i|x_i; θ)[log p(x_i, y_i| z_i; θ) + log p(z_i) - log q(z_i|x_i; θ) ]
    
    =     _q(z_i|x_i; θ)[log p(x_i, y_i| z_i; θ)] - D_KL(q(z_i|x_i; θ)||p(z_i)).

We assume that x, z, and y form a Markov chain, i.e., y  x | z, and therefore

    ℒ_1(θ; x_i, y_i)
    =     _q(z_i|x_i; θ_E)[log p(x_i|z_i; θ_D)] + _q(z_i|x_i; θ_E)[log p(y_i|z_i; θ_R) ]
        -D_KL(q(z_i|x_i; θ_E)||p(z_i)),

where θ_E are the parameters of the encoder, θ_D are the parameters of the decoder, and θ_R are the parameters of the regressor. Similarly, we have a variational lower bound for log p(x_i; θ):

    ℒ_2(θ; x_i)
    =  _q(z_i|x_i; θ_E)[log p(x_i|z_i; θ_D)]-D_KL(q(z_i|x_i; θ_E)||p(z_i)).


By substituting Eq. (<ref>) and Eq. (<ref>) into Eq. (<ref>), we obtain a lower bound for the log probability of the data and aim to minimize the negative lower bound:

    𝒥(θ; 𝐱, 𝐲)
    =     - ∑_i=1^N_Lℒ_1(θ; x_i, y_i) - ∑_i=N_L+1^Nℒ_2(θ; x_i) 
    
    =    ∑_i=1^ND_KL(q(z_i|x_i; θ_E)||p(z_i))
    -∑_i=1^N_L_q(z_i|x_i; θ_E)[log p(y_i|z_i; θ_R)]
        -∑_i=1^N_q(z_i|x_i; θ_E)[log p(x_i|z_i; θ_D)].




  
Latent Variable Prior p(z).
We let the latent variable prior p(z) be a multivariate normal distribution, which serves to regularize the latent representation of images. 



  
Latent Representation q(z|x).
We apply the reparameterization trick used in <cit.>. Conditioned on image x_i, the latent representation becomes a multivariate Gaussian variable, z_i|x_i∼𝒩(z_i; μ_i, Λ_i), where μ_i is a D-dimensional vector [μ_ik]_k=1^D and Λ_i is a diagonal covariance matrix represented by its diagonal elements as [λ_ik^2]_k=1^D. Thus, the first term in Eq. (<ref>) becomes:

    𝒥_KL(θ_E; x_i) = -1/2∑_k=1^D(logλ_ik^2 - μ_ik^2 - λ_ik^2 ) + const.

We implement the encoder as a neural network f_E(x; θ_E) that estimates the mean and the variance of z|x. Samples of z can be readily generated from this estimated Gaussian distribution. We use one sample per image for training the model.



  
Ordinal Regression p(y|z).
In radiology reports, pulmonary edema severity is categorized into four groups: no/mild/ moderate/severe. Our goal is to assess the severity of pulmonary edema as a continuous quantity. We employ ordinal representation to capture the ordering of the categorical labels. We use a 3-bit representation y_i=[y_ij]_j=1^3 for the four severity levels. The three bits represent the probability of any edema, of moderate or severe edema, and of severe edema respectively (i.e., “no" is [0,0,0], “mild" is [1,0,0], “moderate" is [1,1,0], and “severe" is [1,1,1]). This encoding yields probabilistic output, i.e., both the estimate of the edema severity and also uncertainty in the estimate. The three bits are assumed to be conditionally independent given the image:

    p(y_i|z_i; θ_R) = ∏_j=1^3 f_R^j(z_i; θ_R)^y_ij(1-f_R^j(z_i; θ_R))^1-y_ij,

where y_ij is a binary label and f_R^j(z_i; θ_R) is interpreted as the conditional probability p(y_ij=1|z_i). f_R(·) is implemented as a neural network. The second term in Eq. (<ref>) becomes the cross entropy:

    𝒥_ℛ(θ_E, θ_R; y_i, z_i) 
    
    =     - ∑_j=1^3y_ijlog f_R^j(z_i; θ_R) 
    - ∑_j=1^3 (1-y_ij)log(1-f_R^j(z_i; θ_R)).




  
Decoding p(x|z).
We assume that image pixels are conditionally independent (Gaussian) given the latent representation. Thus, the third term in Eq. (<ref>) becomes: 

    𝒥_𝒟(θ_E, θ_D; x_i, z_i)
    
        = - log𝒩(x_i; f_D(z_i; θ_D), Σ_i) 
        = 12(x_i-f_D(z_i; θ_D))^TΣ_i^-1(x_i-f_D(z_i; θ_D))+const.,

where f_D(·) is a neural network decoder that generates an image implied by the latent representation z, and Σ_i is a diagonal covariance matrix. 



  
Loss Function.
Combining Eq. (<ref>), Eq. (<ref>) and Eq. (<ref>), we obtain the loss function for training our model:

    𝒥 (θ_E, θ_R, θ_D; 𝐱, 𝐲)
    =    ∑_i=1^N𝒥_KL(θ_E; x_i) + ∑_i=1^N_L𝒥_ℛ(θ_E, θ_R; y_i, z_i) + ∑_i=1^N𝒥_𝒟(θ_E, θ_D; x_i, z_i) 
    
    =     -1/2∑_i=1^N∑_k=1^D(logλ_ik^2 - μ_ik^2 - λ_ik^2 ) 
        - ∑_i=1^N_L( ∑_j=1^3y_ijlog f_R^j(z_i; θ_R) + ∑_j=1^3 (1-y_ij)log(1-f_R^j(z_i; θ_R)) ) 
        + 12∑_i=1^N(x_i-f_D(z_i; θ_D))^TΣ_i^-1(x_i-f_D(z_i; θ_D)).

We employ the stochastic gradient-based optimization procedure Adam <cit.> to minimize the loss function. Our training procedure is outlined in the supplementary materiel. The pulmonary edema severity category extracted from radiology reports is a discrete approximation of the actual continuous severity level. To capture this, we compute the expected severity:

    ŷ
     =0 × (1-ŷ_1) + 1× (ŷ_1-ŷ_2)+2× (ŷ_2-ŷ_3) +3 ×ŷ_3 =ŷ_1+ŷ_2+ŷ_3.




§ IMPLEMENTATION DETAILS

The size of the chest x-ray images in our dataset varies and is around 3000×3000 pixels. We randomly rotate and translate the images (differently at each epoch) on the fly during training and crop them to 2048×2048 pixels as part of data augmentation. We maintain the original image resolution to preserve subtle differences between different levels of pulmonary edema severity.

The encoder is implemented as a series of residual blocks <cit.>. The decoder is implemented as a series of transposed convolutional layers, to build an output image of the same size as the input image (2048×2048). The regressor is implemented as a series of residual blocks with an averaging pooling layer followed by two fully connected layers. The regressor output ŷ has 3 channels. The latent representation z has a size of 128×128. During training, one sample is drawn from z per image. The KL-loss (Eq. (<ref>)) and the image reconstruction error (Eq. (<ref>)) in the loss function are divided by the latent feature size and the image size respectively. The variances in Eq. (<ref>) are set to 10, which gives a weight of 0.1 to the image reconstruction error. The learning rate for the Adam optimizer training is 0.001 and the minibatch size is 4. The model is trained on a training dataset and evaluated on a separate validation dataset every few epochs during training. The model checkpoint with the lowest error on the validation dataset is used for testing. The neural network architecture is provided in the supplementary material.



§ EXPERIMENTS




  
Data.
Approximately 330,000 frontal view x-ray images and their associated radiology reports were collected as part of routine clinical care in the emergency department of Beth Israel Deaconess Medical Center and subsequent in-hospital stay. A subset of the image set has been released <cit.>.



In this work, we extracted the pulmonary edema severity labels from the reports by searching for keywords that are highly correlated with a specific stage of pulmonary edema. Due to the high variability of wording in radiology reports, the same keywords can mean different clinical findings in varying disease context. For example, perihilar infiltrate means moderate pulmonary edema for a heart failure patient, but means pneumonia in a patient with a fever. To extract meaningful labels from the reports using the keywords, we limited our label extraction to a CHF cohort. This cohort selection yielded close to 30,000 images, of which 5,771 images could be labeled via our keyword matching. Representative images of each severity level are shown in Fig. <ref>. The data details are summarized in the supplementary material. 



  
Evaluation.
We randomly split the images into training (4,537 labeled images, 334,664 unlabeled images), validation (628 labeled images), and test (606 labeled images) image sets. There is no patient overlap between the sets. Unlabeled images of the patients in the validation and test sets are excluded from training. The labeled data split is 80%/10%/10% into training/validation/test respectively. 

We evaluated four methods: (i) Supervised: purely supervised training that uses labeled images only; (ii) EM: supervised training with labeled images that imputes labels for unlabeled images and minimizes the entropy of predictions <cit.>; (iii) DGM: Semi-supervised training with deep generative models (multiple VAEs) as described in <cit.>; (iv) VAE_R: Our method that learns probabilistic feature representations from the entire image set with limited labels. For the baseline supervised learning method, we investigated different neural network architectures previously demonstrated for chest x-ray images <cit.> and did not find the network architecture changes the supervised learning results significantly.

We evaluate the methods on the test image set using the root mean squared (RMS) error and Pearson correlation coefficient (CC). 






  
Results.
Fig. <ref> summarizes the prediction performance of the four methods. The method that jointly learns probabilistic feature representations outperforms the other three models.



§ CONCLUSIONS


In this paper, we demonstrated a regression model augmented with a VAE trained on a large image dataset with a limited number of labeled images. Our results suggest that it is difficult for a generative model to learn distinct data clusters for the labels that rely on subtle image features. In contrast, learning compact feature representations jointly from images and limited labels can help inform prediction by capturing structure shared by the image distribution and the conditional distribution of labels given images. 

We demonstrated the first attempt to employ machine learning algorithms to automatically and quantitatively assess the severity of pulmonary edema from chest x-ray images. Our results suggest that granular information about a patient's status captured in medical images can be extracted by machine learning algorithms, which promises to enable clinicians to deliver better care by quantitatively summarizing an individual patient's medical history, for example response to different treatments. This work also promises to enable clinical research studies that require quantitative summarization of patient status.  







splncs04
10

chakko1991clinical
Chakko, S., Woska, D., Martinez, H., De Marchena, E., Futterman, L., Kessler,
  K.M., Myerburg, R.J.: Clinical, radiographic, and hemodynamic correlations in
  chronic congestive heart failure: conflicting results may lead to
  inappropriate care. The American journal of medicine  90(1),
  353–359 (1991)

francis2014heterogeneity
Francis, G.S., Cogswell, R., Thenappan, T.: The heterogeneity of heart failure:
  will enhanced phenotyping be necessary for future clinical trial success?
  (2014)

grandvalet2005semi
Grandvalet, Y., Bengio, Y.: Semi-supervised learning by entropy minimization.
  In: Advances in neural information processing systems. pp. 529–536 (2005)

hammon2014improving
Hammon, M., Dankerl, P., Voit-Höhne, H.L., Sandmair, M., Kammerer, F.J.,
  Uder, M., Janka, R.: Improving diagnostic accuracy in assessing pulmonary
  edema on bedside chest radiographs using a standardized scoring approach. BMC
  anesthesiology  14(1),  94 (2014)

he2016deep
He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image
  recognition. In: Proceedings of the IEEE conference on CVPR. pp. 770–778
  (2016)

johnson2019mimic
Johnson, A.E., Pollard, T.J., Berkowitz, S., Greenbaum, N.R., Lungren, M.P.,
  Deng, C.y., Mark, R.G., Horng, S.: MIMIC-CXR: A large publicly available
  database of labeled chest radiographs. arXiv preprint arXiv:1901.07042
  (2019)

kamnitsas2018semi
Kamnitsas, K., Castro, D., Le Folgoc, L., Walker, I., Tanno, R., Rueckert, D.,
  Glocker, B., Criminisi, A., Nori, A.: Semi-supervised learning via compact
  latent space clustering. In: International Conference on Machine Learning.
  pp. 2464–2473 (2018)

kingma2014adam
Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv
  preprint arXiv:1412.6980  (2014)

kingma2014semi
Kingma, D.P., Mohamed, S., Rezende, D.J., Welling, M.: Semi-supervised learning
  with deep generative models. In: Advances in Neural Information Processing
  Systems. pp. 3581–3589 (2014)

kingma2013auto
Kingma, D.P., Welling, M.: Auto-encoding variational bayes. arXiv preprint
  arXiv:1312.6114  (2013)

mahdyoon1989radiographic
Mahdyoon, H., Klein, R., Eyler, W., Lakier, J.B., Chakko, S., Gheorghiade, M.:
  Radiographic pulmonary congestion in end-stage congestive heart failure.
  American Journal of Cardiology  63(9),  625–627 (1989)

rajpurkar2017chexnet
Rajpurkar, P., Irvin, J., Zhu, K., Yang, B., Mehta, H., Duan, T., Ding, D.,
  Bagul, A., Langlotz, C., Shpanskaya, K., et al.: Chexnet: Radiologist-level
  pneumonia detection on chest x-rays with deep learning. arXiv preprint
  arXiv:1711.05225  (2017)

wang2017chestx
Wang, X., Peng, Y., Lu, L., Lu, Z., Bagheri, M., Summers, R.M.: Chestx-ray8:
  Hospital-scale chest x-ray database and benchmarks on weakly-supervised
  classification and localization of common thorax diseases. In: Proceedings of
  the IEEE conference on CVPR. pp. 3462–3471. IEEE (2017)











§ SUPPLEMENTARY MATERIAL














