11institutetext: Mohamed bin Zayed University of Artificial Intelligence, Abu
Dhabi, UAE
11email<EMAIL_ADDRESS>
22institutetext: Sheikh Shakhbout Medical City, Abu Dhabi, UAE
22email<EMAIL_ADDRESS>33institutetext: Big Data Institute, University
of Oxford, UK
33email<EMAIL_ADDRESS>
# TiBiX: Leveraging Temporal Information for Bidirectional X-ray and Report
Generation

Santosh Sanjeev 11 Fadillah Adamsyah Maani 11 Arsen Abzhanov 11 Vijay Ram
Papineni 22 Ibrahim Almakky 11 Bartłomiej W. Papież 33 Mohammad Yaqub 11

###### Abstract

With the emergence of vision language models in the medical imaging domain,
numerous studies have focused on two dominant research activities: (1) report
generation from Chest X-rays (CXR), and (2) synthetic scan generation from
text or reports. Despite some research incorporating multi-view CXRs into the
generative process, prior patient scans and reports have been generally
disregarded. This can inadvertently lead to the leaving out of important
medical information, thus affecting generation quality. To address this, we
propose TiBiX: Leveraging Temporal information for Bidirectional X-ray and
Report Generation. Considering previous scans, our approach facilitates
bidirectional generation, primarily addressing two challenging problems: (1)
generating the current image from the previous image and current report and
(2) generating the current report based on both the previous and current
images. Moreover, we extract and release a curated temporal benchmark dataset
derived from the MIMIC-CXR dataset, which focuses on temporal data. Our
comprehensive experiments and ablation studies explore the merits of
incorporating prior CXRs and achieve state-of-the-art (SOTA) results on the
report generation task. Furthermore, we attain on-par performance with SOTA
image generation efforts, thus serving as a new baseline in longitudinal
bidirectional CXR-to-report generation. The code is available at
https://github.com/BioMedIA-MBZUAI/TiBiX.

###### Keywords:

Report Generation Chest X-ray generation Multimodal Data Bidirectional
Generation Longitudinal data

##  1 Introduction

Image and text generation advancements have propelled the progress in Chest
X-ray (CXR) and report generation tasks. The specific focus on CXR stems from
the time-consuming and error-prone nature of the manual report composition
task, putting a huge strain on radiologists [12]. Moreover, the quality of
manually composed reports depends on the radiologist’s experience, which makes
accurate report generation a good aid for inexperienced practitioners. To this
effect, some research has been conducted in the domain of CXR image [3, 32,
23] and report [29, 27, 22, 17] generation.

Transformer-based models like those in [5, 4] have shown promising
performance, especially with the introduction of expert tokens in [31].
Strategies such as integrating knowledge graphs [35, 11], addressing class-
level bias [19], utilizing disease tags [28, 26], and applying contrastive
learning [18] aim to enhance report quality. In parallel, advancements in
natural image generation, driven by techniques like Generative Adversarial
Networks (GANs) [8, 34], have sparked interest in generating CXR images from
radiology reports aiming to improve a model’s interpretability, helping build
user trust through visual evidence [15]. However, the translation of radiology
reports into CXR images is fraught with challenges, mainly stemming from the
linguistic complexities and varied structures of radiology reports, hindering
accurate representation in image form [32, 3]. Despite efforts to confront
these obstacles, existing methods often fail to capture the intricate
relationships between textual descriptions and corresponding image features,
resulting in suboptimal performance in CXR image generation. Moreover, these
methods overlook the temporal aspect of patient data which is very important
for medical diagnosis.

There are few works exploring the temporal domain for generation tasks [1, 36,
24]. Unlike the previous works, which focus on only report generation, we
focus on bidirectional CXR and report generation. Furthermore, most works
require both the prior and current scan images, which is generally not the
case in real-life scenarios as patients might not always have prior studies.
Recent advancements in bidirectional approaches [14, 10] have shown promise in
bridging the gap between image-to-text and text-to-image generation tasks. In
the medical domain, UniXGen was proposed for bidirectional CXR-to-report
generation, leveraging multiple CXR views from the same patient during the
same visit to enhance report accuracy [16]. Previous research has focused on
generating reports or images at single time points, overlooking the
longitudinal aspect of patient data [13].

Figure 1: Our temporal bidirectional CXR-to-report generation framework which
can deal with three inputs (i.e. current report, current CXR, and previous
CXR). We utilize an image tokenizer $B_{x}$ and text tokenizer $B_{r}$ to
tokenize CXR image(s) and report, respectively. A transformer-based model
$\mathcal{M}_{\theta}$ with causal attention is implemented to handle the
bidirectional generation task in the auto-regressive manner. The input
sequence of $\mathcal{M}_{\theta}$ consists of images and text tokens, a
temporal token (TT) which encodes the time interval between two consecutive
CXR scans, and a learnable cls (class) token. We assign TT as the first input
sequence and cls as the last sequence. During training, the order of the
current report, current CXR, and previous CXR are shuffled, while the missing
modality is placed at the last part of the input sequence during inference.

In this paper, we propose Chest X-ray and Report Bidirectional Generation
along the temporal domain. Our main contributions are summarized as follows:

  * •

We propose a novel approach to integrate temporal data for the bidirectional
generation task by designing a framework consisting of a transformer-based
model equipped with causal attention layers that provide an effective way to
learn temporal information while also reducing computational complexity. We
also introduce a $cls$ (class) token guidance approach based on the patient’s
overall diagnosis and also a temporal token, considering the varying time
intervals between patient visits.

  * •

We achieve state-of-the-art (SOTA) performance on the benchmark dataset for
the report generation task and achieve comparable performance to the (SOTA) on
the image generation task. We also achieve promising results on the
classification metrics of images and reports. Furthermore, we support the
inclusion of temporal scans by conducting an extensive analysis.

  * •

We extract and release the MIMIC-T dataset, curated from the MIMIC-CXR dataset
and consists of longitudinal data of patients extracted from the MIMIC-CXR
dataset.

To the best of our knowledge, this is the first effort to include
prior(temporal) CXRs and reports for the bidirectional CXR-to-report
generation task.

##  2 Methodology

We formulate the temporal bidirectional CXR-to-report generation task starting
from a set of patients $P=\\{p_{1},p_{2},\dots,p_{n}\\}$. For each patient
$p$, we have a set of CXR images $X_{p}=\\{x_{(i)}\\}_{i={t_{0}}}^{t_{m_{p}}}$
along with a set of radiology reports
$R_{p}=\\{r_{(i)}\\}_{i={t_{0}}}^{t_{m_{p}}}$, where $m_{p}\geq 0$ is the
total number of time steps and $r_{(t_{j})}$ is the corresponding radiology
report for $x_{(t_{j})}$ captured at time step $t_{j}$. Therefore, the total
number of samples available is $M=\sum_{p=1}^{n}{(m_{p}+1)}$. We define the
labels for a patient $p$ at time step $t_{j}$ by
$L_{(t_{j})}=\\{l_{(i)}\\}_{i=1}^{c}$, where $l_{(i)}\in\\{0,1\\}$ and $c$ is
the total number of possible pathologies. We also define the variable time
step $\delta=(t_{j}-t_{j-1})$ as the time interval between 2 consecutive
studies, where $\delta>0$. As such, the objective of the temporal
bidirectional CXR-to-report generation task would be to estimate $r_{(t_{j})}$
given $(x_{(t_{j})},x_{t_{(j-1)}})$, while also estimating $x_{(t_{j})}$ given
$(r_{(t_{j})},x_{t_{(j-1)}})$. To learn this, we employ a transformer-based
model with causal attention, which is described later in this section,
starting with the multi-modal input tokenization and ending with the auto-
regressive training process. The high-level visualization of our framework is
depicted in Fig. 1.

###  2.1 Multimodal Tokenization

This work considers three modalities: current report, current CXR, and
previous CXR. To facilitate a simple yet effective and memory-efficient
multimodal interaction, we tokenize every modality accordingly. We use a CXR
image tokenizer $B_{x}(x)=\mathcal{I}^{x}\in\mathbb{Z}^{N_{x}}$ and a report
text tokenizer $B_{r}(r)=\mathcal{I}^{r}\in\mathbb{Z}^{N_{r}}$, where $N_{x}$
and $N_{r}$ are the number of image and text tokens. We adopt $B_{x}$ from
VQGAN [8] tokenizer, which is composed of two primary components: an image
encoder responsible for generating context-rich latent representations and a
codebook used to map the representations to a quantized encoding. The main
component of our bidirectional generation framework $\mathcal{M}_{\theta}$ is
based on a transformer model. We, add unique START and STOP tokens for every
modality, increasing the size of the image and text tokens to
$\bar{\mathcal{I}}^{x}\in\mathbb{Z}^{N_{x}+2}$ and
$\bar{\mathcal{I}}^{r}\in\mathbb{Z}^{N_{r}+2}$, respectively to avoid
confusion between the modalities, as it would treat all input tokens
uniformly. The time interval $\delta$ between two consecutive samples is also
input to $\mathcal{M}_{\theta}$ as a temporal token ($\mathcal{I}^{t}$). This
is added to further improve the model’s prognosis capability by mimicking
clinicians’ approach to leveraging the time gap between different scans.
Finally, we add a learnable token $cls\subset\theta$ to predict the
pathologies existing in the patient’s modalities $L_{(t_{j})}$, allowing
$\mathcal{M}_{\theta}$ to learn from the overall diagnosis.

###  2.2 Multimodal Generation Framework

Transformer-based Model: $\mathcal{M}_{\theta}$ carries out the temporal
bidirectional CXR-to-report generation task as follows:

$\mathcal{M}_{\theta}(\mathcal{I}^{t},\bar{\mathcal{I}}^{m_{1}},\bar{\mathcal{I}}^{m_{2}},\bar{\mathcal{I}}^{m_{3}})=\hat{\mathcal{I}}^{m_{1}},\hat{\mathcal{I}}^{m_{2}},\hat{\mathcal{I}}^{m_{3}},\hat{L}_{t}$
(1)

where $\hat{\mathcal{I}}^{m}\approx\bar{\mathcal{I}}^{m}$, $\hat{L}_{t}$ is
the prediction of the patient’s diagnosis, and
$\\{\bar{\mathcal{I}}^{m_{1}},\bar{\mathcal{I}}^{m_{2}},\bar{\mathcal{I}}^{m_{3}}\\}$
are sampled from set
$\\{\bar{\mathcal{I}}^{x_{(t_{j})}},\bar{\mathcal{I}}^{r_{(t_{j})}},\bar{\mathcal{I}}^{x_{(t_{j-1})}}\\}$
without replacement. Our main training goal is to learn parameters $\theta$ so
that $\mathcal{M}_{\theta}$ is capable of performing Eq. (1) by optimizing the
following objective:

$\min_{\theta}\frac{1}{M}\sum_{i=1}^{M}\Big{(}\mathcal{L}_{gen}\big{(}(\hat{\mathcal{I}}^{m_{1}}_{i},\hat{\mathcal{I}}^{m_{2}}_{i},\hat{\mathcal{I}}^{m_{3}}_{i}),(\bar{\mathcal{I}}^{m_{1}}_{i},\bar{\mathcal{I}}^{m_{2}}_{i},\bar{\mathcal{I}}^{m_{3}}_{i})\big{)}+\lambda\cdot\mathcal{L}_{cls}(\hat{L}_{t_{i}},L_{t_{i}})\Big{)}$

where $\mathcal{L}_{gen}(\cdot,\cdot)$ is the generation cross-entropy loss,
while $\mathcal{L}_{cls}(\cdot,\cdot)$ is the classification cross-entropy
loss based on the $cls$, and $\lambda\in\mathbb{R}$ is a hyperparameter. Some
patients do not have a previous scan, i.e.
$\bar{\mathcal{I}}^{x_{(t_{j-1})}}=\emptyset$. We replace all missing previous
scans with a set of learnable padding tokens to ensure a consistent model
input size during training. We follow a similar approach to Eq. (1) during
inference without providing the modality we wish to generate, i.e., the input
of $\mathcal{M}_{\theta}$ is a patient’s available modalities.

Causal Attention: Inputting multiple modalities raises computational
complexity challenges associated with the standard attention mechanism, which
exhibits $\mathcal{O}(N^{2})$ where $N$ being the number of tokens. As such,
we employ the Performer architecture along with the causal attention mechanism
[6] that reduces the complexity down to $\mathcal{O}(N)$. This is achieved by
approximating the Softmax function, enabling the multiplication of the key and
value matrices to be carried out before the result is multiplied by the query
matrix. Our proposed model $\mathcal{M}_{\theta}$ generates tokens in an auto-
regressive manner, inserting each predicted token into the next step to
generate the next token. To preserve this causality and allow multiprocessing,
we implement unidirectional FAVOR$+$ [6]. The causal constraint prevents a
token from attending to tokens positioned after it in the sequence, thereby
ensuring no information leakage from future tokens. Following our Performer’s
linear projection for the different inputs, we inject sinusoidal positional
embedding to the text embedding, as well as axial positional embedding [9] to
the previous and current image embeddings. We also inject sinusoidal
positional embedding into the full input embeddings before passing it to our
model.

##  3 Datasets and Experiments

Datasets: We extensively test our method, TiBiX, on a selected subset of
MIMIC-CXR [12]. We also introduce MIMIC-T, a new dataset derived from MIMIC-
CXR, containing longitudinal patient data. Our preprocessing involves
discarding studies lacking both impressions and findings sections and
retaining only AP frontal-view scans. This yields approximately $122k$, $1k$,
$2k$ unique studies from $32521$, $255$, $266$ patients in train, validate,
and test splits, respectively. To increase the number of training samples, we
utilize prior scans by considering every consecutive pair for the same subject
as a new sample. Moreover, we also consider single-study cases simulating
real-world scenarios. We prioritize examining previous scans to diagnose
diseases and assess prognosis. The distribution of patients with one or two
studies is outlined in Section A1 in the Appendix.
Implementation Details: We extract impressions and findings from reports and
preprocess them by converting text to lowercase, removing punctuation, and
digits. We use the CXR-BERT [2] tokenizer with 30,522 unique tokens, including
special tokens. AP frontal-view images are resized to 512x512. We employ a
pretrained VQ-GAN from [16], trained on MIMIC-CXR with dz=256 and a codebook
size of 1024, encoding 512x512 images into 32x32=1024 tokens. Our transformer-
based model uses the same configuration of Performer [6]. It handles 1029
image tokens and 30,522 text tokens [8, 6]. The model can process two images
if the previous one is available. Training lasts for 150 epochs with a batch
size of 10, AdamW optimizer, learning rate of 2e-4, weight decay of 1e-2, and
cosine decay schedule. Report and image generation use Top-p sampling (p=0.9,
temperature=0.7). We ensemble five model weights and select the image with the
highest pixel entropy to mitigate stochastic blank image generation.

##  4 Results and Discussion

Table 1: Comparison to the previous methods in report generation. (#) - Subset
of MIMIC-T with $n$ studies where $n\in[1,2]\quad\text{and}\quad
n\in\mathbb{N}$.

Methods | Set (#) | NLG metric | Clinical Efficacy(CE)
---|---|---|---
|  | BLEU-1 | BLEU-2 | BLEU-3 | BLEU-4 | ROUGE | METEOR | Precision | Recall | F1
CNN-RNN [20] | CR$|$CX (2) | 0.190 | 0.105 | 0.066 | 0.045 | 0.200 | 0.096 | 0.303 | 0.221 | 0.230
AdaAtt [21] | CR$|$CX (2) | 0.206 | 0.115 | 0.072 | 0.049 | 0.205 | 0.102 | 0.310 | 0.230 | 0.234
Transformer [25] | CR$|$CX (2) | 0.194 | 0.109 | 0.068 | 0.047 | 0.200 | 0.100 | 0.320 | 0.228 | 0.241
M2Transformer[7] | CR$|$CX (2) | 0.193 | 0.108 | 0.068 | 0.047 | 0.200 | 0.099 | 0.323 | 0.248 | 0.253
R2GenCMN [4] | CR$|$CX (2) | 0.348 | 0.240 | 0.182 | 0.146 | 0.322 | 0.148 | 0.321 | 0.216 | 0.244
XproNet [30] | CR$|$CX (2) | 0.324 | 0.225 | 0.172 | 0.139 | 0.323 | 0.141 | 0.315 | 0.211 | 0.240
| CR$|$CX (1) | 0.234 | 0.156 | 0.116 | 0.091 | 0.272 | 0.119 | 0.340 | 0.190 | 0.230
| CR$|$CX (2) | 0.284 | 0.184 | 0.136 | 0.105 | 0.267 | 0.127 | 0.340 | 0.220 | 0.250
TiBiX(Ours) | CR$|$(PX,CX) (2) | 0.324 | 0.234 | 0.185 | 0.157 | 0.331 | 0.162 | 0.300 | 0.224 | 0.25

Table 2: Comparison to state-of-the-art (SOTA) in regard to image generation
task. (#) - Subset of MIMIC-T with $n$ studies where
$n\in[1,2]\quad\text{and}\quad n\in\mathbb{N}$.

|  | Image Generation metrics | 14 diagnosis classification
---|---|---|---
Method | Set (#) | FID | SSIM | MSSIM | IS | Precision | Recall | AUC | Accuracy | F1
| GT (2) | - | - | - | - | 0.75 | 0.73 | 0.79 | 0.75 | 0.60
Roentgen [3] | CX$|$CR (2) | 81.02 | 0.51 | 0.398 | 1.96±0.08 | 0.64 | 0.64 | 0.84 | 0.77 | 0.60
Cascaded Diffusion[32] | CX$|$CR (2) | 28.78 | 0.53 | 0.428 | 2.05±0.06 | 0.48 | 0.82 | 0.75 | 0.66 | 0.57
| CX$|$CR (1) | 80.48 | 0.50 | 0.442 | 1.70±0.15 | 0.68 | 0.28 | 0.69 | 0.70 | 0.42
| CX$|$CR (2) | 50.40 | 0.50 | 0.438 | 1.76±0.10 | 0.61 | 0.35 | 0.69 | 0.74 | 0.41
TiBiX(Ours) | CX$|$(CR,PX) (2) | 33.85 | 0.52 | 0.472 | 1.82±0.08 | 0.56 | 0.55 | 0.75 | 0.77 | 0.53

###  4.1 SOTA Performance Comparison

To demonstrate the effectiveness of our approach, we compare with several SOTA
methods in the domain of report and CXR image generation. In all the Tables
and Figures below we use the following notations: GT - Ground Truth, CR -
Current Report, CX - Current CXR, PX - Previous CXR.(CX$|$CR) denotes the
current CXR generation given the current report, (CX$|$(CR,PX)) denotes the
current CXR generation given the current report and previous CXR.
Report Generation: We compare our results on the report generation task with
several SOTA image-captioning works such as the CNN-RNN[20], Transformer [25],
AdaAtt [21], M2Transfomer [7] in the natural domain as well as the SOTA works
in the medical domain including R2GenCMN [4] and XproNet [30]. For [33, 20,
25, 7] methods, we rerun them on our datasets, following the official
implementation of [7]. We adhere to the same preprocessing steps as in our
approach, while keeping the other hyperparameters consistent with those
specified in their codebase. Table 1 shows that compared to other works, our
method achieves SOTA results on most of the NLG metrics and also achieves an
on-par performance on the classification scores.
CXR Image Generation: Few works in the medical domain have focused on the
generation of Chest X-ray images. We compare our results with the recent top
works in the field of X-ray generation, which have used Diffusion Models(DM)
[3, 32]. We have implemented these pretrained models on the subset of the test
set having 2 studies, and the results are shown in Table 2. RoentGen is a
medical domain-adapted latent diffusion model based on the Stable Diffusion
pipeline whereas Cascaded Diffusion[32] improves upon this work, by having
cascaded models. The first model in [32] generates a lower resolution image
and the second one generates a higher resolution image, enhancing the first
DM’s output. It is noteworthy that [32] is trained on a huge dataset Massive
Chest X-ray Dataset (MaCheX) around 3.5x the size of MIMIC-T. Our results
outperform Roentgen’s and are comparable to the Cascaded Diffusion approach in
terms of image generation metrics. However, we have found that Roentgen
outperforms all other methods regarding classification scores. It is
counterintuitive that RoentGen’s performance in image generation is poor,
whereas it performs exceptionally well when the generated images are passed
through a classifier. This could be due to the pathology being easily detected
by the model, which is generally not the case in real life.

It is worth noting that all the works we compare are specifically designed for
a single task, either reporting or image generation. In contrast, our
framework excels in bidirectional generation while achieving state-of-the-art
performance across all tasks.

###  4.2 Self-Comparison

We analyse different subsets of the MIMIC-T dataset to understand the
effectiveness and importance of including a prior CXR scan for generation. In
Table 1, we evaluate the quality and clinical efficacy (CE) of the report
generation task on 2 subsets of the test set : (1) patients with 1 study and
(2) patients with 2 studies. Within the second case, we analyse CR $|$ CX as
well as the case of CR $|$ (PX,CX). We observe that when the prior scan is
given along with the current report, the BLEU-4 increases by around 1.5x
times. Furthermore, we observe significant improvement across all the metrics
including Clinical Efficacy. These results, are analogous to the fact that
radiologists when provided with multiple scans, can produce more comprehensive
and accurate diagnostic reports for various diseases. Similarly in Table 2, we
compare the results on the different subsets of MIMIC-T, for the image
generation task using FID, SSIM, MSSIM and IS. Furthermore, we also run a
classifier (MIMIC-CXR pretrained DenseNet121) on the images and get the
classification scores. Given both the prior CXR and CR, the generation quality
improves significantly. Furthermore, it achieves the best results in all the
evaluation metrics including the classification results. We also provide some
ablation studies in Section C1 in the Appendix.

Figure 2: (A) Report Generation Task: The first case describes an example of
(1) current report generation (CR$|$CX) (2) current report generation
(CR$|$(PX,CX)). (B) CXR Image Generation task: The left side of the report is
(CX$|$CR) and the on the right side is (CX$|$(CR,PX)). The GT is present for
comparison.

###  4.3 Qualitative Analysis

The benefits of considering the prior scan when generating a report or a scan
are demonstrated in Fig. 2. Fig. 2-A illustrates that in the case where the
model does not consider the prior scan it produces a report with false
findings that are not in the ground truth (e.g. the report describes
cardiomegaly when it’s not there). Furthermore, it hallucinates and includes
the word "unchanged" even though the previous scan is not provided. On the
other hand, when the model considers the previous scan it generates a report
with findings more aligned with the ground truth (highlighted in different
colors). Fig. 2-B illustrates that when the model considers the previous scan,
it generates a scan with an anatomical structure more similar to the ground
truth. The resultant scan is also closer to the ground truth in terms of
portraying the pathologies. For example, “the opacification of the right lung
base as opposed to the left lung being clear is more accurate when the models
take into account the previous scan”. When the models do not have the prior
scan, the resultant image deviates from the ground truth more. We also provide
extensive qualitative analysis comparing previous methods in Section B1 in the
Appendix.

##  5 Conclusion

In this work, we introduce a novel approach to enhance image and report
generation by integrating temporal data from previous CXR scans. Our framework
includes a transformer-based model with causal attention, facilitating
bidirectional CXR-to-report generation while considering the time difference
between consecutive scans. Experimental results on the MIMIC-T dataset, a
temporal subset of MIMIC-CXR, demonstrate promising performance. Qualitative
analysis and ablation studies justify the inclusion of prior scans and
highlight the importance of each architectural component. While our work lays
a foundational framework, future enhancements could explore incorporating
prior reports and domain expertise through knowledge graphs. Furthermore, we
aim to create a metric capable of capturing the temporal aspect, addressing
the inadequacies of current metrics in this regard.

## References

  * [1] Bannur, S., Hyland, et al.: Learning to exploit temporal structure for biomedical vision-language processing. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 15016–15027 (2023)
  * [2] Boecking, B., Usuyama, et al.: Making the most of text semantics to improve biomedical vision–language processing. In: European conference on computer vision. pp. 1–21. Springer (2022)
  * [3] Chambon, et al.: Roentgen: vision-language foundation model for chest x-ray generation. arXiv preprint arXiv:2211.12737 (2022)
  * [4] Chen, Z., Shen, et al.: Cross-modal memory networks for radiology report generation. arXiv preprint arXiv:2204.13258 (2022)
  * [5] Chen, Z., Song, et al.: Generating radiology reports via memory-driven transformer. arXiv preprint arXiv:2010.16056 (2020)
  * [6] Choromanski, et al.: Rethinking attention with performers. arXiv preprint arXiv:2009.14794 (2020)
  * [7] Cornia, et al.: Meshed-memory transformer for image captioning. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 10578–10587 (2020)
  * [8] Esser, et al.: Taming transformers for high-resolution image synthesis. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 12873–12883 (2021)
  * [9] Ho, J., et al.: Axial attention in multidimensional transformers (2019)
  * [10] Huang, et al.: Du-vlg: Unifying vision-and-language generation via dual sequence-to-sequence pre-training. arXiv preprint arXiv:2203.09052 (2022)
  * [11] Huang, et al.: Kiut: Knowledge-injected u-transformer for radiology report generation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 19809–19818 (2023)
  * [12] Johnson, et al.: MIMIC-CXR, a de-identified publicly available database of chest radiographs with free-text reports. Sci. Data 6(317), 1–8 (Dec 2019)
  * [13] Kayser, et al.: Explaining Chest X-Ray Pathologies in Natural Language. In: MICCAI 2022: 25th International Conference, Singapore, September 18–22, 2022, Proceedings, Part V, pp. 701–713. Springer-Verlag, Berlin, Germany (Sep 2022)
  * [14] Kim, et al.: L-verse: Bidirectional generation between image and text. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 16526–16536 (2022)
  * [15] Lanfredi, et al.: Adversarial Regression Training for Visualizing the Progression of Chronic Obstructive Pulmonary Disease with Chest X-Rays. In: MICCAI 2019 (2019)
  * [16] Lee, H., Kim, et al.: Unified chest x-ray and radiology report generation model with multi-view chest x-rays. arXiv preprint arXiv:2302.12172 (2023)
  * [17] Li, et al.: A self-guided framework for radiology report generation. In: Wang, et al. (eds.) MICCAI. Springer Nature Switzerland (2022)
  * [18] Liu, et al.: Contrastive attention for automatic chest X-ray report generation. In: Zong, et al. (eds.) Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021. pp. 269–280. Association for Computational Linguistics, Online (Aug 2021)
  * [19] Liu, F., Wu, X., Ge, S., Fan, W., Zou, Y.: Exploring and distilling posterior and prior knowledge for radiology report generation. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 13753–13762 (2021)
  * [20] Liu, G., Hsu, et al.: Clinically accurate chest x-ray report generation. In: Machine Learning for Healthcare Conference. pp. 249–269. PMLR (2019)
  * [21] Lu, J., Xiong, et al.: Knowing when to look: Adaptive attention via a visual sentinel for image captioning. arxiv 2016. arXiv preprint arXiv:1612.01887 (2016)
  * [22] Müller, et al.: Radiological reports improve pre-training for localized imaging tasks on chest x-rays. In: MICCAI 2022. pp. 647–657. Springer Nature Switzerland, Cham (2022)
  * [23] Packhäuser, K., et al.: Generation of anonymous chest radiographs using latent diffusion models for training thoracic abnormality classification systems. In: 2023 IEEE 20th International Symposium on Biomedical Imaging (ISBI). pp. 1–5. IEEE (2023)
  * [24] Serra, D., et al.: Controllable chest X-ray report generation from longitudinal representations. In: Bouamor, H., Pino, J., Bali, K. (eds.) Findings of the Association for Computational Linguistics: EMNLP 2023. pp. 4891–4904. Association for Computational Linguistics, Singapore (Dec 2023)
  * [25] Vaswani, et al.: Attention is all you need 30 (2017)
  * [26] Wang, et al.: A self-boosting framework for automated radiographic report generation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 2433–2442 (2021)
  * [27] Wang, et al.: An inclusive task-aware framework for radiology report generation. In: Wang, et al. (eds.) MICCAI. pp. 568–577. Springer Nature Switzerland, Cham (2022)
  * [28] Wang, et al.: A medical semantic-assisted transformer for radiographic report generation. In: MICCAI. pp. 655–664. Springer (2022)
  * [29] Wang, et al.: A medical semantic-assisted transformer for radiographic report generation. In: Wang, et al. (eds.) MICCAI. Springer Nature Switzerland (2022)
  * [30] Wang, J., Bhalerao, A., He, Y.: Cross-modal prototype driven network for radiology report generation. In: European Conference on Computer Vision. pp. 563–579. Springer (2022)
  * [31] Wang, Z., Liu, L., Wang, L., Zhou, L.: Metransformer: Radiology report generation by transformer with multiple learnable expert tokens. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11558–11567 (2023)
  * [32] Weber, et al.: Cascaded latent diffusion models for high-resolution chest x-ray synthesis. In: Pacific-Asia Conference on Knowledge Discovery and Data Mining. pp. 180–191. Springer (2023)
  * [33] Xu, K., Ba, et al.: Show, attend and tell: Neural image caption generation with visual attention. In: International conference on machine learning. pp. 2048–2057. PMLR (2015)
  * [34] Xu, T., et al.: Attngan: Fine-grained text to image generation with attentional generative adversarial networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 1316–1324 (2018)
  * [35] Yang, et al.: Knowledge matters: Chest radiology report generation with general and specific knowledge. Medical image analysis 80, 102510 (2022)
  * [36] Zhu, et al.: Utilizing longitudinal chest x-rays and reports to pre-fill radiology reports. In: MICCAI 2023. pp. 189–198. Springer Nature Switzerland, Cham (2023)

## Appendix

##  Appendix A1 Dataset split

Table 1: Number of patients and studies in the MIMIC-T dataset. #Number of patients | Train | Validation | Test
---|---|---|---
with 1 study | 35498 | 282 | 299
with 2 studies | 112393 | 1088 | 2066
Total studies(scans) | 260284 | 2298 | 4431

##  Appendix B1 Qualitative Comparison of X-ray and report generation tasks
against various methods

Table 1: Comparison for the Report Generation Task with other Report
Generation methods such as R2Gen, R2GenCMN, XProNet. CX - Current CXR, PX-
Previous CXR, CR - Current Report.  Table 2: Comparison for the Image
Generation Task with other Image Generation methods - RoentGen and Cascaded
Diffusion. CX - Current CXR, PX- Previous CXR, CR - Current Report.

##  Appendix C1 Ablation Study

Table 1: Ablation Study on the inclusion of different components of TiBiX.
BASE-Baseline model, CLS - Class token, TT - Temporal token

Method | Image Generation metrics | 14 diagnosis classification
---|---|---
| FID | SSIM | MSSIM | IS | Precision | Recall | AUC | Accuracy | F1
Base (CXRBERT + VQGAN) | 88.51 | 0.383 | 0.396 | 2.42±0.09 | 0.49 | 0.54 | 0.69 | 0.71 | 0.48
Base + CLS | 81.83 | 0.396 | 0.396 | 2.64±0.07 | 0.50 | 0.54 | 0.68 | 0.72 | 0.50
Base(Bi) + CLS + TT | 89.58 | 0.376 | 0.387 | 2.50±0.11 | 0.51 | 0.54 | 0.69 | 0.72 | 0.49
Unidirectional Image (+ CLS + TT) | 70.11 | 0.415 | 0.371 | 2.96±0.18 | 0.44 | 0.45 | 0.60 | 0.68 | 0.42
| Base(Bi) + CLS + TT (w ensemble)
---
33.85 | 0.515 | 0.472 | 1.82±0.08 | 0.56 | 0.55 | 0.75 | 0.77 | 0.53

Table 2: Comparison of unidirectional and bidirectional methods for the report
generation task

Method | NLG metrics | CE
---|---|---
| BLEU-3 | BLEU-4 | P | R | F1
Unidirectional Report | 0.172 | 0.146 | 0.292 | 0.200 | 0.222
Bidirectional(Ours) | 0.185 | 0.157 | 0.300 | 0.224 | 0.250


11institutetext:  Université de Lorraine, CentraleSupélec, CNRS, Inria, LORIA,
F-54000 France 22institutetext: Université de Paris, CNRS, Astroparticule et
Cosmologie, F-75013 Paris, France

# A Bayesian Convolutional Neural Network for Robust Galaxy Ellipticity
Regression

Claire Theobald 11 Bastien Arcelin 22 Frédéric Pennerath 11 Brieuc Conan-Guez
11 Miguel Couceiro 11 Amedeo Napoli 11

###### Abstract

Cosmic shear estimation is an essential scientific goal for large galaxy
surveys. It refers to the coherent distortion of distant galaxy images due to
weak gravitational lensing along the line of sight. It can be used as a tracer
of the matter distribution in the Universe. The unbiased estimation of the
local value of the cosmic shear can be obtained via Bayesian analysis which
relies on robust estimation of the galaxies ellipticity (shape) posterior
distribution. This is not a simple problem as, among other things, the images
may be corrupted with strong background noise. For current and coming surveys,
another central issue in galaxy shape determination is the treatment of
statistically dominant overlapping (blended) objects. We propose a Bayesian
Convolutional Neural Network based on Monte-Carlo Dropout to reliably estimate
the ellipticity of galaxies and the corresponding measurement uncertainties.
We show that while a convolutional network can be trained to correctly
estimate well calibrated aleatoric uncertainty, -the uncertainty due to the
presence of noise in the images- it is unable to generate a trustworthy
ellipticity distribution when exposed to previously unseen data (i.e. here,
blended scenes). By introducing a Bayesian Neural Network, we show how to
reliably estimate the posterior predictive distribution of ellipticities along
with robust estimation of epistemic uncertainties. Experiments also show that
epistemic uncertainty can detect inconsistent predictions due to unknown
blended scenes.

###### Keywords:

Bayesian Neural Networks Convolutional Neural Networks Epistemic uncertainty
Uncertainty calibration Cosmology.

##  1 Introduction

One of the goals of large galaxy surveys such as the _Legacy Survey of Space
and Time_ (LSST, [16]) conducted at the Vera C. Rubin Observatory is to study
_dark energy_. This component of unknown nature was introduced in the current
cosmological standard model to explain the acceleration of the Universe
expansion. One way to probe dark energy is to study the mass distribution
across the Universe. This distribution mostly follows the dark matter
distribution, which does not interact with baryonic matter (i.e. visible
matter) except through gravitation, as dark matter represents around 85% of
the matter in the Universe. Consequently, cosmologists need to use indirect
measurement techniques such as _cosmic shear_ , which measures the coherent
distortion of background galaxies images by foreground matter due to _weak
gravitational lensing_ [17]. In astrophysics, gravitational lensing is the
distortion of the image of an observed source, induced by the bending of
space-time, thus of the light path, generated by the presence of mass along
the line of sight. The mass acts like a lens, in partial analogy with optical
lenses, as illustrated in Fig.1a. The weak gravitational lensing effect is
faint (1% of galaxy shape measurement) and only statistical tools provide a
way to detect a local correlation in the observed galaxies orientations. This
correlation yields a local value at every point of the observable Universe,
defining the cosmic shear field. As pictured in Fig.1b, in an isotropic and
uniform Universe orientations of galaxies are expected to follow a uniform
distribution (left panel). The statistical average of their oriented
elongations, hereafter called _complex ellipticities_ , is expected to be
null. In presence of a lens, a smooth spatial deformation field modifies
coherently the complex ellipticities of neighboring galaxies so that their
mean is no longer zero (right panel).

(a) Gravitational lensing.

Galaxies randomly distributedWith shear: slight bias

(b) Cosmic shear.

Figure 1: LABEL:sub@fig:grav_lens Effect of gravitational lensing: the mass
bends the light and deforms the images of the galaxies.
LABEL:sub@fig:weak_lens Weak lensing: the correlation between orientations and
shapes of neighbour galaxies defines the cosmic shear. In blue: average
ellipticity. Left: the expected ellipticity distribution. Right: the observed
ellipticity distribution. Image: LABEL:sub@fig:grav_lens NASA/ESA

The unbiased measurement of cosmic shear is a major ambition of nowadays
cosmology [21]. One avenue to estimate the cosmic shear locally is to combine
individual galaxy ellipticity measurements. By looking deeper into the sky,
that is to older objects, the next generation of telescopes will allow for the
detection of a very large number of galaxies, potentially leading to very
precise shear measurement and resulting in tight constraints on dark energy
parameters.

Methods already exist to estimate galaxy ellipticities through direct
measurement on images recorded by telescope cameras ([13] for example). This
is a complex problem as, among other things, the shear signal is carried by
faint galaxies which makes it very sensitive to background noise. Another
central issue for current and coming surveys in galaxy shape determination, is
the treatment of statistically dominant overlapping objects, an effect called
_blending_. A current survey projects that 58% of the detected objects will
appear blended [18] and this value is expected to reach around 62% for LSST
[19]. To overcome this issue, solutions exist such as deblending [22, 23, 24]:
the separation of overlapping objects. Yet, they are not perfect and rely on
an accurate detection of blended scenes which is also a complex problem. As
such, in addition to a precise estimation of the complex ellipticities, a
reliable measurement of the uncertainties is crucial in order to discard, or
at least decrease the impact of, unreliable and inaccurate measurements
avoiding as much as possible the introduction of a bias into the shear
estimation.

Classical ellipticity measurement methods usually adopt assumptions about the
shape of the galaxies (for example via the shape of the window function in
[13]) potentially resulting in model bias. In contrast, _convolutional neural
networks_ or CNNs [2] make it possible to learn and recognize complex and
diverse galaxy shapes directly from data without making any other hypothesis
than the representativeness of the training sample. They consequently are
appropriate tools to learn the regression of galaxy ellipticities, even in the
presence of noise and complex distortions. Yet standard CNNs can only measure
the _aleatoric uncertainty_ : the one due to the presence of noise in the
data. They are unable to estimate the _epistemic uncertainty_ , the one due to
the limited number of samples a CNN has been trained with and to the model [1,
9]. This second type of uncertainty is essential to detect outliers from the
training samples, or formulated accordingly to our problem, to distinguish
between reliable or unreliable galaxy ellipticity estimation. It is only
accessible by considering neural network weights as random variables instead
of constants, that is, by adopting a Bayesian approach. Consequently, we have
focused our work on _Bayesian Deep Learning_ [11] using Monte Carlo dropout
(MC dropout) [1] as the mean to apply Bayesian inference to Deep Learning
models.

Foreseeing a Bayesian estimation of the cosmic shear, combining galaxy
ellipticity posteriors estimated directly from images (with blends or not) in
different filters (or bands), this paper focuses on estimating reliable galaxy
ellipticity posteriors from single band images. This is a necessary step to
check that the proposed method efficiently estimates a calibrated aleatoric
uncertainty and is able to minimize the impact of wrongly estimated
ellipticity values due to outliers in the computation of the shear. We compare
two networks trained on isolated galaxy images with or without noise in order
to test for the calibration of aleatoric uncertainty. Regarding outliers,
blended scenes are perfect examples. Note that these are illustrations of
aleatoric or epistemic uncertainty sources. Most of cosmic shear bias sources
such as detection, Point-Spread-Function (PSF) treatment, or selection for
example [17, 21], can fall in one or the other category. The estimation of
galaxy ellipticity posterior from blended scenes in different bands is a
harder problem that we will investigate in further work.

The contributions of this article are 1) to propose a Bayesian Deep Learning
model that solves a complex multivariate regression problem of estimating the
galaxy shape parameters while accurately estimating aleatoric and epistemic
uncertainties; 2) to establish an operational protocol to train such a model
based on multiple incremental learning steps; and 3) to provide experimental
evidences that the proposed method is able to assess whether an ellipticity
measurement is reliable. This is illustrated, in this paper, by the accurate
differentiation between isolated galaxy or blended scenes, considered here as
outliers, and the relationship between epistemic uncertainty and predictive
ellipticity error. We also show that this last result could not be obtained
with a classical, non Bayesian network.

The rest of the paper is organized as follows. In Section 2 we briefly
describe the problem to be solved and comment on some of its peculiarities. We
detail our proposed solution in Section 3. We analyse the results obtained on
the various experiments we performed in Section 4, and we conclude and give
the directions of further research in Section 5.

##  2 Estimating galaxy ellipticity from images

As mentioned previously, it is possible to estimate cosmic shear combining
individual measurements of galaxy shape. This shape information can be
quantified by the complex ellipticity, which can be defined in cosmology as in
Def. 1.

###### Definition 1

Let $E$ be an ellipse with major axis $a$, minor axis $b$, and with $\theta$
as its position angle. The _complex ellipticity_ of $E$ is defined as:

$\epsilon=\epsilon_{1}+\epsilon_{2}\,i=\frac{1-q^{2}}{1+q^{2}}\,e^{2i\theta},$
(1)

where $q=\frac{b}{a}$ is the axis ratio of the ellipse.

An illustration of the ellipticity parameters is shown in Fig. 2a. The complex
ellipticity defines a bijection between the orientation and the elongation of
the ellipse on one side, and the unit disk on the other side, see Fig. 2b.

$x$$y$$a$$b$$\theta$

(a) Ellipticity parameters: major axis $a$, minor axis $b$, position angle
$\theta$.

$x$$y$$0$$0.5$$1$$-0.5$$-1$$0.5i$$i$$-0.5i$$-i$

(b)  Bijective mapping between ellipse shapes and complex ellipticities.

Figure 2: Geometric representation of the complex ellipticity.
LABEL:sub@fig:ellipse_parameters The ellipse parameters.
LABEL:sub@fig:complex_bijection The complex ellipticity defines a bijection
between ellipse shapes and the unit disk. An ellipticity with low magnitude is
close to a circle, while one with a high magnitude is closer to a straight
line. The argument defines the orientation of the ellipse.

However, the process to achieve an unbiased measurement of cosmic shear,
starting with the estimation of ellipticities, is going to be challenging for
several reasons. We test the reliability of our networks prediction on noise
and blending, two of the many possible bias sources in the cosmic shear
estimation. Both of these issues result from the fact that the shear signal is
mostly carried by faint galaxies. By definition, these objects have a low
signal-to-noise ratio. The noise corrupts the galaxy images, making the shape
estimation much harder (see Fig. 3b), and can introduce a bias in shear
measurement [17].

(a) Isolated noiseless galaxy

(b) Isolated noisy galaxy

(c) Blended noisy galaxies

Figure 3: Three different types of image complexity for the same galaxy:
isolated without noise, isolated with noise, blended with noise. Notice how
the noise slightly deforms the galaxy LABEL:sub@fig:galim2 and how the blended
galaxies makes the ellipticity estimation very difficult LABEL:sub@fig:galim3
when compared to a simple isolated galaxy without noise LABEL:sub@fig:galim1.

Also, a large part of these faint objects will appear blended with foreground
galaxies. Even in scenes where objects are only slightly overlapped, the
apparent shape of the detected object does not correspond to a single galaxy
model and an ellipticity measurement on this image could give a completely
wrong result. Again, this work is the first step of a longer-term goal. Here,
we target a reliable estimation of galaxy ellipticity posterior from single
band images. This includes obtaining a well calibrated aleatoric uncertainty,
tested here with and without the addition of Poisson noise on images, and an
epistemic uncertainty allowing for minimization of the impact of untrustworthy
measurement due to outliers (here, blended scenes).

We simulate LSST-like images, allowing us to control the parameters of the
scene, e.g., the number of galaxies, their location on the image, and the
level and type of noise applied. We consider four categories of simulated
data: isolated centered galaxies without noise, isolated centered galaxies
with noise, and blended scene with and without noise. Images are 64x64 pixels
stamps simulated in the brightest of the six bands corresponding to the LSST
filters, each of them selecting a different part of the electromagnetic
spectrum. These images are simulated placing, in their center, a galaxy whose
ellipticity is to be measured.

The image generating process relies on the GalSim library [14] and is based on
a catalog of parametric models fitted to real galaxies for the third
Gravitational Lensing Accuracy Testing (GREAT3) Challenge [20]. It consists in
1) producing an image of a centered noiseless isolated galaxy from a model
sampled randomly from the catalog, with its corresponding physical properties
(size, shape, orientation, PSF, brightness, redshift, etc) 2) measuring the
complex ellipticity of the galaxy with the KSB algorithm [13] on the image and
record it as the image label, 3) possibly adding on random image location
other galaxy images (from 0 to 5) to generate blended scenes 4) possibly
adding Poisson noise (as in [24]). In this study and for sake of
interpretability, we only provide as input to our CNN the reference band (the
brightest) which we use to define the target ellipticity, making our images
two-dimensional. Once again, while using multiple bands is useful for blended
galaxies [23, 24], here we focus only on predicting the ellipticity of a
single centered galaxy with a correctly estimated uncertainty.

##  3 A method to assess uncertainty in ellipticity estimation

###  3.1 Estimation of noise related uncertainty

Our first goal is to reliably estimate the first layer of complexity in the
galaxy images, the noise. Given the nature of the data, we will be using a CNN
[2]. However, training a CNN to solve a standard regression problem with an L2
loss does not allow us to estimate the uncertainty due to the noise.
Therefore, in place of a complex scalar output, we predict a 2D _multivariate
normal distribution_ (MVN) as an output: given an input image $X$, whose
complex ellipticity is denoted $Y$ and given weight parameters $w$, the
network outputs an MVN $Y\sim\mathcal{N}(\mu(X,w),\Sigma(X,w))$. As such, the
model is no longer trained on a simple L2 loss but rather on the log-
likelihood of the MVN. The mean of the distribution $\mu(X,w)$, which is also
the mode, serves as the predicted output, and the covariance matrix
$\Sigma(X,w)$, which is also an output of the network, represents the so-
called aleatoric uncertainty on the input data $X$. The model is therefore
heteroscedastic, as $\Sigma(X,w)$ depends on the input $X$ [9]. This allows
our model to estimate the aleatoric uncertainty for each image individually.
The determinant of $\Sigma(X,w)$, denoted $|\Sigma(X,w)|$, is a scalar measure
of uncertainty, as it is directly related to the differential entropy,
$\ln\left(\sqrt{(2\pi e)^{2}|\Sigma(X,w)|}\right)$, of an MVN.

The architecture of our network is inspired by the work of Dielman, who
proposed a simple model specifically tuned for the Galaxy Zoo challenge,
therefore adapted to our data [12]. Each image is augmented in four different
parts by cropping thumbnails from high resolution images, centered on spatial
modes of light profile. Then each augmented image is fed to the CNN. The
complete architecture is explained in Fig. 4. More details on the training
process are explained in Section 3.3. Results obtained with this model are
given in Section 4.1.

(a) Convolutional layers (b) Fully connected layers

Figure 4: Convolutional neural network architecture. LABEL:sub@fig:nnim1 The
input after augmentation has dimensions 45x45x1. Each convolutional block
starts with a batch normalization layer and has a PReLU activation. The first
convolutional layer is of dimension 45x45x32 with a 5x5 kernel size (in
yellow), followed by a 2x2 Max-Pooling operation (in orange). The second
convolutional layer is 22x22x64 with kernel size 3x3, followed by a 2x2 Max-
Pooling operation. Then, we add two 11x11x128 convolutional layers with a 3x3
kernel that ends with a final 2x2 Max-Pooling operation, and the resulting
feature maps are flattened into a 3200 fully connected layer (in purple).
LABEL:sub@fig:nnim2 Each augmented image gives a 3200 fully connected layer
convolutional output (all augmented images share the same convolutional layers
and filters), which are then concatenated into a 12800 fully connected layer.
The two final layers have 4096 neurons in the case of an MVN regression, 2048
else; with Maxout activation [15] and dropout with a rate of 0.5. The output
layer has 5 neurons in the case of an MVN regression.

###  3.2 Estimation of blend related uncertainty

As seen in Section 2, estimating the uncertainty due to the noise in the data
is only one part of the problem. An estimated 60% of the images represent
blended scenes, for which a direct estimation of ellipticity does not make
sense in the context of this work. The uncertainty related to the blended
images cannot be estimated simply with the variance of the MVN distribution.
Indeed, in the case of a blended scene image, the network is not uncertain
because of the noise but rather because this kind of images is not part of the
training sample. This can be characterized by the epistemic uncertainty.

This uncertainty can be estimated using a Bayesian Neural Network (BNN), which
assumes a probability distribution on the weights $W$ of the network instead
of a single point estimate [11]. Given a prior $p(w)$ on $W$ and a set
$\mathcal{D}=\\{(X_{i},Y_{i})\\}_{i}$ of observations, the resulting posterior
distribution $p(w|\mathcal{D})\propto p(\mathcal{D}|w)\,p(w)$ is analytically
impossible to compute. A variational Bayes optimization method is necessary to
derive an approximate posterior $q_{\theta}(w)$ parameterized by hyper-
parameters $\theta$. In MC dropout [1, 3], the considered search space
includes all approximate posteriors resulting from applying _dropout_ [7], i.e
multiplying every neuron output (of selected layers) by an independent
Bernoulli variable. The dropout rate is set to the conventional value of 0.5,
as this leads to an approximate posterior that can achieve well calibrated
uncertainty estimates [1]. However, there are other ways to define the
posteriors such as dropout rate tuning [5], or ensemble methods by training
many networks [6]. During training, standard stochastic gradient descent
techniques can be used, thanks to the reparameterization trick, to search for
an approximate posterior maximizing locally the ELBO [1]. During testing, the
_posterior predictive distribution_ $p(Y|X,\mathcal{D})$ for some input $X$
can be estimated using Monte Carlo sampling:

$p(Y|X,\mathcal{D})\approx\int
p(Y|X,w)q_{\theta}(w)dw\approx\frac{1}{K}\sum_{k=1}^{K}p(Y|X,w_{k})\,,$ (2)

where $(w_{k})_{k=1}^{K}\sim q_{\theta}(w)$ refer to weights of $K$
independent dropout samples.

In the case of a multivariate regression problem like ours, every distribution
$p(Y|X,w_{k})$ is a MVN so that the resulting posterior predictive
distribution in Eq. 2 is a Gaussian mixture of order $K$. The uncertainty
underlying this mixture can be summarized by its covariance matrix
$\Sigma_{pred.}(X,\mathcal{D})=\mathrm{Cov}(Y|X,\mathcal{D})$. This matrix
accounts for both aleatoric and epistemic uncertainties, whose respective
contributions can actually be separated in a way that generalizes the variance
decomposition described in Depeweg [10]:

$\Sigma_{pred.}(X,\mathcal{D})=\Sigma_{aleat.}(X,\mathcal{D})+\Sigma_{epist.}(X,\mathcal{D})\,,$
(3)

where the first term represents the aleatoric uncertainty and can be computed
as the mean of the covariance matrices for each of the $K$ output samples:

$\displaystyle\Sigma_{aleat.}(X,\mathcal{D})$
$\displaystyle=\mathbb{E}_{W|\mathcal{D}}(\mathrm{Cov}(Y|X,W))\,,$
$\displaystyle\approx\frac{1}{K}\sum_{k=1}^{K}\Sigma(X,w_{k})\,.$ (4)

while the second represents the epistemic uncertainty and is estimated as the
empirical covariance matrix of the $K$ mean vectors produced as outputs:

$\displaystyle\Sigma_{epist.}(X,\mathcal{D})$
$\displaystyle=\mathrm{Cov}_{W|\mathcal{D}}(\mathbb{E}(Y|X,W))\,,$
$\displaystyle\approx\frac{1}{K}\sum_{k=1}^{K}(\mu(X,w_{k})-\mu(X))(\mu(X,w_{k})-\mu(X))^{T}\,.$
(5)

where $\mu(X)=\frac{1}{K}\sum_{k=1}^{K}\mu(X,w_{k})$.

Matrix $\Sigma_{epist.}(X,\mathcal{D})$ defines the epistemic uncertainty as
the covariance matrix of the mean vectors over the posterior. This uncertainty
will be high if the sampled predictions from each model vary considerably with
respect to $W$. This would mean that no consistent answer can be deduced from
the model and therefore it would be highly uncertain.

Finally when the context requires to reduce these uncertainty matrices to
uncertainty levels so that they can be compared, their determinants are used
to define two corresponding scalar quantities:

$\mathcal{U}_{aleat.}(X,\mathcal{D})=|\Sigma_{aleat.}(X,\mathcal{D})|\quad\text{and}\quad\mathcal{U}_{epist.}(X,\mathcal{D})=|\Sigma_{epist.}(X,\mathcal{D})|\,.$

###  3.3 Training protocol

In order to train a BNN with an MVN output, the model needs to learn both the
mean and the covariance matrix. The network’s training diverges when trying to
learn both at the same time, forcing us to separate the training into two
steps. First, we train a simple neural network without an MVN output - we use
only two output neurons representing the mean - using a L2 loss. Then, we
transfer the filters of the convolutional layers into the model with an MVN
output, but reinitialize the fully connected layers. This allows the model to
converge smoothly as the mean of the MVN distribution has already been
learned, allowing the covariance matrix to be calibrated accordingly.

This protocol works well when training on noiseless images of isolated
galaxies but fails when training on noisy images. Indeed, overfitting occurs
during the training of the network without MVN. When transferring the filters
to the MVN model, the mean of the MVN is not well calibrated enough and the
training of the BNN diverges. To fix this, we adjust the protocol for the
model without MVN, adding noise incrementally during training: we first submit
noiseless images, and modify 5% of the sample, switching from noiseless to
noisy images, every 50 epochs for 1000 epochs. This prevents overfitting and
allows the MVN model to converge after the transfer.

##  4 Experiments

###  4.1 Estimation of uncertainty related to noise

In this section we show that using an MVN as an output allows for a reliable
and well calibrated estimation of the aleatoric uncertainty, i.e. uncertainty
related to the noise in the data. In order to show that estimating the
ellipticity of galaxies in the presence of background noise is complex and can
induce incorrect predicted ellipticity values, we first train two simple CNNs
without an MVN output: one on noiseless images and one on noisy images,
accordingly tested on noiseless and noisy images respectively. Figure 5 shows
the images of galaxy with their target complex ellipticity superimposed, as
well as the predicted one.

(a) Predicted ellipticity without noise

(b) Predicted ellipticity with noise

Figure 5: Galaxy images with the predicted ellipticity superimposed on them.
The arrow and the corresponding elliptic shape are rendered in an arbitrary
scale for visualization purposes. In orange: the true ellipticity. In green:
the predicted ellipticity. LABEL:sub@fig:classic_noiseless_galaxy Galaxy
without noise. LABEL:sub@fig:classic_noise_galaxy Galaxy with noise.

The ellipse represents the estimated shape - with a fixed scale adapted for
visualization - and the arrow is the corresponding complex ellipticity -
modified with half its argument in order to be aligned wih the main axis of
the ellipse. On this example, we can qualitatively see that the galaxy
ellipticity on the noisy image is harder to estimate as the noise deforms the
shape of the galaxy. Figure 6 generalizes this observation as it shows a
sample of the predicted ellipticities on the complex plane within the unit
circle, with the target ellipticity and the difference between predicted and
targeted values.

(a) Predicted ellipticities without noise

(b) Predicted ellipticities with noise

Figure 6: Predicted ellipticities on the complex plane. In red: unit circle.
In yellow: predicted ellipticities. In blue: target ellipticites. In green:
difference between true and predicted values.

While the model trained on noiseless data performs really well (Fig. 6a), it
cannot achieve the same level of performance when trained on noisy data,
losing part of its reliability (Fig. 6b). As such, using a simple CNN without
any estimation of aleatoric uncertainty is not satisfying for our application.

We now train two Bayesian Convolutional Neural Networks with an MVN
distribution to estimate both epistemic and aleatoric uncertainties, as seen
in Section 3.2. Like the simple CNN models, we show in Fig. 7, the
ellipticities estimated from the BNNs on the complex plane. We also add the
90% confidence ellipses of both epistemic, aleatoric and predictive
uncertainties.

(a) Predicted ellipticities without noise

(b) Predicted ellipticities with noise

Figure 7: Predicted ellipticities on the complex plane. In red: unit circle.
In yellow: predicted ellipticities. In blue: target ellipticites. In light
green: difference between true and predicted values. In pink: 90% epistemic
confidence ellipse. In dark green: 90% aleatoric confidence ellipse. In grey:
90% predictive confidence ellipse.

We observe that in both cases, the epistemic uncertainty is low if not
negligible, meaning that the model is confident in its predictions. Put
another way, all $K$ pairs of outputs $\mu(X,w_{k})$ and $\Sigma(X,w_{k})$ are
roughly equal to their mean, respectively $\mu(X,w)$ and
$\Sigma_{aleat.}(X,w)$, so that, according to Eq. 3 and Eq. 5 ,
$\Sigma_{epist.}(X,w)\approx 0$ and
$\Sigma_{pred.}(X,w)\approx\Sigma_{aleat.}(X,w)$. The aleatoric uncertainty is
low for noiseless images but higher for noisy ones, confirming that the noise
corrupting galaxy images makes it more difficult for the model to consistently
give an accurate ellipticity estimation.

Finally, in order to see if the MVN distribution is well calibrated, we
standardize the output and check if the resulting distribution follows the
standard distribution. More precisely, if we define:

$Z(X,w)=\Sigma_{pred.}(X,w)^{-\frac{1}{2}}(Y-\mu(X,w))\quad,$ (6)

then the distributions of its two independent components $z_{1}\sim
Z(X,w)_{1}$ and $z_{2}\sim Z(X,w)_{2}$ should be equivalent to the standard
distribution $\mathcal{N}(0,1)$. Note that this is true only because all $K$
output MVNs are confounded. Figure 8 shows that the standardized distributions
for the model trained on noisy images are indeed well calibrated and therefore
the model is neither overestimating nor underestimating the predictive
uncertainty.

(a) Distribution of $z_{1}$ (b) Distribution of $z_{2}$

Figure 8: Histogram of the standardized distributions on the model trained
with noisy images. In red: standard bell curve. In blue: histogram of the
standardized distribution with the smoothed curve.

###  4.2 Estimation of uncertainty related to blending

In the previous part we showed that our BNNs are well calibrated. Here we
submit outliers to the networks in order to study the impact on epistemic
uncertainty and whether it can be used to detect them. Our models have only
been trained on images of isolated galaxies, but astrophysical images can
contain multiple overlapped galaxies. In that case, asking the model to
measure a single ellipticity does not make sense. If the epistemic uncertainty
behaves as expected, then its measurement would allow us to detect when a
predicted ellipticity is incorrect due to the presence of multiple galaxies in
the image. We fed images of blended scenes to the two models trained on
isolated galaxies (with or without noise), adding noise to the blended scenes
only for the model trained on noisy images.

Results shown in Fig. 9 demonstrate that in both cases the predictions are
particularly inexact when compared to the target ellipticity of the central
galaxy. Also, and as expected, the epistemic uncertainty is much higher for
these blended scenes than for isolated galaxy images. However, the aleatoric
uncertainty gives incoherent values as the model has not been trained to
evaluate it on blended images: notice how the aleatoric ellipses are more
flattened with a lower area. Figure 10 permits to visualise the behavior of
the epistemic uncertainty. It shows how the ellipticities sampled with dropout
slightly diverge compared to the mean prediction. Here the model cannot give a
consistent answer and therefore its prediction should be deemed untrustworthy.

(a) Predicted ellipticities, without noise

(b) Predicted ellipticities, with noise

Figure 9: Predicted ellipticities on the complex plane for blended galaxies
images. In red: unit circle. In yellow: predicted ellipticities. In blue:
target ellipticites (label of the centered galaxy). In light green: difference
between true and predicted values. In pink: 90% epistemic confidence ellipse.
In dark green: 90% aleatoric confidence ellipse. In grey: 90% predictive
confidence ellipse.

(a) Blended galaxies without noise

(b) Blended galaxies with noise

Figure 10: Blended galaxies images with the predicted ellipticity superimposed
on them. The arrow and the corresponding elliptic shape are rendered in an
arbitrary scale for visualization purposes. In orange: the true ellipticity
(label for the galaxy in the center). In green: the predicted ellipticity. In
pink: the individual MC dropout predicted ellipses. The green ellipticity is
therefore the mean of the pink ones. On both images the prediction is
uncertain as the individual MC samples slightly diverge from the mean.

To quantify the quality of the epistemic uncertainty when it comes to
detecting incoherent predictions due to outliers, we computed the ROC curves
for each uncertainty type. More precisely, we reduce each covariance matrix
(aleatoric, epistemic and predictive) to a scalar by computing its
determinant. We interpret these estimates as a scoring function to assess
whether an image is an outlier, i.e. a blended image: the higher the score,
the more likely the image contains a blend. Finally, we compute for each of
these scoring functions its ROC curve. We repeat that process for both
networks trained with noisy and noiseless data. The results are shown in Fig.
11.

(a) ROC curve, model without noise

(b) ROC curve, model with noise

Figure 11: ROC curves for detecting outliers for aleatoric, epistemic and
predictive uncertainty. LABEL:sub@fig:roc_curve_noiseless ROC curve, model
without noise. LABEL:sub@fig:roc_curve_noise ROC curve model with noise. Since
the aleatoric ROC curve gives incoherent answers on outliers (see Fig. 9), we
also plot the complementary classifier as a dashed line.

These ROC curves are also summarized by their associated Area Under Curve
(AUC) on Tab. 1. The epistemic uncertainty clearly appears as the most
consistent “metric” to detect outliers and therefore to give useful
information about the confidence in the model predictions. Even the predictive
uncertainty performs worse than the epistemic one. This is especially true in
the presence of noise since the aleatoric uncertainty then occupies a more
important part of the predictive one compared to the noiseless case. Notice
that the aleatoric ROC curve is mostly below the diagonal with an AUC below
0.5, meaning it performs worse than a random classifier. This is due to the
fact that the model has not been trained to evaluate aleatoric uncertainty on
blended scenes. As seen in Fig. 9, the aleatoric ellipses are more flattened
in the blended cases, meaning its determinant is lower. Thus the aleatoric
uncertainty is on average lower on blended scenes when compared to isolated
ones. To compensate, results of the complementary classifier for the aleatoric
uncertainty are shown. It is still not as satisfying as the epistemic
uncertainty. While using epistemic uncertainty to identify inconsistent
predictions due to a lack of knowledge is highly effective, we note that few
blended images still have low epistemic uncertainty due, for instance, to a
large galaxy that obstructs all of the other ones, making the image actually
closer to an isolated galaxy image.

Finally, we evaluate how each type of uncertainty is a reliable representation
of the risk of error in ellipticity prediction. Unfortunately, in the presence
of blended images, the predictive distribution is no longer a simple MVN but a
mixture of K well separated Gaussian distributions. The normalization process
that allowed us to obtain the results presented in Fig. 8 is no longer
applicable here. It is still possible to study the relationship between the
uncertainty and the ellipticity prediction error testing a trivial rule: the
higher the uncertainty, the more important we expect the error to be. To do
so, we do three sorting of the images according to each uncertainty type, from
the lowest uncertainty to the highest, on a scale from $0$ to $0.4$ for
isolated objects, and from $0.4$ to $1$ for blended scenes. We then compute
the mean ellipticity error considering the proportion of the sorted data from
0 to 1. For blended scenes the ellipticity prediction error is computed w.r.t.
the ellipticity of the centered galaxy. We repeat this experiment twice, for
networks trained on noiseless and noisy data. Finally we add an ”oracle” curve
where the data is sorted directly according to the ellipticity prediction
error which represents a perfect sorting. Results are shown in Fig. 12.

Uncertainty | AUC noiseless | AUC noise
---|---|---
Epistemic | 0.956 | 0.969
Aleatoric | 0.394 | 0.306
Aleatoric (inverse) | 0.606 | 0.694
Predictive | 0.856 | 0.594

Table 1: AUC values for all uncertainties, for the model with and without
noise. Here, the epistemic uncertainty is clearly the best to detect outliers,
as its AUC value is close to 1 in both noisy and noiseless datasets.

(a) Mean error curve, model without noise

(b) Mean error curve, model with noise

Figure 12: Mean error curves w.r.t. data proportion for aleatoric, epistemic
and predictive uncertainty. LABEL:sub@fig:error_curve_noiseless Mean error
curve without noise. LABEL:sub@fig:error_curve_noise Mean error curve with
noise. In black the threshold between the proportion of isolated galaxies:
$[0,0.4]$ and blended galaxies: $[0.4,1]$. In pink the oracle curve, where the
data is sorted by the predictive error. The closest a curve is to the oracle
the better.

Once again, epistemic uncertainty proves to be best suited to anticipate
ellipticity predictive error. The samples with the lowest epistemic
uncertainty have the lowest mean ellipticity error and conversely, while
samples with low aleatoric uncertainty can already have high mean error.
Consequently, on real astrophysical data, when the predictive ellipticity
error is obviously unknown, relying on the epistemic uncertainty to reject, or
minimize the impact of, a sample because of its probable predictive error is
the best way to go.

##  5 Conclusion

We developed a Bayesian approach to estimate the posterior distribution of
galaxy shape parameters using convolutional neural networks and MC-Dropout. In
addition to a precise measurement of the ellipticities, this approach provides
a calibrated estimation of the aleatoric uncertainty as well as an estimation
of the epistemic uncertainty. We showed that the latter is behaving according
to expectations when applied to different kind of galaxy images, and is well-
suited to identify outliers and to anticipate high predictive ellipticity
error. These results confirm the suitability of Bayesian neural networks for
galaxy shape estimation and incite us to continue exploring their use to go
from ellipticity posterior distributions, estimated from multi-band galaxy
images, to cosmic shear estimation.

####  5.0.1 Acknowledgements

The first author is preparing a PhD thesis at the LORIA Lab in the context of
the AstroDeep Research Project
(https://astrodeep.pages.in2p3.fr/website/projects/) funded by ANR under the
grant ANR-19-CE23-0024.

## References

  * [1] Gal, Y., Ghahramani, Z.: Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In: ICML Proceedings 2016, JMLR Workshop and Conference Proceedings, vol. 48, pp. 1050–1059.
  * [2] LeCun, Y., Bengio, Y. and Hinton, G.: Deep learning. Nature, 521(7553), pp. 436-444 (2015).
  * [3] Gal, Y., 2016, Uncertainty in Deep Learning, PhD thesis, University of Cambridge
  * [4] Gal, Y., Hron, J., Kendall, A.: Concrete Dropout. In: NIPS Proceedings 2017, pp. 3581-3590.
  * [5] Lakshminarayanan, B., Pritzel, A., Blundell, A.: Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. In: NIPS Proceedings 2017, pp. 6402-6413.
  * [6] Kendall, A., Gal, Y.: What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision? In: NIPS Proceedings 2017, pp. 5574–5584.
  * [7] Srivastava, N., et al.: Dropout: a simple way to prevent neural networks from overfitting. In: J. Mach. Learn. Res., vol. 15, pp. 1929-1958 (2014).
  * [8] Le, V. Q., Smola, J. A., Canu, S.: Heteroscedastic Gaussian process regression. In: ICML Proceedings 2005, vol. 119, pp. 489–496
  * [9] Hüllermeier, E., Waegeman, W.: Aleatoric and Epistemic Uncertainty in Machine Learning: A Tutorial Introduction. CoRR, http://arxiv.org/abs/1910.09457 (2019)
  * [10] Depeweg, S., et al.: Decomposition of uncertainty in bayesian deep learning for efficient and risk-sensitive learning. In: ICML Proceedings 2018, vol. 80, pp. 1192–1201.
  * [11] Hinton, G.E., van Camp, D.: Keeping the neural networks simple by minimizing the description length of the weights. In: Proceedings of the Sixth Annual ACM Conference on Computational Learning Theory, COLT 1993. pp. 5–13.
  * [12] Dielman, S.: My solution for the Galaxy Zoo challenge (2014) https://benanne.github.io/2014/04/05/galaxy-zoo.html
  * [13] Kaiser N., Squires G., Broadhurst T.: A Method for Weak Lensing Observations, In: Astrophys.J. 449:460-475 (1995)
  * [14] Rowe B., et al.: GalSim: The modular galaxy image simulation toolkit, https://github.com/GalSim-developers/GalSim (2015)
  * [15] Goodfellow, I.J., et al.: Maxout Networks. In: Proceedings of the 30th International Conference on Machine Learning, vol. 28, pp. 1319-1327 (2013)
  * [16] LSST Science Collaboration, Abell P. A.,et al.: LSST Science Book. 2009, arXiv, arXiv:0912.0201
  * [17] Kilbinger, M.: Cosmology with cosmic shear observations: a review. 2015, Reports on Progress in Physics 78, 086901.
  * [18] Bosch, J., et al.: The Hyper Suprime-Cam Software Pipeline. 2018, Publications of the Astronomical Society of Japan 70, S5.
  * [19] Sanchez, J., Mendoza, I., Kirkby, D.P., and Burchat: Effects of overlapping sources on cosmic shear estimation: Statistical sensitivity and pixel-noise bias. 2021, arXiv e-prints, arXiv:2103.02078.
  * [20] Mandelbaum, R., et al.: The Third Gravitational Lensing Accuracy Testing (GREAT3) Challenge Handbook. 2014, The Astrophysical Journal Supplement Series 212, 5.
  * [21] Mandelbaum, R.: Weak lensing for precision cosmology. 2018, Annual Review of Astronomy and Astrophysics 56, 393.
  * [22] Bertin, E. and Arnouts, S.: SExtractor: Software for source extraction. 1996, Astronomy and Astrophysics Supplement Series 117, 393.
  * [23] Melchior, P., Moolekamp, F., Jerdee, M., Armstrong, R., Sun, A.-L., Bosch, J., and Lupton, R.: SCARLET: Source separation in multi-band images by Constrained Matrix Factorization. 2018, Astronomy and Computing 24, 129.
  * [24] Arcelin, B., Doux, C., Aubourg, E., Roucelle, C., and LSST Dark Energy Science Collaboration: Deblending galaxies with Variational Autoencoders: a joint multi-band, multi-instrument approach. 2021, Monthly Notices of the Royal Astronomical Society 500, 531.