11institutetext: The University of Tokyo, Tokyo, Japan
11email<EMAIL_ADDRESS>
22institutetext: Preferred Networks, Tokyo, Japan
22email<EMAIL_ADDRESS>
# Label-Efficient Multi-Task Segmentation
using Contrastive Learning

Junichiro Iwasawa This work was done when J.I. worked at Preferred Networks as
an intern and part-time researcher11 Yuichiro Hirano 22 Yohei Sugawara 22

###### Abstract

Obtaining annotations for 3D medical images is expensive and time-consuming,
despite its importance for automating segmentation tasks. Although multi-task
learning is considered an effective method for training segmentation models
using small amounts of annotated data, a systematic understanding of various
subtasks is still lacking. In this study, we propose a multi-task segmentation
model with a contrastive learning based subtask and compare its performance
with other multi-task models, varying the number of labeled data for training.
We further extend our model so that it can utilize unlabeled data through the
regularization branch in a semi-supervised manner. We experimentally show that
our proposed method outperforms other multi-task methods including the state-
of-the-art fully supervised model when the amount of annotated data is
limited.

###### Keywords:

Multi-task learning Brain tumor segmentation Semi-supervised learning.

##  1 Introduction

For precision medicine, it is imperative that interpretation and
classification of medical images is done quickly and efficiently; however, it
is becoming a major hurdle due to the shortage of clinical specialists who can
provide informed clinical diagnoses. Automated segmentation can not only save
physicians’ time but can also provide accurate and reproducible results for
medical analysis. Recent advances in convolutional neural networks (CNN) have
yielded state-of-the-art segmentation results for both 2D and 3D medical
images [15, 11, 12], a significant step toward fully automated segmentation.
However, this level of performance is only possible when sufficient amount of
labeled data is available. Furthermore, obtaining annotations from medical
experts is both expensive and time-consuming, despite its importance for
training CNNs. Thus, methods that utilize small labeled datasets have been
explored extensively. Specifically, multi-task learning has been considered as
an efficient method for small data, since parameter sharing for both the main
segmentation task and regularization subtask could reduce the risk of
overfitting [16, 12]. Although subtasks in multi-task learning have been
extensively investigated, we still lack a systematic understanding of the
impact of subtasks on the main segmentation model, especially in the low
labeled data regime.

Contrastive learning based approaches have recently shown state-of-the-art
performance in image classification with small amount of labels [19, 7, 5, 3].
In this work, we integrated a contrastive learning based approach as a subtask
of a multi-task segmentation model. However, its applicability in segmentation
tasks, especially medical image segmentation is yet to be explored [5]. Here,
we systematically assess the performance of a brain tumor segmentation problem
for three different multi-task models including our proposed method. The main
contributions can be summarized as follows111Our implementation is available
at github.com/pfnet-research/label-efficient-brain-tumor-segmentation.:

  * •

We propose a novel method for tumor segmentation by utilizing contrastive
learning as a subtask for the main segmentation model.

  * •

We experimentally show that our proposed method, combined with a semi-
supervised approach which utilizes unlabeled data, could enhance segmentation
performance when the amount of labeled data is small.

Figure 1: Schematic image of the encoder-decoder based network combined with a
regularization branch. The green block represents a ResBlock.

##  2 Methods

###  2.1 Encoder-decoder network with regularization branches

Fig. 1 is the schematic image of our model. The model constitutes an encoder-
decoder architecture with a regularization branch starting from the encoder
output. The encoder and decoder are composed of ResNet [6]-like blocks
(ResBlock) with skip-connections between the encoder and decoder [15, 11, 12].
We adopted the encoder-decoder architecture from Myronenko (2018) [12], and
each ResBlock consists of two $3\times 3\times 3$ convolution layers with
group normalization [21].

The loss function of the model consists of two terms: $\mathcal{L}_{\rm
total}=\mathcal{L}_{\rm Dice}+\mathcal{L}_{\rm branch}$, where
$\mathcal{L}_{\rm Dice}$ is a softmax Dice loss applied to the decoder output
$p_{\rm pred}$, compelling it to match the ground truth label $p_{\rm true}$.
The Dice loss [11] is given by,

$\mathcal{L}_{\rm Dice}\left(p_{\rm true},p_{\rm
pred}\right)=1-\frac{2\times\sum_{i}p_{{\rm true},i}p_{{\rm
pred},i}}{\sum_{i}(p_{{\rm true},i}^{2}+p_{{\rm pred},i}^{2})+\epsilon},$ (1)

where $\epsilon=10^{-7}$ is a constant preventing zero division.
$\mathcal{L}_{\rm branch}$ is the loss applied to the output of the
regularization branch and is dependent on the architecture of the branch.

To explore the regularization effects of the multi-task learning, we compared
the performance of different subtasks for the regularization branch on the
encoder endpoint. Based on the type of information it uses to calculate the
loss function, subtasks for multi-task learning can be categorized into the
following classes:

  1. 1.

Subtasks that use the input $X$ itself, without using any additional
information, such as a decoder-like branch attempting to reconstruct the
original input $X$ [12].

  2. 2.

Subtasks that attempt to predict a transformed feature of the label $y$, such
as boundary-aware networks that predict the boundary of the given label [4,
13].

  3. 3.

Subtasks that compel the encoder to obtain certain representations of the
input by predicting low or high-level features from the input $X$, such as
tasks predicting the angles of rotated images [14]. Our proposed method using
contrastive predictive coding (CPC) [19, 7] would also be classified in this
class.

To investigate a wide spectrum of subtasks, we implemented three different
types of regularization branches that use either a variational autoencoder
(VAE) loss, a boundary loss, or a CPC loss. The description of each branch and
loss is provided subsequently (the architectures are shown in supplementary
materials).

####  2.1.1 Variational autoencoder branch

The purpose of the VAE branch is to guide the encoder to extract an efficient
low-dimensional representation for the input [12]. The encoder output is first
fed to a linear layer to reduce its dimension to 256 (128 to represent the
mean $\mu$, and 128 to represent the standard deviation (SD) $\sigma$).
Accordingly, a sample is drawn from a 128-dimension Gaussian distribution with
the given mean and SD. This sample is processed by several upsizing layers,
where the number of features is reduced by a factor of two using $1\times
1\times 1$ convolution, and the spatial dimension is doubled using 3D bilinear
sampling, so that the final output size matches the input size. The output is
fed to the VAE loss ($\mathcal{L}_{\rm VAE}$) which is given by
$\mathcal{L}_{\rm branch}=\mathcal{L}_{\rm
VAE}=0.1\times\left(\mathcal{L}_{\rm rec}+\mathcal{L}_{\rm KL}\right)$, where
$\mathcal{L}_{\rm rec}$ is the L2 loss between the output of the VAE branch
and the input image, and $\mathcal{L}_{\rm KL}$ is the Kullback-Leibler
divergence between the two normal distributions, $\mathcal{N}(\mu,\sigma^{2})$
and $\mathcal{N}(0,1)$.

####  2.1.2 Boundary attention branch

The boundary attention branch aims to regularize the encoder by extracting
information for predicting the boundaries of the given labels [20, 4, 13]. We
prepared the boundary labels by applying a 3D Laplacian kernel to the ground
truth binary labels. The attention layer first upsizes the output of the
encoder, and concatenates it with the feature map from the encoder with the
same spatial dimensions. Accordingly, a $1\times 1\times 1$ convolution is
applied to the concatenated sample, followed by a sigmoid function that yields
an attention map. The output of each attention layer is given by an element-
wise multiplication of the attention map and the input to the layer. This
operation is repeated until the spatial dimension matches that of the model
input.

The loss for the boundary attention branch is given by $\mathcal{L}_{\rm
branch}=\mathcal{L}_{\rm boundary}=\mathcal{L}_{\rm Dice}\left(b_{\rm
pred},b_{\rm true}\right)+\mathcal{L}_{\rm edge}$, where $\mathcal{L}_{\rm
Dice}\left(b_{\rm pred},b_{\rm true}\right)$ is the Dice loss between the
branch output ($b_{\rm pred}$) and boundaries of the ground truth label
($b_{\rm true}$). $\mathcal{L}_{\rm edge}$ is given by a weighted binary cross
entropy loss:

$\mathcal{L}_{\rm edge}=-\beta\sum_{j\in y_{+}}\log{P\left(b_{{\rm
pred},j}=1|X,\theta\right)}-(1-\beta)\sum_{j\in y_{-}}\log{P\left(b_{{\rm
pred},j}=0|X,\theta\right)},$ (2)

where $X$, $\theta$, $y_{+}$ and $y_{-}$ denote the input, model parameters,
and the boundary and non-boundary voxels, respectively. $\beta$, the ratio of
the non-boundary voxels to the total number of voxels is introduced to handle
the imbalance of the boundary and non-boundary voxels.

####  2.1.3 Contrastive predictive coding branch

Self-supervised learning, where a representation is learned from unlabeled
data by predicting missing input data based on other parts of the input, has
become a promising method for learning representations; it is useful for
downstream tasks such as classification [9, 7] and segmentation [22].
Recently, CPC has been proposed as a self-supervised method that can
outperform the fully supervised methods in ImageNet classification tasks in
the small labeled data regime [7]. Despite its performance in classification
tasks, the effectiveness of CPC in segmentation tasks is yet to be explored.
Here, we incorporated the CPC architecture into the encoder-decoder structure
as a regularization branch and investigated its performance with regard to
medical image segmentation.

First, we divided the input image into $32\times 32\times 32$ overlapping
patches with a 16-voxel overlap (resulting in an $8\times 8\times 7$ grid for
the brain tumor dataset). Each divided patch is individually encoded to a
latent representation using the encoder in the encoder-decoder architecture
and spatially mean-pooled into a single feature vector $z_{i,j,k}$. Here, it
should be noted that the visual field of the encoder when using the CPC branch
is $32\times 32\times 32$ pixels, which is smaller than that with the other
regularization branches due to the initial image division. Accordingly, eight
layered ResBlocks $f_{\rm res8}$ and a linear layer $W$ were applied to the
upper half of the feature vectors, $\hat{z}_{i,j,k_{\rm low}}=W\left(f_{\rm
res8}\left(z_{i,j,k_{\rm up}}\right)\right)$, to predict the lower half of the
feature vectors $z_{i,j,k_{\rm low}}$. The predictions were evaluated based on
the CPC loss (i.e. the InfoNCE [19]),

$\mathcal{L}_{\rm branch}=\mathcal{L}_{\rm CPC}=-\sum_{i,j,k_{\rm
low}}\log{\frac{\exp{\left(\hat{z}^{T}_{i,j,k_{\rm low}}z_{i,j,k_{\rm
low}}\right)}}{\exp{\left(\hat{z}^{T}_{i,j,k_{\rm low}}z_{i,j,k_{\rm
low}}\right)}+\sum_{l}\exp{\left(\hat{z}^{T}_{i,j,k_{\rm
low}}z_{l}\right)}}},$ (3)

where the negative samples $\\{z_{l}\\}$ are randomly taken from the other
feature vectors which are encoded from different patches of the same image.

Figure 2: Comparison of the Dice scores of the test predictions by the
proposed models with different training data sizes. For the ssVAEseg and
ssCPCseg, the training data size refers to the number of labeled images.

##  3 Experiments and Results

####  3.0.1 Brain tumor dataset

We used the brain tumor dataset provided by the Medical Segmentation Decathlon
[17], which is a subset of the data used in the 2016 and 2017 Brain Tumor
Image Segmentation (BraTS) challenges [10, 1, 2]. This dataset includes
multimodal 3D magnetic resonance imaging (MRI) scans from patients diagnosed
with either glioblastoma or lower-grade glioma with an image size of
$240\times 240\times 155$. For the experiments, we randomly split the 484
labeled data into the training (387), validation (48), and test set (49).

The dataset has three different labels corresponding to different tumor sub-
regions (i.e., necrotic and non-enhancing parts of the tumor (NCR & NET), the
peritumoral edema (ED), and the enhancing tumor (ET)). Following the BraTS
challenge, we evaluated the model’s prediction accuracy using the three nested
structures of these sub-regions: ET, tumor core (TC: ET+NCR+NET), and the
whole tumor (WT: TC+ED). The output of the network was set to three channels
to predict each of the three nested structures above.

####  3.0.2 Preprocessing and augmentation

All the input images were normalized to have zero mean and unit SD.
Accordingly, a random scale (0.9, 1.1), random intensity shift (-0.1, 0.1 of
SD), and random axis mirror flip (all axes, probability 0.5) were applied. All
the inputs were randomly cropped to $(160,192,128)$, except when utilizing the
CPC branch ($(144,144,128)$ in this case), due to the graphics processing unit
(GPU) memory limitation. Note that the random cropping was performed without
respect to the location of the brain region.

Table 1: Performance of EncDec, VAEseg, Boundseg, CPCseg (trained with six labeled data), and ssVAEseg, ssCPCseg (trained with six labeled and 381 unlabeled data). Evaluation was done using the mean Dice and $95^{\rm th}$ percentile Hausdorff distance. The performance of VAEseg trained with 387 labels is shown for comparison. | Dice | Hausdorff distance (mm)
---|---|---
Test data |  ET |  TC |  WT |  ET |  TC | WT
VAEseg (387 labels) |  0.9077 |  0.9323 |  0.9536 |  3.6034 |  10.2344 | 8.3895
EncDec (6 labels) |  0.8412 |  0.8383 |  0.9144 |  11.4697 |  20.12 | 24.3726
VAEseg (6 labels) |  0.8234 |  0.8036 |  0.8998 |  14.3467 |  22.4926 | 17.9775
Boundseg (6 labels) |  0.8356 |  0.8378 |  0.9041 |  17.0323 |  27.2128 | 25.8112
CPCseg (6 labels) |  0.8374 |  0.8386 |  0.9057 |  10.2839 |  14.9661 | 15.0633
ssVAEseg (6 labels) |  0.8626 |  0.8425 |  0.9131 |  9.1966 |  12.5302 | 14.8056
ssCPCseg (6 labels) |  0.8873 |  0.8761 |  0.9151 |  8.7092 |  16.0947 |  12.3962

####  3.0.3 Comparison of different regularization branches

We implemented our network in Chainer [18] and trained it on eight Tesla V100
32GB GPUs. We used a batch size of eight, and the Adam optimizer [8] with an
initial learning rate of $\alpha_{0}=10^{-4}$ that was further decreased
according to $\alpha=\alpha_{0}(1-n/N)^{0.9}$, where $n,N$ denotes the current
epoch and the total number of epochs, respectively.

To compare the different regularization effects, we measured the segmentation
performances of four models: the encoder-decoder alone (EncDec), EncDec with a
VAE branch (VAEseg), EncDec with a boundary attention branch (Boundseg), and
EncDec with a CPC branch (CPCseg). The performance was evaluated using the
mean Dice and $95^{\rm th}$-percentile Hausdorff distance of the ET, TC, and
WT. To evaluate the regularization effect with varied amounts of labeled data,
we also evaluated each model’s performance with the training data size reduced
to 3 – 96. The results are given in Fig. 2 and Table 1, S4.

It can first be observed that no regularization branch consistently
outperforms the others. Furthermore, it can be observed that in some cases the
EncDec that has no regularization branch had the highest mean Dice score among
the fully supervised models (Table 1). However, it should also be noted that
multi-task models tended to outperform the EncDec when using all 387 labeled
data (Table S4). These results imply that the regularization branches using
labeled data have a limited effect on the segmentation performance when the
amount of labeled data is small. Interestingly, the VAEseg, the state-of-the-
art model in the BraTS 2018 challenge [12], was not necessarily the best model
for various training data sizes. This was surprising, although the dataset we
used in this study slightly differs from that used in the BraTS 2018
challenge. Our results suggest that the VAE subtask would not always be the
optimal approach to brain tumor segmentation tasks.

####  3.0.4 Semi-supervised multi-task learning

Typically, unlabeled data is readily accessible in large quantities, compared
to annotated data. This naturally leads to the question: is it possible to
utilize unlabeled data to guide the segmentation model at small data regimes?
To answer this question, we focused on the VAEseg and CPCseg, because they do
not require labels to optimize the regularization branch. For training, we
used all the 387 unlabeled data from the training set, and varied the number
of labels used. To utilize the unlabeled data, we devised a semi-supervised
update method wherein the model could be updated using only $\mathcal{L}_{\rm
branch}$ when the image had no label, and by $\mathcal{L}_{\rm
Dice}+\mathcal{L}_{\rm branch}$, otherwise. This update method lets the
encoder (and regularization branch) learn representations from both labeled
and unlabeled data. The segmentation results for the semi-supervised VAEseg
(ssVAEseg) and semi-supervised CPCseg (ssCPCseg) are shown in Fig. 2, 3 and
Table 1. It can be observed that the semi-supervised methods outperform their
fully supervised counterparts. In addition, the ssCPCseg outperformed all the
other regularization methods, including the fully supervised state-of-the-art
model, VAEseg, in the small labeled data regime. This tendency was most
apparent for the ET and TC. For example, the ssCPCseg using six labels
exhibited a 6% decrease in the Dice score for TC compared to the VAEseg using
387 labels, while other methods using six labels exhibited a 10 – 14% decrease
(Table 1). We speculate that this is because the areas of the ET and TC are
smaller, compared to that of the WT, and thus, provide less supervision
signals per sample to the model. Our results imply that difficult and non-
trivial subtasks such as CPC, as well as unlabeled data, can be exploited to
achieve state-of-the-art performance when the amount of annotated data is
limited.

Figure 3: Prediction results of the semi-supervised CPCseg (24 labeled + 363
unlabeled) and the CPCseg (24 labeled). The union of purple, orange, and
yellow corresponds to the WT, orange plus yellow corresponds to the TC, and
yellow corresponds to the ET.

To explore the importance of the number of unlabeled data to the segmentation
performance, we compared the ET Dice score of the ssCPCseg trained using six
labeled and different amounts of unlabeled data (Fig. 4). It can be observed
that the Dice score increases monotonically with the number of unlabeled data.
However, it should be noted that the Dice score seems to increase linearly
with the $\log$ of unlabeled data, indicating that the number of data required
to improve accuracy would increase exponentially. We also investigated the
effect of unlabeled data on the learning behavior of the model. As can be seen
in Fig. 4, the semi-supervised CPCseg not only produced more accurate results
but was also more robust to overfitting. This might be because the encoder
needs to map good representations for both the labeled and unlabeled data to
optimize the CPC objective. Overall, our results imply that utilizing
unlabeled data could both enhance and stabilize the model’s performance.

Figure 4: The left panel shows the dependency of the number of unlabeled data
on the segmentation accuracy for the ssCPCseg when using six labeled data. The
right panel show learning curves for the test data of the brain tumor dataset.
The learning curves for the CPCseg (24 labeled) and ssCPCseg (24 labeled + 363
unlabeled) are shown.

##  4 Discussion and Conclusion

In this work, we systematically investigated the effectiveness of different
regularization subtasks for medical image segmentation. Our experiments on the
brain tumor dataset showed that utilizing unlabeled data through the
regularization branch improved and stabilized the performance of segmentation
models when the number of labeled data was small. Especially, our proposed
ssCPCseg outperformed other methods including the state-of-the-art fully
supervised model in the small labeled data regime. In previous works, CPC has
been used for self-supervised pre-training for image classification tasks with
unlabeled images [19, 7]. Our work is the first to show the effectiveness of
CPC as a regularization subtask for image segmentation by utilizing both
unlabeled and labeled images, providing a novel direction for label efficient
segmentation. It should also be noted that ssCPCseg achieved particularly
higher Dice scores than the other methods for tumor sub-regions with small
area size (i.e. ET and TC) whose structure varies across the data and provides
limited supervision signals to the model. Although our results are based on a
single dataset, we believe that our method could be applicable to various
targets in the field since target labels for medical images are often
relatively small, and have widely varied structures.

It is generally expensive to obtain annotations for 3D medical images. On the
other hand, large number of unlabeled images are often available. Thus, our
semi-supervised method should have wide applicability to medical image
segmentation tasks. However, it should be noted that all the MRI scans in the
brain tumor dataset are normalized using a reference atlas; furthermore, they
have the same size and voxel resolution. Therefore, we have not been able to
evaluate the segmentation performance when the quality of the unlabeled data
varies. An important future work is to evaluate the model’s segmentation
performance when it is fed unlabeled images from different modalities and
domains. We believe that our systematic study provides important designing
principles for segmentation models, leading to more cost-efficient medical
image segmentation.

####  4.0.1 Acknowledgments

J.I. was supported by the Grant-in-Aid for JSPS Fellows JP18J21942.

## References

  * [1] Bakas, S., Akbari, H., Sotiras, A., Bilello, M., Rozycki, M., Kirby, J.S., Freymann, J.B., Farahani, K., Davatzikos, C.: Advancing The Cancer Genome Atlas glioma MRI collections with expert segmentation labels and radiomic features. Scientific Data 4, 170117 (2017)
  * [2] Bakas, S., et al.: Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge. arXiv e-prints arXiv:1811.02629 (2018)
  * [3] Chen, T., Kornblith, S., Norouzi, M., Hinton, G.: A Simple Framework for Contrastive Learning of Visual Representations. arXiv e-prints arXiv:2002.05709 (2020)
  * [4] Hatamizadeh, A., Terzopoulos, D., Myronenko, A.: End-to-End Boundary Aware Networks for Medical Image Segmentation. In: Machine Learning in Medical Imaging - 10th International Workshop, MLMI 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China, 2019, Proceedings. pp. 187–194 (2019)
  * [5] He, K., Fan, H., Wu, Y., Xie, S., Girshick, R.: Momentum Contrast for Unsupervised Visual Representation Learning. In: The IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (2020)
  * [6] He, K., Zhang, X., Ren, S., Sun, J.: Identity Mappings in Deep Residual Networks. In: Computer Vision - ECCV 2016 - 14th European Conference, Amsterdam, The Netherlands, 2016, Proceedings, Part IV. pp. 630–645 (2016)
  * [7] Hénaff, O.J., Srinivas, A., De Fauw, J., Razavi, A., Doersch, C., Eslami, S.M.A., van den Oord, A.: Data-Efficient Image Recognition with Contrastive Predictive Coding. arXiv e-prints arXiv:1905.09272 (2019)
  * [8] Kingma, D.P., Ba, J.: Adam: A Method for Stochastic Optimization. In: 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, 2015, Conference Track Proceedings (2015)
  * [9] Kolesnikov, A., Zhai, X., Beyer, L.: Revisiting Self-Supervised Visual Representation Learning. In: IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, 2019\. pp. 1920–1929 (2019)
  * [10] Menze, B.H., et al.: The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS). IEEE Transactions on Medical Imaging 34(10), 1993–2024 (2015)
  * [11] Milletari, F., Navab, N., Ahmadi, S.: V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation. In: Fourth International Conference on 3D Vision, 3DV 2016, Stanford, CA, USA, 2016. pp. 565–571 (2016)
  * [12] Myronenko, A.: 3D MRI Brain Tumor Segmentation Using Autoencoder Regularization. In: Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries - 4th International Workshop, BrainLes 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, 2018, Revised Selected Papers, Part II. pp. 311–320 (2018)
  * [13] Myronenko, A., Hatamizadeh, A.: 3D Kidneys and Kidney Tumor Semantic Segmentation using Boundary-Aware Networks. arXiv e-prints arXiv:1909.06684 (2019)
  * [14] Noroozi, M., Favaro, P.: Unsupervised Learning of Visual Representations by Solving Jigsaw Puzzles. In: Computer Vision - ECCV 2016 - 14th European Conference, Amsterdam, The Netherlands, 2016, Proceedings, Part VI. pp. 69–84 (2016)
  * [15] Ronneberger, O., Fischer, P., Brox, T.: U-Net: Convolutional Networks for Biomedical Image Segmentation. In: Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015 - 18th International Conference Munich, Germany, 2015, Proceedings, Part III. pp. 234–241 (2015)
  * [16] Ruder, S.: An Overview of Multi-Task Learning in Deep Neural Networks. arXiv e-prints 1706.05098 (2017)
  * [17] Simpson, A.L., Antonelli, M., Bakas, S., Bilello, M., Farahani, K., van Ginneken, B., Kopp-Schneider, A., Landman, B.A., Litjens, G., Menze, B., Ronneberger, O., Summers, R.M., Bilic, P., Christ, P.F., Do, R.K.G., Gollub, M., Golia-Pernicka, J., Heckers, S.H., Jarnagin, W.R., McHugo, M.K., Napel, S., Vorontsov, E., Maier-Hein, L., Cardoso, M.J.: A large annotated medical image dataset for the development and evaluation of segmentation algorithms. arXiv e-prints arXiv:1902.09063 (2019)
  * [18] Tokui, S., Okuta, R., Akiba, T., Niitani, Y., Ogawa, T., Saito, S., Suzuki, S., Uenishi, K., Vogel, B., Vincent, H.Y.: Chainer: A Deep Learning Framework for Accelerating the Research Cycle. In: Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD 2019, Anchorage, AK, USA, 2019. pp. 2002–2011 (2019)
  * [19] van den Oord, A., Li, Y., Vinyals, O.: Representation Learning with Contrastive Predictive Coding. arXiv e-prints arXiv:1807.03748 (2018)
  * [20] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., Polosukhin, I.: Attention is All you Need. In: Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 2017, Long Beach, CA, USA. pp. 5998–6008 (2017)
  * [21] Wu, Y., He, K.: Group Normalization. In: Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, 2018, Proceedings, Part XIII. pp. 3–19 (2018)
  * [22] Zhou, Z., Sodha, V., Siddiquee, M.M.R., Feng, R., Tajbakhsh, N., Gotway, M.B., Liang, J.: Models Genesis: Generic Autodidactic Models for 3D Medical Image Analysis. In: Medical Image Computing and Computer Assisted Intervention - MICCAI 2019 - 22nd International Conference, Shenzhen, China, 2019, Proceedings, Part IV. pp. 384–393 (2019)

## Supplementary Materials

Table S1: Architecture for the VAE regularizing branch. The architecture was adopted from [12]. GN stands for group normalization, Conv: $3\times 3\times 3$ convolution, Conv1: $1\times 1\times 1$ convolution, AddId: addition of identity / skip connection, UpLinear: 3D linear spatial upsampling, Dense: fully connected layer.  Operations | Repeats | Output size
---|---|---
GN, ReLU, Conv(16) stride 2, Dense (256) | 1 | $256\times 1$
sample $\sim$ $\mathcal{N}\left(\mu(128),\sigma^{2}(128)\right)$ | 1 | $128\times 1$
Dense, ReLU, Conv1, UpLinear | 1 | $256\times 20\times 24\times 16$
Conv1, UpLinear | 1 | $128\times 40\times 48\times 32$
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 | $128\times 40\times 48\times 32$
Conv1, UpLinear | 1 | $64\times 80\times 96\times 64$
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 | $64\times 80\times 96\times 64$
Conv1, UpLinear | 1 |  $32\times 160\times 192\times 128$
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 |  $32\times 160\times 192\times 128$
Conv1 | 1 | $4\times 160\times 192\times 128$
Table S2: Architecture for the boundary attention branch. ConcSC stands for concatenate input with skip connection, EM: element-wise multiplication with input and attention layer. Operations | Repeats | Output size
---|---|---
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 |  $128\times 40\times 48\times 32$
ConcSC, Conv1, Sigmoid, EM | 1 | $128\times 40\times 48\times 32$
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 |  $64\times 80\times 96\times 64$
ConcSC, Conv1, Sigmoid, EM | 1 | $64\times 80\times 96\times 64$
GN, ReLU, Conv, GN, ReLU, Conv, AddId | 1 |  $32\times 160\times 192\times 128$
ConcSC, Conv1, Sigmoid, EM | 1 | $32\times 160\times 192\times 128$
Conv1 | 1 | $4\times 160\times 192\times 128$
Table S3: Architecture for the CPC regularization branch. The initial random crop size for the model with the CPC branch is $(144,144,128)$. The input image is first divided into overlapping $32\times 32\times 32$ grids before being fed to the encoder. Operations | Repeats | Output size
---|---|---
Spatial average | 1 |  $1\times 8\times 8\times 7$
GN, ReLU, Conv, GN, ReLU, Conv | 8 |  $1\times 8\times 8\times 7$
Conv1 | 1 | $1\times 8\times 8\times 7$
Table S4: Performance of EncDec, VAEseg, Boundseg, CPCseg (trained with six
and 387 labeled data, respectively), and ssVAEseg, ssCPCseg (trained with six
labeled and 381 unlabeled data). Evaluation was done using the mean Dice
score. using all 387 labels
---
| Dice
Test data |  ET |  TC | WT
EncDec (387 labels) |  0.9081 |  0.9234 | 0.9473
VAEseg (387 labels) |  0.9077 |  0.9323 | 0.9536
Boundseg (387 labels) |  0.9083 |  0.9241 | 0.9568
CPCseg (387 labels) |  0.9116 |  0.9305 | 0.9538
using 6 labels
EncDec (6 labels) |  0.8412 |  0.8383 | 0.9144
VAEseg (6 labels) |  0.8234 |  0.8036 | 0.8998
Boundseg (6 labels) |  0.8356 |  0.8378 | 0.9041
CPCseg (6 labels) |  0.8374 |  0.8386 | 0.9057
ssVAEseg (6 labels + 381 unlabeled) |  0.8626 |  0.8425 | 0.9131
ssCPCseg (6 labels + 381 unlabeled) |  0.8873 |  0.8761 | 0.9151


# Rapid Generation of Kilonova Light Curves Using Conditional Variational
Autoencoder

Surojit Saha Institute of Astronomy,
National Tsing Hua University, Hsinchu City, Taiwan (R.O.C.)  Michael J.
Williams Institute for Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK  Institute of Cosmology and Gravitation, University of
Portsmouth, Portsmouth PO1 3FX, UK  Laurence Datrier Institute for
Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK  Fergus Hayes Institute for Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK  Matt Nicholl Astrophysics Research Centre, School of
Mathematics and Physics, Queens University Belfast, Belfast BT7 1NN, UK
Albert K. H. Kong Institute of Astronomy,
National Tsing Hua University, Hsinchu City, Taiwan (R.O.C.)  Martin Hendry
Institute for Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK  IK Siong Heng Institute for Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK  Gavin P. Lamb Astrophysics Research Institute,
Liverpool John Moores University, IC2 Liverpool Science Park, 146 Brownlow
Hill, Liverpool L3 5RF  En-Tzu Lin Institute of Astronomy,
National Tsing Hua University, Hsinchu City, Taiwan (R.O.C.)  Daniel Williams
Institute for Gravitational Research,
School of Physics and Astronomy, University of Glasgow,
Glasgow, Scotland, UK

###### Abstract

The discovery of the optical counterpart, along with the gravitational waves
from GW170817, of the first binary neutron star merger, opened up a new era
for multi-messenger astrophysics. Combining the GW data with the optical
counterpart, also known as AT2017gfo, classified as a kilonova, has revealed
the nature of compact binary merging systems by extracting enriched
information about the total binary mass, the mass ratio, the system geometry,
and the equation of state. Even though the detection of kilonova brought about
a revolution in the domain of multi-messenger astronomy, since there has been
only one kilonova from a gravitational wave detected binary neutron star
merger event so far, this limits the exact understanding of the origin and
propagation of the kilonova. Here, we use a conditional variational
autoencoder trained on light curve data from two kilonova models having
different temporal lengths, and consequently, generate kilonova light curves
rapidly based on physical parameters of our choice with good accuracy. Once
trained, the time scale for light curve generation is of the order of a few
milliseconds, thus speeding up generating light curves by $1000$ times
compared to the simulation. The mean squared error between the generated and
original light curves is typically $0.015$ with a maximum of $0.08$ for each
set of considered physical parameter; while having a maximum of $\approx 0.6$
error across the whole parameter space. Hence, implementing this technique
provides fast and reliably accurate results.

Kilonova, light curves, conditional variational autoencoder

††software: matplotlib (Hunter, 2007), pandas(pandas development team, 2020),
tensorflow(Abadi et al., 2015), keras(Chollet et al., 2015) MOSFIT
https://github.com/mnicholl/MOSFiT,https://github.com/guillochon/MOSFiT

##  1 Introduction

The concomitant discovery of the gravitational waves (GW) and optical
counterpart electromagnetic waves (EM) (Soares-Santos et al., 2017; Lipunov et
al., 2017; Tanvir et al., 2017; Arcavi et al., 2017; Valenti et al., 2017;
Coulter et al., 2017) from the merger of the binary neutron star GW170817
(Abbott et al., 2017), also accompanied by short gamma-ray burst (GRB)
(Goldstein et al., 2017; Savchenko et al., 2017), advanced the domain of
multi-messenger astronomy. This optical counterpart, designated as kilonova
(KNe) and known as AT2017gfo bolstered previous predictions of the existence
of such electromagnetic transients (Li & Paczyński, 1998; Rosswog et al.,
1999). Prior to the discovery of GW170817, it was predicted that such an event
would be accompanied by an EM counterpart, including short-duration GRBs
(Eichler et al., 1989; Nakar, 2007), emissions ranging from radio to X-rays
pertaining to on or off axis afterglow (van Eerten & MacFadyen, 2012; Coward
et al., 2014; Fong et al., 2015; Lamb & Kobayashi, 2016) and optical to near-
IR resulting from the decay of r-process nuclei. A kilonova is an isotropic
quasi-thermal transient which is powered by the radioactive decay, in the
merger ejecta, of r-process nuclei, having luminosities $10^{40}-10^{42}\,$erg
s-1 (Li & Paczyński, 1998; Rosswog et al., 1999; Metzger et al., 2010; Barnes
& Kasen, 2013). This EM counterpart can provide deeper understanding of the
merger environment and its products. When EM information is further combined
with GW, it leads to a unique platform for an extensive understanding of such
binary events. KNe models consists of one or more radioactive ejecta
components that produce light curves peaking at different timescales and
temperature depending on atomic mass number of the ejecta, and luminosity. The
two-components consists of the blue KNe (Metzger et al., 2010; Roberts et al.,
2011; Metzger & Fernández, 2014) emission having poor lanthanide fraction
$(10^{-5})$ in the merger ejecta peaking at a relatively earlier time scale
and red KNe (Barnes & Kasen, 2013; Kasen et al., 2013; Tanaka & Hotokezaka,
2013) comprising of lanthanide-rich $(10^{-2})$ merger ejecta with peak values
at later days. In the three-component model (Perego et al., 2017) there is an
inclusion of purple KNe that describes the presence of purple KNe with a
lanthanide fraction of $10^{-3}$. Various papers (Cowperthwaite et al., 2017;
Villar et al., 2018) have provided the best-fit parameter for AT2017gfo
related to the blue, purple and red kilonova in terms of ejecta mass, ejecta
velocity and lanthanide fraction. However, since there has been only one GW-
confirmed detection of a KNe from a merger of binary neutron stars, along with
some possible candidates from other sources, it is difficult to understand and
verify the properties of binary merging systems that emit electromagnetic
radiation. An overview on the physical parameters that govern the KNe has been
provided by Metzger (2019).

In recent years, machine learning has been used for various data analysis and
formulating techniques required in astronomy (Vander Plas et al., 2014; Ball &
Brunner, 2010; Ntampaka et al., 2021; Jones et al., 2022). There has been
certain areas where the application of machine learning techniques has
produced a remarkable result, especially focusing on providing a relatively
faster results alternative to the existing methods (Ball, 2011; Sipőcz et al.,
2020). There are plenty of excellent literature available on different topics
of machine learning application in astronomy and astrophysics spanning over
wide range of topics (Li et al., 2022; Gheller & Vazza, 2022; García-Jara et
al., 2022; Sheng et al., 2022). In this paper, we incorporate a method from ML
domain known as autoencoder (Rumelhart et al., 1986), which is based on feed-
forward mechanism, to generate light curves. The primary task of an
autoencoder is to encode the input into a lower dimensional latent
representation and then decode it back to data. In feed-forward mechanism, the
neural network has unidirectional processing of information. The striking
feature for an autoencoder is that the encoder compresses the high-dimensional
input to a low-dimensional latent space and the decoder decompressed to
produce result to match the input.

In this work, we regenerate the KNe light curves by implementing a conditional
variational autoencoder (CVAE) (Kingma & Welling, 2019), which is developed
based on variational autoencoeder (VAE)(Kingma & Welling, 2013). We use this
CVAE to generate light curves based on our choice of physical parameter values
of the binary merging system. For the CVAE, once the training has been
completed, we are more flexible to rapidly generate light curves based on the
physical parameters depending on choice of the data set. Although, KNe data
from only two models for training and producing the results are used in this
work, it is also possible to extend the same idea for data from other models.
Although the data have been produced from the model, this unique data analysis
technique can not only generate light curves for the parameter values that are
not explicitly mentioned in the model but can also replace time-consuming and
resource-draining simulations required for predicting the light curves. The
novelty of this technique is expressed in the fact that once the CVAE is
trained on a KNe model, we can generate numerous light curves for a different
combinations of physical parameters in a very short time. An alternative
approach for such conditional variational autoencoder and KNe work has been
presented in Lukošiūtė et al. (2022) where different method has been studied
for obtaining results. One of the fundamental differences between Lukošiūtė et
al. (2022) and our work is in using spectra rather than light curves during
the training. In our method, we look into the light curve evolution with
respect to different sets of physical parameters aiming to complete the
process in a relatively shorter time-scale compared to the existing simulation
methods. This kind of rapid generation technique for KNe light curves can be
useful when used as template for rapid parameter estimation of KNe.

In the following sections, we show the gradual implementation of our idea, and
the results are discussed thereafter. In Section 3, we discuss the CVAE
architecture implemented in our work. Section 4 puts forward the differences
in the data and their physical parameters used for training and generation of
the KNe light curves. Section 5 provides the detailed discussion of results
obtained after implementation of CVAE. In Section 6 we summarize our approach
and present some important features of this technique. Additional results for
references are included in the appendix.

##  2 Kilonova Models

Kilonova emission results from the mass ejection in the NS mergers
(Freiburghaus et al., 1999; Ruffert, M. & Janka, H.-Th., 2001; Hotokezaka et
al., 2013). The properties of the ejecta such as the ejecta mass, ejecta
velocity and opacity dominates the peak luminosity and the time of the peak
luminosity (Li & Paczyński, 1998; Kasen et al., 2013, 2015; Barnes & Kasen,
2013; Tanaka et al., 2017). This luminosity is sourced from the radioactive
decay of the r-process elements synthesized in the merger. This work primarily
focuses on two KNe models having different physical parameters that determine
the light curves. Below we provide an outline of the two models used in this
work. In Kasen et al. (2017),the light curves are dependent upon the peak
magnitude and decay time with respect to the ejecta mass, velocity and
lanthanide fraction is provided. Quantitatively, peak luminosity to first
order is denoted by $L\propto M_{\rm ej}$ and the width of the light curve is
$\tau\propto(\kappa M_{\rm ej}/v_{\rm ej})^{1/2}$, pointing to lanthanide-rich
ejecta with higher $\kappa$, where $M_{\rm ej}$, $v_{\rm ej}$ and $\kappa$ are
the mass, velocity and opacity of the ejecta. Hence heavier ejecta from the
merger leads to higher peak magnitude with relatively longer duration KNe
whereas ejecta with higher velocities have short duration with bright KNe. We
see that lanthanide fraction plays a major role on the light curves where the
ejecta with lower lanthanide fraction decays on a relatively shorter timescale
compared to lanthanide-rich ejecta decaying over weeks. From the literature in
Nicholl et al. (2021), the study presents an approach where light curves are
predicted from the chirp mass, mass ratio, orbital inclination also including
the properties related to the nuclear equation of state. Here we see the
nature of light curves based on these physical parameters. It is important to
note that, (Kasen et al., 2017) is based on radiative transfer model, where as
(Nicholl et al., 2021) is a semi-analytic model. While KNe light curves can be
generated with CVAE, we choose these models due to their different
parametrizations and the data availability. However, this does not put any
limitations on the use of CVAE, since data from any other available KNe models
can be equally used.

##  3 Autoencoder

With the advancement of new machine learning (ML) methods and their implicit
application in data analysis, has opened a new era where ML techniques can be
implemented to obtain relatively faster results. Time-consuming and resource-
draining simulations can be completed on a reasonable time scale (Carleo et
al., 2019). There are many available ML techniques that can be used for data
analysis with specific techniques built to obtain certain results. We use an
autoencoder, an ML technique that is based on a feed-forward mechanism, and
its function is to reproduce the input. As is the case with any standard
autoencoder, the CVAE consists of three sections, an encoder, a latent space,
and a decoder, as shown in Fig. 1. The encoder $(Q(\phi))$ compresses the
high-dimensional input data into a lower-dimensional latent space while
capturing the features of the input data. Latent space(_Z_), which is also
referred to as the output of encoder, has a lower dimensional representation.
Generally, a well-trained CVAE has the entire high-dimensional input data
smoothly distributed over the latent space. To generate light curves of our
choice, we draw samples from this latent space and pass it through the decoder
where the decoder $(P(\theta))$ takes the compressed representation and
generates the required light curves. This facilitates using CVAE as a tool,
not only limited to a generative model but also quickly look into the
parameter dependency of KNe light curves. In CVAE, the training is regularised
in order to avoid over-fitting. Hence, the latent space is well distributed to
enable generative process Even though there are other generative models like
Generative Adversarial Networks (GAN)(Goodfellow et al., 2014) but since in
our work we want to sample new data from the probability distribution of the
input data, to look at the model parameters crucial for the light curves, we
prefer to utilize the strength of an autoencoder because of its ability to
have a probability distribution of input.

Even though GAN and CVAE are both generative models and belongs to the
category of unsupervised learning, one of the primary differences lies in
there architecture where GAN has a generator and discriminator and CVAE has
encoder, latent space and decoder. Along with this, the loss functions used in
GAN and CVAE are different from each other. It is important to note that for
generating new data, sampling a hidden state in GAN takes place from a
predefined distributions and is then fed into the discriminator in contrast to
that of CVAE where sampling of hidden state, to be fed into the decoder, is
done from selecting a prior distribution related to the actual data. Here, to
achieve our goal, we take advantage of the conditional variational
autoencoder, where we train it on the light curves by conditioning them on
physical parameters and generate new light curves based on the physical
parameters of our choice. In this case, we have complete control over the
generated data set. In this work, we take advantage of this feature of CVAE to
generate new data. In our CVAE architecture, the loss function is the
combination of the reconstruction-loss and the Kullback-Leibler divergence
(Asperti & Trentin, 2020). Training is carried out with $batch\\_size=50$ for
$epochs=1000$ and having a learning rate of $0.0001$. Commencing from the
training till generation of the light curves for the different physical
parameter values, it takes $\simeq$ 20 minutes for the process to complete
using $3.9$ GHz $8$-core Intel Core i$9$ processor with $32$ GB memory. In our
training and regeneration of the light curves, we have not used the GPU of the
system and have completely relied on the CPU. Once the model is trained and
saved, it takes only a few milliseconds to generate the light curves. In this
section, we outline the CVAE without delving into the details, however, for
more detailed variational encoder readers can refer to (Kingma & Welling,
2013) and (Kingma & Welling, 2019). In Section 5 more insight into the results
obtained from CVAE are shown.

Figure 1: Schematic diagram of the conditional variational autoencoder,
consisting of the encoder, latent space and the decoder, from input to data
generation is shown. In the figure, we have the probabilistic encoder
$Q(\phi)$ and the probabilistic decoder $P(\theta)$. _Z_ represents the
compressed input or encoded representation. Here, the encoder transforms the
high-dimensional input data, into the low-dimensional latent space which
represents the compressed form of the high-dimensional input data while the
decoder maps this compressed representation in the low-dimensional latent
space leading to the reconstruction of the original high-dimensional data. We
take a set of data _x_ , which is the KNe light curves and condition it on a
given parameter _y_ , which in our case are the physical parameters from the
model, and then the encoder compresses this representation in the latent
space. The decoder decompresses the representation and provides the relevant
output data, KNe light curves, based on the physical parameter values of our
choice. This is an overview of the CVAE architecture that has been used in
this work to train and generate the KNe light curves. Figure 2: (a) In this
plot, we represent only $50$ KNe light curves of the g-band of $D_{1}$ which
are a part of the data set for training. These light curves consist of the
mass of the ejecta, the velocity of the ejecta, and the fraction of lanthanide
as physical parameters with values ranging between
($0.001-0.1M_{\odot}$),$(0.03-0.05c)$ , $(10^{-9}-10^{-4})$ respectively.
(b)In this plot, we show $50$ light curves for g-band of $D_{2}$ which are a
part of the training data set. These light curves consists of chirp mass, mass
ratio, fraction of the remnant disk with values ranging between
($0.7-1.10M_{\odot}$), $(0.50-1.0)$, $(0.15-0.45)$ respectively and having a
viewing angle of $0\degree$ as their physical parameters. This figure gives an
outline of the KNe light curves and their decay time, which is used for
training the CVAE. For the both data, the KNe light curves are in absolute
magnitude. These light curves are scaled between [$0$-$1$] while feeding them
into the CVAE for training, but, while representing the generated results,
these light curves are scaled back to the absolute magnitude.

##  4 Data

In the ML domain, one needs to be very careful about the data that is fed into
the algorithm. The data have to be pre-processed and prepared accordingly. The
data used for training, validation and test set are categorized into two types
as mentioned in Section 2. We adopt the data
(https://github.com/dnkasen/Kasen_Kilonova_Models_2017), hereafter $D_{1}$
(Kasen et al., 2017) to prepare our data set to be fed into the CVAE. The
physical parameters in this data set are ejecta mass ($0.001-0.1M_{\odot}$),
ejecta velocity($0.1-0.3c$) and lanthanide fraction(10-4-10-9). There are
$329$ light curves with duration of $\approx 25$ days. Each light curve has
different values of the physical parameter within the range mentioned above.
The light curves are available for different filter bands (u,g,r,y,z) within
the same physical parameter range. Using these data, $241$ light curves are
used as training sets, and the rest are equally divided into the test set and
the validation set. In our analysis, we have trained the CVAE on each filter
band data separately, hence there are five trained CVAE’s, and carried out
light curves generation; however, in the main sections, only the g-band
results are shown; the rest are added in appendix.

The second type of data, hereafter $D_{2}$, used here consists of simulated
light curves at same the filter bands but having different physical
parameters. The physical parameters of these data are the chirp mass of the
binary system, the mass ratio, the fraction of the remnant disk that is
ejected, the viewing angle in degrees from the pole, and the opening angle of
the cocoon shock (Nicholl et al., 2021). The range of values are
$(0.7M_{\odot}-2.0M_{\odot})$ for the chirp mass, $(0.5-1.0)$ for the mass
ratio, $(0.15-0.45)$ for the fraction of the remnant disk, respectively. For
viewing angle the data corresponds to the values of $0\degree$, $60\degree$
and $90\degree$. There are $529$ light curves having temporal length of $30$
days, out of which $401$ light curves are included in the training data and
the rest are equally divided into test and validation set. The difference in
the two data set is appears in the physical parameters used in the respective
KNe models. These data are used to verify the pliability of the CVAE. Since
the physical parameters in the two data sets used are entirely different, this
gives an extra edge in verifying our approach, and the results are detailed in
Section 5. Since both data sets contain light curves that have different
absolute magnitude peak values, while feeding them for training, we have
scaled the light curve data and the physical parameters between $0$ and $1$,
as this would be a more effective approach. Later, after training and
generating light curves, we rescale the light curves and plot accordingly.
Throughout the paper, we will follow a particular format to represent the
physical parameters in the text and in the legends of the plots. For the texts
and plots relevant to $D_{1}$, we will use the format where $a$,$b$ and $c$
will be the values of ejecta mass in $M_{\odot}$, velocity of the ejecta in
units of velocity of the light and lanthanide fraction respectively. Similarly
for the $D_{2}$, we will use the format of $[w,x,y,z]$ where $w$ is the chirp
mass in units of $M_{\odot}$, $x$ is the mass ratio, $y$ gives the fraction of
the remnant disk that is ejected and $z$ is the viewing angle in degrees from
the pole. We will use the above format throughout the paper wherever required,
without further mentioning the respective units. Fig. 2(a) shows $50$ out of
$329$ original light curves from $D_{1}$ in the g-band for physical parameters
corresponding to ejecta mass, ejecta velocity and lanthanide fraction with
values ranging between $0.001-0.1M_{\odot}$, $0.03-0.05c$ and
$10^{-4}-10^{-9}$ respectively whereas in Fig 2(b), we show $50$ out of $529$
light curves corresponding to $D_{2}$ having physical parameters that of chirp
mass ($0.7-1.10M_{\odot}$), mass ratio ($0.5-1.0$), fraction of the remnant
disk ($0.15-0.45$) and having a viewing angle of $0\degree\hskip 1.42271pt$
from the pole as mentioned in Section 4. In Fig. 2, we show the input light
curves corresponding to $D_{1}$ and $D_{2}$ which is fed into the CVAE. During
training and validation, scaled values of the KNe light curves and the
conditional physical parameters are used for both $D_{1}$ and $D_{2}$, and the
obtained results are shown in absolute magnitude values. The input light
curves for training from $D_{1}$ consist of $241$ light curves and from
$D_{2}$ has $401$ light curves from the g-band. As previously mentioned, each
light curve has different physical parameters corresponding to ejecta mass,
ejecta velocity and lanthanide fraction for $D_{1}$ and chirp mass, mass
ratio, fraction of the remnant disk and the viewing angle from the pole for
$D_{2}$. Throughout the text, $D_{1}$ and $D_{2}$ will be referred to
accordingly, keeping the above training, test, and validation set unchanged.
In Table 1, we tabulate the physical parameters and the data ranges used in
this work. In this study, while defining the training, test, and validation
set we do not restrict the physical parameters to the different regions of the
parameter space. Instead, the physical parameters present in the training,
test, and validation sets cover the entire parameter space. But, at the same
time, we ensure that none of the parameter combinations, i.e ejecta mass,
ejecta velocity, and lanthanide fraction for $D_{1}$ and chirp mass, mass
ratio, a fraction of the remnant disk, and viewing angle for $D_{2}$, are
repeated in the training, test, and validation set. Specifically, for $D_{2}$,
apart from the training, test, and validation set, we have kept aside a
separate set of simulated light curves from
MOSFiT(https://github.com/guillochon/MOSFiT), which are utilized to evaluate
the performance of the CVAE. We augmented the existing test set with these
separately simulated light curves, hereafter $D_{2}\dagger$, to provide a
robust performance check. The combination of the physical parameters present
in $D_{2}\dagger$, as shown in Table 2, covers the parameter space but is
previously unseen to the CVAE. The CVAE will be employed with these sets of
physical parameters to generate light curves and hence compare with the true
light curves, thus providing a robust approach to appraise the CVAE
performance. Hence, using the CVAE, we generate light curves over a grid on
the parameter space.

Data | Physical Parameters | Data Features
---|---|---
$D_{1}$ | Ejecta Mass ($M_{\odot}$) |  $0.001$, $0.0025$, $0.005$, $0.01$, $0.02$, $0.025$, $0.03$, $0.04$, $0.05$, $0.075$
Ejecta Velocity(c) |  $0.03$, $0.05$, $0.1$, $0.2$, $0.3$
Lanthanide Fraction |  $10^{-9}$, $10^{-5}$, $10^{-4}$, $10^{-3}$, $10^{-2}$, $10^{-1}$
$D_{2}$ | Chirp Mass ($M_{\odot}$) |  $0.7$, $0.8$, $0.9$, $1.0$, $1.1$, $1.2$, $1.3$, $1.4$, $1.5$, $1.6$, $1.7$, $1.8$
Mass Ratio |  $0.5$, $0.6$, $0.7$, $0.8$, $0.9$, $1.0$
Fraction of the Remnant Disk |  $0.15$, $0.30$, $0.45$
Viewing Angle |  $0\degree$, $60\degree$, $90\degree$

Table 1: This table gives an overview of the physical parameters and range of
their values corresponding to $D_{1}$ and $D_{2}$ used in this work for
training and generation of the KNe light curves. $D_{1}$ is taken from
https://github.com/dnkasen/Kasen_Kilonova_Models_2017 and $D_{2}$ is adopted
from https://github.com/mnicholl/kn-models-nicholl2021. The training, test and
validation data for the corresponding data set consists of these individual
physical parameter.

Data | Physical Parameters | Data Features
---|---|---
$D_{2}^{\dagger}$ | Chirp Mass ($M_{\odot}$) |  $1.0$, $1.2$, $1.4$, $1.6$, $1.8$
Mass Ratio |  $0.7$, $0.75$, $0.8$, $0.85$, $0.9$
Fraction of the Remnant Disk |  $0.15$, $0.20$, $0.25$, $0.30$, $0.35$, $0.40$
Viewing Angle |  $45\degree$ ,$60\degree$, $75\degree$, $90\degree$

Table 2: This table contains the interpolated physical parameters that were
used to simulate light curves using MOSFIT and compare those with the CVAE-
generated light curves. Light curves from these physical parameters, alongside
from the test set, are employed to evaluate the performance of the CVAE over
the parameter space.

##  5 Results

In this section, we present the results after the implementation of CVAE.
After training, samples are drawn from the latent space to reconstruct light
curves based on the required physical parameters. This section points out the
comparison between the generated light curves from the latent space and the
original light curves from the simulation. This technique allows us to
generate as many light curves as we require over a wide range of physical
parameters within the range provided by the model.

Figure 3: (a)This figure show light curves comparison corresponding to g-band
after training and generation are implemented on $D_{1}$. We display $100$
light curves for physical parameter of ejecta mass, ejecta velocity and
lanthanide fraction having values $[0.001M_{\odot},0.03c,10^{-5}]$ and compare
it along side the light curve from the test data for the same physical
parameter. The original light curve is shown in diamond marker and is within
distribution of the generated light curves shown by dashed lines. The
deviations seen in the generated light curves arise due to the variations in
the latent space. (b)This plot corresponds to the learning curve after
implementing the CVAE. The gap between the two curves is small indicating a
good-fit

In Fig.3(a), we show the generated light curves of g-band after implementing
CVAE for the physical parameter of $[0.001M_{\odot},0.03c,10^{-5}]$ which are
the ejecta mass, ejecta velocity, and lanthanide fraction, respectively as
mentioned in Section 4. The loss plot for the CVAE is shown in Fig 3(b). We
find that the validation loss and the training loss has decreased to a point
of stability and there is small gap between the two curves. Although, CVAE has
the ability to generate as many light curves we want for one or many physical
parameters, but in order to avoid congestion in plots, we limit ourselves to
$100$ generated light curves based on the above physical parameter. The $100$
light curves are closely spaced and the apparent variations at the later days
comes from the variations of the probability distribution in the latent space.
These variations are expected since, to regenerate light curves, we randomly
draw samples from the latent space. This randomness allows the CVAE to produce
variations in the light curves even for the same input. Thus, we see these
variations in the CVAE-generated light curves for the physical parameters.
Besides, the light curves are more sensitive to the changes in the physical
parameters in the later days. However, since we are more interested in the
earlier KNe evolution, we restrict the plots to $14$ days. This parameter is
chosen from the simulated data in order to verify the robustness of the CVAE.
It is clear that the input light curve, shown by diamond markers, used from
the validation data of the CVAE is well within the limit of the regenerated
light curves.

This assures that the CVAE is well trained. In Fig. 3(a), we see some
overlapping in the light curves, since the $100$ generated light curves have
relatively similar values. While generating these light curves, samples were
drawn from different points in latent space. It is possible to generate as
many light curves as desired and observe the possible variations from the
latent representation. Since the encoder maps the input light curves to a
probability distribution over the latent space. Hence, even though in the
respective KNe models, there is one light curve associated with a single value
of the physical parameter, implementing CVAE provides a comparatively wide
range of light curve distribution for the same value of physical parameter
arising from the latent representation. This variation of the generated KNe
light curves are particularly interesting when the light curves are generated
for an entirely new physical parameter set which the CVAE has not come across
during training and validation.

Figure 4: In this figure, we compare $100$ light curves (solid lines)
generated from CVAE with original the light curves (diamond marker) for
physical parameter values having ejecta masses of $0.02M_{\odot}$(_upper
left_),$0.1M_{\odot}$(_upper right_) and $0.03M_{\odot}$(_bottom left_)
keeping the ejecta velocity $0.03c$ and lanthanide fraction$(10^{-9})$
constant for all the generated light curves. These values of the physical
parameter were chosen from the test set. In the bottom right panel, $10$ light
curves for an arbitrary physical parameter $[0.025M_{\odot}$, $0.035c$,
$10^{-5}]$ is shown. All the generated and the original light curves
corresponds to the g-band. Figure 5: Comparison for the true and generated
light curves corresponding to $D_{2}$. These $100$ light curves are
particularly from r-band generated from the CVAE where the comparison is shown
with the respective physical parameter of chirp mass, mass ratio, fraction of
the remnant disk and viewing angle from the pole having values of
$[1.8M_{\odot},0.9,0.15,90\degree]$ respectively. The diamond marker
corresponds to the true light curve, whereas the solid lines corresponds to
the CVAE-generated light curves. The physical parameter for the true light
curve is taken from the test set for verifying with the generated ones.

Kilonova light curves are extremely sensitive to ejecta mass(Kasen et al.,
2017). With the change in the ejecta mass, the peak value of the light curves
increases and shifts to later days, which is explicitly shown in the top left,
top right and bottom left panel in Fig.4. This result is expected in
accordance to the KNe evolution. In this figure, the comparative analysis for
the different values of ejecta masses $(0.02M_{\odot},0.1M_{\odot}$ and
$0.03M_{\odot})$ are shown keeping the ejecta velocity $(0.03c)$ and
lanthanide fraction $(10^{-9})$ unchanged. $100$ light curves are generated
from the CVAE comparing with the original light curve for the above physical
parameters. In the bottom right panel, light curve for an arbitrary parameter
$(0.025M_{\odot}$,$0.035c$,$10^{-5})$ is shown. This physical parameter value
configuration in not present in $D_{1}$ therefore we do not have any benchmark
to validate it, however, with more data from the respective KNe model, this
can be verified. Similar comparative analysis, for all other filter bands,
showing the evolution of the light curves with the change in ejecta mass,
ejecta velocities and lanthanide fractions are shown in appendix in Fig 8. We
next take the second data set, $D_{2}$, with physical parameters of chirp
mass, mass ratio, fraction of the remnant disk and viewing angle from the
pole. We run the same analysis as above. As mentioned in Section 4, this data
set has a total of $529$ light curves, with the above physical parameters
having values given as mentioned in Section 4. The training data has $401$
light curves, and the remaining are equally divided into a test set and a
validation set. In Fig 5, we demonstrate the application of the CVAE on
$D_{2}$ where we have compared $100$ light curves generated from CVAE with the
original light curve for the physical parameter of
$[1.8M_{\odot},0.9,0.15,90\degree]$ correlated to the chirp mass, the mass
ratio, the fraction of the remnant disk and the viewing angle from the pole,
respectively. This set of physical parameters is taken from the test data set

The input light curve is within the distribution of the generated light curve
which indicates that the CVAE is performing well and the training and
regeneration of the light curves are successful. Fig. 6, shows the confidence
plot for a set of physical parameter where it is important to mention that the
neural network has not seen this set of physical parameter and hence from the
similarity between the generated and true light curves we can substantiate the
performance of the CVAE. Other relevant plots for different combination of
physical parameter and the changes in the light curves corresponding to those
for other filter bands are shown in Fig 9 in appendix. To measure the
predictive accuracy of the trained CVAE, we use mean absolute error (MAE),
where lower value corresponds to more accurate prediction, calculated using

$MAE=\frac{1}{n}\sum_{i=1}^{n}|y_{i}-\hat{y}_{i}|$ (1)

, where $y_{i}$ and $\hat{y}_{i}$ are the true values and predicted values of
the light curves of the test data and $n$ is the total number of points drawn
for each light curve from latent space. Alongside the above mean absolute
error, we also calculate and plot the mean squared error (MSE) between the
true and predicted light curve values using the equation

$MSE=\frac{1}{n}\sum_{i=1}^{n}(z_{i}-\hat{z}_{i})^{2}$ (2)

where, $n$ being the total number of data points and $z_{i}$ and $\hat{z}_{i}$
are the respective values of true and CVAE-predicted light curves.

Figure 6: This figure corresponds to the $90\%$ confidence plot for the CVAE-
generated and true light curves in the g,r,z,y,i and u bands that have the
physical parameter of $1.2M_{\odot}$ chirp mass, $0.7$ mass ratio, $0.15$
fraction of the remnant disk and a viewing angle of $60\degree$. True values
of the light curves for the above parameter are represented by the solid lines
in each color-filled region. We find quite satisfactory agreement between the
true and CVAE-generated light curves. The light curves corresponding to the
above parameters were taken from the $D_{2}\dagger$ data set and thus are
entirely new to the CVAE for prediction and generation. Figure 7: In this
figure, we present the performance of the CVAE spanning over the whole
parameter space of chirp mass $(1.0M_{\odot}-1.80M_{\odot})$,viewing
angle$(45\degree-90\degree)$,fraction of remnant disk$(0.15-0.40)$ and mass
ratio $(0.70-0.90)$ in the _u,g,i,r,z_ and _y_ filter bands represented by the
different colors bars grouped accordingly. The _y-axis_ represents the mean
squared values whereas in the _x-axis_ we have the respective physical
parameter values for which we have calculated the overall mean squared error
in the different filter bands shown with different color bars. For the chirp
mass frame(_upper left_), the mean squared errors for different filer bands
are grouped by the chirp mass values,
$1.0M_{\odot},1.2M_{\odot},1.4M_{\odot},1.6M_{\odot},1.8M_{\odot}$, while
covering parameter space of viewing angle, mass ratio and fraction of the
remnant disk. For the viewing angle frame(_upper right_), we show the error
for $45\degree$, $60\degree$, $75\degree$ and $90\degree$ viewing angles
calculated over the entire parameter values of chirp mass, mass ratio and
fraction of the remnant disk. In the remnant disk frame(_lower left_), for the
values of $0.15$,$0.20$,$0.25$,$0.30$,$0.35$ and $0.40$ fraction of remnant
disk, errors are grouped in different filter bands encompassing the chirp
mass, viewing angle and mass ratio. For the mass ratio frame(_lower right_),
errors corresponding to mass ratio values of $0.70$,$0.75$,$0.80$,$0.85$,and
$0.90$ are shown as grouped bar plots across the parameter space of chirp
mass, viewing angle and fraction of the remnant disk for the different filter
bands. In the insets for all the frames, we show the CVAE performance over the
entire parameter space, as discussed above, obtained from CVAE generated _g-
band_ light curves. Each dot in the inset corresponds to the calculated value
of the mean squared error of a light curve for the relevant set of physical
parameters. For each of the histograms corresponding to the different
parameter space, error bars have been shown. The histograms without the error
bar correspond to a single light curve that is available for calculating the
mean squared error between the true and CVAE generated light curves.

##  6 Discussion

Since the CVAE depends on the input data, light curves from different KNe
models can be used for training and the corresponding trained CVAE is saved
for generating light curves.The input data contains both training data and
physical parameters. This versatility of the technique provide a unique
opportunity to test this on other data set. Using eq(1), we have the mean
absolute error to be $0.0995$ and $0.0339$ for $D_{1}$ and $D_{2}$
respectively. The mean squared error for $D_{1}$ and $D_{2}$ calculated on the
test data using eq(2) are found to be $0.00241$ and $0.00153$ respectively.

Compared to $D_{2}$, the MAE for $D_{1}$ is higher, as the training data is
relatively lower for $D_{1}$. We not only successfully reconstructed the light
curves for known values of physical parameters but also generated light curves
for the physical parameters that are not present in the data set but are well
within the range of the physical parameters values provided in the respective
models. In Fig. 3, we presented the reconstructed light curves and the
corresponding loss values for the CVAE. Consequently, based on these results,
we proceed to generate light curves for different combinations of physical
parameters and verify those as presented in Fig. 4. The obtained results are
consistent with the KNe evolution presented in the model. Comparing the plots
in Fig. 3(a) and Fig. 4, as mentioned in Section 5, we can see more variations
in the latter, and for the latter figure the light curves has been truncated
to $14$ days while in the former, CVAE-generated light curves are extended
till the end. For the plots generated from $D_{2}$, since there has been no
significant change related to astrophysical importance, we have truncated the
light curves to $14$ days. Model evaluation is performed with the test data,
in the current work, we move a step ahead to evaluate the performance via the
confidence interval (Fig. 6) and mean squared error (Fig.7) between the CVAE-
generated and true light curves by taking physical parameters from
$D_{2}\dagger$ which contains the physical parameters from the test set as
well as the entirely new physical parameters that were absent in the training,
test and validation set. Therefore, as previously mentioned in Section. 4, we
use the physical parameters combination from Table 2 to produce Fig. 7.
Therefore, when evaluating the performance on these sets of previously unseen
parameters, simulates the real scenario akin to observing a new kilonova. In
Fig. 6, we represent the figure for the $90\%$ confidence interval between the
generated and true light curves in different filter bands of u,g,r,i,z and y
having the chirp mass of $1.2M_{\odot}$, mass ratio of $0.7$, $0.15$ fraction
of the remnant disk having a viewing angle of $60\degree$. The true light
curves for this set of parameters were not included in the training, test, or
validation data set, thus this data is unseen to the network. To plot the
confidence interval, we take the mean light curves from the $2000$ CVAE-
generated light curve for the above physical parameter values
[$1.2M_{\odot}$,$0.7$,$0.15$,$60\degree$] in each filter bands. Each result is
highlighted with color-shaded region for each filter band including the true
light curve. Here we see a good agreement between the true and generated light
curves which provides supporting evidence in favour of the performance of
CVAE. In addition to the above, we have calculated the mean squared error
between the generated and true light curves where we have have obtained the
error to be always between $(0.015-0.08)$ for the set of physical parameters.
To provide further evidence of CVAE performance, we illustrate the mean
squared error calculated over the entire parameter space while considering the
following values of physical
parameters,$1.0M_{\odot}$,$1.2M_{\odot}$,$1.4M_{\odot}$,$1.6M_{\odot}$, and
$1.8M_{\odot}$ for the chirp mass region, $45\degree$, $60\degree$,
$75\degree$ and $90\degree$ for viewing angle region, $0.15$, $0.20$, $0.25$,
$0.30$, $0.35$ and $0.40$ for the fraction of remnant disk parameter and for
the $0.70$, $0.75$, $0.80$, $0.85$ and $0.90$ in mass ratio parameter for the
different filter bands as shown in Fig.7. To calculate the mean squared error,
we have used the mean of $2000$ CVAE-generated light curves and the true light
curves for the respective sets of physical parameters.

Fig.7 consists of four frames, each having an inset, which show the error
calculated for the respective value of physical parameter (along _x-axis_)
across the parameter space of the other remaining physical parameters. For the
chirp mass frame (_upper left_),error is calculated for the different chirp
masses
($1.0M_{\odot}$,$1.2M_{\odot}$,$1.4M_{\odot}$,$1.6M_{\odot}$,$1.8M_{\odot}$)
across the parameter space of viewing angle, mass ratio and the fraction of
the remnant disk, grouped together for each value of chirp mass for the filter
bands. For the viewing angle frame (_upper right_), for the values of
$45\degree$,$60\degree$,$75\degree$ and $90\degree$ error is calculated across
the chirp mass, mass ratio and fraction of the remnant disk and grouped
accordingly. For calculating the mean squared error for the fraction of
remnant disk, we have considered $0.15$,$0.20$,$0.25$,$0.30$,$0.35$,and $0.40$
across the parameter space of chirp mass, viewing angle and mass ratio grouped
in parameter values in different filter bands. In the mass ratio frame (_lower
right_), error is calculated for the mass ratio values of
$0.70$,$0.75$,$0.80$,$0.85$, and $0.90$, across the parameter space of chirp
mass, viewing angle and fraction of the remnant disk grouped accordingly. For
certain set of physical parameter in certain filter bands, tend to have
comparatively high mean squared error value, thus adding more to the overall
error. For instance, the above case can be seen in the viewing angle frame in
Fig.7 for $45\degree$ viewing angle in _i-band_. However, the overall error
are considerably low. In the insets of Fig.7, _g-band_ results are shown where
each dot refers to the mean squared error of a single light curve for the
respective set of physical parameter across the parameter space. Here also, we
find that for certain set of physical parameter, the error value is
comparatively higher as evident from outliers in $1.8M_{\odot}$ chirp mass,
$60\degree$ viewing angle, $0.15$ fraction of remnant disk and $0.80$ mass
ratio in the insets. In comparison to one-one mean squared error, the overall
error is relatively higher as it includes all the parameter sets across the
individual parameter considered. From the above result, we find the regions in
parameter space for different filter bands where the CVAE has comparatively
unsatisfactory results. In addition to the above, the network summary, Fig 10
has been added in the appendix for reference. While comparing Fig. 3 and Fig.
4 for $D_{1}$ with results from $D_{2}$ in Fig .5, one will find similar
deviations in the tail of the light curves after $10$ days for the different
sets of physical parameters shown in plots presented in the appendix(Fig 9).
In the main text, only one such plot (Fig. 5) has been shown for a single set
of physical parameter, where the deviations are comparatively less. One of the
main reasons for such different deviations is distribution of the latent space
owing to the difference in the number of training data, since $D_{1}$ has
comparatively less training data than $D_{2}$.

##  7 Conclusion

In this paper, we look into rapid generation of KNe light curves based on
different physical parameters values while implementing CVAE. We present a
methodological approach to our idea. In the initial stages, the performance
check was carried out while inspecting the loss curves during the training and
validation. We also evaluated the CVAE performance on different physical
parameter from test data, hence perform rapid interpolation of light curves
and subsequently provide evidence that the results are quite in agreement with
the true light curves. However, in certain cases, while generating the light
curves, we find that for $D_{1}$, as compared to the original light curves,
the generated light curves tends to deviate after $8$ days but for $D_{2}$ we
do not see such deviations in the later days. Here, we have used publicly
available KNe data to provide a proof of our concept, but this does not put
any limitation on this technique. This technique can also be used for similar
kind of data analysis comprising of data related to other domain in astronomy
and astrophysics. The striking point for such an approach is that it reduces
the time which is required by simulations to rerun and reproduce similar
results for different physical parameters every time. We train the data
available and produce the desired results rapidly rather than adjusting the
simulation code each time and the saved CVAE model can be used to generate the
required KNe light curves.In this work, by speeding up the light curve
generation by $1000$ times, we have achieved rapid results. Besides the
current application for KNe, the CVAE technique demonstrated above has
prospective applications in a similar procedure where data are available for
training, test and, validation. In the current work, since the detailed
calculations of the simulation are not incorporated into the CVAE, we do not
expect any new results, but at the same time, the trained-CVAE produces
results without looking into the particulars of the simulation.An alternative
approach for generating simulation results, with the help of machine learning
tools, without actually re-executing those is presented here.Additionally,
this method can accommodate other KNe models where data are available to be
fed into the network. This kind of CVAE approach has also the potential to be
utilized for rapid parameter estimation not only limiting to KNe but also for
other astrophysical sources where rapid data analysis is required.

##  8 Acknowledgements

This work is supported by the National Science and Technology Council of
Taiwan under the grants 110-2628-M-007-005 and 111-2112-M-007-020, and a joint
grant of the National Science and Technology Council and the Royal Society of
Edinburgh through 110-2927-I-007-513. MJW is supported by the Science and
Technology Facilities Council [2285031,ST/V005634/1,ST/V005715/1]. ISH is
supported by the Science and Technology Research Council [ST/ L000946/1]. ISH
and MJW are also supported by the European Cooperation in Science and
Technology (COST) action [CA17137]. MN is supported by the European Research
Council (ERC) under the European Union’s Horizon 2020 research and innovation
programme (grant agreement No. 948381) and by funding from the UK Space
Agency. The authors are grateful to the valuable suggestions from He-Feng
Hsieh, John Veitch, and Nicola De Lillo.

## References

  * Abadi et al. (2015) Abadi, M., Agarwal, A., Barham, P., et al. 2015, TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems.  https://www.tensorflow.org/
  * Abbott et al. (2017) Abbott, B. P., Abbott, R., Abbott, T. D., et al. 2017, Phys. Rev. Lett., 119, 161101, doi: 10.1103/PhysRevLett.119.161101
  * Arcavi et al. (2017) Arcavi, I., McCully, C., Hosseinzadeh, G., et al. 2017, The Astrophysical Journal, 848, L33, doi: 10.3847/2041-8213/aa910f
  * Asperti & Trentin (2020) Asperti, A., & Trentin, M. 2020, arXiv e-prints, arXiv:2002.07514.  https://arxiv.org/abs/2002.07514
  * Ball (2011) Ball, N. M. 2011, arXiv e-prints, arXiv:1110.5688.  https://arxiv.org/abs/1110.5688
  * Ball & Brunner (2010) Ball, N. M., & Brunner, R. J. 2010, International Journal of Modern Physics D, 19, 1049, doi: 10.1142/S0218271810017160
  * Barnes & Kasen (2013) Barnes, J., & Kasen, D. 2013, ApJ, 775, 18, doi: 10.1088/0004-637X/775/1/18
  * Carleo et al. (2019) Carleo, G., Cirac, I., Cranmer, K., et al. 2019, Rev. Mod. Phys., 91, 045002, doi: 10.1103/RevModPhys.91.045002
  * Chollet et al. (2015) Chollet, F., et al. 2015, Keras, GitHub.  https://github.com/fchollet/keras
  * Coulter et al. (2017) Coulter, D. A., Foley, R. J., Kilpatrick, C. D., et al. 2017, Science, 358, 1556, doi: 10.1126/science.aap9811
  * Coward et al. (2014) Coward, D. M., Branchesi, M., Howell, E. J., Lasky, P. D., & Böer, M. 2014, MNRAS, 445, 3575, doi: 10.1093/mnras/stu1863
  * Cowperthwaite et al. (2017) Cowperthwaite, P. S., Berger, E., Villar, V. A., et al. 2017, ApJ, 848, L17, doi: 10.3847/2041-8213/aa8fc7
  * Eichler et al. (1989) Eichler, D., Livio, M., Piran, T., & Schramm, D. N. 1989, Nature, 340, 126, doi: 10.1038/340126a0
  * Fong et al. (2015) Fong, W., Berger, E., Margutti, R., & Zauderer, B. A. 2015, ApJ, 815, 102, doi: 10.1088/0004-637X/815/2/102
  * Freiburghaus et al. (1999) Freiburghaus, C., Rosswog, S., & Thielemann, F. K. 1999, ApJ, 525, L121, doi: 10.1086/312343
  * García-Jara et al. (2022) García-Jara, G., Protopapas, P., & Estévez, P. A. 2022, arXiv e-prints, arXiv:2205.06758.  https://arxiv.org/abs/2205.06758
  * Gheller & Vazza (2022) Gheller, C., & Vazza, F. 2022, MNRAS, 509, 990, doi: 10.1093/mnras/stab3044
  * Goldstein et al. (2017) Goldstein, A., Veres, P., Burns, E., et al. 2017, ApJ, 848, L14, doi: 10.3847/2041-8213/aa8f41
  * Gompertz et al. (2018) Gompertz, B. P., Levan, A. J., Tanvir, N. R., et al. 2018, ApJ, 860, 62, doi: 10.3847/1538-4357/aac206
  * Goodfellow et al. (2014) Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., et al. 2014, Generative Adversarial Networks.  https://arxiv.org/abs/1406.2661
  * Hotokezaka et al. (2013) Hotokezaka, K., Kiuchi, K., Kyutoku, K., et al. 2013, Phys. Rev. D, 87, 024001, doi: 10.1103/PhysRevD.87.024001
  * Hunter (2007) Hunter, J. D. 2007, Computing in Science & Engineering, 9, 90, doi: 10.1109/MCSE.2007.55
  * Jones et al. (2022) Jones, E., Do, T., Boscoe, B., et al. 2022, Photometric Redshifts for Cosmology: Improving Accuracy and Uncertainty Estimates Using Bayesian Neural Networks, arXiv, doi: 10.48550/ARXIV.2202.07121
  * Kasen et al. (2013) Kasen, D., Badnell, N. R., & Barnes, J. 2013, The Astrophysical Journal, 774, 25, doi: 10.1088/0004-637x/774/1/25
  * Kasen et al. (2015) Kasen, D., FernÃ¡ndez, R., & Metzger, B. D. 2015, Monthly Notices of the Royal Astronomical Society, 450, 1777, doi: 10.1093/mnras/stv721
  * Kasen et al. (2017) Kasen, D., Metzger, B., Barnes, J., Quataert, E., & Ramirez-Ruiz, E. 2017, Nature, 551, 80, doi: 10.1038/nature24453
  * Kawaguchi et al. (2016) Kawaguchi, K., Kyutoku, K., Shibata, M., & Tanaka, M. 2016, ApJ, 825, 52, doi: 10.3847/0004-637X/825/1/52
  * Kawaguchi et al. (2020) Kawaguchi, K., Shibata, M., & Tanaka, M. 2020, ApJ, 889, 171, doi: 10.3847/1538-4357/ab61f6
  * Kingma & Welling (2013) Kingma, D. P., & Welling, M. 2013, arXiv e-prints, arXiv:1312.6114.  https://arxiv.org/abs/1312.6114
  * Kingma & Welling (2019) Kingma, D. P., & Welling, M. 2019, Foundations and Trends® in Machine Learning, 12, 307–392, doi: 10.1561/2200000056
  * Lamb & Kobayashi (2016) Lamb, G. P., & Kobayashi, S. 2016, The Astrophysical Journal, 829, 112, doi: 10.3847/0004-637x/829/2/112
  * Li et al. (2022) Li, C., Zhang, Y., Cui, C., et al. 2022, MNRAS, 509, 2289, doi: 10.1093/mnras/stab3165
  * Li & Paczyński (1998) Li, L.-X., & Paczyński, B. 1998, The Astrophysical Journal, 507, L59, doi: 10.1086/311680
  * LIGO Scientific Collaboration et al. (2017) LIGO Scientific Collaboration, Virgo Collaboration, Fermi GBM, et al. 2017, Astrophysical Journal Letters, 848, doi: 10.3847/2041-8213/aa91c9
  * Lipunov et al. (2017) Lipunov, V. M., Gorbovskoy, E., Kornilov, V. G., et al. 2017, The Astrophysical Journal, 850, L1, doi: 10.3847/2041-8213/aa92c0
  * Lukošiūtė et al. (2022) Lukošiūtė, K., Raaijmakers, G., Doctor, Z., Soares-Santos, M., & Nord, B. 2022, KilonovaNet: Surrogate Models of Kilonova Spectra with Conditional Variational Autoencoders, arXiv, doi: 10.48550/ARXIV.2204.00285
  * Metzger (2019) Metzger, B. D. 2019, Living Reviews in Relativity, 23, doi: 10.1007/s41114-019-0024-0
  * Metzger & Fernández (2014) Metzger, B. D., & Fernández, R. 2014, MNRAS, 441, 3444, doi: 10.1093/mnras/stu802
  * Metzger et al. (2010) Metzger, B. D., Martínez-Pinedo, G., Darbha, S., et al. 2010, MNRAS, 406, 2650, doi: 10.1111/j.1365-2966.2010.16864.x
  * Nakar (2007) Nakar, E. 2007, Phys. Rep., 442, 166, doi: 10.1016/j.physrep.2007.02.005
  * Nicholl et al. (2021) Nicholl, M., Margalit, B., Schmidt, P., et al. 2021, MNRAS, 505, 3016, doi: 10.1093/mnras/stab1523
  * Ntampaka et al. (2021) Ntampaka, M., Ho, M., & Nord, B. 2021, arXiv e-prints, arXiv:2111.14566.  https://arxiv.org/abs/2111.14566
  * pandas development team (2020) pandas development team, T. 2020, pandas-dev/pandas: Pandas, latest, Zenodo, doi: 10.5281/zenodo.3509134
  * Perego et al. (2017) Perego, A., Radice, D., & Bernuzzi, S. 2017, 850, L37, doi: 10.3847/2041-8213/aa9ab9
  * Roberts et al. (2011) Roberts, L. F., Kasen, D., Lee, W. H., & Ramirez-Ruiz, E. 2011, ApJ, 736, L21, doi: 10.1088/2041-8205/736/1/L21
  * Rosswog et al. (1999) Rosswog, S., Liebendörfer, M., Thielemann, F. K., et al. 1999, A&A, 341, 499.  https://arxiv.org/abs/astro-ph/9811367
  * Rosswog et al. (2013) Rosswog, S., Piran, T., & Nakar, E. 2013, MNRAS, 430, 2585, doi: 10.1093/mnras/sts708
  * Ruffert, M. & Janka, H.-Th. (2001) Ruffert, M., & Janka, H.-Th. 2001, A&A, 380, 544, doi: 10.1051/0004-6361:20011453
  * Rumelhart et al. (1986) Rumelhart, D. E., Hinton, G. E., & Williams, R. J. 1986, Nature, 323, 533, doi: 10.1038/323533a0
  * Savchenko et al. (2017) Savchenko, V., Ferrigno, C., Kuulkers, E., et al. 2017, The Astrophysical Journal, 848, L15, doi: 10.3847/2041-8213/aa8f94
  * Scolnic et al. (2018) Scolnic, D., Kessler, R., Brout, D., et al. 2018, The Astrophysical Journal Letters, 852, L3, doi: 10.3847/2041-8213/aa9d82
  * Shen et al. (2019) Shen, H., George, D., Huerta, E. A., & Zhao, Z. 2019, ICASSP 2019 - 2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), doi: 10.1109/icassp.2019.8683061
  * Sheng et al. (2022) Sheng, X., Ross, N., & Nicholl, M. 2022, MNRAS, 512, 5580, doi: 10.1093/mnras/stac803
  * Sipőcz et al. (2020) Sipőcz, B. M., Connolly, A., Ivezić, Ž., Vander Plas, J., & Gray, A. 2020, in American Astronomical Society Meeting Abstracts, Vol. 235, American Astronomical Society Meeting Abstracts #235, 109.15
  * Soares-Santos et al. (2017) Soares-Santos, M., Holz, D. E., Annis, J., et al. 2017, ApJ, 848, L16, doi: 10.3847/2041-8213/aa9059
  * Tanaka & Hotokezaka (2013) Tanaka, M., & Hotokezaka, K. 2013, The Astrophysical Journal, 775, 113, doi: 10.1088/0004-637x/775/2/113
  * Tanaka et al. (2017) Tanaka, M., Utsumi, Y., Mazzali, P. A., et al. 2017, PASJ, 69, 102, doi: 10.1093/pasj/psx121
  * Tanvir et al. (2017) Tanvir, N. R., Levan, A. J., González-Fernández, C., et al. 2017, The Astrophysical Journal, 848, L27, doi: 10.3847/2041-8213/aa90b6
  * Valenti et al. (2017) Valenti, S., Sand, D. J., Yang, S., et al. 2017, ApJ, 848, L24, doi: 10.3847/2041-8213/aa8edf
  * van Eerten & MacFadyen (2012) van Eerten, H. J., & MacFadyen, A. I. 2012, The Astrophysical Journal, 751, 155, doi: 10.1088/0004-637x/751/2/155
  * Vander Plas et al. (2014) Vander Plas, J., Connolly, A. J., & Ivezic, Z. 2014, in American Astronomical Society Meeting Abstracts, Vol. 223, American Astronomical Society Meeting Abstracts #223, 253.01
  * Villar et al. (2018) Villar, V. A., Cowperthwaite, P. S., Berger, E., et al. 2018, ApJ, 862, L11, doi: 10.3847/2041-8213/aad281

##  Appendix A Additional Results from $D_{1}$

Figure 8: This figure encompasses the CVAE-generated light curves for
different values of physical parameters for the respective filter bands. In
each row of the figure, light curves corresponds to different ejecta mass
having values of $(0.1M_{\odot},0.01M_{\odot},0.001M_{\odot})$keeping ejecta
velocity $(0.03c)$ and lanthanide fraction $(10^{-5})$ constant; different
ejecta velocity having values $(0.01c,0.03c,0.05c)$ keeping ejecta
mass$(0.025M_{\odot})$ and lanthanide fraction $(10^{-5})$ to be constant;
different lanthanide fraction having values $(10^{-4},10^{-5},10^{-9})$,
keeping ejecta mass $(0.025M_{\odot})$ and ejecta velocity $(0.03c)$ to be
constant for _i,r,y_ and _z_ filter bands respectively.

##  Appendix B Additional Results from $D_{2}$

Figure 9: This figure demonstrates the evolution of light curves for different
physical parameters of chirp mass, viewing angle, mass ratio and fraction of
the remnant disk. In each row of the plot, we bring together the CVAE-
generated light curves for different chirp mass values
$(1.0M_{\odot})$,$(1.2M_{\odot})$,$(1.4M_{\odot})$,$(1.6M_{\odot})$ and
$(1.8M_{\odot})$ while keeping the viewing angle($60\degree$), mass
ratio($0.9$) and fraction of the remnant disk($0.15$) to be constant; for
different viewing angles having values
$(10\degree,30\degree,45\degree,60\degree,75\degree,90\degree)$ from the pole
keeping the chirp mass $(1.2M_{\odot})$, mass ratio$(0.9)$ and the fraction of
the remnant disk ($0.15$) constant; for different fraction of the remnant disk
having values $(0.15,0.20,0.25,0.30,0.35,0.40)$, keeping the chirp
mass$(1.2M_{\odot})$, mass ratio $(0.9)$ and the viewing angle $(90\degree)$
constant for the _i,r,y,z_ and _u_ filter bands respectively.

##  Appendix C The CVAE Summary.

Figure 10: This is the summary of the CVAE implemented in the current work