GSOC 25 with ML4SCI | Foundation Model for Gravitational Lensing

I worked with the Machine Learning for Science (ML4Sci) team as a Google Summer of Code contributor. My project, “Foundation Model for Gravitational Lensing,” was conducted in collaboration with the DeepLense team. My work focused on developing new self supervised framework for gravitaional lensing.

Image credits: Google Summer of Code, ML4Sci

🌌 Lens-JEPA: Building a Foundation Model for Gravitational Lensing

A physics-informed journey toward next-generation AI for astrophysics

🚀 Introduction: Why Gravitational Lensing Needs a Foundation Model ?

Gravitational lensing is one of the most powerful tools in modern astrophysics capable of revealing the invisible universe, mapping dark matter substructure, and probing galaxies in the early cosmos. Yet despite its scientific potential, one challenge remains persistent:

Astrophysical imaging still lacks a unified, general-purpose AI system, what the broader machine learning community calls a foundation model.

While vision transformers, large language models, and multimodal architectures have revolutionized natural image understanding, astrophysics continues to rely on highly task-specific pipelines:

1. One model for lens classification
2. Another for mass modeling
3. Specialized setups for super-resolution or anomaly detection
4. Custom deep learning architectures for dark matter substructure analysis

This fragmentation leads to duplicated effort, large labeled datasets, and poor cross-task generalization.

Lens-JEPA a new physics-informed foundation model design aims to change this landscape.

🔭 What Is Lens-JEPA?

Lens-JEPA introduces a Physics-Informed Joint Embedding Predictive Architecture that merges:

I-JEPA self-supervised representation learning , Vision Transformers and explicit gravitational lensing physics (via the lens equation) into a single cohesive model.

This hybrid solution takes inspiration from the emerging class of foundation models systems that learn general-purpose representations once, then adapt rapidly to multiple downstream tasks.

While traditional JEPA models focus on semantic image features, lensing images require something more nuanced:

They encode distortions governed not just by visual features, but by physics specifically, the Einstein deflection of light by mass.

Lens-JEPA is the first JEPA-style architecture that:

• Embeds gravitational lens physics directly into the encoder
• Uses learnable potential functions based on the Singular Isothermal Sphere (SIS) model
• Incorporates the lens equation to reconstruct the source galaxy
• Uses transformer-driven feature extraction afterward

This results in representations that are both visually rich and physically meaningful.

🛰 Why Existing Models Fall Short

Most machine learning approaches in lensing fall into two categories:

1. Purely data-driven models (e.g., CNNs, ViTs, ResNets)

These learn image features but ignore lensing physics.

2. Physics-informed networks (e.g., PINNs, Lensformer)

These use physical priors but lack general-purpose pretraining frameworks.

Neither provides a pathway toward a multi-task and data-efficient foundation model.

Lens-JEPA fills this gap by combining the scalability of transformers and JEPA training with the constraints of astrophysical theory.

🪐 Dataset Setup: Simulating the Universe

To evaluate Lens-JEPA, We generated strong lensing images using Lenstronomy, simulating two types of Data:

Model A: generic telescope (Gaussian PSF ~0.05") Model B: Euclid-like telescope setup

The simulated galaxies include:

• Lensed background galaxies with Sérsic profiles

• Foreground dark matter halos via the SIS model

• Multiple dark matter substructure models:

• CDM with truncated NFW subhalos

• Axion-like dark matter with wave-like features

• A no-substructure baseline

Datasets:

10,000 images from Model A for pretraining 3,000 per class (Axion, CDM, No-substructure) for downstream classification

Each image: 150×150 pixels, SNR ≈ 25.

⚙️ Architecture: A Physics-Guided JEPA

Lens-JEPA introduces a major modification to the standard I-JEPA encoder.

Overview of Lens JEPA Architecture

📌 Key idea:

Before extracting transformer features, the model applies the gravitational lens equation using a learnable potential.

1️⃣ The Lens Equation

2️⃣ SIS Potential Approximation

The potential is approximated as:

Here, k(x, y) is learned by a Vision Transformer for Small Datasets (ViTSD). This lets the encoder undo the lensing distortion before tokenization.

3️⃣ Hybrid Encoder

Once the source-plane reconstruction is produced:

• It is tokenized using Shifted Patch Tokenization (SPT)
• Passed to a transformer with Local Multi-Head Attention (LMA)

This blend of physics + transformer is what makes Lens-JEPA uniquely effective.

🧪 Experiments: Testing the Model as a Foundation

Lens-JEPA is pre-trained via I-JEPA’s self-supervised objective on Model A, then fine-tuned on Model B for three class classification:

• Axion dark matter
• CDM substructure
• No substructure

Baselines included:

• ResNet18
• Vision Transformer (ViT)
• CaiT
• CvT
• ViTSD
• Lensformer
• Standard I-JEPA

All models were trained with identical:

• Parameters
• Optimizer (AdamW)
• Learning rate (1e-5)
• Cross entropy loss
• 50 epochs

📊 Results: Physics + JEPA = Significant Gains

Accuracy on Model B

This is consistently higher than both supervised ViTs and even the JEPA baseline.

Why does it work so well?

Because the model is not just learning texture or morphology:

It is learning representations shaped by the physics of gravitational lensing itself.

🌠 Why Lens-JEPA Is a Path Toward a Foundation Model

Even though this study only evaluates classification, the underlying architecture is designed for:

• Lens Detection
• Mass Modeling
• Super-Resolution
• Substructure Characterization
• Inverse Modeling
• Physics-Aware Embedding for Dark Matter Studies

JEPA’s latent-space prediction makes the model naturally suitable for:

• Low-label regimes
• Large-scale pretraining
• Cross-task transfer
• Physically meaningful embeddings

This is the first step toward a universal representation model for gravitational lensing.

🧭 Limitations & Future Directions

The natural limitations are:

1. Training foundation models is computationally intensive
2. Physics-informed modules make training costlier than plain ViTs.
3. Real observations introduce domain shifts
4. Current experiments rely on simulated data. Future versions need domain adaptation.
5. Only classification was demonstrated

The real power of Lens-JEPA emerges when extended to:

• Inverse modeling
• Cosmological parameter inference
• Strong lens mass modeling
• High-resolution reconstruction
• Subhalo mass function forecasting

Ongoing work aims to unify all these tasks into a single pretrained model, similar to how vision foundation models unify segmentation, detection, and classification.

🌌 Conclusion: The Beginning of Physics-Informed Foundation Models

Lens-JEPA is more than another architecture it represents a shift in how machine learning can be integrated with astrophysical theory.

It demonstrates that:

• Physics-informed transformers are feasible
• JEPA objectives work extremely well for lensing
• Embedding physical priors boosts generalization
• A foundation model for lensing is within reach

As future telescopes like Euclid, LSST, and Roman deliver unprecedented data volumes, models like Lens-JEPA will be crucial for:

• Efficient discovery
• High-precision inference
• Dark matter characterization
• Automated scientific analysis at scale

Lens-JEPA is an early but promising step toward a future where astrophysical AI is:

unified, physics-aware, scalable, and capable of powering the next era of cosmological discovery.

This work was submitted to NeurIPS ML4PS workshop and got accepted.

Paper Title: Lens-JEPA: Physics Informed Joint Embedding Predictive Architecture for Gravitational Lensing

Authors: J Rishi, Pranath Reddy, Michael W. Toomey, Sergei Gleyzer

Code Link: https://github.com/RishiJinka/LensJEPA