DeepLense GSoC 2026: Tasks Solution

Author: Yash Gupta
Affiliation: 3rd Year Undergraduate, Majoring in Data Science & Artificial Intelligence (DSAI) at IIT Roorkee
Project Title: Test I. Multi-Class Classification, Specific Test VI. Image Super-resolution

---

2. Executive Summary

This pipeline is build to bridge the gap between "seeing" a galaxy and "understanding" its physics. Instead of a standard black-box model, this project follows the actual lifecycle of an astrophysical observation:

1. Task 1

Before we can reconstruct a lens, we have to find it. A Two-Stage LP-FT classifier is implemented for this that hits a 0.9936 Macro AUC.

2. Task 6A

To master the geometry of lensing, developed the PI-HAF (Physics-Informed Hybrid Attention Flow) backbone. It reached 38.4 dB PSNR in simulation by learning the deterministic "flow" of mass maps rather than just memorizing pixels.

3. Task 6B

Real telescope data is messy. To handle the Sim2Real Gap, used LoRA (Low-Rank Adaptation) to perform a "surgical transfer."

• The trick: Froze the simulated physics and only trained tiny "adapters" to learn real-world noise from HST/HSC.
• The result: Stable reconstruction on real data (26.6 dB with TTA) without "hallucinating" non-physical features.

Note: You'll notice the PSNR drop from 38dB to 26dB—that’s the real-world physics talking. By using LoRA + PI-HAF, We are prioritizing scientific integrity over artificial smoothness.

3. Phase I: Task 1

The first stage of the pipeline is a multi-class classification model. It classifies images into three types: No Substructure, Subhalo Substructure, and Vortex Substructure.

The LP-FT Strategy

To train the model effectively, a two-step approach is used:

1. Linear Probing (LP):
   The main model (backbone) is kept frozen, and only the final classification layer is trained for 5 epochs. This helps the model adapt to the lensing task without changing the pre-trained features too much.

2. Full Fine-Tuning (FT):
   Then, the entire model is trained for 30 epochs with different learning rates for different parts. This helps the model learn fine details like small patterns and distortions in the images.

Final Validation Metrics

The model demonstrates exceptional discriminative power across all three classes (see output folder for more details):

Class	AUC (OVR)	Precision	Recall
No Substructure	0.9959	0.956	0.984
Subhalo (Sphere)	0.9905	0.958	0.935
Vortex	0.9943	0.970	0.964
Macro Average	0.9936	0.961	0.961

---

Note: The model has very high recall (0.984) for the "No Substructure" class. This means it rarely misses important lensing cases, ensuring that useful data is not lost before the Super-Resolution stage.

4. The PI-HAF Backbone

The Physics-Informed Hybrid Attention Flow (PI-HAF) is the heart of this project. While standard UNets often "hallucinate" pixel-level details, PI-HAF is built to respect the actual geometry of gravitational lensing.

Why it’s a Novel Approach:

• FlowHAB (Hybrid Attention): Instead of just convolutions, I used Flow-based Hybrid Attention Blocks. This allows the model to capture local pixel details and the global curvature of Einstein rings simultaneously.
• Cross-Attention Guidance: The Low-Resolution image acts as a "physical anchor." Through Cross-Attention, the model is forced to generate high-resolution features that stay strictly aligned with the observed lensing arcs.
• Rectified Flow Logic: Rather than a simple pixel-to-pixel map, PI-HAF learns a velocity field to transform noise into a clean galaxy distribution. This makes the reconstruction much smoother and mathematically grounded.
• LoRA-Ready Design: The modular nature of these attention blocks is what makes the Task 6B Surgical Transfer possible. It allows us to "patch" the model for real-world noise without breaking the physics it learned in simulation.

Note: I chose this over a standard UNet because astrophysical data isn't just a picture, it’s a coordinate-based physics problem. PI-HAF treats it like one.

5. Phase II: Task 6A

Before addressing the complexities of real telescope data, the PI-HAF backbone was rigorously trained on a given dataset of strong lensing images. This phase establishes the theoretical upper bound of the architecture's performance in a noise-controlled environment.

Benchmarking Synthetic Excellence

In Task 6A, the model learns the deterministic mapping between Low-Resolution (LR) and High-Resolution (HR) lensing arcs. By leveraging the Rectified Flow objective, the model achieves near-perfect reconstruction of the galactic mass distribution.

Key Performance Metrics (Validation):

• Peak PSNR: 38.40 dB
• SSIM: 0.9117
• Sampler Steps: 50 (ODE Integration)
• Evaluation Samples: 1,500 pairs

Note: The 38.4 dB result is significantly higher than the standard 30-32 dB seen in generic SR models. This is due to the PI-HAF's Cross-Attention mechanism, which uses the LR input as a hard physical constraint, preventing the model from hallucinating non-physical galactic features.

6. Phase III: Task 6B

The most challenging phase of the project involves bridging the Sim2Real Gap: applying the physics learned from 10,000+ simulations (Task 6A) to a sparse, noisy dataset of 300 real HST/HSC telescope images.

The LoRA Mechanism: Surgical Adaptation

Instead of standard fine-tuning, which would lead to "catastrophic forgetting" of the lensing physics, I implemented a Surgical Transfer using Low-Rank Adaptation (LoRA).

• Key Idea: The Task 6A backbone (3.85M params) is frozen to preserve lensing geometry, while lightweight LoRA adapters (Rank=16) are added on top using Task 6B weights to learn real-world noise and instrument effects.

Sim2Real Success Metrics

On the real telescope dataset, the model achieved stable, high-fidelity results.

Metric	Baseline (Standard)	TTA Enhanced (Final)	Delta (Boost)
PSNR (dB)	26.00 dB	26.60 dB	+0.60 dB
SSIM	0.5251	0.5946	+0.0695

Analyzing the Sim2Real Gap

The transition from 38.4 dB (Sim) to 26.6 dB (Real) is a documented scientific reality. Real telescope images contain irreducible aleatoric uncertainty from cosmic rays, atmospheric turbulence, and sensor bleeding.

Key Achievement: Achieving an SSIM of nearly 0.60 on real data proves that the model is recovering true structural galactic features rather than merely smoothing pixels to "cheat" the PSNR metric.

7. Advanced Engineering & Optimization

To push the boundaries of the Sim2Real transfer, I implemented two key engineering optimizations that stabilize the training of LoRA adapters and maximize the perceptual quality of the final reconstructions.

Test-Time Augmentation (TTA)

• Key Idea: TTA creates multiple transformed versions of each image and averages their predictions, improving robustness and giving a consistent boost in PSNR and SSIM, while an OOM-safe wrapper ensures it runs efficiently on 8GB GPUs.

Stochastic Weight Averaging (SWA)

• Key Idea: SWA averages model weights over the final training epochs to reduce overfitting and improve generalization, leading to more stable performance on real-world data.

8. Installation, Environment & Reproducibility

To ensure full reproducibility of the results across Task 1 and Task 6, this repository follows a strict configuration-driven design. All dependencies are managed via a dedicated Conda environment.

1. Environment Setup (Recommended)

The fastest way to reproduce this environment is to use the provided environment.yml file:

# Clone and enter repo
git clone https://github.com/Yash-g2310/ml4sci.git
cd ml4sci

# Create the environment from the file
conda env create -f environment.yml
conda activate deeplense

# Additional: Create a weights folder if you want to verify them or can run cells in nb from start
mkdir weights

# Now download and add weights from google drive

2. Weight Layering Logic (Critical)

The weights are added in google drive. The Super-Resolution weights follow a hierarchical loading strategy. You must have both Task 6A and Task 6B weights to run the real-world inference:

• Task 6A Base: weights/task6a_reflow_best.pth (Full PI-HAF Backbone)
• Task 6B Adapters: weights/task6b_swa_final.pth (LoRA Adapters)

Note: When running Task6B_Real_SR.ipynb, the script automatically loads the 6A backbone first, then "injects" the 6B LoRA adapters into the attention layers. Attempting to load 6B weights without the 6A backbone will result in a shape mismatch.

---

📂 9. Repository Organization & File Mapping

The repository is organized into modular components.

Directory Structure

.
├── configs/                # Centralized YAML configurations for all tasks
├── notebooks/              # Executable research environments (Task 1, 6A, 6B)
├── src/                    # Core source code
│   ├── data_utils/         # Custom Dataset and Transform pipelines
│   ├── engine/             # Training loops, ODE samplers, and TTA logic
│   └── models/             # PI-HAF backbone and LoRA surgery scripts
├── weights/                # Model checkpoints (Hierarchical loading)
├── outputs/                # Generated reports, ROC curves, and SR grids
├── requirements.txt        # Python dependency list
└── README.md               # Project documentation

---

Final Note to GSoC 2026 Maintainers: All notebooks are self-contained and pre-configured to point to the correct relative paths within this structure. For any technical inquiries regarding the PI-HAF implementation or the Sim2Real transfer strategy, please refer to the detailed docstrings within the src/ directory.