RepsNet: Combining Vision with Language for Automated Medical Reports

The encoder has four constituent parts: 1) image encoder to extract visual features, 2) text encoder to tokenize and contextualize natural language questions and answers features, 3) bilinear attention network to fuse the image and the question, and 4) contrastive alignment of visual features and textual answers.

Image Encoder: We use the ResNeXt-101 [37] architecture as the base image encoder. We remove the last linear and pooling layer and add a 2D adaptive average pooling layer to resize the input image to a fixed feature space of $14\times 14\times 2048$ that preserves the correspondence between the visual features and the input image ($n_{\mathrm{x}}=196,d_{\mathrm{x}}=2048$). Moreover, we add image transformations, namely color jittering, normalization, random erasing, to augment the training data distribution within each batch before extracting the visual features.

Text Encoder: We adapt the BERT [9] model for the text encoder that is pre-trained to predict masked words locally based on the context provided by other non-masked words in the sequence. We filter out the punctuation marks and tokenize the text using WordPiece algorithm [9], before extracting the textual features.

Bilinear Attention Network (BAN): We use a BAN to fuse the cross-modal encoded question and image features [17]. The outer product or the bilinear product exhaustively combines the multi-modal features at the cost of higher computational complexity; in comparison to naive concatenation or inner product between the features. Compared to other co-attention mechanisms, BAN exploits bilinear interaction maps where each feature is pooled by low-rank bilinear approximations. Residual learning on top combines multiple bilinear attention maps for effective joint representation of question and image features.

For the sake of brevity and a slight abuse of notation, we denote ${\boldsymbol{\mathrm{\bar{X}}}}$ for both the image and the combined (image and question) features in describing the rest of the encoder and the decoder sections.

Contrastive Vision and Text Learning: We align images with natural language descriptions via bidirectional contrastive learning [5], that pulls together a given image-answer pair, while pushing away observations that correspond to different image-answer pairs.

Given the encoded image (and question) ${\boldsymbol{\mathrm{\bar{X}}}}$ and the natural language answer features ${\boldsymbol{\mathrm{\bar{Y}}}}\in\mathbb{R}^{n_{\mathrm{y}}\times d_{\mathrm{y}}}$ with $n_{\mathrm{y}}$ tokens of dimension $d_{\mathrm{y}}$, we first project them to a $d$-dimensional space with a linear transformation to ${\boldsymbol{\mathrm{\hat{X}}}}\in\mathbb{R}^{d}$ and ${\boldsymbol{\mathrm{\hat{Y}}}}\in\mathbb{R}^{d}$. During training, the loss operates on a mini-batch of $N_{T}$ image-text pairs $\{{\boldsymbol{\mathrm{\hat{x}_{i}}}},{\boldsymbol{\mathrm{\hat{y}_{i}}}}\}_{i=1}^{N_{T}}$, where each pair is in turn taken as a positive sample to maximize agreement against all other negative samples, i.e.,

where $\langle{\boldsymbol{\mathrm{\hat{x}}}},{\boldsymbol{\mathrm{\hat{y}}}}\rangle=\frac{{\boldsymbol{\mathrm{\hat{x}}}}^{{\!\scriptscriptstyle\top}}{\boldsymbol{\mathrm{\hat{y}}}}}{\|{\boldsymbol{\mathrm{\hat{x}}}}\|\|{\boldsymbol{\mathrm{\hat{y}}}}\|}$ represents the cosine similarity distance and $\tau\in\mathbb{R}^{+}$ represents the temperature parameter to scale the similarity metric. Similar to the image-to-text loss in Eq. (2), we also define the text-to-image loss $\mathcal{L}_{{\boldsymbol{\mathrm{\hat{y}}}}\rightarrow{\boldsymbol{\mathrm{\hat{x}}}}}$ to account for the asymmetry with respect to each input modality as in [43, 27]. Overall bidirectional encoder loss $\mathcal{L}_{\mathrm{enc}}$ is the sum of the two constituent contrastive losses weighted by constant $\alpha_{l}\in\mathbb{R}^{+}$,

Prior Context Knowledge: We store the normalized natural language answers of the train set ${\boldsymbol{\mathrm{\hat{Y}_{train}}}}$ during model training. We then compute topk nearest neighbours ${\boldsymbol{\mathrm{\bar{C}}}}$ that maximize the cosine similarity between a given encoded image ${\boldsymbol{\mathrm{\hat{X}}}}$ and the stored natural language answers ${\boldsymbol{\mathrm{\hat{Y}_{train}}}}$. We use the FAISS library for scalable nearest neighbour search [16]. The prior context aids the decoder to attend to longer horizon dependencies and get additional case-based details for controlled text generation. This is particularly relevant in describing medical images with specific terminologies, writing style and class imbalanced abnormalities, i.e.,

The probability distribution of generating the output text sequence ${\boldsymbol{\mathrm{Y_{1:t}}}}$ conditioned on the contextualized encoding sequence $\mathcal{P}_{{\boldsymbol{\theta_{\mathrm{dec}}}}}\left({\boldsymbol{\mathrm{Y_{1:t}}}}|{\boldsymbol{\mathrm{\bar{X},\bar{C}}}}\right)$ can be decomposed into a product of conditional distributions using Bayes’ rule,

where ${\boldsymbol{\mathrm{y_{0}}}}=\langle{\boldsymbol{\mathrm{BOS}}}\rangle$ is a special token reserved for the beginning of a sentence. We model the conditional language generation with a stack of transformer-based blocks, using the GPT-2 model as the base pretrained language decoder [28]. We introduce modifications to the GPT-2 model for conditioning on image and prior context features by directly adding their attention outputs to the pretrained self-attention layers of the model, similar to  [44, 2], thereby adding the attention outputs for different conditional inputs with a parsimonious increase in the number of parameters only (see supplementary materials for details). The conditional probability distribution in Eq. (5) is maximized by optimizing the cross-entropy loss on the ground-truth and the predicted sequences.

Overall Approach: During training, we adapt the pretrained language and vision models in an end-to-end manner for contrastive encoding and conditional decoding with a small amount of image-text pairs. The overall training loss function comprises of the contrastive loss and the cross-entropy loss. During natural language generation, we predict the output sequence in an auto-regressive manner with greedy or beam search decoding, and stop generating the sequence once we predict a special end of text $\langle{\boldsymbol{\mathrm{EOS}}}\rangle$ token.

VQA-Rad Dataset: We use the VQA-Rad datasets [18] from $2018-2019$ and also introduce an aggregrated dataset, VQA-Rad All, that combines all the VQA-Rad datasets from $2018-2021$. Radiology images in the datasets are taken from open-access MedPix database, and the questions are predominantly posed from categories, such as image plane, imaging modality, organ system involved and image abnormalities. The VQA problem is posed as a multi-class classification over all possible set of answers, and classification accuracy on evaluation set is used as the performance metric. We use the standard training and evaluation splits provided with the datasets (see summary in supplementary materials).
Results: Table 2 shows that RepsNet outperforms all other competing methods across all the datasets. Similar to other methods, RepsNet uses bilinear attention mechanism to fuse the image and the question features. Contrary to other methods, RepsNet does not use fixed Glove word embeddings [26] or RNNs for sentence level representations; instead it learns the entire contextual embeddings using BERT-style transformer with WordPiece tokenization. Ablation study in Table 4 also shows that the performance increases the most with the use of pre-trained models. We observe from Table 2 that simply filtering out instances and class categories with less than $5$ and $10$ instances per class category $M_{o}=\{5,10\}$ proportionally increases the classification accuracy across all datasets, at the cost of reducing the overall number of instances and class categories to mitigate class imbalance in the datasets. Note that we do not take into account unseen class category instances of the evaluation set in computing the classification accuracy.

IU-XRay: The Indiana University x-ray dataset [7] comprises of frontal and lateral views of chest x-ray images that are associated with radiology report sections, namely impressions, findings and manual tags. For brevity, we only report results for populating the findings question in this work, i.e., we associate the same question for all the answers. After omitting the reports without findings section, we randomly split the remaining $3607$ reports into $80\%$ training and $20\%$ evaluation sets. On average, each report instance has $5.7$ sentences, while each sentence has $6.5$ words that describe the image findings. Note that no classification labels are available to detect the anomalies. Max number of tokens for a report section is set to $200$ and the report findings are zero-padded in case its length is less than max number of tokens. We use the sentence level BLEU scores as the performance metric computed using the nltk library that compares $n$-gram similarity between the ground-truth and the generated report where $n$ varies from $1$ to $4$ (whereas for classification accuracy evaluation in VQA-Rad datasets, we compare the predicted and the ground-truth indices of the class categories).

Results: Results are summarized in Table 2. It can be seen that RepsNet performs significantly better than the state-of-the-art report generation methods across all the BLEU scores, suggesting the feasibility of adapting large-scale pretrained language and vision models on a small set of domain-specific medical data. Ablation study in Table 4 reveals that adding visual features, contrastive learning and prior context subsequently boosts the performance of the GPT2-decoder. Fig. 4 provides a qualitative comparison between the ground-truth and the generated report findings, along with the heatmap visualizations using grad-cam [31] for an intuitive understanding of the approach. We observe a strong alignment in generating normal report findings, whereas part of the findings sometimes get omitted and/or added in describing the abnormalities, especially for rare cases (see supplementary materials for video demonstration and other examples). Systematic dealing of rare cases with external domain knowledge and past medical history of patients is a promising direction of our future work. We are also interested in incorporating attention mechanisms for conditional visualization of generated text on image patches as a measure of uncertainty in the prediction. Making these reports self-explainable is critical for its wider adoption. Other areas of interest include reducing the liability of generated report errors, as well as working with medical experts to evaluate the generated reports.