🌌 AION-1: AstronomIcal Omnimodal Network

🎯 Overview

AION-1 is a cutting-edge large omnimodal model specifically designed for astronomical surveys. It seamlessly integrates multiple data modalities, and enables simple adaptation to a wide range of astronomical tasks.

Alpha Testing

AION-1 model weights are hosted on Huggingface behind gates during the alpha testers phase. First, ensure that you have access to the Hugginface model weights. If you don't have access, you can request it directly on the hugginface repo here.

Once you have access, you will need to set up a huggingface token locally. This can be done by first installing hugginface_hub:

pip install huggingface_hub

and then logging in via

huggingface-cli login --token YOUR_HF_TOKEN

All of the ensuing steps should work out of the box after this point.

🚀 Quick Start

Assuming you have PyTorch installed, you can install AION trivially with:

pip install polymathic-aion

Then you can load the pretrained model and start analyzing astronomical data:

import torch
from aion import AION
from aion.codecs import CodecManager
from aion.modalities import LegacySurveyImage

# Load model and codec manager
model = AION.from_pretrained('aion-base').to('cuda')  # or 'aion-large', 'aion-xlarge'
codec_manager = CodecManager(device='cuda')

# Prepare your astronomical data (example: Legacy Survey image)
image = LegacySurveyImage(
    flux=your_image_tensor,  # Shape: [batch, 4, height, width] for g,r,i,z bands
    bands=['DES-G', 'DES-R', 'DES-I', 'DES-Z']
)

# Encode data to tokens
tokens = codec_manager.encode(image)

# Option 1: Extract embeddings for downstream tasks
embeddings = model.encode(tokens, num_encoder_tokens=600)

# Option 2: Generate predictions (e.g., redshift)
from aion.modalities import Z
predictions = model(
    tokens,
    target_mask={'tok_z': torch.zeros(batch_size, 1)},
    num_encoder_tokens=600
)

🔬 Scientific Overview

🧬 Architecture

AION-1 employs a two-stage, transformer-based design:

1. Modality-Specific Tokenizers transform raw inputs into discrete tokens
2. Unified Encoder–Decoder Transformer ingests all token streams via a multimodal masked modeling (4M) objective

🗂️ Supported Modalities

AION-1’s tokenizers cover 39 distinct data types, grouped by survey and data category

📈 Model Variants

Pretraining – Global batch size: 8 192 – Steps: Base (1.5 days on 64 H100), Large (2.5 days on 100 H100), XLarge (3.5 days on 288 H100) – Optimizer: AdamW, peak LR 2 × 10⁻⁴, linear warmup + cosine decay

🔧 Data Preparation

AION uses a typed data system to understand the provenance of each astronomical observation. Each modality must be properly formatted:

Modality Types

from aion.modalities import (
    LegacySurveyImage, HSCImage,           # Images
    DESISpectrum, SDSSSpectrum,            # Spectra
    LegacySurveyFluxG, HSCMagG,            # Photometry
    GaiaParallax, Z,                       # Scalars
    # ... and 30+ more modalities
)

Example: Preparing Legacy Survey Data

import torch
from aion.modalities import LegacySurveyImage, LegacySurveyFluxG

# Format image data (shape: [batch, 4, height, width])
image = LegacySurveyImage(
    flux=torch.tensor(image_data, dtype=torch.float32),
    bands=['DES-G', 'DES-R', 'DES-I', 'DES-Z']
)

# Format scalar photometry
flux_g = LegacySurveyFluxG(value=torch.tensor([flux_values]))

Supported Data Formats

💡 Example Use Cases

🔍 Similarity Search

Find galaxies similar to a query object across different modalities:

# Extract embeddings for similarity search
query_embedding = model.encode(codec_manager.encode(query_image))
all_embeddings = model.encode(codec_manager.encode(*dataset_images))

# Find most similar objects using cosine similarity
from sklearn.metrics.pairwise import cosine_similarity
similarity_scores = cosine_similarity(query_embedding, all_embeddings)
similar_objects = similarity_scores.argsort()[::-1][:10]  # Top 10 similar

📊 Property Prediction

Build lightweight models on AION embeddings:

# Extract embeddings from multiple modalities
embeddings = model.encode(codec_manager.encode(
    image, spectrum, flux_g, flux_r, flux_i, flux_z
), num_encoder_tokens=900)

# Train simple regressor for stellar mass, redshift, etc.
from sklearn.neighbors import KNeighborsRegressor
regressor = KNeighborsRegressor(n_neighbors=5)
regressor.fit(embeddings.mean(axis=1), target_property)

🌌 Generative Modeling

Predict missing astronomical properties:

# Predict redshift from photometry + morphology
predictions = model(
    codec_manager.encode(image, flux_g, flux_r, flux_i, flux_z),
    target_mask={'tok_z': torch.zeros(batch_size, 1)},
    num_encoder_tokens=600
)
redshift_probs = torch.softmax(predictions['tok_z'], dim=-1)

📚 Documentation

🎓 Tutorials

Start with our interactive tutorial:

• Open in Google Colab - Learn AION basics interactively, no local setup required!

For detailed guides, see the online documentation.

📦 Advanced Installation

AION offers flexible installation options to suit your environment and requirements.

To install AION with PyTorch included:

pip install aion[torch]

For contributors and developers:

pip install aion[torch,dev]

This includes testing frameworks, linting tools, and development dependencies.

For specific PyTorch versions (e.g., CUDA support):

# Install PyTorch with CUDA 12.4 support
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124

# Then install AION
pip install aion

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🌟 Acknowledgments

AION is developed by Polymathic AI, advancing the frontier of AI for scientific applications.

📬 Contact

• Issues: GitHub Issues
• Discussions: GitHub Discussions