METransformer: Radiology Report Generation by Transformer with Multiple Learnable Expert Tokens

Our ViT encoder adopts vision transformer [9]. In addition to the common input of visual patches/tokens embedding and position embedding, we further introduce expert tokens embedding and segment embedding.

Visual patches Embedding. Given an image $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$, followed  [9], we reshape $\mathbf{x}$ into a sequence of flattened 2D patches $\mathbf{x}_{p}\in\mathbb{R}^{N\times(P^{2}\dot{C})}$, where $(H,W)$ is the resolution of the original image, $C$ is the number of channels, $(P,P)$ is the resolution of each image patch, and $N=HW/P^{2}$ is the resulting number of patches. We consider $\mathbf{x}_{p}$ as visual tokens and $N$ is the input sequence length.

Expert tokens Embedding. In addition to $N$ visual tokens, we post-pend $M$ learnable embeddings $\mathbf{x}_{e}\in\mathbb{R}^{M\times D}$ which have the same dimension as the visual token embeddings and are called expert tokens. We introduce an Orthogonal Loss as in Eqn. 6 to enforce orthogonality among the expert token embeddings, encouraging different expert tokens to attend different image regions.

Segment Embedding. We introduce two types of the segment, “[0]” and “[1]”, to separate the input tokens from different sources [29], i.e., “[0]” for visual tokens and “[1]” for expert tokens. A learned segment embedding $\mathbf{E}_{seg}$ is added to each input token to indicate its segment type.

Position Embedding. A standard 1D learnable position embedding $\mathbf{E}_{pos}$ is added to each input token to indicate its order in the input sequence.

Model Structure. The expert ViT encoder adopts the same structure as the standard Transformer [30], which consists of alternating layers of multi-headed self-attention (MSA) and Multi-Layer Perceptron blocks (MLP). Layernorm(LN) [2] is applied before every block and residual connections are applied after every block. Mathematically:

where $\mathbf{E}\in\mathbb{R}^{(P^{2}C)\times D}$ is a learnable matrix parameter to map visual tokens to a fixed dimension D. The subscript $l=1\dots L$ and $L$ is the total number of the transformer layers. $\mathbf{E}_{pos}\in\mathbb{R}^{(N+M)\times D}$ is the position embeddings and $\mathbf{E}_{seg}\in\mathbb{R}^{(N+M)\times D}$ is the segment embeddings.

The output of our expert ViT encoder, $\mathbf{z}_{L}\in\mathbb{R}^{(N+M)\times D}$, combines the visual token embeddings $\mathbf{z}_{L}^{v}=\mathbf{z}_{L}\left[:\mathrm{N}\right]$ and the expert token embeddings $\mathbf{z}_{L}^{e}=\mathbf{z}_{L}\left[\mathrm{N}:(\mathrm{N}+\mathrm{M})\right]$. They have attended to each other using multi-head self-attention, which is linear attention. Considering the fine-grained nature of medical images and the effectiveness of high-order attention in fine-grained recognition [20], we further enhance both embeddings by expert bilinear attention.

Expert Bilinear Attention. We followed  [25] to design our expert bilinear attention (EBA) module. As shown in Figure. 1 Left, we take the enhanced expert token embeddings $\mathbf{z}_{L}^{e}$ as query $\mathbf{Q}\in\mathbf{R}^{M\times 1\times D_{q}}$, the enhanced visual tokens embeddings $\mathbf{z}_{L}^{v}$ as key $\mathbf{K}\in\mathbf{R}^{1\times N\times D_{k}}$ and value $\mathbf{V}\in\mathbf{R}^{1\times N\times D_{v}}$, a low-rank bilinear pooling [14] if first performed to obtain the joint bilinear query-key $\mathbf{B}_{k}$ and query-value $\mathbf{B}_{v}$ by $\mathbf{B}_{k}=\sigma(\mathbf{W}_{k}\mathbf{K})\odot\sigma(\mathbf{W}_{q}^{k}\mathbf{Q})$ and $\mathbf{B}_{v}=\sigma(\mathbf{W}_{v}\mathbf{V})\odot\sigma(\mathbf{W}_{q}^{v}\mathbf{Q})$, where $\mathbf{W}_{k}\in\mathbf{R}^{D_{B}\times D_{k}}$, $\mathbf{W}_{v}\in\mathbf{R}^{D_{B}\times D_{v}}$, $\mathbf{W}_{q}^{k}\in\mathbf{R}^{D_{B}\times D_{q}}$ and $\mathbf{W}_{q}^{v}\in\mathbf{R}^{D_{B}\times D_{v}}$ are learnable parameters, resulting $\mathbf{B}_{k}\in\mathbb{R}^{M\times N\times D_{B}}$, $\mathbf{B}_{v}\in\mathbb{R}^{M\times N\times D_{B}}$. $\sigma$ denotes ReLU unit, and $\odot$ represents Hadamard Product. Afterward, we compute attention for input tokens both spatial and channel-wisely. 1) Spatial-wise attention. We use a linear layer to project ${\mathbf{B}_{k}}$ into an intermediate representation ${\mathbf{B}_{mid}}=\sigma(\mathbf{W}_{B}^{k}\mathbf{B}_{k})$, where $\mathbf{W}_{B}^{k}\in\mathbb{R}^{D_{B}\times D_{mid}}$. Then another linear layer is applied to mapping $\mathbf{B}_{mid}\in\mathbb{R}^{M\times N\times D_{mid}}$ from the dimension $\mathbf{D}_{mid}$ to 1 and further followed a softmax function to obtain the spatial-wise attention weight $\alpha_{s}\in\mathbb{R}^{M\times N\times 1}$. 2) Channel-wise attention. A squeeze-excitation operation [10] to ${\mathbf{B}_{mid}}$ is performed to obtain channel-wise attention $\beta_{c}=\mathrm{sigmoid}(\mathbf{W}_{c}\bar{\mathbf{B}_{mid}})$, where $\mathbf{W}_{c}\in\mathbf{R}^{\mathbf{D}_{mid}\times\mathbf{D}_{B}}$ is learnable parameters and $\bar{\mathbf{B}_{mid}}\in\mathbf{R}^{M\times\mathbf{D}_{mid}}$ is the average pooling of ${\mathbf{B}_{mid}}$. The first layer of EBA is formulated as,

Bilinear Encoder Layer. Besides EBA, we also use the “Add & Norm” layer with the residual connection as in a standard transformer in our bilinear encoder layer. The $n$-th layer bilinear encoder is expressed as,

where $n=1\dots N$ and $N$ is the number of bilinear encoder layer, $\mathrm{LN}(\cdot)$ denotes layer normalization [2], $\mathbf{W}_{e}^{n}\in\mathbb{R}^{(D_{q}+D_{v})\times D_{v}}$ is learnable parameters and $[;]$ denotes concatenation. We set $\hat{\mathbf{z}}_{L}^{e(0)}=\mathbf{z}_{L}^{e}$ and $\hat{\mathbf{z}}_{L}^{v(0)}=\mathbf{z}_{L}^{v}$ when n=1.

The output of the last expert bilinear encoder, expert token embeddings $\hat{\mathbf{z}}_{L}^{e(N)}\in\mathbb{R}^{N\times D_{B}}$ and visual tokens embeddings $\hat{\mathbf{z}}_{L}^{v(N)}\in\mathbb{R}^{M\times D_{B}}$, are sent to the decoder for report generation. The bilinear decoder layer comprises an $\mathrm{EBA}_{mask}$ layer to compute attention of the shifted-right reports and an $\mathrm{EBA}_{cross}$ to compute the cross-modal attention and an adjust block to regulate the word and visual tokens embeddings by expert token embeddings. For convenience, we denote $\hat{\mathbf{z}}_{L}^{e(N)}$ and $\hat{\mathbf{z}}_{L}^{v(N)}$ as $\mathbf{f}_{e}$ and $\mathbf{f}_{v}$.

Adjust block. To incorporate the expert tokens into the report generation process, we propose an expert adjustment block that allows each expert token embedding to influence the embedding of the words and visual tokens, thereby generating the report associated with that expert token. Since our expert tokens are trained orthogonally, we can generate discrepant reports by the different expert tokens. To regulate visual tokens embeddings $\mathbf{f}_{v}$ by expert token embeddings $\mathbf{f}_{e}$, the adjust block is calculated as follows,

where $\mathbf{W}_{e}$ and $\mathbf{W}_{v}$ are learnable parameters. $\sigma$ denotes ReLU unit, and $\odot$ represents Hadamard Product.

Bilinear Decoder Layer. We first perform a masked EBA to word embeddings $\mathbf{E}_{r}$ and then followed another EBA block to compute the cross attention of word embeddings and visual tokens embeddings. We also employ residual connections around each EBA block similar to the encoder, followed by layer normalization. Mathematically, the $i$-th layer bilinear decoder can be expressed as,

where $i=1\dots I$ and $I$ represents the total number of decoder layer. $\hat{\mathbf{E}}_{r}^{(i-1)}=\mathrm{F_{adjust}}(\mathbf{f}_{e},\mathbf{E}_{r}^{(i-1)})$ and $\hat{\mathbf{f}}_{v}=\mathrm{F_{adjust}}(\mathbf{f}_{e},\mathbf{f}_{v})$. $\mathbf{W}_{d}^{i}\in\mathbb{R}^{(D_{r}+D_{r})\times D_{r}}$ is learnable parameters and $[;]$ denotes concatenation. Specifically, $E_{r}^{(0)}\in\mathbb{R}^{T\times D_{r}}$ is the original word embeddings where T is the total number of words in the reports and $D_{r}=D_{B}$ is the dimension of word embedding, and we extend it to M replicates corresponding to M expert tokens to compute parallelly. The final output of the decoder is $E_{c}^{I}\in\mathbb{R}^{M\times T\times D_{r}}$, which will be further used to predict word probabilities by a linear projection and the softmax activation function.

Orthogonal Loss. To encourage orthogonality among expert token embeddings, we introduce an orthogonal loss term to the output of expert bilinear encoder $\hat{\mathbf{z}}_{L}^{e}$,

where $\ell(\cdot)$ denotes $L_{2}$ normalization and $I$ represents the identity matrix with dimension M.

Report Generation Loss. We train our model parameters $\theta$ by minimizing the negative log-likelihood of $\mathbf{P}(t)$ given the image features:

where $\mathbf{P}({\mathbf{t}}_{i}^{(m)}|{\mathbf{I}},{\mathbf{t}}_{i-1}^{(m)},\cdots,{\mathbf{t}}_{1}^{(m)})$ represents the probability predicted by the $m$-th expert tokens for the $i$-th word $t_{i}$ based on the image $\mathbf{I}$ and the first $(i-1)$ words.

Our overall objective function is: ${\mathcal{L}}_{all}={\mathcal{L}}_{CE}+\lambda{\mathcal{L}}_{OrL}$. The hyper-parameter $\lambda$ simply balances the two loss terms, and its value is given in Section 4.

For the $M$ diagnostic reports $R=\left[r_{1},r_{2},\dots,r_{M}\right]$ produced by $M$ experts, we design a metric-based expert voting strategy to select the optimal one, where the “metric” means the conventional natural language generation (NLG) metrics, such as BLEU-4 and CIDEr (we use CIDEr in our paper). The voting score $\mathbf{S}_{i}$ for $i$-th expert’s report can be calculated by the following equation:

where $\mathrm{CIDEr(,)}$ denotes the function for computing CIDEr with $\mathbf{r}_{i}$ as candidate and $\mathbf{r}_{j}$ as reference. In this way, each expert’s report can get a vote score, indicating the degree of consistency of the report with that of other experts.The diagnostic report with the highest score is the winner of the voting. As demonstrated in Table. 3, our voting strategy is more effective than the commonly used ensemble/fusion methods [25]. The possible reasons are that i) our method utilizes NLG metrics as reference scores, thereby the voted results are directly related to the final evaluation metrics; ii) since we ultimately select a single result with the highest score, we can produce a more coherent report than fusing multiple results at the word level.

In this experiment, two datasets are used for the performance evaluation, i.e., a widely-used benchmark IU-Xray [8] and the currently largest dataset MIMIC-CXR [13] for medical report generation.

IU-Xray Indiana University Chest X-ray Collection (IU-Xray) [8] is the most widely used publicly accessible dataset in medical report generation tasks. It contains 3,955 fully de-identified radiology reports, each of which is associated with frontal and/or lateral chest X-ray images, and 7,470 chest X-ray images in total. Each report is comprised of several sections: Impression, Findings, Indication, etc. In this work, we adopt the same data set partitioning as [5] for a fair comparison, with a train/test/val set by 7:1:2 of the entire dataset. All evaluations are done on the test set.

MIMIC-CXR The recently released MIMIC-CXR [13] is the largest public dataset containing both chest radiographs and free-text reports. In total, it consists of 377110 chest x-ray images and 227835 reports from 64588 patients of the Beth Israel Deaconess Medical Center examined between 2011 and 2016. In our experiment, we adopt MIMIC-CXR’s official split following the work [5] for a fair comparison, resulting in a total of 222758 samples for training, and 1808 and 3269 samples for validation and test.

Evaluation Metrics  Following the standard evaluation protocol¹¹1https://github.com/tylin/coco-caption, we utilize the most widely used BLEU-4 [26], METEOR [3], ROUGE-L [19], and CIDEr [31] as the metrics to evaluate the quality of the generated diagnostic reports. To measure the accuracy of descriptions for clinical abnormalities, we follow [5, 4, 22] and further report clinical efficacy metrics. For this purpose, the CheXpert [11] is applied to labeling the generated reports and the results are compared with ground truths in 14 different categories related to thoracic diseases and support devices. We use precision, recall, and F1 to evaluate model performance for clinical efficacy metrics.

Implementation Details  For IU-Xray, we use image pairs for report generation as [5]. For both datasets, we use the ”bert-base-uncased” model’s tokenizer in huggingface transformer [36] to tokenize all words in the reports. We utilize a pre-trained vision transformer with a patch size of 32 to initialize our Expert ViT encoder. The number of layers of the expert bilinear encoder and decoder is set as (2, 4) for MIMIC-CXR and (2,2) for IU-Xray to reduce the potential overfitting on IU-Xray due to its relatively small data size. The dimensions of the bilinear query-key representation and the transformed bilinear feature ($D_{B}$ and $D_{mid}$) in the expert bilinear attention block is set as 768 and 384, respectively. The hyper-parameter $\lambda$ is set to 2. Our model is trained using Radam optimizer [23] with a mini-batch size of 16. The learning rate is set to be 0.0001 and the model is trained for a total of 20 epochs. We implement our model using Pytorch [27] and Pytorch-lightning library ²²2https://github.com/Lightning-AI/lightning with two NVIDIA GeForce RTX 3090 GPU cards.

Two types of metrics are used in our evaluation: the conventional natural language generation (NLG) metrics and the clinical efficacy (CE) metrics. The results are reported in Table 1 and Table. 2 ³³3It is noted that clinical efficacy metrics only apply to MIMIC-CXR because the labeling schema of CheXpert is designed for MIMIC-CXR., respectively.

Specifically, we compare METransformer with 5 state-of-the-art (SOTA) image captioning methods, including Show-tell [32], AdaAtt [24], Att2in [1], Transformer [7] and M2transformer [6]. For these methods, we use a publicly released codebase [18] and re-run them on both datasets with the same experimental setting as ours. Moreover, eight SOTA medical report generation models are involved in the comparison, including CoAtt [12], HGRN-Agent [17], R2Gen [5], R2GenCMN [4], PPKED [21], GSK [41] and MSAT [34]. It is noteworthy that except R2Gen [5] and R2GenCMN [4], these methods do not have their source code released. CoAtt [12] and HGRN-Agent [17] only report results on IU-Xray in their original paper. For the others, we cite the results from their respective papers. Please note that since IU-Xray does not provide an official training-test partition, the cited performances of these methods on IU-Xray (except R2Gen [5] and R2GenCMN [4] that use the same partition as ours) are not strictly comparable, and they are provided here only for reference. Differently, on MIMIC-CXR, since all these models follow the MIMIC-CXR official training-test partition, their cited performances are comparable.

As shown in Table 1, on both datasets, our METransformer consistently outperforms those “single-expert” based models, including attention-based baselines (Att2in [37], AdaAtt [24]), memory-augmented methods (R2Gen [5], R2GenCMN [4]) and models introducing external information (PPKED [21], MSAT [34]). On MIMIC-CXR, METransformer is the best performer across all metrics. Especially, our CIDEr score is up to 0.362, which is to date the best performance and makes a significant improvement over the second best score of 0.299 of MSAT [34]. These improvements demonstrate the advantage of our framework which is conceptually ”multiple expert co-diagnosis”. Ours METransformer also surpasses these methods on IU-Xray over most metrics, while it is again worth mentioning that the cited performances on IU-Xray may not be strictly comparable as obscure training-test partitions were used by different methods.