sreshwarprasad/hierarchical-gz2

Purpose and entry points

This is an offline ML research pipeline for probabilistic Galaxy Zoo 2 morphology prediction from galaxy images. It trains image models to emit per-question vote distributions rather than hard labels, then evaluates distribution validity, accuracy, calibration, uncertainty, and attention-based interpretability.

Primary execution surfaces are training in src/train.py and src/train_single.py, experiment orchestration in src/ablation.py, and shell launch integration in scripts/training/run_model.sh.

Data preparation and dataset boundary

Raw Galaxy Zoo/image assets are converted, labeled, and validated before model training. The data boundary is handled by scripts/data/convert_zenodo_images.py, scripts/data/prepare_labels.py, and src/dataset.py; dataset verification is available in scripts/data/verify_dataset.py.

The dataset loader supplies transformed image tensors and hierarchical vote-distribution targets to the training/evaluation stack. External dataset/model hosting is Hugging Face; source imagery is associated with Galaxy Zoo 2.

Hierarchical probabilistic model

The core predictor is a Vision Transformer-oriented regression model in src/model.py. Its architectural contract is to produce probability distributions over answer branches for Galaxy Zoo questions, respecting each question’s simplex constraint rather than independent unconstrained scalar regression.

Alternative modeling and baseline paths are centralized in src/baselines.py, enabling comparison with ResNet, ViT MSE-only, and Dirichlet-style approaches. Experiment variants are selected through YAML configurations such as configs/base.yaml, configs/full_train.yaml, and configs/ablation.yaml.

Optimization and experiment execution

Training consumes batches from src/dataset.py, forwards them through src/model.py, and applies the composite probabilistic objective in src/loss.py: KL divergence plus MSE. src/train.py is the main multi-stage training workflow, while src/train_single.py supports a single configured run.

Ablation control in src/ablation.py varies important training choices, notably loss weighting/configuration. Run configuration is externalized in configs/, keeping experiment parameters separate from model and loss implementation.

Evaluation, validity, and calibration

Full evaluation is orchestrated by src/evaluate_full.py, using shared calculations from src/metrics.py. This boundary measures MAE/RMSE and probability-quality properties, including calibration and simplex violations—critical because valid per-question distributions are a core model guarantee.

The evaluation path compares proposed and baseline models, with generated manuscript artifacts retained under outputs/manuscript/03_baseline_comparison/ and outputs/manuscript/05_calibration/.

Interpretability and uncertainty stages

Attention inspection is a distinct post-training analysis stage in src/attention_viz.py, producing ViT attention rollout/entropy artifacts represented in outputs/manuscript/06_interpretability/.

Uncertainty analysis is implemented in src/uncertainty_analysis.py; its outputs relate predicted uncertainty to prediction error and are published under outputs/manuscript/07_uncertainty/. These stages consume trained-model predictions rather than altering the training path.

Output collection and publication artifacts

scripts/training/collect_outputs.py consolidates experiment results. Analysis scripts such as scripts/analysis/generate_comparison_tables.py and scripts/analysis/plot_training_curves.py transform collected metrics/checkpoints into tables and figures.

outputs/manuscript/ is the committed research-result boundary, with outputs/manuscript/MANIFEST.csv indexing final artifacts. The repository itself does not contain the described web demo; the deployed demo is an external Vercel service.