Spherinator & HiPSter

SPACE Training Webinar

Bernd Doser (HITS)

May 2024

Outline

• FAIR Development Setup
• Spherinator: The Model
• Spherinator: The Training
• HiPSter: The Inference
• HiPSter: The Workflow
• Summary and Outlook

FAIR Principles for Research Software

• Findability

  Software, and its metadata, is easy to find for humans and machines

• Accessibility

  Software, and its metadata, is retrievable via standardized protocols

• Interoperability

  Software interoperates with other software through standards

• Reusability

  Software is usable and reusable

Source: (Barker et al. 2022)

FAIR: Development Setup

• GitHub: Repository and project management
• GitHub Actions: Continuous integration
  • Static code analysis
  • Unit tests
• Poetry: Python dependency management
• Dependabot: Automated security / version updates
• ONNX: Open Standard for machine learning interoperability
• Weights&Biases: Visualizing and tracking machine learning experiments

Spherinator: The Model

• Representation learning by dimensionality reduction
• Convolutional Variational Autoencoder (VAE)
• Rotational invariance by using set of rotated images, where only the one with the lowest reconstruction loss is used for the training
• (Hyper-)spherical latent space

Source: (Polsterer et al. 2024, in press)

Power Spherical Distribution

\[\begin{aligned} p_{X}(x; \mu, \kappa) = N_{X}(\kappa, d)^{-1}(1 + \mu^{\top}x)^{\kappa} \end{aligned}\]

\(d\): Dimension

\(\mu\): Direction

\(\kappa\): Concentration

\(N_{X}\): Normalization factor

Source: (Cao and Aziz 2020)

Spherinator: Loss Function

\[\begin{aligned} L = L_{recon} + \lambda \cdot L_{KL} \end{aligned}\]

\(L_{recon}\): Reconstruction loss by pixel-wise euclidean distance

\(L_{KL}\): Kullback-Leibler divergence

• Euclidean space: Gaussian reference distribution
• Hyper-spherical space: Uniform reference distribution

\(\lambda\): Balancing factor

Spherinator: The Training

PyTorch Lightning Command Line Interface (CLI):

seed_everything: 42

model:
  class_path: models.RotationalVariationalAutoencoderPower
  init_args:
    encoder:
      class_path: models.ConvolutionalEncoder
    decoder:
      class_path: models.ConvolutionalDecoder
    h_dim: 256
    z_dim: 3
    image_size: 363
    rotations: 36
    beta: 1.0e-3

data:
  class_path: data.IllustrisSdssDataModule
  init_args:
    data_directories: ["TNG100", "TNG50", "Illustris"]
    extension: fits
    minsize: 100
    batch_size: 512
    shuffle: True
    num_workers: 16

optimizer:
  class_path: torch.optim.Adam
  init_args:
    lr: 0.001

lr_scheduler:
  class_path: lightning.pytorch.cli.ReduceLROnPlateau
  init_args:
    mode: min
    factor: 0.1
    patience: 5
    cooldown: 5
    min_lr: 1.e-5
    monitor: train_loss
    verbose: True

trainer:
  max_epochs: 200
  accelerator: gpu
  devices: 4
  precision: 32

  logger:
    class_path: lightning.pytorch.loggers.WandbLogger
    init_args:
      project: spherinator
      log_model: True
      entity: ain-space
      tags:
        - illustris

  callbacks:
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        monitor: train_loss
        filename: "{epoch}-{train_loss:.2f}"
        save_top_k: 3
        mode: min
        every_n_epochs: 1
    - class_path: callbacks.LogReconstructionCallback
      init_args:
        num_samples: 6

Monitoring with Weights&Biases

• Tracking experiments
• Collaborative project management (teams, reports)
• Model registry (publish models for staging and production)
• Hyperparameter optimization
• Free for academic research

HiPSter: The Inference

HiPS - Hierarchical Progressive Survey

• The HiPS tiles are reconstructed with the corresponding spherical latent space position.

• The HiPS tilings are visualized with Aladin-Lite.

Source: (Fernique et al. 2015)

HiPSter: The Workflow

Summary and Outlook

• Spherinator & HiPSter promising solution for classifying large datasets
• Easy feature extension due to modern software development techniques
• Evaluation of encoder, decoder, loss function for better reconstruction
• Orchestration of the HiPSter workflow
• Geometric Deep Learning: 3D point clouds instead of images
• Prototype and source code are publicly available

Acknowledgement & Disclaimer

References

Barker, Michelle, Neil P. Chue Hong, Daniel S. Katz, Anna Lena Lamprecht, Carlos Martinez-Ortiz, Fotis Psomopoulos, Jennifer Harrow, et al. 2022. “Introducing the FAIR Principles for Research Software.” Scientific Data 9 (December). https://doi.org/10.1038/s41597-022-01710-x.

Cao, Nicola De, and Wilker Aziz. 2020. “The Power Spherical Distribution.” Proceedings of the 37th International Conference on Machine Learning, INNF+, June. http://arxiv.org/abs/2006.04437.

Fernique, P., M. G. Allen, T. Boch, A. Oberto, F-X. Pineau, D. Durand, C. Bot, et al. 2015. “Hierarchical Progressive Surveys: Multi-Resolution HEALPix Data Structures for Astronomical Images, Catalogues, and 3-Dimensional Data Cubes.” Astronomy &Amp; Astrophysics 578 (June): A114. https://doi.org/10.1051/0004-6361/201526075.

Polsterer, Kai L., Bernd Doser, Andreas Fehlner, and Sebastian Trujillo-Gomez. 2024, in press. “Spherinator and HiPSter: Representation Learning for Unbiased Knowledge Discovery from Simulations.” Astronomical Data Analysis Software and Systems XXXIII, 2024, in press.