GSOC 2022 with ML4SCI | Deep Regression for Exploring Dark Matter
In this blog post, I will share my experience as a Google Summer of Code (GSOC) 2022 student at the Machine Learning for Science (ML4Sci) organization.
I have already worked on this project while participating in GSOC 2021, so I will omit some of the results from the previous year (you can find them here) and focus more on the updates produced this summer.
What is the project about?
Scientists have not yet come to a consensus on what dark matter is. They know that there is something that does not omit (or reflect) light but still interacts with gravity. Thus, a few things are considered to be dark matter candidates. For example, axions, weakly interacting massive particles (WIMPs), cold dark matter (CDM), etc.
DeepLense is a machine learning pipeline that explores the properties of the dark matter candidates, which might bring us closer to the origin of dark matter. Namely, we regress some quantities (like axion mass, the total fraction of mass in the substructure, etc.) using strong gravitational lensing images.
Unfortunately, we do not have enough strong gravitational lensing images of real galaxies, so we simulate them using PyAutoLens. But, in combination with another organization project (you can read more about it in Marcos’ blog), we will be able to use models trained on simulations for the real data.
About Me
I am Yurii Halychankyi, a senior CS student at the University of Washington. Most of my research/work experience is connected to Computer Vision. In my free time, I make music. Thus, now I am actively exploring the signal processing field to be able to combine my hobby and job.
I participated in GSOC 2021 and worked for FellowshipAI and Shin-Etsu.
Datasets
To standardize the results, we introduce three new axion datasets for benchmarking: Model I, Model II, and Model III. The smaller the dataset number, the more difficult it is to learn from it. The lensing signature of the images from the dataset is a vortex, similar to the previous year’s dataset. All of the datasets were generated using lenstronomy.
I will briefly describe the datasets, but you can read more about them here.
Model I
The dataset contains 29k grayscale images of the size of (150,150). The simulations are modeled with a simple Sersic light profile and a Gaussian point spread function. The signal-to-noise ratio (SNR) is around 25.
Most of the lenses in the dataset have a mass close to 0. Since the values are very small, we can apply log10 to the masses to make their distribution look more uniform, facilitating the training process.
Model II
Model II contains 30k grayscale images of the size of (64,64). This dataset is easier to train on because of the bigger SNR. It was modeled after Euclid observation characteristics.
Although the noticeable differences in the lenses, the mass density distribution looks very similar to the one of Model I.
Model III
Model III contains 30k grayscale images of the size of (64,64). This is the easiest among these three datasets. We can even see that the halos on the simulations are more distinct. This has the highest SNR and was modeled after HST observation characteristics.
Axion 3-channel
We generated this dataset to test our model’s performance on 3-channel simulations. It turned out that the model showed similar performance on this dataset in comparison to last year’s axion dataset, which contained only grayscale images. But except for a different number of channels, these simulations also model the lensing effect. Thus we can observe a point in the center of the halo. So, to make this dataset results comparable to the previous year’s one, we have to apply the circular mask to cover the point in the center and normalize every channel separately.
To create a circular mask, we use the following function:
def create_circular_mask(h, w, center=None, radius=None):if center is None: # use the middle of the image
center = (int(w/2), int(h/2))
if radius is None: # use the smallest distance between the center and image walls
radius = min(center[0], center[1], w-center[0], h-center[1])Y, X = np.ogrid[:h, :w]
dist_from_center = np.sqrt((X - center[0])**2 + (Y-center[1])**2)mask = dist_from_center <= radius
return mask
Next, we apply it to every image and normalize the channels:
h, w = images[0].shape[:2]
mask = create_circular_mask(h, w,radius=5)for i in tqdm(range(images.shape[0])):
masked_img = images[i].copy()
# Apply the mask
masked_img[mask] = 0# Normalize every channel
for c in range(3):
images[i,:,:,c] = (masked_img[:,:,c]-masked_img[:,:,c].mean()) / masked_img[:,:,c].std()
This is how the images look after the preprocessing:
CDM
This dataset simulates another dark matter candidate: Cold Dark Matter (CDM). Instead of axion mass, we regress on the total fraction of mass in the substructure (f_subs) and the slope of the distribution function (betas), which are the generator’s parameters.
Get Yurii Halychanskyi’s stories in your inbox
Join Medium for free to get updates from this writer.
The dataset contains 40k grayscale images of the size of (64,64).
We can observe that the f_subs distribution looks similar to the axion mass one, so we will apply log10 before training to make the distribution uniform.
The betas distribution is already uniform, so no preprocessing is needed.
Model Architecture and Training
We use the same XResNetHybrid as the previous year because it yields the best performance. But other models also have been tested. For example, the results of XResNetHybrid and CMT, an architecture combining CNN and a vision transformer (VT), are shown in the table below.
And here is a visual comparison of those two models:
Note that the points do not form a diagonal line because, for specific mass values, the lensing signature becomes so subtle that the model cannot learn, which results in the horizontal line part.
The training pipeline is also the same as the previous year. Still, we obtained a better performance by tuning the following hyperparameters (the best set of hyperparameters is marked with an asterisk):
Also, we have tested different loss functions, and SmoothL1Loss happened to work best with our pipeline.
Results
All the showed results were obtained with XResNetHybrid18 trained with batch_size=64, lr=1e-3, wd=None, moms=(0.95,0.85) for 200 epochs.
We have already seen XResNetHybrid performance on our benchmark datasets (Model I, Model II, Model III) when compared to CMT. So, I will show the results for the other datasets described above.
Axion 3-channel
The MAE of the axion 3-channel test dataset is 0.0546, which is reasonable considering that last year’s axion result was 0.0861 (although it was 1-channel). I think the observed performance increase is due to the narrower range of the masses. The number of channels does not affect the performance much, which was proved by training the model on every channel separately and getting similar results.
CDM
The MAE for CDM dataset trained on betas is 0.1168. We do not have previous results, so we cannot compare them directly. But we can see that the line is much wider than the one for the axion mass, which is expected. Betas represent the slope of the distribution function of axion masses where subhalos are sampled from. On every lens, there are about 25 subhalos generated, but only a couple of them near the arcs are detectable in practice. Axion masses are directly related to the length of the vortex, and it is enough to have a couple of subhalos to have a solid vortex length approximation. Still, it is difficult to approximate the slope distribution having only a couple of samples.
MAE for CDM trained on f_subs is 0.0052. We have not previously regressed on f_subs either. Thus we cannot compare them, but MAE is relatively small, and the line is narrow, which shows a small variance (which is good).
Final Thoughts
A lot of work has been done on this project. We have tried different architectures, datasets, sets of hyperparameters, augmentations, and loss functions. But I think we can still get a performance boost by improving how we handle data. Meaning adding some image preprocessing that would make the lensing signatures more distinct or a model that would classify images before regressing on them so that we can have a separate regression for every case.
I enjoyed working on this project, and I hope I will be able to contribute to it in the future. I want to thank Sergei Gleyzer, Michael Toomey, and Brendan Ames for helping and directing me during the program.
Thank GSoC for such an opportunity!
