Physics-Informed Neural Networks Meet Gravitational Lensing: How D4LensPINN Decodes Dark Matter Structures with Deep Learning

This article introduces the open-source project D4LensPINN, which combines Physics-Informed Neural Networks (PINN) with D4 equivariant deep learning to identify dark matter substructure types from gravitational lensing images. Its core highlights include: 1) Integrating physical laws with deep learning to ensure physical consistency of predictions; 2) Using D4 equivariant convolution to improve model efficiency and symmetry; 3) Outperforming traditional baseline models in classification accuracy; 4) Analyzing the symmetry behavior of internal network representations through mechanistic interpretability studies.

The gravitational lensing effect is a key tool for detecting dark matter distribution, but distinguishing dark matter substructures (smooth distribution, spherical cold dark matter clumps, vortex-like warm dark matter structures) is crucial for understanding the nature of dark matter. Traditional methods rely on complex physical modeling and manual feature extraction, while standard CNNs ignore physical laws, leading to models lacking interpretability and physical consistency.

The core innovations of D4LensPINN include:

1. Differentiable Physics Engine: Implements a parameter-free gravitational lensing equation module (Poisson solver, deflection field calculation, inverse lensing layer) to ensure physical self-consistency;
2. D4 Equivariant Convolution: Uses a D4 equivariant U-Net to ensure symmetric transformation of outputs when inputs are transformed, reducing parameters and learning physically correct representations. The model architecture is a four-stage pipeline: Physical Preprocessing → D4 Equivariant U-Net Convergence Estimation → Differentiable Physics Engine → EfficientNetV2 Classification Head.

Experimental results on a dataset of 30,000 images show:

• D4LensPINN achieves a macro-average AUC of 0.9786 (without TTA) / 0.9809 (with D4-TTA), significantly outperforming the ResNet18 baseline (0.9182);
• Introducing four physical losses (total variation regularization, L1 sparsity, center penalty, Poisson residual) improves generalization ability and physical interpretability.

The model is deeply analyzed through three studies:

1. Activation Patching: Intervening on activation values of key layers to identify layers that play a decisive role in classification;
2. Linear Probing: Discovering that equivariant layers extract high-quality physical features earlier than traditional layers;
3. Symmetry Validation: Confirming that the actual equivariance error of equivariant layers is within the range of numerical precision.

The project provides detailed documentation (ARCHITECTURE.md, TRAINING.md, etc.), with code hosted on GitHub, including notebooks for training, control experiments, and interpretability analysis. The dataset is distributed via Google Drive, and the code supports automatic download and caching, while resolving compatibility issues between escnn and numpy (specifying numpy==1.26.4).

D4LensPINN provides a reliable tool for dark matter research and demonstrates the potential of equivariant neural networks in scientific computing. In the future, this physics-data hybrid modeling approach can be extended to fields such as fluid dynamics and materials science, promoting more 'physics-aware' deep learning models to facilitate scientific discoveries.