Model reference · open weights
aion is an open-weight language model from polymathic-ai. AxForge deploys and operates it for you on dedicated EU-owned hardware — with the licence handled where one is required.
Available as managed deployment — configured and operated for you on dedicated EU hardware, quoted per deployment.
What it is
| Released by | polymathic-ai |
|---|---|
| Type | Language models |
| Task | Omni (any→any) |
| Parameters (lead) | 318M |
| Runs with | aion |
| Released | 2025-05-16 |
| Popularity | 2k downloads / month |
| Licence | Open weights |
About
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/
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.
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}")
AION-Base processes data from major astronomical surveys. Here's an overview of the supported categories:
| Category | Description | Token Name(s) |
|---|---|---|
| Imaging (2) | Legacy Survey, HSC Wide | tok_image_ls, tok_image_hsc |
| Catalog (1) | Legacy Survey catalog entries | catalog |
| Spectra (2) | SDSS, DESI | tok_spectrum_sdss, tok_spectrum_desi |
| Gaia (4) | BP/RP spectra, parallax, sky coords | tok_xp_bp, tok_xp_rp, tok_parallax, tok_ra, tok_dec |
| Gaia Photometry (3) | G/BP/RP flux | tok_flux_g_gaia, tok_flux_bp_gaia, tok_flux_rp_gaia |
| Legacy Survey (9) | g,r,i,z bands & WISE W1–W4 flux, E(B–V) | tok_flux_g,…,tok_flux_w4, tok_ebv |
| Legacy Shape (3) | Ellipticity components & effective radius | tok_shape_e1, tok_shape_e2, tok_shape_r |
| HSC Photometry (5) | g,r,i,z,y magnitudes | tok_mag_g,…,tok_mag_y |
| HSC Extinction (5) | g,r,i,z,y extinctions | tok_a_g,…,tok_a_y |
| HSC Shape (3) | Shape components 11,22,12 | tok_shape11, tok_shape22, tok_shape12 |
| Other (1) | Spectroscopic redshift | tok_z |
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/
From the published model card. Full card on the HuggingFace links in the sidebar.
How it works
Using it via the API
Once AxForge deploys aion for you, it answers on the OpenAI-compatible API — the same base URL and keys as every other model. (aion below is illustrative; you get the exact model name on deployment.)
$ curl -sS https://api.axforge.ai/v1/chat/completions \
-H "Authorization: Bearer $AXFORGE_API_KEY" \
-H "Content-Type: application/json" \
-d '{"model":"aion","messages":[{"role":"user","content":"Hello"}]}'
Create an account — your API key is available in the console. 3M free tokens every 30 days with every new account.