Dia-LLaMA: Towards Large Language Model-driven CT Report Generation

The overall architecture is shown in Figure 1. For introducing LLM to report generation, we utilize combined prompts to aggregate the visual embedding and critical diagnostic information. Our designed prompts consist of two segments: $\mathcal{P}=\{\mathcal{S},\mathcal{D}\}$, where the first segment $\mathcal{S}=\{s_{1},s_{2},\ldots,s_{N}\}$ represents a fixed number of special tokens $s_{n}$ for visual embeddings and the second segment $\mathcal{D}=\{d_{1},d_{2},\ldots,d_{L}\}$ represents diagnostic prompt tokens. Let $\mathcal{R}=\{r_{1},r_{2},\ldots,r_{T}\}$ denotes a generated report, where $r_{t}$ represents the token at timestep $t$ and $T$ is the length of report. The decoding process of the LLM $f_{l}$ is described as follows:

where $\mathcal{R^{-}}$ represents the generated report at timestep $t-1$. The report generation process is optimized by minimizing the language modeling loss $\mathcal{L}_{LM}$:

For extracting visual embeddings, the $i_{th}$ CT volume $V_{i}$ is encoded into patch features by a vision encoder $f_{v}$ and subsequently projected into the embedding space of the LLM by a perceiver $f_{p}$:

where $A_{i}^{m}\in\mathbb{R}^{c}$ represents a patch feature, $X_{i}^{n}\in\mathbb{R}^{d}$ represents visual embedding, $c$ and $d$ denote the feature and embedding dimension, $M$ and $N$ represent the number of patch features and visual embeddings, respectively. The visual embeddings $X_{i}$ are integrated into the embedding layer of the LLM.

For deriving diagnostic information, we first utilize disease-aware attention (Section 2.2) to gather disease-level features $D_{i}$ from patch features $A_{i}$. To provide a typical reference for diagnoses, we construct a disease prototype memory bank (Section 2.3) to capture the common representations of various diseases. The diagnostic results can be obtained by comparing disease features and prototypes, then interpreted into diagnostic text prompts (Section 2.4) for LLM.

Employing average-pooled patch features to diagnose various diseases may lead to unreliable diagnosis due to mixed disease information. To alleviate this issue, we propose a disease-aware attention (DAA) module to extract disease-level features from patch features. Specifically, we assign a learnable attention weight for each disease. The patch features from the vision encoder $f_{v}$ are element-wise multiplied with the attention weights and are then aggregated to form the disease-level features. The process can be expressed as:

where $D_{i}\in\mathbb{R}^{L\times c}$ represents the aggregated disease features, $\mathbf{W}_{D}\in\mathbb{R}^{L\times M\times 1}$ denotes the disease-aware attention weights, and $A_{i}\in\mathbb{R}^{1\times M\times c}$ encapsulates the patch features. The disease features $D_{i}$ are then utilized for disease classification, which requires distinguishing abnormal and normal samples.

Due to the biased nature of certain diseases, some abnormalities are relatively rare. To improve the diagnostic accuracy for infrequent abnormalities, we introduce a disease prototype memory bank (DPM) as a reference during diagnosis. The diagnostic results are obtained by comparing the similarity between disease-level features and a set of learnable prototypes. Specifically, the DPM includes both abnormal prototypes $\mathbf{P}_{1}^{l}$ and normal prototypes $\mathbf{P}_{0}^{l}$ to capture the presence and absence of each disease, respectively. These prototypes are updated through the InfoNCE loss [15], which pulls the positive pairs closer and pushes the negative pairs farther. In our case, the positive case $\mathbf{P}_{y_{i}^{l}}^{l}$ and negative case $\mathbf{P}_{1-y_{i}^{l}}^{l}$ are determined based on the disease label $y_{i}^{l}$. The contrastive disease-prototype loss $\mathcal{L}_{DP}$ is defined as

where $y_{i}^{l}$ represents the label of the $l_{th}$ disease, and $\tau$ is the learnable temperature parameter.

It is essential for medical reports to precisely capture abnormal information [19]. Despite the strong capabilities of LLM, directly recognizing abnormalities from visual embeddings without additional guidance is still challenging, which is validated in Section 3.3.1. Therefore, we introduce diagnostic text prompts (DTP), leveraging the diagnostic results as guidance prompts. Specifically, the diagnostic results are converted into text prompts $\mathcal{D}$, which follows a template description “The {disease name} is [disease state]”. For instance, the diagnostic result $c_{1}$: Present in Figure 1 is interpreted as The enlarged cardio mediastinum is present in this image, where $c_{1}$ represents the enlarged cardio mediastinum disease.

The overall training loss of our model is expressed as the weighted sum of the disease-prototype loss $\mathcal{L}_{DP}$ and the language modeling loss $\mathcal{L}_{LM}$:

where $\lambda$ represents the weight adjustment factor.

We adopted a large-scale CT report dataset (CTRG-Chest-548K [19]) to evaluate our method and the compared methods. This dataset comprises 1,804 CT-report pairs. Adhering to the original split ratio [19], we randomly selected 80% of the data for training and 20% for testing. Following the previous works [8, 23], we utilized a pretrained report labeler called CheXbert [18] to extract labels. Despite being pre-trained on the CXR dataset (MIMIC [9]), CheXbert remains effective in our experiments, attributed to the similar content between the chest CT and CXR report. There are 14 diseases recognized by CheXbert, with each having four states: present, absent, uncertain, and blank. Due to the clinical focus on present diagnosis, we categorized all other states as absent.

For evaluation, both NLG and CE metrics are adopted. NLG metrics include BLEU [16], METEOR [4], and ROUGE-L [11]. Following the CE metrics setting in [14, 8], we assess Precision, Recall, and F1 score with CheXbert [18].

For the compared methods in CXR, all settings are consistent with the original paper. We selected 30 CT scans at specific intervals for each sample. For RadFM [22] and our method, the pre-trained ViT3D [22] is adopted as the vision encoder and each volume is resized to $256\times 256\times 64$ as the input. We selected the LLaMA2-7B [21] as the LLM in all our experiments and utilized LoRA [5] for parameter-efficient fine-tuning, where trainable parameters are only 0.06% of the whole parameters. During training, we utilized AdamW [13] as the optimizer, with an initial learning rate of 5e-5, following a constant learning rate schedule that includes a warmup phase. The model was trained on two RTX 3090 GPUs for about 16 hours, built with PyTorch 2.0. The training involved 2000 steps, with an effective batch size of 16. The factor $\lambda$ was set to 4 to balance the two types of loss. To optimize memory usage, we employed the ZeRO [17] stage 2 training strategy in conjunction with gradient checkpointing [1].

Due to the limited works in CTRG, we compared our method with SOTA methods in chest X-ray (CXR), including R2Gen [3], R2GenCMN [2], M2KT [23], and PromptMRG [8]. In addition, we also compared with a CT report generation work SL-DG [19] and a generalist model RadFM [22]. To ensure a fair comparison, we matched the LLM in RadFM with that used in our experiments.

The Table 1 shows the comparison results on CTRG-Chest-548K [19] dataset. We observed that the proposed method achieves SOTA performance across the three CE metrics and the majority of NLG metrics. For CE metrics, we achieved a 0.372 F1 score, which represents a 7.8% improvement compared to the RadFM. As for the Precision and Recall metrics, we obtained 4.5% and 7.2% improvements compared to the second-best results. In terms of NLG metrics, our method also achieved SOTA performance. Regarding the BLEU-1, BLEU-4, and METEOR metrics, our approach obtained improvements of 7.2%, 20%, and 4.3%, respectively, compared to the inferior methods. However, our method did not attain the highest ROUGE-L score. This may be due to the property of this metric, which evaluates reports based on the Longest Common Subsequence with reference reports. Methods based on memory mechanisms [3, 2] can more easily generate common template sentences, resulting in higher ROUGE-L scores.