Lens-JEPA: Physics Informed Joint Embedding Predictive Architecture for Gravitational Lensing
Abstract
We introduce Lens-JEPA, a novel approach towards building a foundation model for gravitational lensing images that extends I-JEPA to the astrophysical domain. Although recent advances in foundation models have transformed vision and language tasks by enabling generalization and transfer across other applications, astrophysical imaging still lacks such a framework. To address this gap, we develop a key component of Lens-JEPA, which is a Physics Encoder that introduces a transformer guided by the lensing equation, that combines the representational power of Vision Transformers with the rigor of Physics Informed Neural Networks (PINNs). This approach enhances the representation capabilities of gravitational lensing. This paper demonstrates the effectiveness of Lens-JEPA via a classification task, which shows that Lens-JEPA surpasses current top baseline architectures, highlighting the benefits of incorporating physics into transformer models. Although this study focuses on classification, it provides a foundation across other applications such as lens detection, mass modeling, and super-resolution, moving toward a foundation model for gravitational lensing