For ImageCLEFmed caption 2021, teams were tasked with developing systems that could automatically produce coherent captions for the entirety of a medical image [1,2]. To succeed, a system must not only identify medical concepts but also their interplay. A system that can achieve this could improve the efficiency of radiologists' interpretation. An example of an image and its ground truth caption for ImageCLEFmed Caption 2021 is provided in Figure 1. Typical medical image captioning approaches make use of either a sequence-to-sequence (seq2seq) model to generate a caption for a medical image [3], or image retrieval, where it is assumed that similar images have similar captions [4]. For our submissions, a seq2seq model was considered only.
While medical data has many unique characteristics, general-purpose Natural Language Processing (NLP) and Computer Vision (CV) methods have proven effective in many domain-specific medical tasks. In NLP, for example, general-domain self-supervised pre-training strategies-such as Masked Language Modeling (MLM) and Next Sentence Prediction (NSP) used to produce Bidirectional Encoder Representations from Transformers (BERT) [5]-have been successfully adapted to medical text. One instance is PubMedBERT [6]-a Transformer encoder [7] pre-trained on PubMed articles using domain-specific self-supervised pre-training strategies. Another example is the use of Transfer Learning (TL) to significantly improve medical image classification accuracy on small datasets [8]. Here, a portion of a Convolutional Neural Network (CNN) trained on ImageNet 2012 [9] (a general-domain image classification dataset) is fine-tuned on the small amount of data for the medical image classification task.
A number of more recent NLP and CV machine learning techniques have not been investigated on medical data. One such approach for sequence generation is the use of pre-trained Transformer checkpoints to initialise both the encoder and the decoder of a seq2seq model [10]. Another method is the Vision Transformer (ViT)-a pre-trained Transformer checkpoint for image classification, which takes 16x16 patches of the image as input [11]. Building from this, a pre-trained ViT checkpoint was paired with a Transformer decoder to form a seq2seq model for image captioning [12].
Motivated by previous adaptations of general-domain NLP and CV machine learning techniques to medical data and the slew of recent techniques that have not been investigated on medical data, we investigate a seq2seq model for medical image captioning that employs a pre-trained ViT checkpoint as the encoder and the pre-trained PubMedBERT checkpoint as the decoder. We also experiment with various pre-training and fine-tuning strategies, such as additionally pre-training the encoder on a multi-label medical image classification task, as well as fine-tuning the seq2seq model with Self-Critical Sequence Training (SCST) [13].
The focus for ImageCLEFmed Caption 2021 was to use real medical images and have participants develop automated systems to predict natural language captions; evaluation was performed by comparing the predicted captions to the annotations provided by medical doctors (i.e. the ground truth captions). Each example from the dataset consisted of a medical image and its associated ground truth caption, as shown in Figure 1. The training, validation, and test sets comprise of 2,756, 500, and 444 examples, respectively. We refer to the ImageCLEFmed Caption 2021 dataset as the task's dataset henceforth.
Each caption (predicted and ground truth) was pre-processed in the following way: The caption was first converted to lower-case. All punctuation was then removed and the caption was tokenized into its individual words. Stopwords were then removed using NLTK's English stopword list (NLTK v3.2.2). Stemming was next applied using NLTK's Snowball stemmer (NLTK v3.2.2). The score was then calculated as the average score of BLEU-1, BLEU-2, BLEU-3, and BLEU-4 [14]. Note that the caption was always considered as a single sentence, even if it contained several sentences. No smoothing function was used. All scores were summed and averaged over the number of captions, giving the final score. One downside of using BLEU for medical image caption evaluation is that it is a word overlap measure and may not capture clinical correctness, as noted in [4].
The nine submissions for the AEHRC CSIRO team are described in Figure 2. In this section, we describe the model architecture, the pre-training and fine-tuning strategy, as well as the external pre-training datasets for each submission. Figure 2 helps the reader identify the stages of pre-training and fine tuning for the encoder, decoder, and the seq2seq model that were used for each submission, along with the epoch chosen for each stage.
The same model architecture was used for each submission. An overview of the model is shown in Figure 3. In terms of architecture, the encoder is identical to ViT [11, ViT-Base, Table 1] and the decoder is identical to PubMedBERT [6]. Both ViT and PubMedBERT use 12 hidden layers, each with a size of 768, an intermediate size of 3,072 [7, see π π π in Section 3.3], and 12 scaled dot-product attention heads. Next, we describe the medical image pre-processing, followed by the encoder and decoder.
A given medical image π β β πΆΓπ Γπ» (where πΆ, π, and π» denote the number of channels, the width, and the height of the medical image, respectively) is first resized using bilinear interpolation so that its smallest side has 416 pixels. Next, the resized image is cropped to a size of β 3Γ384Γ384 (the size required for ViT), with the crop location random during training and centered during testing. The cropped image is then split into a set of non-overlapping patches-each of size β 3Γ16Γ16 (i.e., 576 non-overlapping patches). Each patch is then flattened into a one-dimensional array of size β 768 . A colour depth of 8-bits was used for the images (where images with a higher colour depth were downsampled to 8-bits).
The set of inputs given to the encoder consist of the projection of each patch and the [ C L S ] embedding. The patch projections are formed by passing each flattened patch through a learnable projection matrix: π πππππππ πππ‘πβ β β 768Γ768 . The [ C L S ] embedding is learnt using matrix π πππππππ ππππ π β β 1Γ768 (where [ C L S ] is the classification token whose corresponding output is fed to a classification head during pre-training). The corresponding output for the [ C L S ] embedding forms an aggregate representation over all patches. Before the set of inputs are given to the first ViT hidden layer, a position embedding is added to each element of the set. There are 577 position embeddings for the encoder, with position "0" reserved for the [ C L S ] embedding and positions "1" to "576" reserved for the patch projections (which provide information about the location of each patch within the medical image). Each position embedding is stored in a learnable matrix: π πππππππ πππ ππ‘πππ β β 577Γ768 . The weights for the encoder (including its embeddings and the patch projection) are initialised using one of the pre-trained ViT checkpoints from [11]. This checkpoint has been pre-trained on ImageNet21k [15] and then subsequently on ImageNet 2012 [9]. We also investigate additionally pre-training this checkpoint on a multi-label medical image classification task; we denote this Medical Image Transformer (MIT). The multi-label medical image classification task is comprised of four datasets, as described in Section 3.2.1. The pre-trained checkpoint of either ViT or MIT is used to initialise the encoder for the submissions in Figure 2, where a tick in the "MIT" column indicates that MIT was used over ViT. Moreover, submission identifiers using the pre-trained ViT checkpoint are labelled in Figure 2 starting with "vit", while those using a pre-trained MIT checkpoint have are labelled starting with "mit".
The weights of the decoder (along with its embeddings) are initialised using the pretrained PubMedBERT checkpoint [6]. We classify PubMedBERT as a Medical Report Transformer (MRT)-a Transformer that has been pre-trained on medical text (in this case medical literature) and is suitable for generating medical reports. The output of the last hidden layer of the encoder is fed to each decoder hidden layer via a randomly initialised multi-head cross-attention module, which is inserted between the masked multi-head self-attention module and the Feedforward Neural Network (FNN) module of each layer [7, Section 3.1, Decoder]. PubMedBERT has a vocabulary size of 30,522, comprising subword units. When feeding a subword unit to the encoder, it is first converted to its corresponding token index, and then subsequently into a token embedding. Each token embedding is stored in learnable matrix π πππππππ π‘ππππ β β 30,522Γ768 . Next, a position and a segment embedding are added to the token embedding. The position embedding indicates the location of the subword within the caption. A maximum of 512 positions are used for PubMedBERT, with each position embedding stored in learnable matrix π πππππππ πππ ππ‘πππ β β 512Γ768 . As only one caption is generated per medical image (even though PubMedBERT is pre-trained using two segments), the embedding for segment "0" will only be used. Each segment embedding is stored in learnable matrix π πππππππ π ππππππ‘ β β 2Γ768 . When generating a caption, the token [ B O S ] (beginning of sentence) is first fed to output the first subword of the caption. Caption generation finishes once the decoder generates the [ E O S ] token. Each submission used PubMedBERT, as shown by the submission identifiers in Figure 2 (i.e. "vit2mrt" and "mit2mrt"), where "mrt" indicates that the decoder is an MRT, where, in this case, PubMedBERT is the MRT. During testing, the maximum amount of subwords that the decoder could generate was set to 128. Beam search was also used, with a beam size of eight. Additionally, all n-grams of size three were only allowed to occur once.
Next, we describe the pre-training and fine-tuning strategies for the submissions. Here, finetuning refers to training on the task's dataset. Pre-training refers to training on other, external datasets we selected; this was done before the fine-tuning stage. Teacher Forcing (TF) with categorical cross entropy loss was used to fine-tune each seq2seq model on the task's dataset [16]. We also investigate additionally fine-tuning the seq2seq models (which have already been fine-tuned using TF) with Self Critical Sequence Training (SCST) [13]. Submissions that were fine-tuned with TF and then SCST have a tick in the "SCST" column of Figure 2.
For pre-training the seq2seq models, we used the Radiology Objects in COntext (ROCO) medical image captioning dataset (described in Section 3.2.1) with TF-before fine-tuning on the task's dataset. A tick in the "ROCO" column of Figure 2 signifies if this stage of pre-training was conducted for a submission.
The AdamW optimiser [17] was used for gradient descent optimisation during pre-training and fine-tuning. A learning rate of 1π β 7 was used for fine-tuning on the task's dataset with SCST. A learning rate of 5π β 5 and a linear warm-up of 10,000 training steps from a learning rate of zero was used when pre-training the MITs, pre-training on ROCO with TF, and when fine-tuning on the task's dataset with TF. All other hyperparameters for AdamW were set to their defaults. For the pre-training strategy in [11], L2 regularisation (with a term of 0.9) helped to prevent overfitting the ViT during pre-training. Motivated by this, we investigated L2 regularisation for pre-training only (i.e., for pre-training the MITs and when pre-training on ROCO using TF). We investigated an L2 regularisation term with two of the submissions, as shown by the legend in Figure 2.
A mini-batch size of 64 was used to pre-train each MIT, to pre-train on ROCO with TF, and to fine-tune on the task's dataset with TF. A mini-batch size of eight was used to fine-tuning on the task's dataset with SCST. For epoch selection at each stage of pre-training and fine tuning in Figure 2, early stopping with a patience of five was used. The validation micro F1 score was the monitored metric for early stopping with each MIT and the validation BLEU score (BLEU is described in Section 2.2) was the monitored metric for early stopping with the seq2seq models. For submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 , the early stopping criteria was not enforced until epoch 50, as the BLEU score did not increase from zero until after this epoch. When fine-tuning on the task's dataset using SCST, the maximum amount of subwords that the decoder could generate was set to 32 due to memory restrictions. Moreover, greedy search was used (i.e., a beam size of 1) when generating the baseline for SCST [13].
A number of external medical image datasets were used to pre-training the MIT, to take advantage of any stored knowledge about medical images when fine-tuning on the task's dataset. Specifically, four medical image multi-label classification datasets were identified; these are shown in Table 1. While CheXpert and MURA included validation sets, test sets were not available with any of the four datasets. We refrain from using these validation sets as the CheXpert validation set has been used as a test set previously [18]. Instead, 5% of each training set was selected and removed to form a validation set. Together, the datasets have β’ PadChest includes 160,828 chest X-rays obtained from 67,000 patients of the San Juan Hospital (Spain) from 2009 to 2017 [19]. Each X-ray has an associated report produced by a radiologist. From these reports, labels were extracted manually by trained physicians (for 27% of the X-rays) and automatically (for 73% of the X-rays) using a supervised method. The labels covered six different position views, 174 different radiographic findings, 19 differential diagnoses, and 104 anatomic locations. The labels were then organised into a hierarchical taxonomy and mapped to Unified Medical Language System (UMLS) Concept Unique Identifier (CUI) codes. For our work, the considered 254 classes were derived from the l a b e l C U I S , L o c a l i z a t i o n s C U I S , M o d a l i t y _ D I C O M , and the V i e w P o s i t i o n _ D I C O M labels described in [19,Table 11], where the Digital Imaging and Communications in Medicine (DICOM) fields were extracted from the X-ray. We found that 33 of the images were corrupt and were thus excluded. [18]. The studies were performed between October 2002 and July 2017.
An automatic system was used to extract 14 observations from the associated radiology reports (no finding, enlarged cardiom, cardiomegaly, lung lesion, lung opacity, edema, consolidation, pneumonia, atelectasis, pneumothorax, pleural effusion, pleural other, fracture, and support devices). Each observation class was rated as either positive, uncertain, or negative, thus resulting in 42 classes. 3β’ ChestX-ray14 contains 86,524 chest X-rays (collected from 1992 to 2015) concerned with common thorax diseases [20]. It consists of 15 disease labels (atelectasis, cardiomegaly, effusion, infiltration, mass, nodule, pneumonia, pneumothorax, consolidation, edema, emphysema, fibrosis, pleural thickening, and hernia). The labels were mined from the associated radiological reports of the X-rays. 4β’ The MURA training set comprises of 36,808 musculoskeletal radiographs from 13,457 studies of 11,184 patients, where each study is manually labelled by radiologists as either normal or abnormal [21]. Each radiograph concerns one of seven sections of the body (elbow, finger, hand, humerus, forearm, shoulder, or wrist), where each is classed as either normal or abnormal, resulting in 14 total classes. MURA is a multi-class classification task unlike the previous datasets. 5The ROCO dataset was used to pre-train the seq2seq models before fine-tuning, as depicted in Figure 2. ROCO comprises of image-caption pairs from PubMed Central articles, where compound, multi-pane, and non-radiology images were removed using an automatic system [22]. The ROCO training and validation sets contain 65,450 and 8,180 examples, respectively. The ROCO dataset contains several medical imaging modalities including computer tomography, ultrasound, X-ray, fluoroscopy, positron emission tomography, mammography, magnetic resonance imaging, and angiography. 6
The BLEU scores for each submission on the validation and test sets of the task's dataset are shown in Table 2. 7 Submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 attained the highest test score. This submission had the simplest configuration-no regularisation, no SCST, no ROCO pre-training, and no MIT pre-training. This indicates that the additional steps considered for the other configurations were not suited to the task's dataset. One possible explanation is the small size of the task's dataset.
Another observation is the large discrepancy between the validation and test scores. This indicates a significant difference between the examples of the two sets or that the submissions were overfitted to the training and/or validation sets. Additionally, the validation score was an inconsistent predictor of which submission would achieve the highest test score. While submission v i t 2 m r t -0 . 1 . 2 _ 2 _ e 4 6 attained the highest test score, it was outperformed by multiple submissions in terms of validation score. In fact, submission v i t 2 m r t -0 . 1 . 3 _ 5 _ e 3 attained the highest validation score, which had a configuration that employed ROCO pre-training and SCST, but no regularisation or MIT pre-training.
Using L2 regularisation during pre-training had a negative impact on performance, where both m i t 2 m r t -0 . 1 . 7 _ 1 _ e 0 and m i t 2 m r t -0 . 1 . 8 _ 1 _ e 1 produced worse validation and test scores than m i t 2 m r t -0 . 1 . 5 _ 1 _ e 1 . A regularisation term of 0.9 (mit2mrt-0.1.8_1_e1) was able to attain higher validation and test scores than a term of 0.5 (mit2mrt-0.1.7_1_e0); however, this could be due to submission m i t 2 m r t -0 . 1 . 8 _ 1 _ e 1 completing more epochs of fine-tuning, as shown in Figure 2.
The impact of the medical image multi-label classification task for MIT pre-training was inconclusive. Comparing submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 to m i t 2 m r t -0 . 1 . 9 _ 1 _ e 1 3 8 , using an MIT produced higher validation scores but lower test scores than ViT. Oppositely, employing an MIT over a ViT resulted in a lower validation score, but a higher test score-when comparing submission v i t 2 m r t -0 . 1 . 3 _ 5 _ e 3 to m i t 2 m r t -0 . 1 . 5 _ 1 _ e 1 . It should be emphasised that the medical images that largely make up the medical image multi-label classification task dataset are chest X-rays-whereas those in the task's dataset, as well as ROCO, are more varied in modality and location. This suggests that the datasets used for the medical image multi-label classification task were not suited to the subsequent stages of pre-training and fine-tuning depicted in Figure 2.
The impact of using ROCO to pre-train the seq2seq models before fine-tuning was also inconclusive. Comparing submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 to v i t 2 m r t -0 . 1 . 2 _ 3 _ e 9 1 , ROCO pre-training increased the validation score and decreased the test score. However, pre-training on ROCO substantially reduced the number of epochs until convergence during fine-tuning. Note that medical image captioning datasets derived from PubMed Central articles have been previously criticised for their significant amount of noise [4]. This could mean that pre-training on ROCO may be harmful to performance.
SCST has been effective for medical image captioning previously [3]. Here, we note that SCST is sensitive and can be difficult to attain stable training. Comparing submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 to v i t 2 m r t -0 . 1 . 4 _ 2 _ e 0 , SCST substantially decreased the validation and test scores. This was likely due to the learning rate being too high for this configuration. Oppositely, SCST improved both validation and test scores when comparing submission v i t 2 m r t -0 . 1 . 2 _ 3 _ e 9 1 to v i t 2 m r t -0 . 1 . 3 _ 5 _ e 3 , indicating that the learning rate was suitable for this configuration.
For ImageCLEFmed Caption 2021, the performance of submission v i t 2 m r t -0 . 1 . 1 _ 5 _ e 1 3 1 placed the AEHRC CSIRO team third-with a score 0.078 worse than the first placed participant. This indicates that utilising pre-trained Transformer checkpoints to initialise the encoder and decoder of a seq2seq model is a promising approach for this task. However, the impact of the selected pre-training data was unclear; pre-training the seq2seq model with ROCO produced inconclusive results. Instead of ROCO, an image-caption dataset derived from real medical images and their associated radiologists' reports-such as MIMIC-CXR [23]-is recommended.
The impact of the medical image multi-label classification task was also inconclusive, where its medical images were likely too dissimilar to those from ROCO and the task's dataset. The impact of SCST and L2 regularisation were clearer, with SCST providing a small improvement when configured correctly and the used L2 regularisation terms resulting in a decrease in performance. In future work, we aim to conduct a more thorough investigation of the proposed approach-to better adapt it to medical image captioning. At the same time, our overall approach has been intentionally kept general so that it might be applied to tasks other than medical image captioning.