Enhancing Gravitational Lens Study with Deep Learning: A Study on Effects of Dropout Regularization

Abstract

Strong gravitational lensing provides valuable insights into the mass distribution of galaxies and the nature of dark matter. However, its modeling is computationally demanding due to the large volume of strong lensing observations. In this work, we explore the application of Convolutional Neural Networks to infer physical parameters from simulated galaxy–galaxy lens systems, described by the Singular Isothermal Ellipsoid (SIE) profile for the galaxy lens. We construct a dataset of 76,396 synthetic lensing images derived from the China Space Station Telescope catalog and employ it to train a modified CNN model, based on the AlexNet architecture, to predict four key SIE parameters, the Einstein radius, the axis ratio and ellipticity components. We analyze the network performance under three distinct dropout configurations to quantify their influence on generalization and parameter inference accuracy. The results indicate that the incorporation of dropout is critical for enhancing the precision and robustness of the estimated parameters as demonstrated using a 4-fold cross-validation procedure. When dropout tools are included, we obtain coefficients of determination up to $R^2\sim 0.96$ for most SIE parameters and mean peak signal-to-noise ratios of up to ∼37 dB. Relative to the configuration without dropout, the use of dropout reduces the relative errors in the inferred SIE parameters by approximately 60–76%, resulting in errors of at most ∼9% at the 90% confidence level for the majority of parameters. These findings highlight the potential of deep learning approaches to enable scalable, computationally efficient, and high-precision modeling of strong gravitational lensing systems.

1. Introduction

A gravitational lens is any massive object, such as a galaxy, galaxy cluster, or quasar, whose gravitational influence results in the curvature of space-time, thereby altering the trajectory of light from more distant objects. When the projected mass within the Einstein radius becomes sufficient to produce multiple images, the observed phenomenon is known as a strong gravitational lens (SGL). When the alignment observer-lens-source is perfect, and the source is extended the resulting image is a ring known as an Einstein ring [1,2,3,4,5]. Such phenomena impose significant constraints on the projected mass of the lensed object, aligning with the predictions of Einstein’s theory of general relativity. Thus, strong lensing functions as a vital tool for probing dark matter within the central regions of halos or for determining the gradient of the inner mass density profile of galaxies [5,6]. Moreover, strong lenses facilitate the examination of sources with high redshift as a result of their ability to magnify up to 30 times the background source [7]. Grillo et al. [8] proposed a methodology to estimate cosmological parameters using Strong Gravitational Systems (SGSs). This method leveraged the correlation between the Einstein radius and the central stellar velocity dispersion, assuming an isothermal distribution of the total density for the lensed galaxy. However, the problem with applying this method is the mass-sheet degeneracy and the deviations from isothermality and orbital isotropy that can cause the observed velocity dispersion to differ by up to $20\%$ [5]. Solving this degeneracy requires spatially resolved kinematics, which might not be readily available for the thousands of lensed systems expected in future surveys. Thus, the method proposed by [8] is a first approximation for such a large sample of systems. Another application is its use to estimate the Hubble constant $H_0$ by analyzing time delays between multiple quasar images of the source [5,9,10]. However, this method requires investing in months or years of photometric follow-up observations.

Advanced and sophisticated telescopes, including Euclid [11,12], the Very Large Telescope [13], the James Webb Space Telescope (JWST) [14,15] are already providing unprecedented data. Furthermore, next-generation facilities such as the Simonyi telescope at the Vera C. Rubin Observatory [16], the Extremely Large Telescope (ELT) [17,18] and the Chinese Survey Space Telescope (CSST) [19] are expected to begin operations in the coming years. Together, these instruments will capture images and collect extensive datasets of approximately ∼10⁵ [20] to enhance our understanding of SGSs. Some known large galaxy–galaxy lens surveys include the Cosmic Lens All-Sky Survey (CLASS) [21,22], the COSMOS Survey [23], and the Sloan Lens ACS Survey (SLACS) [24,25]. Such an extensive data collection, comprising more than 100,000 images, poses a challenge to traditional analytical methods, which may require substantial processing time. Conventional lens modeling techniques attempt to address the non-linear inverse problem of reconstructing the brightness of lensed sources while simultaneously modeling the gravitational potential of lenses.

Subsequent to the detection of lenses, it is imperative to develop a model that accurately depicts the total mass distribution, thereby assisting in the elimination of potential false-positive outcomes. Currently, Monte Carlo Markov Chain (MCMC) techniques have been deployed for system modeling; however, these methods are becoming inadequate given the increasing volume of data [26]. To mitigate this issue, artificial intelligence tools such as Convolutional Neural Networks (CNNs) are highly valued due to their proven success in stellar image classification, thereby becoming the burgeoning technique for lens search and identification. CNNs were first introduced in 1988 when Zhang [27] proposed the first two-dimensional convolutional neural network, known as the shift-invariant artificial neural network. The inaugural application of CNNs in astronomical classification involved spectrum classification in the SDSS as documented by Hála’s [28]. Petrillo et al. [29] pioneered the use of a CNN-based morphological classification method to identify strong gravitational lenses in the 225 Kilo Degree Survey. Hezaveh showcased the proficiency of CNNs in predicting parameter values of mass models using preprocessed high-resolution lens images [30]. This team further expanded their research by presenting a CNN specifically designed for modeling high-resolution lens images. The Highly Optimized Lensing Investigations of Supernovae, Microlensing Objects, and Kinematics of Ellipticals and Spirals (HOLISMOKES) team developed a CNN aimed at modeling strongly lensed galaxy images [31], crafted to predict the five parameters of the SIE mass model. Warren R. Morningstar et al. [32] introduced a machine learning approach to reconstructing undistorted images of background sources in strongly lensed systems. Additionally, Bayesian Neural Networks (Bayesian NNs), a probabilistic variant of deep neural networks, have exhibited great success in the extraction of highly abstract information from complex image data. This capability is exemplified in the work of Ji Won Park et al. [33], who used simulated lens time delays to determine the Hubble constant, and Pearson et al. [34], who trained an approximate Bayesian CNN to predict the parameters of the mass profile. More recently, the Euclid Collaboration [35] designed a Bayesian NN to model gravitational lenses automatically and efficiently. With LEMON, key parameters of the lens mass profile, such as the Einstein radius, are estimated in addition to the parameters of the light distribution of the lens galaxy, along with its uncertainties. LEMON was applied to simulated lenses (with characteristics similar to those captured by the Euclid telescope) and also to real lenses. Further works such as those by R. Parlange et al. [36] dive into transformer models to detect strong gravitational lenses from sets from HOLISMOKES VI and SuGOHI X.

Due to the efficacy of CNNs in analyzing SGS, this study provides a comprehensive examination of CNN performance in predicting the physical parameters associated with strong lensing. Specifically, our research emphasizes the influence of CNN architecture and regularization strategies, with particular attention to the incorporation of dropout layers on the accuracy and robustness of parameter estimation. Dropout, a widely used technique introduced by N. Srivastava et al. [37], plays a crucial role in mitigating overfitting in deep learning models, thereby enhancing their generalization capabilities. Its functionality consists on randomly turning off a set of neurons during each training iteration, preventing the units from depending on each other and enabling the network to learn more complex structures. The dropout technique is more detailed in Section 3. In contexts such as gravitational lensing, where the recovery of model parameters such as ellipticity and external shear is acutely sensitive to minute variations in image features, the optimization of dropout placement and rate is crucial in enhancing prediction quality. Consequently, our objective is to determine the optimal configuration and integration of dropout layers within a CNN model to enable precise and reliable estimation of the parameters of the Singular Isothermal Ellipsoid (SIE) model. Our detailed analysis involves comparing the complete image set against the discrepancies observed between simulated images and those modeled based on CNN-predicted values. The synthetic dataset comprises 76,396 images of simulated galaxy–galaxy lens systems derived from the China Space Station Telescope (CSST) catalog, generated in a noise-free environment. In our analysis, the dataset was partitioned into 70,000 images for the k-fold cross-validation process, in which each fold utilized 75% of the images for training and 25% for validation. Additionally, the 6393 remaining images were used in testing purposes. A model based on the AlexNet architecture [38] was selected because actual deep architectures are excessive and computationally expensive to process the millions of images that the great telescopes are going to get. The modifications made to the network enable a smaller architecture compared to current ones, facilitating rapid predictions. The addition of a convolutional layer to the central block enhances data extraction depth and mitigates linearity. The model predicts four parameters from the SIE lens model: complex ellipticity components ${ϵ}_x$ and ${ϵ}_y$, axial ratio f and Einstein radius ${\theta }_E$ (see Section 2 for their definitions). The predictive accuracy of the model is assessed by measuring the relative error between the actual and predicted values. Through this investigation, our aim is to advance the interpretability and effectiveness of deep learning methodologies in modeling gravitational lenses through the validation of the dropout as a critical tool for the reliability of predictions and the use of a lighter architecture designed for the massive data.

The manuscript is organized as follows. Section 2 and Section 3 introduce the subject matter and articulate the principal aim of the paper. These sections provide a theoretical foundation concerning gravitational lensing and succinctly expound on CNNs. Section 4 explores the synthetic data used in this investigation, elaborating on the factors considered in the methodological approach, and offers a comprehensive overview of the evaluation metrics used to assess the performance of the model and analyze the relative errors. Section 5 provides a summary of the performance of the model, elucidates the relative errors associated with the predictions, and discusses these findings. Finally, Section 6 presents the conclusions derived from the study and outlines recommendations for future research.

2. Gravitational Lensing Background

2.1. Strong Lensing Theory

A representative scenario in gravitational lensing is depicted in Figure 1, where a mass concentration is observed to deflect light rays emanating from a source located at a distance $D_S$. In the absence of additional deflectors near the line of sight, and assuming that the extent of the deflecting mass along the line-of-sight is considerably smaller than both $D_L$ and the distance $D_{LS}$ from the deflector to the source, light rays are considered to be deflected in a plane (thin lens hypothesis). The magnitude and orientation of this curvature are characterized by the deflection angle $\overrightarrow{\widehat{\alpha }}$, which depends on the mass distribution of the deflector and the impact vector of the light ray [1]. The real position of the source relative to the observed position of the images in the sky is therefore delineated by the lens equation, which outlines a transformation from the lens plane to the source plane for any mass distribution. For a source positioned in the source plane $\overrightarrow{\eta }=\overrightarrow{\beta }{D}_S$ with a real angular position $\overrightarrow{\beta }$, the light rays that arrive at the lens plane at a distance $\overrightarrow{\xi }=\overrightarrow{\theta }{D}_L$ from the optical axis (defined by the lens-observer line) are deflected by the angle $\overrightarrow{\widehat{\alpha }}$ before reaching an observer. As a result of the deflection, the observer perceives images of the source at the apparent angular positions $\overrightarrow{{\theta }_i}$, with i representing each image. If $\overrightarrow{\theta }$, $\overrightarrow{\beta }$, and $\overrightarrow{\widehat{\alpha }}$ are small angles (of the order of arcseconds), considering the reduced deflection angle by

$$\overrightarrow{\alpha }(\overrightarrow{\theta })\equiv {\displaystyle \frac{D_{LS}}{D_S}}\overrightarrow{\widehat{\alpha }},$$

(1)

where $D_{LS}$ is the angular diameter distance between the lens and the source, the lens equation can be written as

$$\overrightarrow{\beta }=\overrightarrow{\theta }-\overrightarrow{\alpha }(\overrightarrow{\theta }).$$

(2)

One important consequence of gravitational lensing is image distortion, which can be described by the Jacobian matrix as

$$A\equiv {\displaystyle \frac{\partial \overrightarrow{\beta }}{\partial \overrightarrow{\theta }}}=\left({\delta }_{ij}-{\displaystyle \frac{{\partial }^2\widehat{\mathsf{\Psi }}(\overrightarrow{\theta })}{\partial {\theta }_i\partial {\theta }_j}}\right),$$

(3)

where $\widehat{\mathsf{\Psi }}$ represents the effective lensing potential. By segregating the isotropic component (associated to the mass of the main deflector) from the Jacobian and extracting its trace-less segment, we can derive the shear tensor, which encapsulates the external perturbation, thereby characterizing the distortions of background sources. The expression for the shear tensor is

$$\mathsf{\Gamma }={\left({\displaystyle \frac{1}{2}}\mathrm{tr}A·I-A\right)}_{ij}=\left(\begin{array}{cc}{\gamma }_1& {\gamma }_2\\ {\gamma }_2& -{\gamma }_1\end{array}\right),$$

(4)

where ${\gamma }_1=\gamma cos2\varphi$ and ${\gamma }_2=\gamma sin2\varphi$.

2.2. SIE Model Parametrization

In order to construct realistic models of gravitational lens systems, we must employ mass distributions defined by elliptical isodensity contours. In this research, we focus on the surface density of the Singular Isothermal Ellipsoid (SIE) model [4,39,40,41]. The SIE model can be understood as a natural generalization of the singular isothermal sphere (SIS), which assumes spherical symmetry and an isothermal velocity distribution. In the SIS model, three-dimensional mass density ($\rho (\overrightarrow{r})$) follows:

$$\rho (\overrightarrow{r})={\displaystyle \frac{{\sigma }_v^2}{2\pi G{r}^2}},$$

(5)

where ${\sigma }_v^2$ represents the line-of-sight velocity dispersion of the lens galaxy and r is given by $r=\sqrt{{\xi }^2+z^2}$, which represents the spherical radial coordinate. To account for deviations from spherical symmetry, the SIE model introduces ellipticity in the projected mass distribution by replacing the circular radial coordinate in the lens plane with an elliptical one. Following [39], this is achieved through the transformation $\xi \to \sqrt{{\xi }_1^2+f^2{\xi }_2^2}$. Under this prescription, the projected density profile and the corresponding surface mass density ($\Sigma (\overrightarrow{\xi })$) become

$$\begin{array}{c}\hfill \rho (\overrightarrow{r})={\displaystyle \frac{{\sigma }_v^2}{2\pi G{r}^2}},\end{array}$$

(6)

$$\begin{array}{c}\hfill \Sigma (\overrightarrow{\xi })={\displaystyle \frac{{\sigma }_v^2}{2G}}{\displaystyle \frac{\sqrt{f}}{\sqrt{{\xi }_1^2+f^2{\xi }_2^2}}}.\end{array}$$

(7)

In this context, f refers to the axis ratio of the ellipses, specified within the interval $0<f\le 1$. The SIE model profile is fundamentally defined by the line-sight velocity dispersion of the lens galaxy, ${\sigma }_v$. This motivation is key, as ${\sigma }_v$ serves as an observational input to alternatively estimate the galaxy mass. Consequently, the dimensionless surface mass density or convergence $\kappa (\overrightarrow{\theta })$ is directly related to this dispersion. Given the angular coordinates $\overrightarrow{\theta }=({\theta }_1,{\theta }_2)$, the convergence is given by

$$\kappa ({\theta }_1,{\theta }_2)={\displaystyle \frac{{\theta }_E\sqrt{f}}{2\sqrt{{\theta }_1^2+f^2{\theta }_2^2}}}.$$

(8)

The expression delineates the density of the projected surface mass within the lens. Within this formulation, given a perfect alignment source-lens-observer, ${\theta }_E$ denotes the Einstein radius, which is defined by

$${\theta }_E=4\pi {\displaystyle \frac{{\sigma }_v^2}{c^2}}{\displaystyle \frac{D_{\mathrm{LS}}}{D_{\mathrm{S}}}},$$

(9)

where c denotes the speed of light. Using polar coordinates $(r,\varphi )$, where $\varphi$ is the polar angle measured from the semi-major axis (${\widehat{e}}_1$), and taking into account the relationship delineated in $\Delta \mathsf{\Psi }=2\kappa$, where $\mathsf{\Psi }$ represents the lensing potential for the SIE model, which can be expressed as

$$\mathsf{\Psi }(x,\varphi )=x{\displaystyle \frac{\sqrt{f}}{f^{\prime }}}\left[sin\varphi arcsin(f^{\prime }sin\varphi )+cos\varphi \mathrm{arcsinh}\left({\displaystyle \frac{f^{\prime }}{f}}cos\varphi \right)\right],$$

(10)

where $f^{\prime }=\sqrt{1-f^2}$ and characterizes the shape of equipotential surfaces and tends to unit at the circular boundary [1,4,39]. The deflection angle relevant to SIE model can be expressed as follows:

$$\overrightarrow{\alpha }\left(\overrightarrow{x}\right)={\displaystyle \frac{\sqrt{f}}{f^{\prime }}}\left[\mathrm{arcsinh}\left({\displaystyle \frac{f^{\prime }}{f}}cos\varphi \right){\widehat{e}}_1+arcsin(f^{\prime }sin\varphi ){\widehat{e}}_2\right],$$

(11)

where ${\widehat{e}}_1$ and ${\widehat{e}}_2$ denote the unit vectors along the principle axes and the deflection angle is related to the potential by $\overrightarrow{\nabla }\mathsf{\Psi }=\overrightarrow{\alpha }$. In addition to the free parameters of the model, which include the central coordinates and the position angle $\phi$ between the semi-major axis of the mass ellipse of the lens and the axis x of the lens plane, the ellipticity is typically parameterized by the eccentricity moduli

$$\begin{array}{c}\hfill {ϵ}_{l,x}={\displaystyle \frac{1-f}{1+f}}cos2\phi ,\end{array}$$

(12)

$$\begin{array}{c}\hfill {ϵ}_{l,y}={\displaystyle \frac{1-f}{1+f}}sin2\phi ,\end{array}$$

(13)

which are continuously defined according to $-1<{ϵ}_i<1$, thus removing the periodic boundaries associated with $\varphi$ and addressing the discontinuity in $f=0$ as noted in [42].

2.3. Light Model Parametrization

We assume that the lens and source light follow the elliptical Sérsic model distribution, allowing a flexible but physically motivated description of galaxy morphology while keeping the number of free parameters manageable for lens modeling. The Sérsic profile for the surface brightness is defined by the following expression:

$$I(x,y)=I_{e}exp\left[-b_{n_s}\left\{{\left({\displaystyle \frac{\sqrt{x^2+y^2/f^2}}{R_e}}\right)}^{1/n}-1\right\}\right],$$

(14)

where $I_e$ denotes the surface brightness at the half-light radius $R_e$, and $b_{n_s}$ represents a constant related to the Sérsic index $n_s$ by $\mathsf{\Gamma }\left(2n_s\right)=2\gamma (2n_s,b_{n_s})$ [43,44,45,46].

3. Convolutional Neural Networks Theory

Convolutional neural networks, although analogous to traditional feedforward neural networks, are characterized by their ability to extract hierarchical features through convolutional processes, thus reducing computational expenses [47,48]. CNNs have several advantages over conventional neural networks. Firstly, they reduce the number of parameters and enhance convergence by establishing neuron connections with smaller subsets of neurons from the previous layer. Secondly, the method of weight-sharing decreases parameters by distributing weights across various connections. Lastly, CNNs incorporate downsampling, which reduces the input dimensionality, making them suitable for smaller input sizes.

The CNN architecture comprises four main components—convolutional layers, pooling layers, activation functions, and fully connected layers—as depicted in Figure 2. Convolutional layers utilize filters on the input data to identify patterns. The pooling layers perform simplification by downsampling along spatial dimensions, thus reducing parameter loads in the activation process. Their core function is to reduce the spatial dimensions of the activation maps while preserving critical features, which helps to reduce computational demands, control overfitting, and achieve invariant representations of small translations. Activation functions are integral to the expression of complex features, acting similarly to the neuronal model of the human brain by determining which information propagates to the succeeding neuron. The Rectified Linear Unit (ReLU) is prevalently used among activation functions [38,49,50]; it assists the model in discerning non-linear image patterns. If the input x is negative, its function value is 0, while if x is non-negative, its function value equals x itself. The ReLU function is written as

$$ReLU\left(x\right)=\left\{\begin{array}{cc}0& \mathrm{if}x\le 0\\ x& \mathrm{if}x>0\end{array}\right.=max(0,x)$$

(15)

Although the term CNN was assigned in the 1980s, it was Yann LeCun [51] who pioneered the effective application of the backpropagation method for CNNs aimed at handwritten zip code recognition, coining the term convolution in the process. Dropout is a significant machine learning algorithm introduced in the conference on Neural Information Processing Systems (NIPSs) 2012 [52] to train neural networks, which works by randomly dropping out neurons during training to prevent feature detector co-adaptation1. Dropout facilitates the training of numerous network configurations within reasonable timescales [53], alleviates overfitting, and offers an effective approximation to combine exponentially many unique NN structures [37,54]. In considering an NN with hidden layers L, where ${\overrightarrow{z}}^l$ is the input vector for each layer, ${\overrightarrow{y}}^l$ represents the output vector from layer l, and ${\overrightarrow{w}}^l$ signifies the weights in layer l, the dropout feedforward algorithm undergoes redefinition:

$$\begin{array}{cc}\hfill r_j^l& \sim \mathrm{Bernoulli}\left(p\right),\hfill \end{array}$$

(16)

$$\begin{array}{cc}\hfill {\overline{\overrightarrow{y}}}^l& ={\overrightarrow{r}}^l\odot {\overrightarrow{y}}^l,\hfill \end{array}$$

(17)

$$\begin{array}{cc}\hfill z_i^{(l+1)}& ={\overrightarrow{w}}_i^{(l+1)}·{\overline{\overrightarrow{y}}}_i^l+b_i^{(l+1)},\hfill \end{array}$$

(18)

$$\begin{array}{cc}\hfill y_i^{(l+1)}& =f\left(z_i^{(l+1)}\right),\hfill \end{array}$$

(19)

where ${\overrightarrow{r}}^l$ represents a vector of independent Bernoulli random variables, each associated with probability p. The output of that layer is multiplied in elements by ${\overrightarrow{r}}^l$ to produce thinned outputs ${\overline{\overrightarrow{y}}}^l$ [37]. The probability of dropping out depends on the architecture and size of the dataset, and its study is crucial to obtain promising results and reduce overfitting. Normally, the rate at which the dropout is selected is between 20% and 50%. In the subsequent section, we examine three scenarios concerning the influence of dropout on CNN outputs.

4. Methodology

4.1. Simulated Sample from China Space Station Telescope (CSST)

The dataset used to train, validate and test the CNN model consists of 76,396 images, each with dimensions of $100\times 100$ pixels with 1 channel and a resolution of $0.06$ arcseconds/pixel, based on the configuration of strong lens systems presented by CSST [55], corresponding to the single band Wide-Field sample. The images are generated using the Lenstronomy package, in the Python (3.11.10) environment [56], in which the source is modeled by the Sérsic model characterized by the space parameter ${\mathsf{\Theta }}_s=(R_e,{ϵ}_{s,x},{ϵ}_{s,y},x,y)$, where x and y represent the position of the source with respect to the lens plane and setting the Sérsic index $n_s=2.0$. This value of the Sérsic index was chosen because it represents an intermediate value between distinct galaxy populations ($n_s=1.0$ for disk galaxies and $n_s=4.0$ for elliptical galaxies). Furthermore, S. Mukherjee et al. [57] found that the precise choice of source parameters for lens modeling is of secondary importance and does not significantly bias lens modeling. As shown in Equations (12) and (13), the components of the ellipticities were computed knowing the axis ratio from configurations of the CSST sample [55]. For this work, the position angle (p.a.) used to compute the components of the ellipticities for the lens (${ϵ}_{l,x},{ϵ}_{l,y}$) was arbitrarily set at $\mathrm{p}.\mathrm{a}.=10$ deg from North, while the p.a. used to calculate the ellipticities components for the source were extracted from the CSST sample. The lens is characterized by the SIE model for the mass distribution with parameters ${\mathsf{\Theta }}_l=({\theta }_E,{ϵ}_{l,x},{ϵ}_{l,y})$ and the Sérsic model for the brightness distributions used for the source and lens. For the following sections, the notation (${ϵ}_x$,${ϵ}_y$) will be adopted to refer to the ellipticities of the lens. Figure 3 shows the distribution for each parameter used in training issues. The distributions for ${\theta }_E$ and ellipticity components (${ϵ}_x$, ${ϵ}_y$) exhibit a clear positive skew, while axis ratio (f) shows a concentration toward higher values. For ${\theta }_E$, the values are concentrated between [0.07, 3.65]. For f, the values can be found between [0.20, 1.00]. The ellipticity components (${ϵ}_x$, ${ϵ}_y$) can be found between [0.00, 0.63] and [0.00, 0.23], respectively. Figure 4 shows several random examples of synthetic images used in the analysis. To train the model, the dataset was partitioned into 70,000 images for the k-fold cross-validation process2 [58]) and 6396 images for independent testing. Within the 4-fold cross-validation scheme, each fold utilized 75% of the images (52,500) for training and 25% (17,500) for validation.

4.2. CCN

The primary objective is to train an AlexNet-modified CNN (described later) to predict the four characteristic parameters of the SIE model (${\theta }_E$, f, ${ϵ}_x$, ${ϵ}_y$) for lens modeling, assessing the performance of the CNN model under three conditions in which the dropout layers were altered. As mentioned previously, the dataset was partitioned into 70,000 images for 4-fold cross-validation for training and validation and 6396 images for independent testing. In the 4-fold cross-validation scheme, each fold used 75% of the images for training and 25% for validation. A crucial element of the dataset is the binary format used to store images, namely TFRecords, which efficiently stores data sequences, optimized for use with TensorFlow [59]. This format is ideal for large datasets, as it allows for more compact storage and faster data reading during model training, even if the data do not fit completely in memory. This format facilitates the systematic labeling and organization of images, allowing a single file to be partitioned into training, validation, and test samples. Furthermore, the format’s support for grayscale images streamlines the preprocessing procedures.

The CNN model used in this research is based on the AlexNet architecture, which incorporates modifications in both the number of layers and neurons. This architecture, introduced by Alex Krizhevskii et al. in 2012 [38,60,61], achieved significant success in the ImageNet Large Scale Visual Recognition Challenge of 2012, an effort that evaluates algorithmic performance in object detection and image classification. AlexNet achieved an error rate of 15.3% in the classification of large-scale images. AlexNet expands the foundational concepts of LeNet and applies the core principles of CNNs to a more extensive and deep network. Furthermore, AlexNet integrates the ReLU activation function, dropout, and local response normalization (LRN) for the first time within a CNN. For the purposes of this preliminary study, this model was selected among various available models due to its notable performance as evidenced by works such as those of Perreault et al. [62], which demonstrate its ability to accurately estimate lens parameters.

The architecture of the AlexNet-based CNN model is shown in Figure 5. The model comprises six blocks preceding the fully connected layers. The initial blocks consist of a convolutional layer, followed by batch normalization and maximum-pooling layers. The first convolutional layer applies 128 filters of size $7\times 7$ with a stride of 2, thereby reducing the spatial resolution. The second convolutional layer utilizes 256 filters of size $5\times 5$ with a stride of 3, facilitating the extraction of mid-scale features. The deep convolutional block, consisting of four convolutional layers with 256 filters of size $1\times 1$, captures hierarchical representations and enhances the depth of the network without considerably increasing the number of parameters. This block concludes with a Max-Pooling layer to further reduce spatial dimensionality. The flattened feature maps are processed by two dense layers consisting of 512 and 1024 neurons, respectively. The investigation places significant emphasis on the dropout layers, examining changes in dropout rates across three scenarios. Ultimately, the model concludes with a dense layer possessed of linear activation and yielding four outputs corresponding to the parameters of interest. The primordial improvement that presents the model is the batch normalization layers, which stabilize training, allow higher learning rates, and reduce sensitivity to initialization. Another advantage of the architecture used in this work is the inclusion of an additional convolutional layer. This layer enhances the model’s ability to extract more details from the images. Additionally, the architecture undergoes a significant change with the implementation of filter factorization within the $1\times 1$ convolutional layer block. This provides a better representation of the characteristics while keeping the computational cost. Because the ellipticity components (${ϵ}_x$, ${ϵ}_y$) produce weaker variations in the images than other parameters of the lens model, we observed that the network tended to underadjust them. To mitigate this effect, we use a weighting scheme in the loss function defined by Equation (20), assigning higher weights to the terms associated with ${ϵ}_x$ and ${ϵ}_y$. The weights for the Einstein radius and the axis ratio were 1.0 and for the ellipticity components 3.0, respectively. This strategy increases the contribution of these parameters to the gradient during training, stimulating the model to learn more sensitive representations of its morphological signatures.

The NAdam optimizer, an adaptation of the Adam optimizer that incorporates Nesterov momentum [63], was used for weight refinement. Nesterov’s Accelerated Gradient (NAG) is a first-order optimization technique that ensures a higher convergence rate in certain contexts compared to the traditional gradient descent method. This momentum distinguishes itself from the Stochastic Gradient Descent (SGD) optimizer by first taking a substantial step in the direction of the accumulated gradient, followed by the evaluation of the gradient’s terminal position to perform necessary adjustments. NAG achieves a global convergence rate of $\mathcal{O}(1/T^2)$, marking an advance over the convergence rate of gradient descent $\mathcal{O}(1/T)$ [64]. This refinement not only accelerates the convergence process but also enhances precision in minimizing the loss function. To mitigate overfitting, a ReduceLROnPlateau callback was integrated into the model, systematically reducing the learning rate and thus improving the training metrics. This technique allows a dynamic adaptation of the learning rate: if the loss of validation does not show a significant improvement after a threshold of n times (patience), the learning rate is reduced by a factor $\lambda$ [65,66]. In this work, the factor for ReduceLROnPlateau was fixed in $\lambda =0.5$. The chosen minimum learning rate was $ϵ={10}^{-8}$. A significant experimental analysis involved modifying the dropout rates of two dropout layers, which were systematically tested under three conditions: initially using dropout rates of 20% and 30%, then uniformly applying a dropout rate of 20% to both layers, and finally, disabling the dropout layers entirely. For further evaluation of the model, a test sample consisting of 6396 images was selected, with the relative errors related to each prediction visualized.

4.3. Evaluation Metrics

To assess and compare the performance of the model, the weighted Mean Squared Error (MSE) was selected as the loss function, setting the weights for the Einstein radius and axial ratio in 1.0 and the weights for the ellipticity components of the lens in 3.0, following the Multi-Task Learning approach proposed by Caruana (1997) [67], while the Mean Absolute Error (MAE) served as the evaluation metric [68,69]. Furthermore, the coefficient of determination $R^2$ was calculated, providing an indication of the extent to which the variance in the observed values is explicable by the predictions of the model, expecting a value near 1 [70]. These metrics are articulated in the following expressions:

$$\begin{array}{cc}\hfill \mathrm{MSE}=& {\displaystyle \frac{1}{n}}\sum _{i=1}^n\sum _{j=1}^k{w}_j{(y_{i,j}-{\overline{y}}_{i,j})}^2,\hfill \end{array}$$

(20)

$$\begin{array}{cc}\hfill \mathrm{MAE}=& {\displaystyle \frac{1}{n}}\sum _{i=1}^n|y_i-{\overline{y}}_i|,\hfill \end{array}$$

(21)

$$\begin{array}{cc}\hfill R^2=& 1-{\displaystyle \frac{{\sum }_{i=1}^n{(y_i-{\overline{y}}_i)}^2}{{\sum }_{i=1}^n{(y_i-\widehat{y})}^2}},\hfill \end{array}$$

(22)

in which $y_i$ denotes the actual value, ${\overline{y}}_i$ represents the predicted value, and $\widehat{y}$ signifies the mean of the actual values. Additionally, we assess the relative error inherent in the predictions by determining the residuals between the forecasted and actual values and subsequently plotting the actual values against the forecasted values [71]. The relative error was evaluated employing the elementary definition of relative error

$${ϵ}_r=\left|{\displaystyle \frac{\overline{y}-y}{\overline{y}}}\right|.$$

(23)

In order to evaluate the accuracy of the predictions the models on the test images, we computed the bias defined as

$$\mu =\mathrm{median}(\overrightarrow{\overline{y}}-\overrightarrow{y}),$$

(24)

where $\overrightarrow{\overline{y}}$ corresponds to the predicted values array and $\overrightarrow{y}$ to the true values array. The scatter of the distribution can be obtained by computing using the 16th and the 84th percentiles of the distribution by the next definitions:

$$\begin{array}{cc}\hfill {\sigma }^{-}& =\mathrm{median}(\overrightarrow{\overline{y}}-\overrightarrow{y})-P_{16}(\overrightarrow{\overline{y}}-\overrightarrow{y}),\hfill \end{array}$$

(25)

$$\begin{array}{cc}\hfill {\sigma }^{+}& =P_{84}(\overrightarrow{\overline{y}}-\overrightarrow{y})-\mathrm{median}(\overrightarrow{\overline{y}}-\overrightarrow{y}).\hfill \end{array}$$

(26)

Furthermore, we determined the Normalized Median Absolute Deviation (NMAD), which is a robust statistic used to measure data dispersion, defined as the Median Absolute Deviation (MAD) scaled by a factor of approximately $k=1.48$ [72]. NMAD is defined as

$$\mathrm{NMAD}=1.48\mathrm{median}\left[\right|(\overrightarrow{\overline{y}}-\overrightarrow{y})-\mu \left|\right].$$

(27)

A critical examination to reinforce the findings involves simulating images with the predicted outputs of the CNN model in order to facilitate a comparative analysis using the PSNR and MSE between the images, allowing an objective assessment of the quality of the results. Taking into account a reference image X and a test image Y of size $M\times N$, the PSNR, measured in dB, is determined by ${\mathrm{MSE}}_{\mathrm{img}}$ between them using the following expression:

$$\mathrm{PSNR}(X,Y)=10{log}_{10}({\mathrm{MAX}}_{\mathrm{img}}^2/{\mathrm{MSE}}_{\mathrm{img}}(X,Y)),$$

(28)

where

$${\mathrm{MSE}}_{\mathrm{img}}={\displaystyle \frac{1}{MN}}\sum _{i=1}^M\sum _{j=1}^N{(X_{ij}-Y_{ij})}^2,$$

(29)

and ${\mathrm{MAX}}_{\mathrm{img}}$ is the maximum value that a pixel can have in the image. A good value of PSNR is in a range between 30 and 40 dB. A PSNR value above 40 dB means a high-quality reconstruction of the original image [73,74].

5. Results

This section is devoted to analyzing the performance of the CNN model in predicting four parameters of the SIE model (lens galaxy). We note that the reported results correspond to the average values and their standard deviation obtained from a cross-validation procedure employing a 4-fold partitioning scheme. Specifically, the dataset was divided into four independent subsets. Figure 6 illustrates the evolution of loss and MAE functions with respect to the epochs, for which Models 1 and 2 exhibit consistent behavior. Although both functions of Model 3 display greater instability during the initial training phase, they eventually converge. Nevertheless, based on the minimum values attained by these metrics, as shown in

Table 1. Results for $R^2$ and median ${ϵ}_r$ of each parameter for each model. The effect of dropout on generalization is evident. Models 1 and 2 show consistent performance in terms of $R^2$, while the absence of dropout in Model 3 leads to overfitting and reduced generalization capability. Each relative error includes its standard deviation according to the 4-fold-cross procedure.

Model	Metrics		${\mathit{\theta }}_{\mathit{E}}$		f		${\mathit{ϵ}}_{\mathit{x}}$		${\mathit{ϵ}}_{\mathit{y}}$
	Loss	MAE	${\mathit{R}}^{\mathbf{2}}$	${\mathit{ϵ}}_{\mathit{r}}$ (%)	${\mathit{R}}^{\mathbf{2}}$	${\mathit{ϵ}}_{\mathit{r}}$ (%)	${\mathit{R}}^{\mathbf{2}}$	${\mathit{ϵ}}_{\mathit{r}}$ (%)	${\mathit{R}}^{\mathbf{2}}$	${\mathit{ϵ}}_{\mathit{r}}$ (%)
1	$0.008\pm 0.001$	$0.026\pm 0.001$	$0.96\pm 0.01$	$3.{65}_{-2.44}^{+3.31}$	$0.95\pm 0.01$	$2.{56}_{-1.75}^{+2.52}$	$0.96\pm 0.01$	$4.{92}_{-3.40}^{+5.46}$	$0.96\pm 0.01$	$5.{12}_{-3.58}^{+5.49}$
2	$0.006\pm 0.001$	$0.022\pm 0.001$	$0.96\pm 0.01$	$4.{00}_{-2.80}^{+3.68}$	$0.95\pm 0.01$	$2.{44}_{-1.72}^{+2.68}$	$0.97\pm 0.01$	$4.{86}_{-3.43}^{+5.42}$	$0.96\pm 0.01$	$5.{11}_{-3.59}^{+5.23}$
3	$0.092\pm 0.016$	$0.085\pm 0.005$	$0.91\pm 0.05$	$9.{32}_{-6.68}^{+11.54}$	$0.61\pm 0.05$	$7.{38}_{-5.21}^{+9.26}$	$0.56\pm 0.60$	$21.{01}_{-14.75}^{+27.84}$	$0.64\pm 0.15$	$17.{90}_{-12.40}^{+20.96}$

, Model 3 achieves the highest (i.e., worst) loss and MAE values among the three CNN models.

In the following, the results presented are obtained using the testing sample corresponding to 6396 images. Additionally, $R^2$ indicates that models that incorporate dropout exhibit superior predictive performance for SIE parameters (see Table 1), achieving values in the range of 0.95 to 0.97, while Model 3 achieves values comparatively lower (>0.56). Furthermore, Figure 7 shows the distribution of the PSNR for the three models, in which the median value is denoted by a dotted red line. Based on this quantity, we find a good reconstruction of the images of the lensed source for Model 1 and Model 2 compared to Model 3, obtaining similar results (median values of ≈37) with dropout being considered. Figure 8 displays a random example of a galaxy–galaxy lens system reconstructed using the true values (left panel) and the predicted values (middle panel). The right panel shows the residual distribution obtained as the difference between the true image and the predicted. As the intensity of the residual images is very small, we use a Lupton-type asinh transformation [75] that highlight the details of bright and weak areas by transforming weak signals linearly and strong signal logarithmically. We observe that the worst predicted reconstruction occurs when Model 3 is used.

Figure 9 presents the scatter distribution between the true and predicted values of the SIE parameters. Exact predictions lie on the 45-degree diagonal (solid red line), and ideal predictions correspond to a distribution tightly aligned with this line. As anticipated, Models 1 and 2 show broadly consistent results: they exhibit a higher predictive accuracy for low values of ${\theta }_E$, while a systematic bias becomes apparent for ${\theta }_E\gtrsim 1.2$. This can be associated with the values of ${\theta }_E$ utilized for training, as depicted in Figure 3. This figure illustrates a significant concentration of values within the range of $(0,1.2)$. Similarly, the parameter f displays a noticeable bias for $f\gtrsim 0.8$, where the models systematically underestimate the true values. For components ${\epsilon }_x$ and ${\epsilon }_y$, the bias is even more pronounced in the regimes ${\epsilon }_x\gtrsim 0.4$ and ${\epsilon }_y\gtrsim 0.12$, respectively. In contrast, Model 3 shows the largest dispersion and bias among the three, clearly revealing the significant impact of dropout on the quality of the predictions. To assess the uncertainties of the models, we also computed the bias, the standard deviation and the NMAD, which serve as indicators of the accuracy of the predictions. The results of these metrics reveal that Models 1 and 2 exhibit a negligible systematic bias ($\mu \approx -0.02$) but extremely low NMAD values between 0.01 and 0.04, confirming that predictions are highly certain and present minimal dispersion with respect to the line of truth. In contrast, while Model 3 maintains a low average bias ($\mu \approx 0.00$), its high residual dispersion (NMAD ≈ 0.07–0.10) indicates a regime of high variance that limits its reliability for individual parameter inference.

Finally, the relative errors of the SIE parameters for the three CNN architectures are presented in Figure 10. The values shown at the top of each boxplot correspond to the median relative error of the associated parameter. The best predictive performance is obtained for the parameter f, with median relative errors as low as $2.6\%$ for the CNN models that incorporate dropout layers; this median error increases by approximately a factor of three when dropout is disabled. The poorest performance is observed for the ellipticity parameters, which exhibit median relative errors up to $5.12\%$ for Models 1 and 2, and up to $21\%$ for Model 3. These median values are also summarized in Table 1.

6. Discussion and Conclusions

The AlexNet-based CNN architecture was implemented to predict the parameters of the SIE lens model from 76,396 simulated galaxy–galaxy lens systems from the CSST sample [55]. The study compared three model configurations: Model 1 added dropout rates of 20% and 30% for the first and second dense NN layers; Model 2 used dropout rates of 20% in both layers; and Model 3 turned off dropout layers to assess their impact on generalization and predictive capabilities. Additionally, the reported uncertainties of the results account for a cross-validation procedure implemented using a 4-fold partitioning scheme.

First, although Model 3 did not exhibit overfitting, both MSE and MAE error metrics converged to local minima, with values of 0.092 and 0.085, respectively. These are considerably higher than the corresponding minimum values obtained for Models 1 and 2, which are approximately 15 times and 4 times lower for MSE and MAE, respectively. This limitation also affected other performance indicators, such as $R^2$ and PSNR, suggesting that models that incorporate dropout produce higher-quality predictions. We consistently observed that error prediction is substantially improved, i.e., significantly reduced in magnitude, as reported in Table 1. Additionally, we note that Model 1 differs from Model 2 by a 50% increase in the dropout rate in one of the dense layers. We observe a similar behavior during the training phase and a comparable impact on the final results, which supports the robustness of the proposed technique, and achieving $R^2$ values greater than $0.95$ for the predicted parameters, in contrast to $R^2$ values in the range 0.56–0.91 obtained with Model 3. Regarding image reconstruction performance, Figure 7 reports the PSNR distributions, with median values of 36.9 dB, 36.8 dB, and 29.2 dB for Models 1, 2, and 3, respectively. When dropout is incorporated into the models, these values place the reconstructed images in the lower range of the high-quality reconstruction regime. The residual dispersion analysis exposed in Figure 9 shows that Models 1 and 2 present negligible systematic bias ($\mu \approx -0.02$) and low NMAD values (0.01–0.04), indicating highly certain and minimal dispersion with a remarkable stability across all morphological parameters. In contrast, while Model 3 has a low average bias ($\mu \approx 0.00$), its high residual dispersion (NMAD 0.07–0.10) limits its reliability for individual parameter inference, indicating a structural inability to capture the data variance.

All three models demonstrated proficiency in predicting the Einstein radius ${\theta }_E$ and the axial ratio f, both of which consistently exhibited the lowest levels of relative errors (Figure 10). Models 1 and 2 achieved relative errors with an upper limit around 5–9% in the $90\%$ CL for these two parameters. The parameters related to ellipticity (${ϵ}_x$, ${ϵ}_y$) proved more difficulty to predict, presenting larger and more dispersed relative errors. This difficulty can be attributed to the degeneracy of the lens equation, where different parameter combinations yield similar images, and the potential confusion of ellipticity components or source asymmetries. Despite these challenges, the overall performance was remarkable, with relative errors around 5–12% at the $90\%$ CL for most SIE parameters. Furthermore, the model struggled to correctly learn the high values for ${ϵ}_y$, suggesting a scarcity of training samples with high values or the complexity of interpreting the ellipticity.

The implementation of dropout to enhance gravitational lens studies aligns with previous work, such as Perreault Levasseur et al. [62], who optimized this hyperparameter to obtain expected coverage probabilities. Consistent with their findings, this work observed that networks with optimized dropout (Model 2) provided accurate predictions, while the absence of dropout (Model 3) resulted in higher and less reliable predictions. Similarly, Morningstar et al. [76] demonstrated that changing dropout rates affects predictions, achieving good calibration with a dropout rate $1\%$ in a recurrent model. More recently, the LEMON project by Gentile et al. [77] and the Euclid Collaboration [35] used a Bayesian NN with a dropout rate $3\%$ (keep rate $97\%$) to model lenses for HST and Euclid. Their results, achieving $R^2\sim 0.94$ and MAE $\sim 0.20$, are comparable to the results obtained in this work ($R^2$ up to 0.97). However, we would remark that these works are performed with different samples and include noise. The relevance of the results obtained in Model 2 lies in its balance between computational efficiency and analytical accuracy. While previous lens modeling work using CNN used dense architectures that require significant resources, our use of an optimized version of AlexNet shows that it is possible to achieve a determination coefficient of $R^2=0.97$ with a reduced computational load. Since Einstein’s radius (${\theta }_E$) is the main indicator of the projected mass within the critical radius, reaching an uncertainty of just 9% (at 90% CL) for this parameter has a direct impact on the estimation of the mass enclosed within the ${\theta }_E$. This precision is vital for cosmology studies to restrict dark matter profiles [78,79]. Therefore, the model not only delivers a quick prediction but ensures that the derived scientific parameters maintain the integrity necessary for large-scale analysis in missions such as CSST and Euclid.

Finally, this study confirms that CNNs are promising tools for reducing prediction errors in gravitational lensing modeling. Regularization through dropout is crucial to reduce the variability of these errors and prevent overfitting. The predictions of Models 1 and 2 offer high-quality image reconstruction with a minimal MSE around $2.40\times {10}^{-2}$. Furthermore, this method is significantly faster than traditional Markov Chain Monte Carlo (MCMC) techniques, efficiently predicting SIE parameters using a single GPU. This efficiency facilitates future analyses of large datasets from modern telescopes such as the Chinese Southern Sky Telescope (CSST) and the Euclid Space Telescope [19,55]. Future research should investigate alternative network architectures—such as ResNet [80], ConvNeXt, [81] and U-Net [82]—or suitably modified variants thereof, and should incorporate more advanced data augmentation strategies. This would enable a more rigorous characterization of distortions induced by convergence and shear components, particularly in scenarios where the presence of background noise and the finite spatial resolution of the detector could produce image degradation that impacts the model estimation and significantly complicates their analysis.