Lens-LeJEPA: Self-Supervised Learning for Dark Matter Substructure

Simulated strong lenses from the DeepLense Model III test set.

Program

Google Summer of Code 2026

Organisation

ML4SCI · DeepLense

Code

github.com/arnesh2212/Lens-FM-DeepLense

01Introduction

When a massive galaxy lies between us and a more distant one, its gravity bends the light from the far galaxy into arcs, sometimes into a full ring. This is strong gravitational lensing. The shape of these arcs depends on how the mass inside the lens is distributed, including small clumps of dark matter that we can't see directly. So lens images are one of the few places where we can test dark matter on small scales.

Different dark matter models leave slightly different marks. In cold dark matter (CDM), the main halo is full of small subhaloes. If dark matter is an ultra-light axion, it behaves like a wave and creates vortex-like patterns. The DeepLense project tries to tell these apart from images alone.

Supervised models already do this well on simulated data, but they need labels, and every label comes from a simulation that assumed a dark matter model in the first place. Surveys like Euclid will give us a huge number of lens images with no labels. That is where self-supervised learning (SSL) makes sense: learn from unlabelled images first, then fine-tune with a few labels.

The catch is that most SSL methods learn by comparing two augmented copies of an image, made with random crops, zooms and colour changes. For lenses, zooming is not harmless. It changes the apparent Einstein radius \(\theta_E\) (the radius of the ring), and the mass inside the ring grows as \(\theta_E^2\). So a model trained to ignore zoom is also trained to ignore mass. The LenSiam paper pointed this out, and it is where my project started.

02Data

I used the public DeepLense simulations, made with lenstronomy. Every image has a lens galaxy in the centre and a lensed background galaxy around it. There are three classes: axion, CDM and no substructure. The classes differ only by small changes in the arcs.

Table 1. Datasets. All results in this post are on the held-out test sets.

Two things about this data mattered a lot for my design. The lens is always centred and the blur is the same in every direction, so rotating an image by 90° or flipping it gives another valid image of the same system, with the same class and the same mass. And the useful signal is in the thin bright arcs; most of the image is empty sky. Axion images also come with the simulated axion mass (between \(10^{-24}\) and \(10^{-22}\) eV), which I used for a regression task.

03Starting point

I started by setting up a benchmark. I pretrained a few SSL methods on Model I for 300 epochs and tested them on Model II with 100, 500, 1,000, 5,000 or all labels per class. The methods were I-JEPA, Lens-JEPA (an earlier DeepLense model built on I-JEPA) and LeJEPA, a recent method that doesn't need a teacher network. For comparison I trained supervised models: ResNet18, ViT, ViT-SD, Lensiformer and LensPINN.

With a frozen encoder and a single linear layer on top, the best SSL model only got 72% on full Model II data. The features had useful information, but a linear layer could not get it out. Switching to LoRA adapters fixed most of this (on I-JEPA, 0.600 went up to 0.978), so I used adapters for everything after that.

04What did not work

Before arriving at the final method, I tried a few other ideas.

These pushed me towards something simpler: keep the model and training setup plain, and add the physics only through the loss.

05Lens-LeJEPA

Base: LeJEPA

I built on LeJEPA because it is simple. It has no teacher network and no stop-gradient. Both views go through the same ViT, the patch features are averaged and projected, and the loss does two things: it pulls the two views together, and a term called SIGReg keeps the embeddings spread out like a Gaussian, so the model can't collapse everything to a single point.

This loss only looks at the averaged feature, so it loses where things are in the image. My two extra terms work on the individual patch features instead.

Augmentations

I only used changes that could come from a real telescope looking at the same system: a small blur (up to about one pixel) and a bit of noise. No crop, zoom or resize, because those change the Einstein radius, and no shifting, because both of my extra terms assume the lens is in the centre.

Prior 1: matching rotated patches

Rotating or flipping the image is safe, so before making the second view I apply one of the 8 rotations and flips of a square (the group \(D_4\)). The image is 160 × 160 pixels and each patch is 16 × 16, so the patches form a 10 × 10 grid. When the image is rotated, every patch lands exactly on another patch position. That means I know which patch in the rotated image came from which patch in the original, and I can compare them one by one:

Here \(\pi_g\) is that patch mapping. This is the part my earlier global rotation loss was missing.

Prior 2: matching rings around the lens

The first term only makes sure information ends up in the right place. It says nothing about the shape of the arcs, which is where the dark matter signal is. So I divide the patch grid into three rings around the centre and average the patch features inside each ring. Each patch is weighted by how much arc it contains, measured with the image gradient and Laplacian, so empty sky counts for less. Rotations don't change how far a patch is from the centre, so these ring features should be the same for both views. The full loss is:

None of this needs labels or a lens model. If I turn both terms off, I get back plain LeJEPA, so I could run a clean comparison: the same ViT-S/16, trained for 100 epochs on Model I with the same settings, with and without the two terms.

06Adapting the encoder

After pretraining, I froze the encoder and added LoRA adapters to all 48 linear layers. I picked the LoRA rank using only the Model II validation set, and used rank 32 for every run after that.

Table 2. Choosing the LoRA rank on the Model II validation set.

I also tried seven types of adapters and full fine-tuning. Rank-stabilized LoRA (rsLoRA) worked best. The only difference from normal LoRA is that it scales the update by \(\alpha/\sqrt{r}\) instead of \(\alpha/r\). At rank 32 it trains 2.36M parameters.

The big benefit is storage: one shared encoder, and a small 2.36M-parameter adapter for each task.

07Results

Classification

To measure the effect of the two terms, I compared Lens-LeJEPA and LeJEPA with exactly the same adapter:

Table 3. Test accuracy by number of labelled images per class. The gain row is in accuracy points.

The two terms help at every label budget and on both datasets, and the gain is biggest exactly where self-supervised pretraining matters most: +11 points with only 100 labels per class. With rsLoRA the frozen encoder reaches 97.7% on Model II and 97.4% on Model III, and it is 13 to 30 points ahead of Lens-JEPA, the earlier DeepLense self-supervised model.

The improvement is where it should be. The "no substructure" class is already classified correctly every time, so almost all of the gain comes from telling axion and CDM apart, the two classes that both add small structure to the arcs.

The features on their own

How much of this comes from the pretrained features, and how much from the adapter? To check, I froze each encoder completely and trained only a logistic regression on its average feature, choosing the regularisation on the validation set.

Table 4. A linear classifier on the frozen, pooled features. Model II columns are labels per class.

Measured directly on the features, the effect of the priors is even bigger: +17.6 points with 100 labels and still +7 with all of them. The lensing terms make the dark matter class much easier to read out of the representation, before any fine-tuning happens. With very few labels this simple linear classifier is also the best way to use Lens-LeJEPA: at 100 labels per class it reaches 56%, higher than any adapter.

Axion mass regression

Table 5. Predicting \(\log_{10}\) of the axion mass, which spans two orders of magnitude.

The same frozen encoder that classifies the lenses also predicts the axion mass, to within 0.1 dex on Model II and 0.09 dex on Model III (R² = 0.937 on Model II). Only the small adapter and a linear head are trained for this task.

2× super-resolution

For this task I shrink each 64 × 64 image to 32 × 32 and train a small decoder on top of the encoder to rebuild it. Since the dataset has no real low/high-resolution pairs, this measures image reconstruction, not a physical resolution gain.

Table 6. Super-resolution quality with the encoder fine-tuned.

Lens-LeJEPA reconstructs the images about 10 dB better than bicubic interpolation, and slightly better than LeJEPA on both datasets. So one pretrained encoder serves all three tasks: classification, mass regression and image reconstruction.

08A closer look at the features

The tables say what happened, not why. To understand what the two loss terms actually changed, I looked at the frozen encoders directly, before any adapter or label was involved.

Is the encoder really equivariant?

I took 3,000 test images, applied each of the 7 non-trivial rotations and flips, and compared the tokens. For Lens-LeJEPA, every token of the rotated image matches the token of the patch it came from almost perfectly (the error is below 10⁻⁵), and the average feature of the image does not change at all. LeJEPA never learned this: rotating the image changes its average feature a lot (cosine similarity 0.35), even though a rotated lens is the same lens.

This is a nice result: a plain ViT, with no special architecture, learned an exact symmetry of the physics just from the loss.

How it became equivariant

I split every token into two parts: the average token at that position over all images, which only depends on where the patch is, and the rest, which depends on the image. In LeJEPA only 7% of the token variance is this position-only part. In Lens-LeJEPA it is 47%, and 99.8% of it is symmetric under rotations and flips.

So the encoder built itself a coordinate system centred on the lens: part of each token says how far the patch is from the centre, which is exactly the quantity that matters for a ring-shaped arc. This pattern is the same under every rotation and constant inside each ring, which is how both loss terms are satisfied, and the rest of the token then describes the arc itself.

09Discussion

The main lesson for me was that physics can go into the loss instead of the architecture. Two small terms, each with weight 0.1, turned a standard LeJEPA into an encoder that is exactly rotation-equivariant, organises its features around the lens centre, and is much better at recognising dark matter substructure, especially when labels are scarce. They also replace the usual crop and zoom augmentations, which change the Einstein radius and so would teach the model to ignore mass.

Training was also easy. Both lensing terms went down steadily, nothing collapsed, and the extra cost over LeJEPA is small.

10Future work

• Everything here is simulated. The next step is to pretrain on real survey images and test on known real lenses.
• Real images are never perfectly centred and have some lens-galaxy light left over. Fine-tuning with random shifts of a couple of pixels already handles the first; the second needs training with realistic residuals.
• Testing each lensing term on its own would show how much each one contributes.
• Adding a rotation-equivariant backbone on top of the priors could push the results further.
• Running regression and super-resolution with fewer labels, where pretraining should matter most.

11Code

Acknowledgements

Thanks to my mentors at ML4SCI and the DeepLense team for their guidance and feedback during the summer, and to Google Summer of Code for the opportunity. This work builds on earlier DeepLense projects, especially Lens-JEPA, Lensiformer and LensPINN.

References

1. R. Balestriero, Y. LeCun. LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics. 2025. arXiv:2511.08544
2. P.-W. Chang, K.-W. Huang, J. Fagin et al. LenSiam: Self-Supervised Learning on Strong Gravitational Lens Images. NeurIPS 2023 AI for Science Workshop. arXiv:2311.10100
3. J. Rishi, P. R. Kumbam, M. W. Toomey, S. Gleyzer. Lens-JEPA: Physics Informed Joint Embedding Predictive Architecture. NeurIPS 2025 ML4PS Workshop. PDF
4. M. Assran et al. Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023. arXiv:2301.08243
5. L. J. Velôso, M. W. Toomey, S. Gleyzer. Lensformer: A Physics-Informed Vision Transformer for Gravitational Lensing. NeurIPS 2023 ML4PS Workshop. PDF
6. A. Ojha, S. Gleyzer, M. W. Toomey, P. R. Kumbam. LensPINN: Physics Informed Neural Network for Learning Dark Matter Morphology in Lensing. NeurIPS 2024 ML4PS Workshop. PDF
7. D. Kalajdzievski. A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA. 2023. arXiv:2312.03732
8. E. J. Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. ICLR 2022. arXiv:2106.09685
9. S. Alexander, S. Gleyzer, E. McDonough, M. W. Toomey et al. Deep Learning the Morphology of Dark Matter Substructure. The Astrophysical Journal, 2020. arXiv:1909.07346
10. S. Birrer, A. Amara. lenstronomy: Multi-purpose gravitational lens modelling software package. Physics of the Dark Universe, 2018. arXiv:1803.09746