Improving Chest X-Ray Report Generation by Leveraging Warm Starting

In plain words

A chest X-ray is a picture that doctors read to write a short report. A computer can try to write that report by itself. The usual design has two parts: one part that looks at the picture and one part that writes the sentences. Training both from nothing needs a lot of labelled data, and the medical datasets for this task are small.

This paper asks a practical question. Instead of starting from random numbers, can each part start from a model that someone already trained on other data? Starting from trained weights is called warm starting. The authors try several public picture models and several public text models, and they combine them in different pairs. They then measure how close the written reports are to the reports radiologists wrote, and how often the reports mention the right medical findings.

The best pair in their tests is a picture model with some built-in convolution, plus a small, distilled version of GPT-2 for the text. That pair scores higher than an earlier system on the measures the authors use. The authors also say the reports are still not accurate enough for clinical use.

TL;DR

• The best single checkpoint for each side is CvT-21 (ImageNet-21K) for the encoder and DistilGPT2 for the decoder, in the authors' one-at-a-time comparisons §8.4 §8.5. The final model combines them.
• Author-reported on the MIMIC-CXR test set with the Chen et al. splits, this model beats M2 Transformer Progressive on example-based CE F-1, BLEU-4, ROUGE-L and METEOR tab. 5 tab. 6.
• The method can be re-implemented from this card. The official code is GPL-3.0, and several paper and code values disagree (see Paper vs code).
• Diagnostic accuracy stays low for rare findings; the authors call for more work before clinical trials §9.

Problem

Radiologists write a report for each chest X-ray. Automatic report generation could ease that workload, but the authors report that current generators lack the diagnostic accuracy needed for clinical use §1 §2. The usual architecture is an encoder-to-decoder model: an encoder turns the image into visual features, and a decoder writes the report one subword at a time, conditioned on those features §1 fig. 1.

Two pieces are underexplored. Public CV and NLP checkpoints beyond ResNet, DenseNet and the first GPT-2-based decoder have rarely been tried for this task §3. Also, earlier warm-starting work did not compare against a randomly initialised decoder §3. Medical datasets such as MIMIC-CXR and IU X-Ray are small and of lower quality than general image-captioning data §1 §5.

Method

The paper is a systematic comparison of pretrained checkpoints, with one proposed final combination. It is based on the arXiv v2 LaTeX source. Table and figure numbers below follow source order and were not checked against a compiled PDF.

Task and data. Each example is one report with its image or images §6. MIMIC-CXR uses one image per example. IU X-Ray uses two views (frontal and lateral), which are both encoded and then combined §5 §6. The paper uses the splits and labels of Chen et al. 2020 for both datasets §5 tab. 2. Reports are cleaned by lowercasing, removing special characters, capping each report at 60 words, and mapping words seen fewer than three times in the training set to an unknown token §7.4.

Encoder. The encoder is a CV checkpoint whose last hidden state gives S spatial positions, each with F features §6. Where a checkpoint stores features along the other axis, the state is transposed §6. A linear projection maps each position's F features to the decoder's hidden size. The projection is randomly initialised §6.

Decoder. The decoder is an NLP checkpoint with a language-model head. Each decoder layer gets an extra attention block, with random starting weights, placed after its self-attention and before its feed-forward part §6. Visual features enter through this cross-attention §6. The beginning token [BOS] starts generation fig. 5.

Final model (CvT-21²DistilGPT2). The encoder is CvT-21 with ImageNet-21K weights at 384×384 input. The decoder is DistilGPT2, cased, with 50K WebText vocabulary tab. 3 tab. 4 fig. 5. The figure shows CvT-21 with a convolutional token embedding and, in its last stage, convolutional layers in place of the linear query, key and value projections fig. 5.

Pre-processing and augmentation §7.3:

• The image is resized with bilinear interpolation. Its shorter edge becomes W+64, where W is the encoder's input width, and the aspect ratio is kept.
• It is cropped to W×H (square). The crop is random during training and centred at test time.
• During training only, the image is rotated about its centre by an angle drawn uniformly from −5° to +5°.
• Each channel is standardised with the mean and standard deviation of its checkpoint.
• Greyscale images are repeated to three identical channels, and the pixel depth is 8 bit §7.3 tab. 2.

Training objective and optimiser §7.5:

• Teacher forcing with categorical cross-entropy over report subwords.
• AdamW with a learning rate of 1e-5 for the encoder and 1e-4 for all other parameters, following Chen et al. The other AdamW settings are the defaults.
• Batch size 16, mixed precision, PyTorch 1.9.0, four 16 GB P100 GPUs.
• Model selection: the epoch with the highest validation CIDEr. Early stopping has patience 10 epochs and minimum delta 1e-4 on that same metric.

# Training step, teacher-forced  [§6] [§7.5]
# Input:  images (B, 3, 384, 384); report_ids (B, T) as subword ids
# Output: scalar loss

def training_step(images, report_ids):
    # 1. Encode each CXR with the warm-started CvT-21  [tab. 3]
    feature_map = encoder(images)               # (B, 384, H', W')
    # 2. Flatten the spatial grid into S positions  [§6]
    positions = flatten_spatial(feature_map)    # (B, S, 384)
    # 3. Project each position to the decoder width  [§6]
    memory = projection(positions)              # (B, S, 768)
    # 4. Decoder input is every token but the last  [dataset.py#L106-L109]
    decoder_input = report_ids[:, :-1]          # (B, T-1)
    # 5. Labels are every token but the first  [dataset.py#L106-L109]
    targets = report_ids[:, 1:]                 # (B, T-1)
    logits = decoder(decoder_input, cross_attend_to=memory)
    # 6. Cross-entropy, ignoring padding  [§7.5]
    loss = cross_entropy(logits, targets, ignore_pad=True)
    return loss

Decoding. Validation uses greedy search (beam size 1). Test uses beam search with beam size 4. Generation is capped at 128 subwords §7.4.

Warm-start study design §8.4 §8.5 §8.6:

• Encoder comparison: each CV checkpoint is paired with DistilGPT2, eight training runs each, on a 50K-example subset of the MIMIC-CXR training set.
• Decoder comparison: each NLP checkpoint is paired with a ResNet-101 ImageNet-1K encoder, eight runs each, on the same subset.
• The paper assumes the best encoder does not depend on the decoder, and the reverse, and does not test combinations §8.4 §8.5 §9.
• The 50K subset is the smallest size tested that did not differ significantly from the full training set (§8.6, using ResNet-101 with DistilBERT).

Evaluation §7.2 §7.6: NLG metrics (BLEU-1 to BLEU-4, METEOR, ROUGE-L, CIDEr), and clinical efficacy (CE) metrics. CE labels come from CheXbert, which assigns each of 14 CheXpert observations as positive, negative, uncertain or not mentioned. Only "positive" counts as positive. The other three classes count as negative §7.2. Confidence intervals come from 1,000 bootstrap resamples of the test set. Group comparisons use Welch's ANOVA, then Games-Howell post hoc tests, at p = 0.05 §7.6.

Checkpoint findings (the paper's conclusions, not separate method steps):

• Encoder: CvT-13 has the highest mean BLEU-4; CvT-21 has the highest mean METEOR and CIDEr; BEiT has the highest mean ROUGE-L; ResNet-101 has the highest macro precision §8.4. CvT-21 and DeiT could not be separated on the NLG metrics, so CvT-21 was chosen for its smaller size (31.6M vs 86M parameters) §8.4.
• Decoder: DistilGPT2 has the highest mean BLEU-4, METEOR, ROUGE-L and CIDEr. GPT2 and DistilGPT2 are both significantly ahead of the other checkpoints on METEOR and CIDEr, with no significant difference between them. DistilGPT2 was chosen for being smaller §8.5.

Hyperparameters

Values as stated in the paper. Where the code differs, it is noted in Paper vs code.

Results

Author-reported. Setup: MIMIC-CXR test set (3,858 examples) with the labels and splits of Chen et al. 2020; IU X-Ray test set (590 examples); 5-run means with the paper's ± values as printed. The ± values come from the bootstrap §7.6; the paper does not separately define them beyond that. The baseline is M2 Transformer Progressive, as cited by the paper tab. 5 tab. 6.

MIMIC-CXR, NLG metrics tab. 5 (the * row is the single run with the highest validation CIDEr):

BLEU-1 for the n=5 mean is 0.3918 and for the * run 0.3928; the M2 baseline has 0.3780 tab. 5.

MIMIC-CXR, example-based CE tab. 6:

MIMIC-CXR, label-based CE, averaged over the 14 observations tab. 8: macro precision 0.256, recall 0.382, F-1 0.307; micro precision 0.398, recall 0.497, F-1 0.442. The table does not say which run these come from. Support devices, pleural effusion and cardiomegaly score well. Consolidation, enlarged cardiomediastinum, lung lesion, fracture and pleural other score poorly §8.2 tab. 8.

IU X-Ray, NLG metrics tab. 7:

Headline claims (abstract and conclusion, author-reported): CE F-1 +8.3%, BLEU-4 +1.8%, ROUGE-L +1.6%, METEOR +1.0% over M2 Transformer Progressive §10. These do not reproduce with one convention; see Pitfalls.

Paper vs code

Checked against the official repository at c9edfb5 (GPL-3.0). Code lines are cited as they appear in that commit.

• Learning rates differ. The paper states 1e-5 for the encoder and 1e-4 for other parameters §7.5. The MIMIC-CXR and IU X-Ray configs set encoder_lr: 5e-5 and decoder_lr: 5e-4 train_mimic_cxr_chen_cvt2distilgpt2.yaml train_iu_x_ray_chen_cvt2distilgpt2.yaml. The optimiser puts the projection and the decoder in the decoder group cvt2distilgpt2_mimic_cxr_chen.py.
• Report length. The paper caps reports at 60 words §7.4. The code keeps 60 token ids including a boundary token at each end and the period tokens that separate sentences, then decodes from the second id. In effect a report is cut at about 59 tokens, and periods count cvt2distilgpt2_mimic_cxr_chen.py mimic_cxr_chen_tokenizer.py.
• Rare-word threshold matches: a word is kept if it occurs at least three times in the training reports mimic_cxr_chen_tokenizer.py §7.4.
• Projection has no bias. The code uses a bias-free linear layer from 384 to 768 encoder_projection.py. The paper does not say whether the projection has a bias.
• Special tokens. The code adds [BOS] and [PAD] to the GPT-2 tokenizer and resizes the decoder's embeddings to the vocabulary size plus two cvt2distilgpt2_mimic_cxr_chen.py. The paper does not mention these tokens beyond the [BOS] in the figure fig. 5.
• Cross-attention is added and randomly initialised, as the paper says cvt2distilgpt2_mimic_cxr_chen.py.
• Image pipeline matches: resize to 384+64, random crop and ±5° rotation during training only, centre crop at test time cvt2distilgpt2_mimic_cxr_chen.py. The code normalises with the ImageNet mean and standard deviation cvt2distilgpt2_mimic_cxr_chen.py.
• Decoding matches: greedy in validation cvt2distilgpt2_mimic_cxr_chen.py; beam size 4 at test cvt2distilgpt2_mimic_cxr_chen.py train_mimic_cxr_chen_cvt2distilgpt2.yaml; maximum 128 subwords train_mimic_cxr_chen_cvt2distilgpt2.yaml.
• Early stopping matches: patience 10, minimum delta 1e-4, monitoring validation CIDEr train_mimic_cxr_chen_cvt2distilgpt2.yaml train_mimic_cxr_chen_cvt2distilgpt2.yaml.
• Epoch limit: the repository sets no maximum number of epochs; the trainer is built by an external framework stages.py. The paper does not state one either.
• Seeds: each run is seeded with its trial number stages.py. The paper does not state how seeds were chosen.
• Test numbers differ slightly from the paper. The README reports test metrics for one released checkpoint, epoch 8 with validation CIDEr 0.425 README.md. Its test CIDEr is 0.3899 README.md, close to the paper's * row (0.3892), but not identical.
• Example-based precision and recall look swapped. The README lists example-based precision as 0.4183 and recall as 0.3670 README.md README.md. Table 6 gives precision 0.3670 and recall 0.4184 for the * run tab. 6. The paper's F-1 (0.3910) is the harmonic mean of its own values; the README F-1 is 0.3666 README.md, which is not the harmonic mean of its own precision and recall. The definition of F-1 is not stated in the paper.
• IU X-Ray configs are identical to MIMIC-CXR apart from the module name train_iu_x_ray_chen_cvt2distilgpt2.yaml. The two-view handling in the IU module was not reviewed in this compile.
• Freezing exists as an option in the CvT wrapper, off by default cvt.py cvt.py. The paper does not use it.
• Decoder inputs and labels: the decoder input is the report without its last token, and the label is the report without its first token, with [BOS] prepended and the end token appended dataset.py dataset.py.
• Loss: cross-entropy ignoring the padding token cvt2distilgpt2_mimic_cxr_chen.py.

Pitfalls

• Learning rates. Use the paper's values only if you match its results, which the code does not do. Either choice must be stated in your own implementation. The config values are 5× and 5× larger than the paper's §7.5 train_mimic_cxr_chen_cvt2distilgpt2.yaml.
• Report-length cap. Implement the code's behaviour (about 59 tokens, periods included) if you want to reproduce the reported scores. Implement 60 words if you follow the text. The two are not the same.
• Headline percentages do not follow from the tables. Taking the * run in tab. 5 and tab. 6 as absolute differences from M2 Transformer Progressive: example CE F-1 +0.083, BLEU-4 +0.020, ROUGE-L +0.014, METEOR +0.010. The paper states +8.3%, +1.8%, +1.6% and +1.0%. The n=5 mean gives BLEU-4 +0.017 and ROUGE-L +0.013. No single convention reproduces all four. Use the table values.
• IU X-Ray claim. The text says CvT-21²DistilGPT2 has the highest mean BLEU-3, BLEU-4, METEOR and CIDEr on IU X-Ray §8.1. Only the * run beats M2 on BLEU-3 and BLEU-4. The n=5 mean is below M2 on BLEU-3 (0.2141 vs 0.2240) and BLEU-4 (0.1650 vs 0.1690) tab. 7.
• Confidence intervals look very narrow. Some printed ± values are below 0.001 (for example BLEU-1 ± 0.00008). Check whether they describe the variation across the five runs, the bootstrap, or something else before using them to judge significance.
• Final model's training set. The paper says the 50K subset was used for the checkpoint comparisons. It does not state the training-set size of the final model in tab. 5 §8.6.
• Mean versus selected run. The n=5 row is the mean of five runs. The * row is the single run with the highest validation CIDEr tab. 5. The * row is a selected run, so it is not a comparable estimate of the method's average.
• Clinical accuracy is low for rare findings. Several observations score under 0.2 F-1 and the authors name class imbalance as a likely cause §8.2 tab. 8. The authors also state that ROC curves were not included §9.
• Missing inputs. The code uses the CheXbert labeller with the checkpoint path stanford/chexbert/chexbert.pth and bert-base-uncased cvt2distilgpt2_mimic_cxr_chen.py. The checkpoint is not in the repository, and its licence was not checked here.
• Resolution and format. Images are downsampled to 384×384 from 8-bit JPEG and PNG rather than 12-bit DICOM. The authors list this as a limitation §9.
• Licence of the code. The repository is GPL-3.0. Reading it to understand the method is fine, but copying code into a non-GPL project would bring GPL obligations. Re-implement from the description above and the cited lines.
• Joint choice untested. Each side was picked with the other side fixed to one checkpoint §8.4 §8.5. The authors say this assumption may be wrong §9. The best encoder and decoder together were not tested.

Lineage

• Builds on: Chen et al. 2020 for the MIMIC-CXR and IU X-Ray splits, labels and report cleaning §5 §7.4. The paper's footnote links those subsets to the R2Gen repository §5.
• Builds on: an earlier study that used DistilGPT2 to warm start the decoder without a randomly initialised comparison §3.
• Compared against: R2Gen, CMN, Contrastive Attention, PPKED and M² Transformer Progressive, using the numbers reported in their papers tab. 5 tab. 6 tab. 7 tab. 8.
• Superseded by: not known from this wiki.