aion

AION-base is a 300M parameter large omnimodal model specifically designed for astronomical surveys, presented in the paper AION-1: Omnimodal Foundation Model for Astronomical Sciences. It integrates 39 distinct astronomical data types and enables adaptation to a wide range of astronomical tasks through multimodal masked modeling.

Project Homepage: https://polymathic-ai.org/

Model Details

• Architecture: Encoder-Decoder Transformer (12 blocks each, 768 dim, 12 heads)
• Parameters: 300M
• Training: Multimodal Masked Modeling (4M) on astronomical survey data
• Modalities: 39 data types including imaging, spectra, catalogs, and photometry

Installation

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

pip install polymathic-aion

For advanced installation options, including specific PyTorch versions or developer installations, refer to the GitHub repository.

Usage

After installation, 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, Z

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

# Example: Prepare your astronomical data (e.g., a dummy Legacy Survey image)
# In a real scenario, 'your_image_tensor' would come from your dataset.
your_image_tensor = torch.randn(1, 4, 96, 96) # Example: batch_size=1, 4 bands, 96x96 resolution
image = LegacySurveyImage(
    flux=your_image_tensor,
    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)
print(f"Extracted embeddings shape: {embeddings.shape}")

# Option 2: Generate predictions (e.g., redshift)
# For this example, we predict redshift (Z) from the image.
# The target_mask tells the model which modality to generate.
preds = model(
    codec_manager.encode(image),
    target_modality=Z,
)
print(f"Predicted redshift logits shape: {preds['tok_z'].shape}")

Supported Data Types

AION-Base processes data from major astronomical surveys. Here's an overview of the supported categories:

Resources

• GitHub Repository: https://github.com/PolymathicAI/AION

License

This project is licensed under the MIT License. See the LICENSE file in the GitHub repository for full details.

Built with ❤️ for the astronomical community by https://polymathic-ai.org/