Multi-modality Regional Alignment Network for Covid X-Ray Survival Prediction and Report Generation

Survival analysis is a widely applied statistical method for assessing the time until a specific event occurs, with applications spanning various fields. Traditional approaches, such as Random Survival Forest (RSF) [21] and the Cox Proportional Hazards (CPH) model [22], are commonly used for estimating survival risks. DeepSurv [23], an extension of the CPH model, employs a deep learning framework with the Cox negative logarithm partial likelihood loss (CoxPHloss) to capture intricate, nonlinear relationships between survival risks and covariates [23]. Recent advancements emphasize the synergistic use of multiple modalities, including clinical, imaging, genetic, and time series data, to enhance disease survival analysis. Multimodal survival learning proves more robust than unimodal approaches, benefiting from the extraction of diverse information and cross-modal survival predictions [10, 11].

In the realm of radiology report generation, the task is considerably more challenging compared to natural image captioning [7]. Despite focusing on comprehension and generating descriptive sentences, image captioning struggles to fully convey rich feature information in medical images. Radiology report generation extends beyond a single sentence, involving crafting a comprehensive paragraph with structured sentences dedicated to specific medical observations in distinct regions of radiological images [24].

The existing approaches to medical report generation predominantly adopt encoder-decoder frameworks, evolving from LSTM-based CNN-RNN architectures to the more recent Transformer structures, benefitting from substantial datasets [7, 25]. Despite the Transformer’s effectiveness in processing long sequences and contextual information from extensive data, its performance is suboptimal in the context of automatically generating CXR reports due to data imbalances within the medical domain [8, 26]. To address this limitation, various enhancements have been proposed, such as the incorporation of relation memory units or cross-modal memory matrices into the standard Transformer architecture [27, 28]. Many methods introduce additional medical information to better capture anomalous text through auxiliary task learning, e.g., predicting disease keywords or tags [29], incorporating a topic matching mechanism [30], aligning disease tags with visual features [31], and introducing an additional medical knowledge graph [32]. In contrast, region-based generation presents an alternative approach to enhance region-focused attention through bounding box annotations. Wang et al. [33] introduced a region-level extractor using a selective search algorithm, departing from global features. The RGRG [34] method leverages Faster R-CNN [35] and a learnable selector for salient anatomical region object detection, enabling the independent generation of sentences that visually anchor to anatomical structures. Inspired by the working patterns of radiologists, our work aims to overcome these challenges with the multi-modal knowledge of large-scaled models. We evaluate the external survival risk to mitigate textual deviation as part of the attention mechanism for the anatomical regions. This region-aware method emphasizes learning text descriptions for specific diseases under high-risk conditions, ultimately generating well-structured radiology reports.

For report generation, our model closely follows the typical workflow of a radiologist. Given a CXR image $x_{I}$, first, a radiologist identifies the anomalies with in anatomical areas. For each sentence, radiologists pay more attention on region with the high risk evidence and carefully summary with expertise on the anatomical regions. Finally, the structured reports includes sentences of Finding in the 4 critical regions and a global summary as Impression, represented as $x_{T}^{i}\,(i\in[1,...,5])$.

The overall framework is shown in Fig. 2. First, an object detector identifies 29 distinct anatomical regions on $x_{I}$. The Multi-scale Region-feature Encoder (MRE) extracts and aggregates the regions belonging to predefined sentences on the report. The Survival-guided Sentence-feature Encoder (SSE) embeds the imaging survival-attention for the regional sentences, to capture the risk-related textual expression through image-to-text transferring. Employing LLMs-Alignment, an asymmetric encoder-decoder structure constrains for text-to-text and image-to-text alignment in the shared text space. The embedded sentence features are decoded with a transformer-based language generation model. Finally, a two-stage multimodal risk feature integration is employed, leveraging three modalities additionally including clinical variables $x_{C}$, to predict survival outcomes through a survival predictor.

For the anatomical region detector, we use Faster R-CNN [35] trained on Chest ImaGenome dataset [36], which predicts 29 unique anatomical regions for each case. Faster R-CNN consists of a region proposal network (RPN), which generates bounding boxes of potential anatomical regions on the image feature maps. A region of interest (ROI) pooling layer maps each object proposal onto the feature maps, extracting small feature maps of uniform spatial extent. These ROI feature maps are each classified into one of the 29 anatomical region classes (and the background class) following standard procedure in Faster R-CNN. The object proposal is well-detected which should has highest probability score among the 29 class, as well as the top-1 scores for all proposals. Otherwise, the region class does not achieve the highest score for at least one proposal, it is considered as undetected.

A compact and specialized network as shown in Fig. 4, Region Completer is designed to rectify bounding boxes for undetected regions. The missing bounding boxes can be estimated by leveraging the initial predictions made by the region detector. The Completer is a regression network comprising a three-hidden-layer multilayer perceptron (MLPs), which is trained in the way of ”Masked Autoencoders” [37]. Specifically, it is trained by predicting the coordinates of several missing regions which are artificially randomly masked, thereby learning the spatial distribution of the 29 region coordinates. During the inference phase of region detection, the previously predicted coordinates utilized as fixed coordinates, which may contain undetected regions. The relative distribution correlations between known and undetected regions are inferred to predict the missing bounding boxes, thereby completing the bounding boxes for all 29 regions $b_{j}\,(j\in[1,...,29])$.

The image is first passed through the image encoding model to produce the visual features. The image encoder is ResNet-50 model pre-trained with conditional reconstruction task for both vision and language representation [38]. Given an image $x_{I}$, the multi-scale visual features are extract from the 5 ResLayers of image backbone and denoted as $f_{I}^{l}=E_{I}(x_{I})$, $\in R^{{C_{l}}\times{H_{l}}\times{W_{l}}}$, $l$ represents the feature scale $l\in[1,...,5]$.

To obtain the region visual features for each sentence of radiology report, MRE module extracts and aggregates the features across multiple scales on the detected regions. Within region $b_{j}$, ROI align operation of $l$ scale maps local features to the corresponding ROI resolution $\in R^{{C_{l}}\times{H_{r}}\times{W_{r}}}$. After flattening into a one-dimensional vector, the linear transformation layer ”$\textrm{MLPs}\;l$” is applied to $l$-scale feature to reduced the dimension from ${C_{l}}\times{H_{r}}\times{W_{r}}$ of flatten vector to 256. Then, the 256-dimensional local features of 5 scales are subsequently concatenated together into $f_{ROI}^{j}\in R^{1280}$. These 29 extracted ROI features represent the visual characteristics of the 29 anatomical regions. Considering the global expressions present in radiology reports, the feature $f_{I}^{5}$ from the last ResLayer is utilized to extract the global visual feature $f_{G}$ thought average pooling and a fully connected layer. This global representation is then incorporated into the set of region features, $f_{ROI}^{30}=f_{G}$, therefore $j$ belongs to the range $[1,...,29,30]$.

Normally, the sentence in a medical report identifies an observation of the anatomical regions. The MRE selects and aggregates the 30 visual features corresponding to proposal anatomical regions that may be mentioned in sentence, enabling the model to provide spatial discriminability in sentence descriptions. For a given sentence $i$, that may describe within one or more regions. These regions are categorized into the group $i$, for example, $\textbf{ROI group}\,1=\left\{j:1,3,4,...\right\}$. The Region Selector extracts the features belonging to the region group $i$ from the set of region features$f_{ROI}^{j}$ and performs a concatenation operation to obtain sentence-wise visual feature:

where $\operatorname*{\scalerel*{\|}{\sum}}$ stands for concatenation operation. $f_{MRE}$ represents the concatenated visual feature for the sentences.

SSE module is responsible for embedding the survival feature from image into sentence features. Since the deviation of descriptions for normal and abnormal findings in medical reports, SSE employs survival prediction to enhance the language generation model’s risk sensitivity, guiding model to capture specific and abnormal statement and enhancing predictive accuracy for rare and high-risk case.

In SSE, the Survival-attention Module $E_{SA}$ is used to extract the risk features from the image branch, which is implemented upon the Multi-Head Attention (MHA) [39]. The MHA consists of $n$ parallel heads and each head is defined as a scaled dot-product attention, which is formulated as:

where $X$ and $Y$ denote the input matrix for Query and Key/Value. $W^{Q}_{h}$, $W^{K}_{h}$, $W^{V}_{h}$ are their learnable parameters in $h$-th head $Att_{h}$. $W^{O}$ is the weight of output linear layer. $d_{k}=d/n$ is the scaling factors and $d$ is dimension of embedded feature.

We utilize Eq. 4 to construct a self-attention pooling structure, where $E_{SA}$ extracts survival features using a self-attention encoder on the feature of last ResLayer $f_{I}^{5}\in R^{2048\times 7\times 7}$. After spatial-flattening and transposition, visual sequence features $f_{v}\in R^{49\times 2048}$ and their channel-wise mean vectors ${f_{v}}^{\prime}\in R^{1\times 2048}$, augmented with learnable positional embedding, are input into the MHA to extract survival attributes within the image features.

where $Pos\in\mathbb{R}^{49\times 2048}$ and $Pos^{\prime}\in\mathbb{R}^{1\times 2048}$ are the learnable spatial positional embeddings. The acquired the survival feature vectors $f_{sur}\in\mathbb{R}^{1\times 2048}$ are utilized for survival prediction, to provide risk attention for sentence-wise feature embedding. For $i$-th sentence, the local visual embedder $E^{i}_{L}$ aggregates the sentence’s visual features $f_{MRE}^{i}$. The global risk embedder $E^{i}_{G}$ extracts from $f_{sur}$ to associate with the context of sentence $i$ in the global view of survival prediction. The region-to-sentence embedder $E_{R2S}$ combines the survival-guided features, to obtain visual representation $\mathit{v}^{i}_{I}$.

In this way, the SSE can not only obtain sentence feature on the anatomical regions, but also align the risk feature indicating abnormal regions with the relevant descriptions, which imitates the working patterns of radiologists to assign the risk of disease to the abnormal regions when examining the abnormal regions.

We use the GPT-2 model [13] as language decoder $D_{T}$, which is an auto-regressive neural network based on self-attention. To inject the visual features with token embeddings in the MHA of GPT-2, we follow [40] in using Pseudo Self-Attention (PSA) [41], which be formulated as:

where $X$ represents the transferred sentence-wise visual features and $Y$ represents the token embeddings. The new parameters $U^{K}_{h}$ and $U^{V}_{h}$ are the corresponding key and value projection parameters of self attention-visual projecting. Text Decoder $D_{T}$ which is a GPT decoder structure, comprises a total of 24 cascaded Transformer modules including PSA layers, feedforward, residual and normalization layers. In each module, the tokens in a sequence are conditioned on previous tokens and encoded with shared visual feature $\mathit{v}^{i}_{T}$. The predicted logits are obtained in a linear layer followed by a softmax activation. Sentence predictions $\hat{y}_{T}^{i}$ are generated in an iterative process of selecting the most likely next word from the vocabulary, which based on continuously preceding tokens in the previous predicted sequence. The visual features are transferred with an embedder $E_{I2T}$ from image to text space for report generation:

The report generation is optimized by minimizing the sum of the sentence-wise cross-entropy loss:

where $l_{i}$ denotes the length of words in $x_{T}^{i}$ of the $i$-th sentence. The cross-entropy loss for each sentence is calculated word by word until reaching the end of $x_{T}^{i}$.

To enhance the generative capacity of our model, we introduce additional knowledge as semantic constraints, incorporating a self-supervised learning paradigm based on text encoding and decoding during training. In this LLMs-Alignment procedure, GatorTron [12], a BERT (Bidirectional Encoder Representations from Transformers) model, is employed as the Text Encoder $E_{T}$ from another perspective. Renowned for its exceptional capabilities in capturing intricate semantic nuances within biomedical text, the encoder can be regarded as an experienced radiologist, encapsulating the detailed information present in medical reports, and supervising for Text Decoder in the synthesis of report features into coherent textual narratives. The $E_{T2T}$ along with the $E_{i2T}$ embeds sentence features to the shared textual space for report decoding. The text encoder-decoder and the overall language generation loss are formulated as follows:

During training, the reports are attractively generated on textual features $\mathit{u}^{i}_{T}$ and visual features $\mathit{v}^{i}_{T}$, with both branches sharing the Text decoder. In their shared textual space upon the GatorTron’s expertise, the LLMs-Alignment guides image-to-text tranferring with clinical information perception. The alignment refines the visual generated sentences with authentic textual content within the biomedical and health domain.

The MRANet model uses the multi-modal features for survival prediction. The concatenated image-to-text features $\operatorname*{\scalerel*{\|}{\sum}}_{i=1}^{5}\mathit{v}^{i}_{T}$, $f_{sur}$ from the survival attention module, and the clinical variable $x_{c}$ are separately encoded through text, image, and clinical survival encoders, which are MLPs with two hidden layers, mapping multi-modal features to 2048-dimensional survival features $s_{T}$, $s_{I}$, $s_{C}$. Subsequently, a fusion module consisting of a concatenation and linear layer is applied to integrate the features of image $s_{T}$ and text $s_{I}$, followed by the 2nd fusion module to further integrate with clinical feature $s_{C}$. This two-stage fused features are then utilized to predict survival outcomes $\hat{y}_{sur}$ through a survival predictor $C_{sur}$.

The survival prediction model is trained to differentiate between patients with critical disease and those without using CoxPHloss [23], which can handle right-censored survival data and is formulated as:

where $\hat{y}^{i}_{sur}$ is the predicted risk of disease progression for the $i$-th patient, $y^{i}_{t}$ and $y^{i}_{e}$ are survival labels donated as the survival time and the censorship flag, respectively. $y^{i}_{t}$ is defined as the days of diagnosis to death (for $y^{i}_{e}=1$) or patient censored (for $y^{i}_{e}=0$), $N_{Y_{e}=1}$ is the number of patients with an observable event.

The Chest ImaGenome dataset [36] contains automatically constructed scene graphs for the MIMIC-CXR [42] dataset, consisting the images with corresponding free-text radiology reports. Each scene graph describes one frontal CXR and contains bounding box coordinates for 29 unique anatomical regions, as well as sentences describing each region if they exist in the corresponding radiology report. The clinical variables includes age, sex, temperature, oxygen saturation on room air, white blood cell count, lymphocyte count, creatinine, C-reactive protein, and comorbidities (cardiovascular disease, hypertension, chronic obstructive pulmonary disease, diabetes, chronic liver disease, chronic kidney disease, cancer, and HIV status) [43]. We use the split provided by the dataset: 166,512 training, 23,952 validation, and 47,389 test images.

A retrospective review was conducted to identify COVID-19 patients with imaging between March 2020 and July 2020 at Brown Affiliated Hospitals (Brown) and the Hospital of the University of Pennsylvania (Penn). The Brown-COVID dataset contains 1021 front-view CXR and corresponding reports. Different from the free-text reports from MIMIC-CXR dataset, our structured reports include 5 sentences as shown in Fig. 5; each are well defined for 4 Findings which included details for their defined regions and the Impression, which serves as a summary of the report. The concatenation was used as the target report in experiments. The survival labels include mortality event and the length of time from COVID-19 test to date of follow-up or mortality. The Penn-COVID dataset contains 2879 images and survival labels, but lacks radiology reports. We conducted our model evaluations in a multi-center setting to ensure robustness and generalizability. The Brown-COVID dataset was divided into three sets with 70% for training, 10% for validation, and 20% for internal testing. The Penn-COVID dataset was used external validation of survival prediction. The study and collection of dataset have been approved by Lifespan Institutional Review Board, Providence, RI, USA (IRB board protocol: [1584610-18]; approval date: 10/09/2023).

The task-specific training protocol is designed to cater to the unique target of individual modules.

We conducted experiments on survival prediction, comparing our proposed method with widely used baseline approaches. The first three methods listed in Table 1 include three separate survival predictors that are trained to estimate risk scores based on image, ground-truth text report, and clinical variables, respectively. These models utilize image encoders and text encoders in MRANet for feature extraction; their respective survival predictors are RSF [21], CPH model [22], and Vanilla [23]. Subsequently, the risk scores from each modality are combined using a half-parametric linear analysis with the CoxPH model [22] to obtain the fused risk score. The Multi-Fuse model in Table 1 is based on the multi-modal survival prediction structure of MRANet that utilizes one survival model for predicting the fused risk scores on the fused survival feature. Given the absence of text reports in the Penn dataset, there are missing values.

As presented in Table.1, MRANet achieves C-index values of 0.813 and 0.739 for Brown-COVID and Penn-COVID datasets, respectively, surpassing competing methods. The models fused with CoxPH demonstrated the presence of complementary survival-related information between different modalities. Across different modalities, RSF and Cox show relatively better performance as survival predictors. Meanwhile, Vanilla outperforms traditional survival prediction models in clinical experiments (0.739 and 0.703 for Brown and Penn) while maintaining a more stable overall performance with superior generalization on external validation sets. Compared to the single-modal survival prediction then late fusion strategy employed by Vanilla-CoxPH, Multi-Fuse is more effective in analyzing multi-modal features through learnable fusion. Although our proposed MRANet does not rely on real text and employs image-to-text features for survival prediction, it surpasses the performance of existing methods. This highlights that semantic information contained within language can be harnessed to improve survival analysis in generated reports.

Table 3 shows survival prediction performance on single modalities using extracted features from images, reports, and clinical variables as inputs to a single-modal MLP model. We compare image-based performance against various ResNet50 checkpoints while conditioning frozen parameters of $E_{I}$. As demonstrated in the table, PRIOR from large-scale pretraining on MIMIC-CXR data—outperforms ImageNet-pretrained models for survival prediction. Furthermore, when image encoder parameters are fixed, superior performance is achieved. This highlights that, training with smaller datasets leads to overfitting of the survival model, leading to reduced generalization on two test sets. The last three rows of Table 3 correspond to Vanilla-CoxPH’s unimodal survival prediction performance. Merely combining these modalities does improve for survival prediction.

To verify the survival prediction effectiveness of each component, we perform individual experiments on our the ablated variations. Table 4 shows that the steady improvement on each component demonstrates the effectiveness, while highlighting the overall coherence of the module composition within the proposed model.

The experimental results of the report generation on the Brown-COVID datasets are presented in Table 2, including a comparison with state-of-the-art (SOTA) models and ablation studies for MRANet. We used the models trained on MIMIC-CXR as baselines and directly applied them on images of Brown-COVID. In generating radiology reports, our proposed model is either superior or comparable to previous models for NLG and CE metrics. Specifically, our methods achieves the goal of structure report generation when fine-tuned locally on the small-scale dataset, surpassing all large-scaled models optimized for free-text report generation.

Regarding BLEU scores, MRANet surpasses recent state-of-the-art models [34], achieving superior performance on report summarization. Notably, RGRG leverages abnormal region detection to capture rare descriptions in reports, which enhances regional detail description abilities and increases BLEU scores compared to other baselines. Compared to the region-based baseline, MRANet’s base model exhibits a remarkable 70% improvement in BLEU-4 score. This significant difference stems from our method’s utilization of specified regional cluster features, leading to higher anatomy sensitivity and enhanced structural predictability within generated reports. Our final model also excels on the METEOR metric and ROUGE scores for summarization and stemming/synonymy matching as compared to previously reported best results in medical report generation.

MRANet achieves an average 50% increase in CE metrics compared to baseline methods. Global description approaches outperform local baselines due to RGRG’s limited descriptions of only selected areas, which may not be sufficiently comprehensive. If high-risk regions go undetected or are not chosen for description, crucial image features can be overlooked, diminishing disease description performance.

Our method generates sentences based on the aggregation of multiple presumptively defining regions, incorporating all possible descriptive areas for comprehensive learning about region relationships and their importance through survival attention. This approach allows our model to focus on rare details while providing a more complete description compared to the region-based baseline. Furthermore, integrating GatorTron knowledge constraints with LLMs-Align significantly improves Precision (8.5%), Recall (5.9%), and F1 score (6.2%) for medical report generation tasks, highlighting the positive impact of additional clinical knowledge in enhancing the model’s ability to capture clinically relevant details within generated reports.

The proposed MRANet constitutes an integrated three-task model, harmoniously merging anatomical region detection, radiology report generation, and survival prediction. This section highlights the multi-task outcomes for lung regions and conduct a qualitative analysis to assess their correlations. As the results of Brown-COVID texting set demonstrated in Fig. 6, the predicted bounding boxes accurately localize lung and subregions across diverse gray-level images, underscoring the object detector’s robust performance. The model generates accurate descriptions on the disease-affected regions, which can correctly describe the location and diseases. The missing disease grounding descriptions can result from having limited training data, which may not provide enough location context for the language model to specific region.

Risk attention scores are obtained utilizing Grad-CAM tools [51] to evaluate each region feature’s contribution towards survival outcomes within the local visual embedder encoding process. Calculating mean gradient values of predicted multimodal risk scores as weights for region features, regional attention scores result from regional-wise sum of weighted feature map and ReLU activation application. By re-weighted with global survival score, we observed that abnormal areas consistently yield higher risk attention scores compared to normal ones, offering diverse reference to text descriptions. For instance, the 3rd example despite an omission in spatial positioning within generated sentences, elevated risk attention scores can still assist in corroborating the accuracy of ”worse on the right than on the left” descriptions.