[DeepLense] Foundation Model for Gravitational Lensing - WaveLens-JEPA
Strong gravitational lensing is a powerful probe of dark matter structure and cosmology, yet no foundation model exists that is purpose-built for lensing data with theoretical training guarantees. This proposal introduces WaveLens-JEPA, a self-supervised vision foundation model for gravitational lensing that extends LeJEPA (Balestriero & LeCun, 2025) to the astrophysical domain. The core idea is to decompose lensing images using a 2D Discrete Wavelet Transform before encoding, separating the physically meaningful frequency scales — Einstein ring morphology, subhalo perturbations, and ALP vortex structure — into orthogonal channels. A physics cross-attention module injects convergence-map priors from the lensing equation at every ViT block, while Einstein-Ring-Aware Masking (ERAM) targets the physically informative annular region during pre-training. Fourier phase regularization preserves substructure-discriminative spatial information that standard encoders discard. The training objective inherits provable collapse-free guarantees via Sketched Isotropic Gaussian Regularization (SIGReg), unlike all prior lensing SSL methods. The pre-trained backbone is evaluated on six downstream tasks: two substructure classification setups, lens finding, mass density regression, super-resolution, and a novel interpretability task via Wavelet Attribution Maps (WAM). Deliverables include: (1) open-source modular codebase, (2) pre-trained backbone weights on Hugging Face Hub, (3) a six-task evaluation benchmark, (4) a sim-to-real domain adaptation protocol targeting COSMOS HST data, and (5) a full research paper.
Project details
Technologies
Not listed in the archive