README.md

---
license: apache-2.0
tags:
- vision-transformer
- image-classification
- pytorch
- timm
- cait
- gravitational-lensing
- strong-lensing
- astronomy
- astrophysics
datasets:
- parlange/gravit-j24
metrics:
- accuracy
- auc
- f1
paper:
- title: "GraViT: A Gravitational Lens Discovery Toolkit with Vision Transformers"
  url: "https://arxiv.org/abs/2509.00226"
  authors: "Parlange et al."
model-index:
- name: CaiT-b3
  results:
  - task:
      type: image-classification
      name: Strong Gravitational Lens Discovery
    dataset:
      type: common-test-sample
      name: Common Test Sample (More et al. 2024)
    metrics:
    - type: accuracy
      value: 0.8430
      name: Average Accuracy
    - type: auc
      value: 0.8265
      name: Average AUC-ROC
    - type: f1
      value: 0.5613
      name: Average F1-Score
---

# 🌌 cait-gravit-b3

🔭 This model is part of **GraViT**: Transfer Learning with Vision Transformers and MLP-Mixer for Strong Gravitational Lens Discovery

🔗 **GitHub Repository**: [https://github.com/parlange/gravit](https://github.com/parlange/gravit)

## 🛰️ Model Details

- **🤖 Model Type**: CaiT
- **🧪 Experiment**: B3 - J24-all-blocks
- **🌌 Dataset**: J24
- **🪐 Fine-tuning Strategy**: all-blocks



## 💻 Quick Start

```python
import torch
import timm

# Load the model directly from the Hub
model = timm.create_model(
    'hf-hub:parlange/cait-gravit-b3',
    pretrained=True
)
model.eval()

# Example inference
dummy_input = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    output = model(dummy_input)
    predictions = torch.softmax(output, dim=1)
print(f"Lens probability: {predictions[0][1]:.4f}")
```

## ⚡️ Training Configuration

**Training Dataset:** J24 (Jaelani et al. 2024)  
**Fine-tuning Strategy:** all-blocks


| 🔧 Parameter | 📝 Value |
|--------------|----------|
| Batch Size | 192 |
| Learning Rate | AdamW with ReduceLROnPlateau |
| Epochs | 100 |
| Patience | 10 |
| Optimizer | AdamW |
| Scheduler | ReduceLROnPlateau |
| Image Size | 224x224 |
| Fine Tune Mode | all_blocks |
| Stochastic Depth Probability | 0.1 |


## 📈 Training Curves

![Combined Training Metrics](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/training_curves/CaiT_combined_metrics.png)


## 🏁 Final Epoch Training Metrics

| Metric | Training | Validation |
|:---------:|:-----------:|:-------------:|
| 📉 Loss | 0.0188 | 0.0494 |
| 🎯 Accuracy | 0.9935 | 0.9882 |
| 📊 AUC-ROC | 0.9997 | 0.9981 |
| ⚖️ F1 Score | 0.9934 | 0.9882 |


## ☑️ Evaluation Results

### ROC Curves and Confusion Matrices

Performance across all test datasets (a through l) in the Common Test Sample (More et al. 2024):

![ROC + Confusion Matrix - Dataset A](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_a.png)
![ROC + Confusion Matrix - Dataset B](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_b.png)
![ROC + Confusion Matrix - Dataset C](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_c.png)
![ROC + Confusion Matrix - Dataset D](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_d.png)
![ROC + Confusion Matrix - Dataset E](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_e.png)
![ROC + Confusion Matrix - Dataset F](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_f.png)
![ROC + Confusion Matrix - Dataset G](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_g.png)
![ROC + Confusion Matrix - Dataset H](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_h.png)
![ROC + Confusion Matrix - Dataset I](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_i.png)
![ROC + Confusion Matrix - Dataset J](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_j.png)
![ROC + Confusion Matrix - Dataset K](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_k.png)
![ROC + Confusion Matrix - Dataset L](https://huggingface.co/parlange/cait-gravit-b3/resolve/main/roc_confusion_matrix/CaiT_roc_confusion_matrix_l.png)

### 📋 Performance Summary

Average performance across 12 test datasets from the Common Test Sample (More et al. 2024):

| Metric | Value |
|-----------|----------|
| 🎯 Average Accuracy | 0.8430 |
| 📈 Average AUC-ROC | 0.8265 |
| ⚖️ Average F1-Score | 0.5613 |


## 📘 Citation

If you use this model in your research, please cite:

```bibtex
@misc{parlange2025gravit,
      title={GraViT: Transfer Learning with Vision Transformers and MLP-Mixer for Strong Gravitational Lens Discovery}, 
      author={René Parlange and Juan C. Cuevas-Tello and Octavio Valenzuela and Omar de J. Cabrera-Rosas and Tomás Verdugo and Anupreeta More and Anton T. Jaelani},
      year={2025},
      eprint={2509.00226},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2509.00226}, 
}
```

---


## Model Card Contact

For questions about this model, please contact the author through: https://github.com/parlange/