GraViT: Transfer Learning with Vision Transformers and MLP-Mixer for Strong Gravitational Lens Discovery

2025-08-29

AI Summary

Automated strong gravitational lens detection in the massive image datasets expected from next-decade surveys (e.g., LSST) remains challenging due to limited labeled examples and high annotation costs. Method: We propose a transfer learning framework leveraging Vision Transformers and MLP-Mixers—two state-of-the-art vision architectures—and systematically evaluate their adaptability to lens classification. Our approach integrates multi-source data augmentation, fine-grained fine-tuning, model ensembling, and optimizer hyperparameter tuning to enhance few-shot generalization. Results: Fine-tuned on HOLISMOKES VI and SuGOHI X, ten transformer- and mixer-based models consistently outperform conventional CNN baselines in classification accuracy. Crucially, they achieve favorable trade-offs between inference speed and model complexity, demonstrating scalability and robustness for large-scale survey operations. This work establishes a new, extensible, and highly robust paradigm for automated lens detection—directly supporting dark matter mapping and precision cosmological parameter inference.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Gravitational lensing offers a powerful probe into the properties of dark matter and is crucial to infer cosmological parameters. The Legacy Survey of Space and Time (LSST) is predicted to find O(10^5) gravitational lenses over the next decade, demanding automated classifiers. In this work, we introduce GraViT, a PyTorch pipeline for gravitational lens detection that leverages extensive pretraining of state-of-the-art Vision Transformer (ViT) models and MLP-Mixer. We assess the impact of transfer learning on classification performance by examining data quality (source and sample size), model architecture (selection and fine-tuning), training strategies (augmentation, normalization, and optimization), and ensemble predictions. This study reproduces the experiments in a previous systematic comparison of neural networks and provides insights into the detectability of strong gravitational lenses on that common test sample. We fine-tune ten architectures using datasets from HOLISMOKES VI and SuGOHI X, and benchmark them against convolutional baselines, discussing complexity and inference-time analysis.

Problem

• Automating detection of gravitational lenses in large surveys
• Leveraging transfer learning with Vision Transformers and MLP-Mixer
• Evaluating classification performance across architectures and datasets

Innovation

• Vision Transformer transfer learning
• MLP-Mixer architecture integration
• Fine-tuning multiple pretrained models