Domain adaptation in application to gravitational lens finding

Gravitational lensing, the distortion of light from a distant source due to the presence of a massive object between that source and the observer, is one of the most important phenomena predicted by the general relativity theory and a powerful tool to study the universe. For example, the sensitivity of lensing observables to the distribution of mass in the lens provides a way to study the dark matter profiles of individual galaxies, galaxy groups, and galaxy clusters (Koopmans et al., 2009; Barnabè et al., 2011; Sonnenfeld et al., 2013; Newman et al., 2015; Li et al., 2018; Etherington et al., 2023). It also allows to probe dark matter substructures on subgalactic scales (Daylan et al., 2018; Brehmer et al., 2019; Diaz Rivero & Dvorkin, 2020; Bayer et al., 2023), and to test various dark matter theories (Li et al., 2016; Gilman et al., 2018; Alexander et al., 2020; Şengül & Dvorkin, 2022; Vegetti et al., 2023; Anau Montel et al., 2023; Gilman et al., 2023). Time-delay measurements of gravitationally lensed time-variable sources, such as supernovae and quasars, allow for measuring absolute distances and constraining the Hubble constant and other cosmological parameters (Refsdal, 1964; Shajib et al., 2020; Wong et al., 2020; Treu et al., 2022; Birrer et al., 2022; Suyu et al., 2024). Cluster-scale lenses provide magnification that enables spectroscopic studies of high-redshift galaxies, helping to uncover star formation processes in the early universe (Rigby et al., 2018; Sharon et al., 2022; Zhang et al., 2023; Mainali et al., 2023; Keerthi Vasan G. et al., 2024).

Strong gravitational lensing, characterized by the occurence of multiple images, arcs, and rings, requires precise alignment between foreground and background sources. Therefore, these systems are very rare, and their discovery demands extensive search efforts. Upcoming wide-field surveys, such as Euclid Wide Survey (Laureijs et al., 2011; Euclid Collaboration et al., 2022) and Legacy Survey of Space and Time (LSST; LSST Science Collaboration et al., 2009) are expected to deliver on the order of $10^{5}$ strong lenses (Collett, 2015). However, finding them will require searching through hundreds of millions of sources, which is infeasible for human inspection. Over the last 20 years, various methods have been employed to discover gravitational lenses, such as citizen science (e.g. Marshall et al., 2016), spectroscopic selection of lens candidates (Bolton et al., 2004, 2006; Shu et al., 2017), arcs and ring-like feature finders (More et al., 2012; Gavazzi et al., 2014; Sonnenfeld et al., 2018). Nevertheless, in recent years, the majority of lens discoveries have come from deep learning algorithms.

Typically, the lens finding problem is treated as an image classification problem, which is well-suited for Convolutional Neural Networks (CNN) (Lecun et al., 1998). Supervised CNN-based algorithms have demonstrated superior performance in the gravitational lens finding challenge (Metcalf et al., 2019) and have recently been applied to various wide-field surveys, yielding a few thousand lens candidates (e.g. Petrillo et al., 2017; Pourrahmani et al., 2018; Canameras et al., 2021; Cañameras et al., 2020; Li et al., 2021; Rojas et al., 2022; Shu et al., 2022a; Storfer et al., 2022; Huang et al., 2020, 2021; Jacobs et al., 2019). To ensure that the model, trained in a supervised way, is able to learn relevant features and generalize well to unseen data, the training dataset typically requires a large number of labeled samples (on the order of $10^{3}-10^{4}$), which exceeds the number of known lenses. For that reason, it is common to use simulated lenses at the training stage (however, see e.g. Huang et al., 2020, 2021 for supervised training only with observational data).

Despite efforts to increase the realism of simulations, differences between real observations and mock images are unavoidable and can significantly degrade the performance of the model (Ćiprijanović et al., 2020, 2022). Among the ways to tackle this issue are, for example, implementing unsupervised methods and train a model purely on observational data, as in Cheng et al. (2020) and Stein et al. (2022), or applying domain adaptation techniques to reduce the gap between the simulated and observational domains.

Domain adaptation (DA) is a class of methods applied in situations when the training and test datasets come from different, but related domains, namely the source and the target domains. Based on the availability of labels in the target dataset, DA methods can be divided into three categories: unsupervised domain adaptation (UDA), when ground truth for the target data is unknown; semi-supervised domain adaptation (SSDA), where labels are available for a fraction of the target data; and supervised domain adaptation (SDA), where labels for the entire dataset are accessible. In astrophysics, domain adaptation has found various applications, including, but not limited to: improving the performance of galactic morphology classifiers across different surveys (Xu et al., 2023; Ćiprijanović et al., 2023), identifying merging galaxies in simulated and observational datasets (Ćiprijanović et al., 2021), and classifying strong gravitational lenses with different dark matter substructure simulated for surveys with varying observational characteristics (Alexander et al., 2023).

In this paper, we assess how much domain adaptation techniques can enhance the performance of lens-finding algorithms trained in a supervised way on simulated or real data. Building upon the work in Alexander et al. (2023), we examine three domain adaptation methods: Adversarial Discriminative Domain Adaptation (ADDA), Wasserstein Distance Guided Representation Learning (WDGRL), and Supervised Domain Adaptation (SDA). The paper is structured as follows: Section 2 describes the algorithms we employ; Section 3 presents the datasets used in the experiments; and Section 4 outlines the neural network architecture and training process. We discuss our results in Section 5 and conclude in Section 6.

In the context of gravitational lens finding, our goal is to train a model on a source dataset containing simulated lenses and adapt it to a smaller unlabeled target dataset with real lenses, using either the Adversarial Discriminative Domain Adaptation (ADDA) or the Wasserstein Distance Guided Representation Learning (WDGRL) method. We will also explore the supervised setting for domain adaptation and compare the supervised classifier trained purely on real data with the Supervised Domain Adaptation (SDA) algorithm trained on a labeled target dataset. A brief description of each domain adaptation algorithm follows. For the remainder of the section, $(\mathbf{X}^{s},Y^{s})$ denotes images and labels of the source dataset, while $(\mathbf{X}^{t},Y^{t})$ represents images and labels of the target dataset.

As a source dataset, we use a mixture of simulated lenses and real non-lensed galaxies. For the target dataset, we use purely observational data, combining previously discovered lens candidates with non-lensed galaxies. We describe the construction of these datasets below.

We begin by selecting a parent dataset from Hyper Suprime-Cam Subaru Strategic Program (HSC-SSP) PDR2 Wide field (Aihara et al., 2019), following the criteria in Shu et al. (2022b) and Cañameras et al. (2021). The selection targets the extended sources brighter than 26 mag in the g, r, and i bands with the following color cuts: $0.6<\texttt{g\_cmodel\_mag}-\texttt{r\_cmodel\_mag}<3$, $2<\texttt{g\_cmodel\_mag}-\texttt{i\_cmodel\_mag}<5$. The color cuts serve to narrow the sample to preferentially red galaxies, as massive red ellipticals are expected to dominate the lens population. We cross-matched the parent dataset with the list of known lens candidates and removed matches with separation less than 40 arcsec to reduce possible contamination among the non-lensed sources. We then used this cleaned parent dataset to draw negative examples (non-lenses) for both the source and target dataset. While there is a possibility of undiscovered lenses in the parent sample, we expect the fraction of such serendipitous lenses to be negligible.

We compare the performance of three domain adaptation methods described in Sec. 2. While these methods differ in their training approaches, they share common components: an encoder that maps input images to lower-dimensional representations, and a classifier that outputs binary classes for these representations.

For the encoder, we explored two different architectures: the residual neural network (ResNet, He et al., 2016) and the equivariant neural network (ENN, Weiler & Cesa, 2019). ResNet is a type of convolutional neural networks that was developed to overcome the problem of degrading accuracy with an increasing number of layers in earlier models such as AlexNet or VGG. Its key component is the residual block with skip connections. In a residual block, the input is propagated through the layers without being changed, so the block is learning the difference between the original function and an identity transformation. Since their introduction, ResNets have gained popularity in image recognition tasks; specifically in the field of gravitational lens detection, a ResNet-based model won the gravitational lens finding challenge (Metcalf et al., 2019; Lanusse et al., 2018).

We based the encoder on the ResNet-18 architecture. The encoder’s core consists of four residual blocks, each containing two convolutional layers with batch normalization and ReLU activations. We modified the standard ResNet-18 architecture by replacing its final fully connected layer with a custom module: a dropout layer for regularization, followed by a linear layer that outputs a 256-dimensional embedding vector of the input image.

In addition to ResNet, we explored equivariant neural networks (ENN), which are designed to preserve inherent symmetries in the data through their architecture. For gravitational lenses, which often display rotational and reflectional symmetries, the use of ENNs can lead to more efficient learning, better generalization, and a reduced need for data augmentation. For example, in Alexander et al. (2023) domain adaptation with an ENN-based encoder showed superior performance in the task of classifying signatures of dark matter substructures in simulated strong gravitational lenses.

We used the e2cnn (Cesa et al., 2022) package and implemented an ENN with the dihedral group D4, which includes the identity transformation, rotations by $\pm\pi/2$ and $\pi$, and horizontal/vertical reflections. Our model comprises six equivariant convolutional blocks, each composed of a convolutional layer, a batch normalization layer, and a ReLU activation function. The final layer also outputs a 256-dimensional representation.

For the classification task, both encoders are followed by a simple classifier network, consisting of two fully connected layers for the binary classification.

We used data augmentation with rotations and horizontal/vertical flips. While augmentation is not required for ENNs due to symmetries being incorporated in the structure of the network, it turned out to be crucial for the successful training of the ResNet-based model.

We used the Adam optimizer (Kingma & Ba, 2014) to minimize losses and 1-cycle scheduler to adjust the learning rate. We trained both ENN and ResNet-based networks for 100 epochs with a patience of 5 epochs, so the training stops if the validation loss of the model does not improve for 5 consecutive epochs.

After the source encoder and the classifier are trained on the source dataset in a supervised way, we continue with the domain adaptation training.

In the ADDA method, the training step is split into two parts. First, we update a discriminator which operates on the embeddings and define whether they come from a source or target dataset. In our case, the discriminator is implemented as a simple three-layer network. Next we update a target encoder, initialized with the same weights as the source encoder, using inverted labels. For the ResNet-based model, we update encoder 5 times for each discriminator step for more stable training.

For the WDGRL approach, we train a critic network and update the pre-trained encoder and classifier. The goal of the critic network is to estimate the Wasserstein distance between the representations of source and target images. For each mini-batch, we first train the critic for n_critic steps and then train the encoder and classifier with cross-entropy loss for the classifier and Wasserstein distance as the loss for the encoder.

For the supervised domain adaptation, we also update the pre-trained classifier and encoder; however we use a contrastive semantic alignment (CSA) loss. In this case, the most important hyperparameters are the balancing coefficient $\alpha$, which determines the relative contribution of classifier loss and CSA loss, and the margin $m$, which sets the minimal separation distance for the examples from different domains and different class labels.

In all cases we use the Adam optimizer to optimize the losses. All hyperparameters were fine-tuned via grid search and are listed in Table 1. The models were trained on the high-performance computing cluster at The University of Alabama.

To compare the efficiency of domain adaptation techniques for lens finding, we use the Receiver Operation characteristic Curve (ROC) curve as our main metric. The ROC curve illustrates how the false positive rate (FPR) and true positive rate (TPR) change with the variation of the classification threshold. TPR, also known as sensitivity or recall, is defined as the number of correctly classified lenses relative to the total number of lenses in the test dataset:

FPR is the ratio of incorrectly classified non-lensed objects to the total number of non-lenses:

For an ideal classifier, FPR is close to 0 when the TPR is close to 1. For a random classifier, FPR equals TPR for all thresholds. The model that performs better has higher area under the ROC curve (AUROC).

Besides the AUROC, it is also particularly interesting to compare the behaviour of the ROC curve at low FPRs; we quantify it by measuring $\mathrm{TPR}_{1\%}$ and $\mathrm{TPR}_{0.1\%}$ – true positive rates at the thresholds that provide false positive rate of 1% and 0.1% respectively. In practice, all objects classified as positive examples by automated methods are visually validated. For the purpose of saving time in human inspection, the ideal algorithm should have a very low contamination fraction while maintaining a high number of identified lenses.

Fig. 3 displays ROC curves for unsupervised domain adaptation methods with a ResNet-based encoder on the left panel and an ENN-based encoder on the right, evaluated on the imbalanced test dataset with 200 lenses and 20 000 non-lenses. For the ResNet, both ADDA and WDGRL result in an overall higher AUROC; however, the increase in performance is rather small, and at low false positive rates the recall of all three algorithms is similar (see Table 2). Compared to ResNet, the ENN-based model works much better even without domain adaptation (AUROC score 0.921 vs 0.858) and improvement from application of domain adaptation is more prominent, especially at low false positive rates: $\mathrm{TPR}_{1\%}=0.354$ for the naive approach, 0.462 for ADDA, and 0.528 for WDGRL. In both cases, the WDGRL method outperforms ADDA.

Fig. 4 demonstrates the results for the supervised domain adaptation. In this setting, we compared the adapted model (blue curve) with a model trained on an unbalanced dataset containing only observational images of lensed and non-lensed sources (red curve). The supervised domain adaptation in the case of ENN-based encoder gives some improvement, however it is rather small (AUROC score 0.987 vs 0.991). For the ResNet-based classifier, SDA performed even worse than naive approach. One interesting result is that the model, trained on real sources, significantly outperform the classifiers trained on simulations (compare red curves on Fig. 3 and 4), with AUROC for ENN improving from 0.921 to 0.987 and for ResNet – from 0.858 to 0.978, even though the number of real lenses is $\sim 15$ times smaller than the number of mock lenses. While the unbalanced proportion of the training dataset is in better agreement with the unbalanced nature of the test dataset and of real world data, it is usually considered a good practice to balance training dataset to make sure that the model is able to see a sufficient number of examples from all classes and learn relevant features. The fact, that the model trained on a smaller set of real data performs better than the model trained on a larger dataset with simulations highlights that simulations are not perfect representation of real data. Another explanation might come from the fact that majority of real lens candidates that consitute training and test set were found with machine learning algorithms and might represent subsample of objects with similar properties, for example, in terms of relative brightness and size of lens features. Therefore, in future work it would be interesting to apply both models to the larger sample of unseen data and compare properties of the identified lens candidates.

In this work, we investigate whether domain adaptation is able to enhance the performance and reliability of CNN-based gravitational lens detection algorithms. Due to the limited sample of known lenses, deep learning models are typically trained on simulated strong lenses, which can lead to degraded performance on real data. The shift between training and test domains can be alleviated with domain adaptation techniques.

We used simulated mock lenses as positive examples in the source dataset and observations of lens candidates from HSC SSP PDR2 Wide layer as constituents in the target dataset. We explored two neural network architectures, ResNet-18 and Equivariant Neural Network, for the encoder and found that ENNs are more successful in identifying lenses.

For unsupervised domain adaptation, we implemented two methods, ADDA and WDGRL. The WDGRL approach resulted in the largest increase of ROC score compared to the naive inference of a classifier trained with mock lenses on real data. For the supervised domain adaptation, we compared the performance of the adapted algorithm with a model trained on an unbalanced dataset containing purely observational data. Somewhat surprisingly, despite the relatively small number of lenses in the training sample (1754 real lenses vs 31754 mock lenses in the source dataset), the model showed superior performance. Supervised domain adaptation provided only marginal improvement in the case of the ENN-based algorithm, and did not yield any improvement for the ResNet-based model.

In the common setting for lens searches, algorithms are applied to large datasets ($10^{6}-10^{7}$ objects), and objects with a classification score above a predefined threshold are visually inspected by a team of experts. A higher threshold leads to increased sample purity at the cost of reduced completeness, and vice versa. This trade-off means that human experts must examine a large number of lens candidates. For example, in Jaelani et al. (2023), 20 241 sources with a score higher than 0.9 were visually inspected, first reduced to 1522 cutouts, and then to 43 definite and 269 probable lenses. Upcoming surveys are expected to contain $\mathcal{O}(10^{5})$ lenses, which makes the task of reducing false positives particularly important. In our experiment we saw that the improvement in recall from the combination of ENN-based algorithm and WDGRL domain adaptation method is particularly noticeable for low false positive rates.

One of the challenges of the lens finding task is the correct classification between lenses and similar-looking contaminants, such as ring galaxies or spiral galaxies. In the supervised domain adaptation setting, we show that an algorithm trained on a target dataset that includes a relatively small number of spirals is able to enhance the model’s ability to distinguish between lenses and spiral galaxies. The model’s performance increases with the number of spiral galaxies in the target dataset. However, adding spirals only to the source dataset leads to performance degradation.

Domain adaptation techniques can be seamlessly integrated into lens finding pipelines and easily combined with other methods. Our results emphasize that DA methods should be considered an essential component in future lens finding campaigns.

We acknowledge useful conversations with Stephon Alexander. H. P. and S. G. were supported in part by U.S. National Science Foundation award No. 2108645. P. R. was a participant in the Google Summer of Code 2023 program. This material is based upon work supported by the U.S. Department of Energy, Office of Science, Office of High Energy Physics of U.S. Department of Energy under grant Contract Number DE-SC0012567. M. W. T. acknowledges financial support from the Simons Foundation (Grant Number 929255).