The Galaxy's Guide to the Tokenizer: A Benchmark for Scientific Foundation Models

Technical Summary: The Galaxy's Guide to the Tokenizer

Problem Statement
While transformer-based foundation models have become a leading paradigm for extracting scientific insight from astronomical data, the specific impact of tokenization strategies on learned representations remains poorly understood. In scientific domains like astronomy, where observations are governed by known physical processes and accompanied by independently measured ground truth (e.g., redshift, stellar mass), there is a critical need to evaluate how converting raw observations into token sequences shapes the model's ability to perform scientific inference. Current literature often focuses on architectural choices, leaving a gap in understanding whether reconstruction fidelity correlates with the accessibility of physical information within the latent space.

Methodology
The authors conduct a systematic comparison of four distinct tokenization strategies within a unified, controlled framework using the AstroPT backbone, a decoder-only transformer trained autoregressively on astronomical data. The study utilizes 640,000 galaxy images from the DESI Legacy Survey (DR8), processed into 256×256 pixel postage stamps.

The four tokenization strategies evaluated are:

1. Affine: A minimal baseline using a simple linear projection from image patches to the embedding space.
2. AIM (MLP-based): Replaces the linear projection with a Multi-Layer Perceptron (MLP) to introduce non-linearity, following the original AstroPT framework.
3. JetFormer: A flow-based approach using an invertible normalizing flow to produce continuous tokens. It is trained end-to-end to maximize image likelihood without a separate pre-trained encoder.
4. VQ-VAE: A discrete approach that quantizes image patches into a learned codebook of visual tokens. The encoder is pre-trained on in-distribution galaxy imagery and frozen during the main AstroPT training.

Evaluation Protocol
The study employs two primary evaluation axes:

• Physical Property Prediction: Linear and MLP probes are trained on embeddings extracted from the center layer of the AstroPT models to predict 167,000 held-out galaxies' physical properties. These properties include photometric magnitudes, colors, photometric and spectroscopic redshifts, specific star formation rate (sSFR), stellar mass, and morphological indicators.
• Reconstruction Fidelity: Image reconstruction quality is assessed using SSIM and PSNR on 5,000 random test set galaxies.
• Information Content Verification: To distinguish between information loss and inaccessible information, the authors probe the reconstructed images themselves using ImageNet-pretrained ResNet-50 features, isolating the information content of the pixel data from the embedding geometry.

Key Results
The results reveal a significant decoupling between reconstruction fidelity and the quality of representations for physical inference:

• JetFormer: Achieves the highest reconstruction fidelity (Mean PSNR ~31.11 dB, SSIM ~0.76), preserving fine spatial details and low surface brightness components. However, its embeddings are the least linearly accessible for physical property prediction. Probing the reconstructed images confirms that JetFormer retains nearly all original information, but this information is "scattered" in the latent space in a way that simple probes cannot easily recover.
• VQ-VAE: Yields the strongest performance for predicting physical properties, particularly with linear probes. Its discrete bottleneck organizes physical information in a more linearly accessible manner. However, it achieves lower reconstruction fidelity (Mean PSNR ~23.57 dB) and systematically loses fine-grained details (e.g., morphology, absolute photometry), though it preserves color indices.
• Affine and AIM: These methods perform nearly identically across almost all properties, suggesting that the MLP head in AIM adds negligible benefit over a simple linear projection when the transformer backbone is sufficiently expressive. They preserve localized morphological information better than JetFormer but generally underperform VQ-VAE on regression tasks.
• Trade-offs: The study identifies a fundamental trade-off: soft, dense, information-complete representations (JetFormer) favor high-fidelity reconstruction, while hard, patch-based tokenization strategies (VQ-VAE) promote representations that are more directly interpretable and accessible for physical inference.

Significance and Claims
The paper argues that tokenization is not a neutral preprocessing step but a critical design choice that dictates what information is learned and how it is organized. The authors claim that reconstruction quality is not a reliable proxy for the utility of a representation for scientific inference.

By grounding the evaluation in independently measured physical quantities, the study serves as a demonstration of how scientific data can provide objective, non-circular benchmarks for foundation models. The authors posit that the choice of tokenizer should be driven by the specific downstream task: VQ-VAE is preferable for photometric or spectroscopic inference, while JetFormer is superior for image generation or reconstruction tasks. The work highlights the potential of astronomy as a setting for constructing interpretable benchmarks, noting that similar frameworks could be valuable in other domains where physical ground truth is available.

Limitations
The authors acknowledge that the comparison is not entirely ablative due to the necessity of different training procedures for each tokenizer (e.g., pre-training VQ-VAE vs. end-to-end training for JetFormer). Additionally, the reliance on linear and MLP probes may understate the information present in the representations, although probing reconstructed images helps mitigate this for the extreme cases. The study does not explore scaling with dataset size or model capacity.