Original: Deconvolution of Astronomical Images with Deep Neural Networks
Revised: Deconvolution of Astronomical Images with Deep Neural Networks

Original: Anonymous Author(s)

Original: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: Optical astronomical images are strongly affected by the point spread function(PSF) of the optical system and the atmosphere (seeing) which blurs the observedimage.
Revised: Optical astronomical images are strongly affected by the point spread function (PSF) of the optical system and the atmosphere (seeing) which blurs the observed image.

Original: The amount of blurring depends on both the observed band, and morecrucially, on the atmospheric conditions during observation.
Revised: The amount of blurring depends on both the observed band, and more crucially, on the atmospheric conditions during observation.

Original: A typical astronomicalimage will therefore have a unique PSF that is non-circular and different in differentbands.
Revised: A typical astronomical image will therefore have a unique PSF that is non-circular and different in different bands.

Original: Observations of known stars give us an estimation of this PSF.
Revised: Observations of known stars give us an estimation of this PSF.

Original: Any seriouscandidate for production analysis of astronomical images must take the known PSF into account during image analysis.
Revised: Any serious candidate for production analysis of astronomical images must take the known PSF into account during image analysis.

Original: So far the majority of applications ofneural networks (NN) to astronomical image analysis have ignored this problemby assuming a fixed PSF in training and validation.
Revised: So far the majority of applications of neural networks (NN) to astronomical image analysis have ignored this problem by assuming a fixed PSF in training and validation.

Original: We present a neural networkarchitecture based on Deep Wiener Deconvolution Network (DWDN) that takesthe PSFshape into account when performing deconvolution, a possible approachof leveraging PSF information in neural networks.
Revised: We present a neural network architecture based on Deep Wiener Deconvolution Network (DWDN) that takes the PSF into account when performing deconvolution, a possible approach of leveraging PSF information in neural networks.

Original: We study the performance ofthis algorithm under realistic observational conditions.
Revised: We study the performance of this algorithm under realistic observational conditions.

Original: We employ two regularization schemes and study custom loss functions that are optimized for quantitiesof interest to astronomers.
Revised: We employ two regularization schemes and study custom loss functions that are optimized for quantities of interest to astronomers.

Original: We show that our algorithm can successfully recoverunbiased image properties such as colors, ellipticities and orientations for sufficiently high signal-to-noise.
Revised: We show that our algorithm can successfully recover unbiased image properties such as colors, ellipticities and orientations for sufficiently high signal- to-noise.

Original: This study represents a comprehensive application of AI in astronomy, where the experimental design, model construction, optimizationcriteria, error estimation and metrics of benchmarks are all meticulously tailored tothe domain problem.
Revised: This study represents a comprehensive application of AI in astronomy, where the experimental design, model construction, optimization criteria, error estimation and metrics of benchmarks are all meticulously tailored to the domain problem.

Original:  Introduction
Revised:  Introduction

Original: The advent of large telescopes and big data sets are bringing about a transformable era for astronomicalsurvey science.
Revised: The advent of large telescopes and big data sets are bringing about a transformable era for astronomical survey science.

Original: New datasets obtained from upcoming ground-based and space-based observingfacilities will contain data from large volumes of the observable sky at unprecedented depths andcadences [10, 1, 7, 13, 6].
Revised: New datasets obtained from upcoming ground-based and space-based observing facilities will contain data from large volumes of the observable sky at unprecedented depths and cadences [11, 1, 7, 14, 6].

Original: While the resolution of space-based telescopes is usually diffractionlimited, the large apertures of ground based survey telescopes are limited by seeing .
Revised: While the resolution of space-based telescopes is usually diffraction limited, the large apertures of ground based survey telescopes are limited by seeing .

Original: Seeing is adistortion of the image caused by the perturbation of an optical wavefront as it passes through theturbulent atmosphere.
Revised: Seeing is a distortion of the image caused by the perturbation of an optical wavefront as it passes through the turbulent atmosphere.

Original: This results in a finite point-spread function (PSF).
Revised: This results in a finite point-spread function (PSF).

Original: The ground based surveysoften do not employ adaptive optics due to the large field of view, resulting in an arcsecond sizedpoint spread function.
Revised: The ground based surveys often do not employ adaptive optics due to the large field of view, resulting in an arcsecond sized point spread function.

Original: The two main features of the PSF is that it is a variable and that it is known in each exposure.
Revised: The two main features of the PSF is that it is a variable and that it is known in each exposure.

Original: The
Revised: The

Original: PSF is a variable because the atmosphere above the telescope fluctuates.
Revised: PSF is a variable because the atmosphere above the telescope fluctuates.

Original: It varies in size, shape andorientation.
Revised: It varies in size, shape and orientation.

Original: It is typically asymmetric and different in different bands.
Revised: It is typically asymmetric and different in different bands.

Original: On the other hand, it is also known in each exposure, because invariably every exposure contains stars that are excellent pointsources and therefore can directly be used to probe the PSF.
Revised: On the other hand, it is also known in each exposure, because invariably every exposure contains stars that are excellent point sources and therefore can directly be used to probe the PSF.

Original: Recovery of an object’s shape, position, flux and other intrinsic parameters in the presence of noise andvarying PSF is goal of astronomical image analysis.
Revised: Recovery of an object’s shape, position, flux and other intrinsic parameters in the presence of noise and varying PSF is goal of astronomical image analysis.

Original: Machine learning techniques using convolutionaloperations have demonstrated promising avenues for speeding up the astronomical processing.
Revised: Machine learning techniques using convolutional operations have demonstrated promising avenues for speeding up the astronomical processing.

Original: Majority of these grapple with the galaxy deblending problem such as [14, 3, 2, 16].However, mostof the current neural-network based galaxy deblenders have a number of simplifying assumptions,including a constant PSF.
Revised: Majority of these grapple with the galaxy deblending problem such as [15, 3, 2, 18].However, most of the current neural-network based galaxy deblenders have a number of simplifying assumptions, including a constant PSF.

Original: This problem is usually “swept under the rug” by assuming that severaltraining sets can be provided, one for each PSF.
Revised: This problem is usually “swept under the rug” by assuming that several training sets can be provided, one for each PSF.

Original: However, this is likely impractical for reasons wewill cover in the following sections.
Revised: However, this is likely impractical for reasons we will cover in the following sections.

Original: It is also clear that simply ignoring this problem will lead tobiased and sub-optimally inferred parameters for individual objects.
Revised: It is also clear that simply ignoring this problem will lead to biased and sub-optimally inferred parameters for individual objects.

Original: Therefore the proper treatmentof PSF is one of the issues that need to be addressed before neural network based approaches can belifted from the realm of toy problems and benchmarking datasets to real observations from telescopes.
Revised: Therefore the proper treatment of PSF is one of the issues that need to be addressed before neural network based approaches can be lifted from the realm of toy problems and benchmarking datasets to real observations from telescopes.

Original: Other issues include blending, masking and modeling various detector artifacts.
Revised: Other issues include blending, masking and modeling various detector artifacts.

Original: In this work wespecifically address the issue of leveraging PSF information and build a network that deconvolves anoisy, convolved input image to produce the true image of the galaxy at the resolution supposed bythe PSF.
Revised: In this work we specifically address the issue of leveraging PSF information and build a network that deconvolves a noisy, convolved input image to produce the true image of the galaxy at the resolution supposed by the PSF.

Original:  Description of the problem
Revised:  Description of the problem

Original: We use a linear model for the observed astronomical image, a convolution with a point spread functionfollowed by addition of observational noise:
Revised: We use a linear model for the observed astronomical image, a convolution with a point spread function followed by addition of observational noise:

Original: [Equation]
Revised: [Equation]

Original: where PSF i is the PSF for the channel i , N i is the noise, assumed to be homogeneous and normallydistributed and index gt refers to ground truth.
Revised: where PSF i is the PSF for the channel i , N i is the noise, assumed to be homogeneous and normally distributed and index gt refers to ground truth.

Original: It is important to note that PSF is band dependentand that the amplitude of noise is also band-dependent.
Revised: It is important to note that PSF is band dependent and that the amplitude of noise is also band-dependent.

Original: This model is an approximation, because thebrighter-fatter effects, amplifier non-linearities, and other detector phenomena make the transformation weakly non-linear.
Revised: This model is an approximation, because the brighter-fatter effects, amplifier non-linearities, and other detector phenomena make the transformation weakly non-linear.

Original: The shape of PSF also varies across the focal plane.
Revised: The shape of PSF also varies across the focal plane.

Original: However, in this work weassumed that the region of interest is always small enough that the the PSF can be assumed constant.
Revised: However, in this work we assumed that the region of interest is always small enough that the the PSF can be assumed constant.

Original: One of the classical regularization techniques of deconvolution is the Wiener Deconvolution formulated in Fourier space as (dropping the band subscripts for clarity):
Revised: One of the classical regularization techniques of deconvolution is the Wiener Deconvolution formulated in Fourier space as (dropping the band subscripts for clarity):

Original: [Equation]
Revised: [Equation]

Original: where P s ( k ) and P n ( k ) are spectral densities of signal and noise respectively.
Revised: where P s ( k ) and P n ( k ) are spectral densities of signal and noise respectively.

Original: A neural networkextension of the above formulation is the Deep Wiener Deconvolution Network [5].
Revised: A neural network extension of the above formulation is the Deep Wiener Deconvolution Network [5].

Original:  Band-wise Deep Wiener Deconvolution Network
Revised:  Band-wise Deep Wiener Deconvolution Network

Revised: Works using neural networks for deblurring have emerged recently [9, 17], among which the Deep

Original: The Deep Wiener Deconvolution Network (DWDN, [5]) provides a brand new solution for thenon-blind image deblurring problems.
Revised: Wiener Deconvolution Network (DWDN, [5]) provides a brand new solution for the non-blind image deblurring problems.

Original: Instead of deconvolution in the image space like the existingmethods do, DWDN applies Wiener Deconvolution explicitly in the feature space.
Revised: Instead of deconvolution in the image space like the existing methods do, DWDN applies Wiener Deconvolution explicitly in the feature space.

Original: As a simplebut effective integration of classical Wiener Deconvolution and deep learning, DWDN achievesoutstanding performance with fewer artifacts.
Revised: As a simple but effective integration of classical Wiener Deconvolution and deep learning, DWDN achieves outstanding performance with fewer artifacts.

Original: DWDN contains two components, the feature-based Wiener deconvolution module and a multi-scale feature refinement module, where the former carriesout the deconvolution process and the latter restore high-quality images using features from theprevious module.
Revised: DWDN contains two components, the feature-based Wiener deconvolution module and a multi-scale feature refinement module, where the former carries out the deconvolution process and the latter restore high-quality images using features from the previous module.

Original: A restriction that limits the direct use of DWDN on astronomical images is that typically, for galaxyimages, different bands have different PSFs.
Revised: A restriction that limits the direct use of DWDN on astronomical images is that typically, for galaxy images, different bands have different PSFs.

Original: However, in DWDN, the feature maps extracted fromthe input have already mixed up and no longer have distinguishable bands as the image does.
Revised: However, in DWDN, the feature maps extracted from the input have already mixed up and no longer have distinguishable bands as the image does.

Original: Thus,we propose an enhanced version of DWDN that can work on band-dependent PSFs.
Revised: Thus, we propose an enhanced version of DWDN that can work on band-dependent PSFs.

Original: Figure 1 shows the architecture of our model.
Revised: Figure 1 shows the architecture of our model.

Original: In order to apply feature-based deconvolution on different input bandsseparately, we parallel three deconvolution modules.
Revised: In order to apply feature-based deconvolution on different input bands separately, we parallel three deconvolution modules.

Original: This number varies based on the number ofbands in the input image and PSF.
Revised: This number varies based on the number of bands in the input image and PSF.

Original: Each band of the input image with the corresponding PSF bandwill go through one of those modules to complete the feature-based deconvolution.
Revised: Each band of the input image with the corresponding PSF band will go through one of those modules to complete the feature-based deconvolution.

Original: Eq.
Revised: Eq.

Original: 3 expressesthe deconvolution for l -th feature channel in i -th deconvolution module
Revised: 3 expresses the deconvolution for l -th feature channel in i -th deconvolution module

Original: [Equation]
Revised: [Equation]

Original: where h li is the deconvolved feature of the l th channel, and L is the total number of feature channels.
Revised: where h li is the deconvolved feature of the l th channel, and L is the total number of feature channels.

Original: F refers to discrete Fourier transformation and PSF
Revised: F refers to discrete Fourier transformation and PSF

Original: i ( k ) is the complex conjugate of PSF i ( k ) .
Revised: i ( k ) is the complex conjugate of PSF i ( k ) .

Original: s ni,land s xi,l are the variance and standard deviation of this channel for noise and input image separatelywhich can be estimated by blurred features and mean-filtered features.
Revised: s ni,l and s xi,l are the variance and standard deviation of this channel for noise and input image separately which can be estimated by blurred features and mean-filtered features.

Original: After concatenation as Eq. 4,we get output from i -th deconvolution module H i , that is, deconvolved features for i -th band of theinput image and PSF.
Revised: After concatenation as Eq. 4, we get output from i -th deconvolution module H i , that is, deconvolved features for i -th band of the input image and PSF.

Original: H i will then be passed to the second module  in Eq.
Revised: H i will then be passed to the second module R as in Eq.

Original: 5 with C denoting thenumber of deconvolution modules.
Revised: 5 with C denoting the number of deconvolution modules.

Original: [Equation]
Revised: [Equation]

Original: We keep the refining module R the same except the number of channels of its input feature is Ctimes as original DWDN because of the concatenation.
Revised: We keep the refining module R the same except the number of channels of its input feature is C times as original DWDN because of the concatenation.

Original: In this way, different bands of the input imagewill execute their own deep Wiener deconvolution and the deconvolved features are utilized to restorethe clean images.
Revised: In this way, different bands of the input image will execute their own deep Wiener deconvolution and the deconvolved features are utilized to restore the clean images.

Original: More details of the architecture can be found in [5].
Revised: More details of the architecture can be found in [5].

Original: feature refinement module contains convolution-based encoder-decoder network for multiple levels.

Original: The kernel size is 5 × 5 for all convolutional layers.

Original:  Custom loss function
Revised:  Custom loss function

Original: In majority of image analysis problems, the networks are trained with one of the standard loss- functions that encode some intuitive notion of similarity between the true and reconstructed image.
Revised: In majority of image analysis problems, the networks are trained with one of the standard lossfunctions that encode some intuitive notion of similarity between the true and reconstructed image.

Original: Typical example includes ℓ 1 , ℓ 2 loss functions and the cross-entropy functions.
Revised: Typical example includes ℓ 1 , ℓ 2 loss functions and the cross-entropy functions.

Original: However, in astronomical image analysis we are interested in particular properties of the output image.
Revised: However, in astronomical image analysis we are interested in particular properties of the output image.

Original: In fact the first few image moments described in the next section naturally describe the astronomical quantities of interest: object brightness, its position and its shape.
Revised: In fact the first few image moments described in the next section naturally describe the astronomical quantities of interest: object brightness, its position and its shape.

Original: It is therefore natural to include them into the loss function minimized during training – this way we ask neural network to focus explicitly on the quantities of interest.
Revised: It is therefore natural to include them into the loss function minimized during training – this way we ask neural network to focus explicitly on the quantities of interest.

Original: We will describe this non-standard training procedure next.
Revised: We will describe this non-standard training procedure next.

Original: 4.1 Quantities of interest
Revised: 4.1 Quantities of interest

Original: Let I ij be flux value in pixel ( i, j ) .
Revised: Let I ij be flux value in pixel ( i, j ) .

Original: We define raw moments as
Revised: We define raw moments as

Original: [Equation]
Revised: [Equation]

Original: The following quantities are of interest:
Revised: The following quantities are of interest:

Original: [Equation]
Revised: [Equation]

Original: Second central

Original: [Equation]

Original: These quantities can be trivially calculated from a noise-less image in a way that enable their use inloss function.
Revised: These quantities can be trivially calculated from a noise-less image in a way that enable their use in loss function.

Original: 4.2 Training procedure and loss function
Revised: 4.2 Training procedure and loss function

Original: The model is trained in an end-to-end manner with two phases.
Revised: The model is trained in an end-to-end manner with two phases.

Original: In the first phase, the model ispre-trained to acquire the ability to deblur and produce reasonable prediction as output, while in thesecond phase, the model is fine-tuned with a more complex loss including quantities listed in Section4.1 for more accurate description of the galaxy.
Revised: In the first phase, the model is pre-trained to acquire the ability to deblur and produce reasonable prediction as output, while in the second phase, the model is fine-tuned with a more complex loss including quantities listed in Section 4.1 for more accurate description of the galaxy.

Original: The loss function for the first phase is formulated in Eq.
Revised: The loss function for the first phase is formulated in Eq.

Original: It contains ℓ 1 -norm between the outputand the ground truth as shown in Eq.
Revised: It contains ℓ 1 -norm between the output and the ground truth as shown in Eq.

Original: 8 and the difference of the total flux as Eq.
Revised: 8 and the difference of the total flux as Eq.

Original: In Eq.
Revised: In Eq.

Original: 8, f is thenetwork, I and I gt refer to the input and true image.
Revised: 8, f is the network, I and I gt refer to the input and true image.

Original: λ 0 controls the trade-off between the two parts.
Revised: λ 0 controls the trade-off between the two parts.

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation] In the second phase, in addition to L ℓ 1 and L M 00 , the astrometric position ⟨ x ⟩ and ⟨ y ⟩ , the secondorder central moments µ 11 , µ 20 , µ 02 are utilized to fine-tune the model.
Revised: In the second phase, in addition to L ℓ 1 and L M 00 , the astrometric position ⟨ x ⟩ and ⟨ y ⟩ , the secondorder central moments µ 11 , µ 20 , µ 02 are utilized to fine-tune the model.

Original: The quantities are defined in
Revised: The quantities are defined in

Original: Section 4.1 and the loss for these quantities has the same form as Eq. 9.
Revised: Section 4.1 and the loss for these quantities has the same form as Eq. 9.

Original: Thus, the total loss for thesecond phase can be written as Eq.
Revised: Thus, the total loss for the second phase can be written as Eq.

Original: λ ’s are the weights for different terms and manually adjustedso that those terms are subdominant, but not negligible compared to the dominant ℓ 1 loss.
Revised: λ ’s are the weights for different terms and manually adjusted so that those terms are subdominant, but not negligible compared to the dominant ℓ 1 loss.

Original: [Equation]
Revised: [Equation]

Original:  Training and testing data
Revised:  Training and testing data

Original: The dataset used in this work was generated using WeakLensingDeblending package 1 developedin order to test the shear measurement of weak lensing by LSST Data Management.
Revised: The dataset used in this work was generated using WeakLensingDeblending package 1 developed in order to test the shear measurement of weak lensing by LSST Data Management.

Original: From thegalaxy catalog class FixedGalaxyCatalog , composite exponential galaxy images of single galaxyare simulated in g, r
Revised: From the galaxy catalog class FixedGalaxyCatalog , composite exponential galaxy images of single galaxy

Original:  and i bands with single exposures per band and Gaussian PSFs that varyacross the different bands (using industry standard galsim , [15]).
Revised: [Table] are simulated in g, r and i bands with single exposures per band and Gaussian PSFs that vary across the different bands (using industry standard galsim , [16]).

Original: For each galaxy, a noisy image, a truth image and a PSF image are generated as a stack of the three aforementioned bands.
Revised: For each galaxy, a noisy image, a truth image and a PSF image are generated as a stack of the three aforementioned bands.

Original: Since ourobjective is to ascertain the influence of incorporating PSF in neural networks for denoising puprosesof astronomical images, we have set shear to be zero.
Revised: Since our objective is to ascertain the influence of incorporating PSF in neural networks for denoising puproses of astronomical images, we have set shear to be zero.

Original: Moreover, we incorporated a signal-to-noise(SNR) cut of > 10 .
Revised: Moreover, we incorporated a signal-to-noise (SNR) cut of > 10 .

Original: The 100 , 000 images thus generated for each of the three categories (truth, noisy and PSF) havedimensions of 35 × 35 × 3 .
Revised: The 100 , 000 images thus generated for each of the three categories (truth, noisy and PSF) have dimensions of 35 × 35 × 3 .

Original: The entire dataset is divided into training set with 90 , 000 images and testsets with 10 , 000 images.
Revised: The entire dataset is divided into training set with 90 , 000 images and test sets with 10 , 000 images.

Original: Each image is sized 32 × 32 × 3 and normalized to [0 , 1] by dividing itsoverall maximum pixel value so that colors are preserved.
Revised: Each image is sized 32 × 32 × 3 and normalized to [0 , 1] by dividing its overall maximum pixel value so that colors are preserved.

Original: During testing, in order to get the accurateperformance of the model, the prediction from the network are rescaled back by exactly same factorwhich allows a direct comparison with the raw ground truth.
Revised: During testing, in order to get the accurate performance of the model, the prediction from the network are rescaled back by exactly same factor which allows a direct comparison with the raw ground truth.

Original:  Results 
Revised:  Results 6.1 Implementation and comparison models

Original: 6.1 Implementation and comparison models

Original: For experiments, we adopt Adam optimizer [8] and the batch size is 128 .
Revised: For experiments, we adopt Adam optimizer [8] and the batch size is 128 .

Original: During the first phase, themodel is trained for 80 epochs with learning rate as 10 − 4 .
Revised: During the first phase, the model is trained for 80 epochs with learning rate as 10 − 4 .

Original: In the second phase, we fine-tune themodel for 10 epochs and the learning rate is set to 10 − 5 .
Revised: In the second phase, we fine-tune the model for 10 epochs and the learning rate is set to 10 − 5 .

Original: The weights for moment-related loss are asfollows:
Revised: The weights for moment-related loss are as follows: λ 0 , 1 , 2 = 0 .

Original:  λ 2 = 0 .λ
Revised: 01 , λ 3 = 0 .

Original: = 002 and λ 4 , 5 = 0 .
Revised: 002 and λ 4 , 5 = 0 .

Original: In order to evaluate the performance of our model, we train 4 different models:
Revised: In order to evaluate the performance of our model, we train 4 different models:

Original: • The fiducial model as explained in the text; • The Average PSF model, where the training and testing datasets remain the same, but insteadof fitting actual per-image PSF to the network, we freeze the PSF to its average shape.This
Revised: • The fiducial model as explained in the text; • The Average PSF model, where the training and testing datasets remain the same, but instead of fitting actual per-image PSF to the network, we freeze the PSF to its average shape.

Original:  allows us to explore the information gain coming from knowledge of the actual PSF; • The ℓ 1 loss only model, where training is not fine-tuned with custom loss function;• The Gaussian regularization model where Wiener deconvolution formula is replaced with a differeng regularization as described in Appendix.
Revised: This allows us to explore the information gain coming from knowledge of the actual PSF; • The ℓ 1 loss only model, where training is not fine-tuned with custom loss function; • The Gaussian regularization model where Wiener deconvolution formula is replaced with a different regularization as described in Appendix.

Original: 6.2 Evaluation
Revised: 6.2 Evaluation

Original: In most tests we split the sample into three categories by the value of M 00 as a proxy for signalto-noise.
Revised: In most tests we split the sample into three categories by the value of M 00 as a proxy for signal-tonoise.

Original: Fig. 2 visualizes the performance of the fiducial model on the testing set.
Revised: Figure 2 visualizes the performance of the fiducial model on the testing set.

Original: We select oneexample from samples with low, medium and high M 00 separately.
Revised: We select one example from samples with low, medium and high M 00 separately.

Original: From the results, our deblurringmodel can recover the morphological and color information for images under different conditions,even for images with low fluxes.
Revised: From the results, our deblurring model can recover the morphological and color information for images under different conditions, even for images with low fluxes.

Original: It can also be observed that the model performs better on imageswith higher fluxes.
Revised: It can also be observed that the model performs better on images with higher fluxes.

Original: We will analyse in detail in the following sections.
Revised: We will analyse in detail in the following sections.

Original: We first apply some standard metrics such as the peak signal-to-noise ratio (PSNR) and the structuralsimilarity index (SSIM) [17] to evaluate the quality of images recovered from our model comparedwith the ground truth images.
Revised: We first apply some standard metrics such as the peak signal-to-noise ratio (PSNR) and the structural similarity index (SSIM) [19] to evaluate the quality of images recovered from our model compared with the ground truth images.

Original: The results are listed in Table 1, including the mean and medianamong the testing set.
Revised: The results are listed in Table 1, including the mean and median among the testing set.

Original: The best results are from the model with accurate PSF and moment loss.
Revised: The best results are from the model with accurate PSF and moment loss.

Original: Specifically, by comparison between the first and third columns in the table, we can conclude thatthose moment-related loss terms help improve the performance, reflected by a higher PSNR and
Revised: Specifically, by comparison between the first and third columns in the table, we can conclude that those moment-related loss terms help improve the performance, reflected by a higher PSNR and

Original: SSIM value, but the improvement is modest.
Revised: SSIM value, but the improvement is modest.

Original: The results from models with accurate PSF and average
Revised: The results from models with accurate PSF and average

Original: PSF also demonstrate that the quality of the prediction will get degraded if the accurate PSF is notprovided for each image again with surprisingly modest improvements. 
Revised: PSF also demonstrate that the quality of the prediction will get degraded if the accurate PSF is not provided for each image again with surprisingly modest improvements. [Figure]

Original: 6.3 Recovery of astronomical quantities
Revised: 6.3 Recovery of astronomical quantities

Original: Table 2 in Appendix presents the difference between output from the model and true image forastronomical quantities of interest.
Revised: Table 2 in Appendix presents the difference between output from the model and true image for astronomical quantities of interest.

Original: These values are averaged over the three available bands.
Revised: These values are averaged over the three available bands.

Original: Forour fiducial network, the relevant scatter plots can also be found in Figure 3.
Revised: For our fiducial network, the relevant scatter plots can also be found in Figure 3.

Original: We find that across[Table] [Figure] the noise levels, the networks correctly recover the quantities of interest.
Revised: We find that across the noise levels, the networks correctly recover the quantities of interest.

Original: The results nicely scatteraround x = y line with no obvious bias for the total flux, first moments (astrometric position) andsecond moments.
Revised: The results nicely scatter around x = y line with no obvious bias for the total flux, first moments (astrometric position) and second moments.

Original: Ellipticity suffers from small bias, where the network tends to preferrably makeobjects less round.
Revised: Ellipticity suffers from small bias, where the network tends to preferrably make objects less round.

Original: This effect is observed for both e 1 and e 2 .
Revised: This effect is observed for both e 1 and e 2 .

Original: To understand this effect better we havesplit objects by signal-to-noise and replotted the histograms for e 1 in the Figure 4a.
Revised: To understand this effect better we have split objects by signal-to-noise and replotted the histograms for e 1 in the Figure 4a.

Original: Here we see thathigh signal-to-noise objects have their shapes recovered without bias, while there is an emergence of"multiplicative" bias for low signal-to-noise where measured ellipticites are systematically low.
Revised: Here we see that high signal-to-noise objects have their shapes recovered without bias, while there is an emergence of "multiplicative" bias for low signal-to-noise where measured ellipticites are systematically low.

Original: Next we want to study how much information we are getting by using a per-band known PSF.
Revised: Next we want to study how much information we are getting by using a per-band known PSF.

Original: Figure4b plots the histograms for model with average PSF.
Revised: Figure 4b plots the histograms for model with average PSF.

Original: We see essentially the same structure, but thescatter around the mean is somewhat increased.
Revised: We see essentially the same structure, but the scatter around the mean is somewhat increased.

Original: Similar observations for e 2 can be found in Figure5 in the Appendix.
Revised: Similar observations for e 2 can be found in Figure 5 in the Appendix.

Original: This is confirmed in the Table 2, where rms errors when using average PSF aresystematically larger by about 5-10% compared to those of fiducial network for ellipticity and secondmoments.
Revised: This is confirmed in the Table 2, where rms errors when using average PSF are systematically larger by about 5-10% compared to those of fiducial network for ellipticity and second moments.

Original: We also find that the error on the astrometries are not affected by knowing or not the true
Revised: We also find that the error on the astrometries are not affected by knowing or not the true

Original: PSF – this makes sense because our PSFs are Gaussian with zero mean and therefore using a wrong
Revised: PSF – this makes sense because our PSFs are Gaussian with zero mean and therefore using a wrong

Original: PSF to deconvolve the image is not going to affects its central position.
Revised: PSF to deconvolve the image is not going to affects its central position.

Original: Somewhat more surprising isthat the error on the object flux also seems to be unchanged between the two cases.
Revised: Somewhat more surprising is that the error on the object flux also seems to be unchanged between the two cases.

Original: More analysis onthe recovery of shapes and its biases can be found in Appendix.[Figure] [Figure]
Revised: More analysis on the recovery of shapes and its biases can be found in Appendix.

Original: 6.4 Uncertainty quantification (UQ)
Revised: 6.4 Uncertainty quantification (UQ)

Original: Like majority of the deep learning frameworks, our model is also a point-prediction method, i.e.,the predictions in the network are not associated with error estimates.
Revised: Like majority of the deep learning frameworks, our model is also a point-prediction method, i.e., the predictions in the network are not associated with error estimates.

Original: Large number of uncertaintyquantification approaches are currently applied in the context of AI applications in physics [4, 12, 9],some of which may be computationally expensive in our problem, where the individual model trainingrequires 25.5 hours on one Nvidia V100 GPU with 32GB memory.
Revised: Large number of uncertainty quantification approaches are currently applied in the context of AI applications in physics [4, 13, 10], [Figure] [Figure] some of which may be computationally expensive in our problem, where the individual model training requires 25.5 hours on one Nvidia V100 GPU with 32GB memory.

Original: In this paper, we perform a two-fold treatment of uncertainty.
Revised: In this paper, we perform a two-fold treatment of uncertainty.

Original: The first one aims at estimating themodel uncertainties, where we quantify the fluctuations in the predicted images due to imperfect model training.
Revised: The first one aims at estimating the model uncertainties, where we quantify the fluctuations in the predicted images due to imperfect model training.

Original: Second, we quantify the uncertainty propagated due to the noisy input data duringtesting.
Revised: Second, we quantify the uncertainty propagated due to the noisy input data during testing.

Original: Both these treatments are approximations aimed at obtaining a rudimentary understanding ofthe confidence intervals of our methods, within reasonable computational budgets.
Revised: Both these treatments are approximations aimed at obtaining a rudimentary understanding of the confidence intervals of our methods, within reasonable computational budgets.

Original: In applicationsinvolving real observations, several systematic effects may have to be considered in a greater detail.
Revised: In applications involving real observations, several systematic effects may have to be considered in a greater detail.

Original: Model uncertainty using Monte Carlo methods: For model uncertainties, we choose the UQtechnique of Monte Carlo (MC) Dropout, that provides approximate errors with little computationaloverhead and minimal changes to the architecture.
Revised: Model uncertainty using Monte Carlo methods: For model uncertainties, we choose the UQ technique of Monte Carlo (MC) Dropout, that provides approximate errors with little computational overhead and minimal changes to the architecture.

Original: The implementation of MC Dropout is: First,we train a model without any UQ considerations to obtain a fiducial model.
Revised: The implementation of MC Dropout is: First, we train a model without any UQ considerations to obtain a fiducial model.

Original: Next, we consider anensemble of N tot models with the same architecture, only differing from each other by a fractionof trained neurons re-initialized to a random value.
Revised: Next, we consider an ensemble of N tot models with the same architecture, only differing from each other by a fraction of trained neurons re-initialized to a random value.

Original: We obtain this ensemble of networks and thefraction of randomized (or “dropped-out”) weights is dictated by the dropout rate d .
Revised: We obtain this ensemble of networks and the fraction of randomized (or “dropped-out”) weights is dictated by the dropout rate d .

Original: Since the Deep
Revised: Since the Deep

Original: Neural Networks are over-parametrized, a small number of randomized weights will only result insmall fluctuations in the output images.
Revised: Neural Networks are over-parametrized, a small number of randomized weights will only result in small fluctuations in the output images.

Original: Each image in the testing set will be evaluated by those N totmodels and obtain N tot outputs.
Revised: Each image in the testing set will be evaluated by those N tot models and obtain N tot outputs.

Original: The mean and variance of these outputs will be considered as theuncertainty quantified prediction from the ensemble.
Revised: The mean and variance of these outputs will be considered as the uncertainty quantified prediction from the ensemble.

Original: Statistical uncertainty due to noisy image inputs: In addition to model uncertainty which inphysicist language corresponds to systematic error, any observation also has noise or random uncertainty.
Revised: Statistical uncertainty due to noisy image inputs: In addition to model uncertainty which in physicist language corresponds to systematic error, any observation also has noise or random uncertainty.

Original: Since neural network results are point estimates, one would ideally process the sameimage with multiple noise realizations to estimate the noise in any derived quantities.
Revised: Since neural network results are point estimates, one would ideally process the same image with multiple noise realizations to estimate the noise in any derived quantities.

Original: Since this isimpossible, we instead generate N tot images with additional noise drawn from the same distributionas the real noise and run those through the network.
Revised: Since this is impossible, we instead generate N tot images with additional noise drawn from the same distribution as the real noise and run those through the network.

Original: These form an ensemble from which a noiseestimate on the derived quantities can be formed.
Revised: These form an ensemble from which a noise estimate on the derived quantities can be formed.

Original: Since these images now have double the noisevariance, the result is strictly speaking an overestimate of measurement uncertainty.
Revised: Since these images now have double the noise variance, the result is strictly speaking an overestimate of measurement uncertainty.

Original: In the experiment, N tot = 200 for both approaches and dropout rate d = 0 . 2 .
Revised: In the experiment, N tot = 200 for both approaches and dropout rate d = 0 . 2 .

Original: Results of these twotechniques are presented in the right panel of Fig.
Revised: Results of these two techniques are presented in the right panel of Figure 2.

Original: We see that the the truth is always within theestimated uncertainty region, whose size is likely overestimated.
Revised: We see that the the truth is always within the estimated uncertainty region, whose size is likely overestimated.

Original: We also see that the size of thestatistical uncertainty decreases with increasing SNR, but interestingly the size of the systematicuncertainty also somewhat shrinks with an increasing signal to noise.
Revised: We also see that the size of the statistical uncertainty decreases with increasing SNR, but interestingly the size of the systematic uncertainty also somewhat shrinks with an increasing signal to noise.

Original: In other words, as the SNRincrease, the network seems to internally produce more consistent results.
Revised: In other words, as the SNR increase, the network seems to internally produce more consistent results.

Original:  Discussion & Conclusions
Revised:  Discussion & Conclusions

Original: Despite numerous applications of neural networks to the astronomical image analysis, the problemof using the information from the known, band dependent and relatively large PSF has not beenaddressed directly.
Revised: Despite numerous applications of neural networks to the astronomical image analysis, the problem of using the information from the known, band dependent and relatively large PSF has not been addressed directly.

Original: In this work, we propose a neural network based deconvolution algorithm toaddress this crucial missing link.
Revised: In this work, we propose a neural network based deconvolution algorithm to address this crucial missing link.

Original: It derives from a network that deploys Wiener Deconvolution inthe feature space.
Revised: It derives from a network that deploys Wiener Deconvolution in the feature space.

Original: We have two notable improvements aimed at the properties of astronomical tasks.
Revised: We have two notable improvements aimed at the properties of astronomical tasks.

Original: Firstly, instead of a simplifying assumption of constant PSF, we employ a band-wise deconvolutionmodule.
Revised: Firstly, instead of a simplifying assumption of constant PSF, we employ a band-wise deconvolution module.

Original: This enables the network to work with each band separately, rendering is compatiblewith the PDFs of the imaging systems in astronomy.
Revised: This enables the network to work with each band separately, rendering is compatible with the PDFs of the imaging systems in astronomy.

Original: Secondly, we engage quantities representingparticular properties of astronomical images into our custom loss function.
Revised: Secondly, we engage quantities representing particular properties of astronomical images into our custom loss function.

Original: The model can have betterinvestigation on those properties by directly training on those quantities.
Revised: The model can have better investigation on those properties by directly training on those quantities.

Original: Extensive experimentsand analysis demonstrate that the model can acquire more better recovery of shape than comparisonmodels, especially for images with higher SNR.
Revised: Extensive experiments and analysis demonstrate that the model can acquire more better recovery of shape than comparison models, especially for images with higher SNR.

Original: This work is the start of incorporating band-separated
Revised: This work is the start of incorporating band-separated

Original: PSFs information in astronomical image processing.
Revised: PSFs information in astronomical image processing.

Original: More broadly, this study motivates the need ofrobust, domain-focussed architectures in AI-for-Astronomy problems.
Revised: More broadly, this study motivates the need of robust, domain-focussed architectures in AI-for-Astronomy problems.

Revised: Acknowledgements

Revised: We thank Erin Sheldon for useful discussions.

Revised: This material is partly based upon the work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing

Revised: Research and Office of High Energy Physics, Scientific Discovery through Advanced Computing (SciDAC) program on "Accelerating HEP Science: Inference and Machine Learning at Extreme

Original: A Additional results
Revised: Scales". NR’s work at Argonne National Laboratory was supported by the U.S. Department of Energy, Office of High Energy Physics. Argonne, a U.S. Department of Energy Office of Science Laboratory, is operated by UChicago Argonne LLC under contract no.

Revised: DE-AC02-06CH11357.

Revised: A Additional results

Original: Difference of moments between predicted images and ground truth: Table 2 shows the differencebetween the model’s output and the true image for astronomical quantities of interest.
Revised: Difference of moments between predicted images and ground truth: Table 2 shows the difference between the model’s output and the true image for astronomical quantities of interest.

Original: All fourcomparison models in the text are included.
Revised: All four comparison models in the text are included.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: Recovery of ellipticity: Figure 5 is the same as Figure 4 but for e 2 .
Revised: Recovery of ellipticity: Figure 5 is the same as Figure 4 but for e 2 .

Original: Similar conclusion can be drawnthat the model recovers the shape high signal-to-noise objects without bias and model trained withaverage PSF lost some information with a increased scatter around the mean compared with thefiducial model.
Revised: Similar conclusion can be drawn that the model recovers the shape high signal-to-noise objects without bias and model trained with average PSF lost some information with a increased scatter around the mean compared with the fiducial model.

Original: Recovery of shapes and its biases: We also study the performance of our network in removing thesubtle biases that remain in the prediction.
Revised: Recovery of shapes and its biases: We also study the performance of our network in removing the subtle biases that remain in the prediction.

Original: In particular, we want to study the remaining contaminationby the PSF shape.
Revised: In particular, we want to study the remaining contamination by the PSF shape.

Original: Our main motivation is to understand whether such methods could ever be used forweak lensing application, where residual contamination with the PSF leads to to the so-called additivebias (see [11] and references there-in for an in-depth review).
Revised: Our main motivation is to understand whether such methods could ever be used for weak lensing application, where residual contamination with the PSF leads to to the so-called additive bias (see [12] and references there-in for an in-depth review).

Original: In short if the shape of galaxy remainscorrelated with the PSF after the convolution, then this PSF shape will leak into shear estimation andlead to shape correlations that can mimic real weak gravitational signal.
Revised: In short if the shape of galaxy remains correlated with the PSF after the convolution, then this PSF shape will leak into shear estimation and lead to shape correlations that can mimic real weak gravitational signal.

Original: The true galaxy ellipticities are uncorrelated with PSF shapes, but since the observed image is aconvolution of the PSF with true shape, the resulting predicted image ellipticies might still correlatewith the PSF shape.
Revised: The true galaxy ellipticities are uncorrelated with PSF shapes, but since the observed image is a convolution of the PSF with true shape, the resulting predicted image ellipticies might still correlate with the PSF shape.

Original: On the other hand, if a measurement is perfect, the predicted ellipticities shouldcorrelated perfectly with true ones.
Revised: On the other hand, if a measurement is perfect, the predicted ellipticities should correlated perfectly with true ones.

Original: To quantify this effect, we define the following correlations
Revised: To quantify this effect, we define the following correlations

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: where averages are calculated by simply averaging data in the testing dataset.
Revised: where averages are calculated by simply averaging data in the testing dataset.

Original: Note that since invidiualellipticites of either PSF, true galaxy image or predicted galaxy image can reach zero, we average thenumerator and denominator separately.
Revised: Note that since invidiual ellipticites of either PSF, true galaxy image or predicted galaxy image can reach zero, we average the [Figure] numerator and denominator separately.

Original: In ideal case, we expect the predicted image to remove anytrace of PSF from the resulting image, so ξ PSF = 0 .
Revised: In ideal case, we expect the predicted image to remove any trace of PSF from the resulting image, so ξ PSF = 0 .

Original: However, this could also be achieved by makingall objects round (i.e. e pred = 0 for every objects).
Revised: However, this could also be achieved by making all objects round (i.e. e pred = 0 for every objects).

Original: Therefore we also look at the correlation with thetrue value, which should be, in the ideal case perfect, i.e ξ true = 1 .
Revised: Therefore we also look at the correlation with the true value, which should be, in the ideal case perfect, i.e ξ true = 1 .

Original: Results of this exercise are plotted in the Figure 6.
Revised: Results of this exercise are plotted in the Figure 6.

Original: We see that our fiducial case (denoted "Accurate
Revised: We see that our fiducial case (denoted "Accurate

Original: PSF" and plotted red) performs bettter than the "Average PSF" case across all signal-to-noiselevels.
Revised: PSF" and plotted red) performs bettter than the "Average PSF" case across all signal-to-noise levels.

Original: The cross-correlation with PSF is significantly lower and the cross-correlation with thetruth significantly higher.
Revised: The cross-correlation with PSF is significantly lower and the cross-correlation with the truth significantly higher.

Original: We conclude that the deconvolution is taking place, albeit imperfectly:although the network does successfuly decorrelate PSF in the high SNR, but it still contains significantcontamination in the lower SNR region.
Revised: We conclude that the deconvolution is taking place, albeit imperfectly: although the network does successfuly decorrelate PSF in the high SNR, but it still contains significant contamination in the lower SNR region.

Original: In order to improve upon this, we explore another version of deconvolution in which we replace the
Revised: In order to improve upon this, we explore another version of deconvolution in which we replace the

Original: Winer filter Equation 2 with a modified term, which we refer to as "Gaussaian Regularization":
Revised: Winer filter Equation 2 with a modified term, which we refer to as "Gaussian Regularization":

Original: [Equation]
Revised: [Equation]

Original: In other words, the PSF is linearly deconvolved from the noisy image, but the resulting image isthen regularized with a circular Gaussian roughly the size of the PSF.
Revised: In other words, the PSF is linearly deconvolved from the noisy image, but the resulting image is then regularized with a circular Gaussian roughly the size of the PSF.

Original: In the original formulation ofthe Wiener filter, modes are anisotropically regularized, taking into account that anistropic PSF hasdestroyed more information in some directions than the others.
Revised: In the original formulation of the Wiener filter, modes are anisotropically regularized, taking into account that anistropic PSF has destroyed more information in some directions than the others.

Original: However, the result is that the PSFshape can sneak back into the resulting image.
Revised: However, the result is that the PSF shape can sneak back into the resulting image.

Original: The Guassian regularization tries to prevent this at theexpense of formally decreased optimality, i.e. the resulting process is now manifestly worse thanminimum variance solution.
Revised: The Guassian regularization tries to prevent this at the expense of formally decreased optimality, i.e. the resulting process is now manifestly worse than minimum variance solution.

Original: We choose a Gaussian whose size is close to typical PSF size and haveexplicitly demonstrated that the results are largely insensitive to precise choice of this kernel, if thenetwork is retrained.
Revised: We choose a Gaussian whose size is close to typical PSF size and have explicitly demonstrated that the results are largely insensitive to precise choice of this kernel, if the network is retrained.

Original: In our tests, however, this methods seems to perform pretty well.
Revised: In our tests, however, this methods seems to perform pretty well.

Original: In Figure 6, we see that the Gaussianregularization model is performs signficiantly better than our fiducial network with no evidenceof PSF correlation at no loss in correlation with the truth.
Revised: In Figure 6, we see that the Gaussian regularization model is performs signficiantly better than our fiducial network with no evidence of PSF correlation at no loss in correlation with the truth.

Original: In Tables 2, we see that the Gaussian regularization performs essentially as well as the fiducial model within the noise and that whileindividual RMS values scatter favouring one or the other method, there is no clear winner.
Revised: In Tables 2, we see that the Gaussian regularization performs essentially as well as the fiducial model within the noise and that while individual RMS values scatter favouring one or the other method, there is no clear winner.