Hyrax: An Extensible Framework for Rapid ML Experimentation and Unsupervised Discovery in the Era of Rubin, Roman, and Euclid

Abstract: The NSF-DOE Vera C. Rubin Observatory, Roman Space Telescope, Euclid, and other next-generation surveys will deliver imaging, spectroscopic, and time-domain data at scales that increasingly shift the bottleneck in astronomical ML projects from model design to infrastructure. We present Hyrax, an open-source, modular, GPU-enabled Python framework that supports the full ML lifecycle in astronomy: from data acquisition and training to inference and experiment comparison, with capabilities including multimodal dataset support, integrated vector databases for similarity search, and interactive two- and three-dimensional latent-space exploration for unsupervised discovery. We demonstrate Hyrax's versatility through five representative applications on real survey data: (i) unsupervised representation learning on $\sim 4\times10^5$ Rubin Legacy Survey of Space and Time (LSST) Data Preview 1 (DP1) galaxies, surfacing new merger and low-surface-brightness candidates missing from reference Euclid and Dark Energy Survey catalogs, while also isolating imaging artifacts -- all without labeled training data; (ii) hybrid density-based clustering for identifying cluster-scale gravitational lens candidates in DP1 data; (iii) multimodal early-time transient classification in the Zwicky Transient Facility leveraging light curves, spectra, images, and metadata; (iv) supervised false-positive filtering in shift-and-stack searches for distant solar system objects in the Dark Energy Camera Ecliptic Exploration Project survey; and (v) supervised detection of semi-resolved dwarf galaxies in Hyper Suprime-Cam and LSST-like imaging using synthetic source injection. Together, these results demonstrate that Hyrax provides astronomy-specific ML infrastructure that enables systematic discovery and rapid methodological iteration across next-generation astronomical surveys.

• The paper introduces Hyrax, a modular Python framework that streamlines survey-scale ML workflows through standardized data handling, experiment management, and interactive latent space visualization.
• It achieves up to 10x training speedup via efficient data caching and accelerates unsupervised discovery with astronomy-aware similarity search for uncovering rare objects.
• The framework’s integration with PyTorch, MLflow, and Optuna offers a reproducible, scalable platform that addresses critical technical challenges in large-scale astronomical research.

Hyrax: An Extensible Framework for Survey-Scale ML and Unsupervised Discovery in Astronomy

Introduction and Motivation

The exponential growth in data from modern astronomical surveys such as the Vera C. Rubin Observatory (LSST), Roman Space Telescope, and Euclid mission has rendered existing, project-specific ML pipelines insufficient for scalable, systematic astronomical analysis. The paper introduces Hyrax, an open-source, modular Python framework designed to standardize and streamline the ML lifecycle for astronomical research, specifically optimized for both supervised and unsupervised workflows targeting vast and heterogeneous data modalities.

Hyrax directly addresses the reality that the ML model itself represents only a fraction of the technical stack required for successful scientific deployment. Surrounding infrastructure—including scalable data acquisition, high-throughput data ingestion, experiment tracking, similarity search, and 2D/3D latent space visualization—is a persistent bottleneck for both individual investigators and survey-scale collaborations. Hyrax's design explicitly targets this technical debt (Figure 1).

Figure 2: The ML model constitutes a small fraction of a typical ML astronomy workflow; critical infrastructure is required for data, monitoring, caching, and scientific visualization.

Framework Design and Architecture

Hyrax implements a verb-oriented, modular API that abstracts core workflow stages—data acquisition (hyrax.download()), dataset construction, training (hyrax.train()), inference (hyrax.infer()), similarity search, and interactive visualization (hyrax.visualize()). The architecture utilizes established open-source components for ML (PyTorch), experiment tracking (MLflow, Optuna), and vector databases (LanceDB, ChromaDB), overlaying these with domain-specific support for FITS, survey formats, and multimodal data.

Figure 4: Architectural schematic of Hyrax, highlighting modular stages from parallelized data acquisition to 3D latent space visualization.

Dataset acquisition leverages native support for major astronomical archives (e.g., Rubin Butler), importing and constructing on-disk caches of survey patches and dynamically generating image/spectra cutouts for GPU-resident model training. The LSSTDataset abstraction achieves high training throughput and reproducibility by caching at the patch—rather than cutout—level and enabling streaming ingestion via multi-threaded pipelines (Figure 5).

Figure 1: Efficient data acquisition and patch-level caching in LSST-scale training using Hyrax's dataset abstraction.

Data ingestion is further decoupled from acquisition logic via the DataProvider architecture, which supports heterogeneous, multimodal fusion (images, spectra, time series) and runtime-optimized memory use, demonstrated to speed training epochs by up to an order of magnitude (Figure 6).

Figure 3: Data flow orchestration during training/inference with flexible caching and fusion.

Training, Hyperparameter Optimization, and Experiment Management

Hyrax centralizes model training via a configuration-driven interface that enables rapid prototyping and scaling to HPC/cloud environments. Models and datasets are modular, with full user control—custom PyTorch models and dataset definitions are natively supported. Automated hyperparameter optimization via Optuna is integrated, as is systematic experiment tracking using MLflow (Figure 7), enabling visualization, filtering, and quantitative comparison across diverse experiments and hyperparameter sweeps.

Figure 5: Native experiment management and metric tracking through MLflow integration.

Unsupervised Latent Space Analysis, Similarity Search, and Visualization

The framework makes a significant methodological contribution by introducing scalable, astronomy-aware similarity search and interactive latent space exploration. Model outputs are indexed in state-of-the-art vector databases, enabling approximate nearest neighbor queries in $\mathcal{O}(\log n)$ time for datasets of tens of millions of sources. The scientific value of this is demonstrated by the ability to surface rare or scientifically distinct populations, such as mergers and low-surface-brightness candidates, without relying on labeled data (see Figure 8).

Figure 6: Example of similarity search in latent space, retrieving morphologically consistent nearest neighbors.

For interpretability and validation, Hyrax provides interactive latent-space visualizers operating both within Jupyter notebooks (2D; Figure 9) and browser-based (3D; Figure 10). These tools enable users to select regions of the embedding, retrieve multimodal data, and correlate latent structure with astrophysical properties—a capability hitherto lacking in general-purpose ML astronomy tools.

Figure 7: 2D Latent Space Explorer supports interactive, in-notebook inspection and scientific labeling.

Figure 8: 3D Latent Space Explorer for interactive navigation of high-dimensional embeddings.

Empirical Applications and Science Demonstrations

Hyrax is empirically validated on multiple survey-scale tasks drawn from diverse astronomical domains:

• Unsupervised representation learning ($N \sim 4\times10^5$ galaxies, Rubin DP1): Emergent clusters corresponding to scientifically meaningful populations (mergers, LSB galaxies) are recovered without labels. A critical finding is that training loss does not monotonically correlate with scientific utility; interactive inspection is essential for model selection (Figure 11, Figure 12).
• Hybrid unsupervised+manual lensing arc searches: Combination of UMAP embedding and DBSCAN clustering enables scalable identification of gravitational lens candidates.
• Supervised multimodal transient classification (AppleCiDEr): Integrated light curves, spectra, images, and metadata, showing enhanced performance over unimodal approaches and revealing the limitations imposed by class imbalance in rare-event regimes.
• Distant solar system object search (KBMOD integration): Hyrax's dataset and tracking infrastructure enables filtering with a 72.2% rejection rate for candidate tracklets, while recovering 95.2% of synthetic injected signals at scale.
• Semi-resolved dwarf galaxy search: Supervised CNNs trained with synthetic source injection achieve a false positive rate of 0.2%, and the method is prepared for deployment on future LSST data.

Numerical and Methodological Highlights

• Training speed-up: DataProvider architecture yields an observed $\sim10\times$ speedup in training epochs through aggressive on-disk/memory caching.
• Unsupervised discovery: The latent space visualizers, coupled with similarity search, can iteratively “peel off” artifacts and systematics, achieving near-artifact-free galaxy samples without labeled data (Figure 11).
• Experiment management: Automated configuration and hyperparameter tracking via MLflow allow transparent and reproducible large-scale experimentation.
• Rare object recovery: In task-specific pipelines, unsupervised discovery recovers populations missing from prior visual catalogs, and in the case of TNO detection, filtering efficiency and completeness are both quantitatively reported.

  Figure 9: Successive rounds of unsupervised peeling isolate imaging artifacts, leaving a purified sample of astrophysical sources.

  

  Figure 10: UMAP projections illustrate that validation loss is a poor proxy for scientifically coherent clustering; visual tools are essential for meaningful model selection.

Implications, Limitations, and Future Directions

Practically, Hyrax lowers the technical barrier for research groups entering survey-scale ML by condensing the infrastructure stack into a distributable, community-maintained framework. This enables more of the methodological burden to be absorbed by astronomers, thereby facilitating reproducibility and interoperability as the cadence and heterogeneity of data increase.

Theoretically, the explicit coupling of unsupervised ML with scalable similarity search and programmatic visualization establishes new workflows for unsupervised discovery. The finding that scalar metrics (e.g., loss) are not reliable guides for representation learning outcomes will influence evaluation protocols for upcoming foundation models and domain-specific SSL approaches.

Hyrax's architecture is positioned for extension to real-time alert ingestion (e.g., transient surveys), collaborative model/dataset sharing, and tight integration with survey science platforms and cloud-native appliance deployment—directions that will be pivotal as the LSST and Roman alert rates come online.

Conclusion

Hyrax constitutes a formal convergence of modern ML infrastructure, domain-specific astronomy tooling, and scalable unsupervised discovery. By abstracting infrastructure complexity, providing reproducible and efficient pipelines, and enabling interactive latent space workflows, Hyrax accelerates methodological iteration and deployability for the next generation of survey astronomy. Its empirical validation across multiple astrophysical domains and its strong experiment-tracking and interpretive toolsets establish it as a core platform for forthcoming foundation model ecosystems in astronomy.

Reference: Hyrax: An Extensible Framework for Rapid ML Experimentation and Unsupervised Discovery in the Era of Rubin, Roman, and Euclid (2605.18959).