Source-linked AI summary
The Galaxy's Guide to the Tokenizer: A Benchmark for Scientific Foundation Models
Sogol Sanjaripour, Michael J. Smith, Manuel Pérez-Carrasco, Juan Rafael Martínez-Galarza, Bahram Mobasher, Gabriela Canalizo
TL;DR
Tokenization’s effect on learned scientific representations remains poorly understood. This study compares four strategies in a shared astronomical-imaging transformer framework and finds that reconstruction and physical-property representation quality are decoupled, with different tokenizers excelling on different tasks.
Problem
Tokenization’s influence on learned representations in scientific foundation models remains poorly understood.
Method
The study compares four tokenization strategies within a shared AstroPT transformer, evaluating reconstruction and physical-property prediction using galaxy images and probes.
Results
Reconstruction fidelity and accessible physical information are decoupled: JetFormer reconstructs best, VQ-VAE produces the most informative embeddings, and Affine and AIM perform well on morphology-driven features.
Takeaways & Limitations
Tokenizer choice should follow the intended downstream task rather than reconstruction metrics alone, while physics-grounded benchmarks provide objective evaluation of representations.
Takeaways & Limitations
The comparison is not purely ablative because tokenizer training procedures differ, and VQ-VAE’s performance may partly reflect its pretrained encoder.
Abstract
from arXiv · showhide
Tokenization is central to adapting scientific data for transformer-based foundation models, yet its impact on learned representations remains poorly understood. We compare four tokenization strategies, Affine, AIM, JetFormer, and VQ-VAE, within a unified transformer framework for astronomical imaging. Using 640,000 galaxy images from the DESI Legacy Survey and a shared AstroPT backbone, we evaluate each method on reconstruction fidelity and prediction of physical properties. Our results reveal trade-offs across approaches. The flow-based JetFormer achieves higher reconstruction quality, while VQ-VAE yields strong probe performance for galaxy physical properties. Affine and AIM better preserve localized morphological information. We find that reconstruction and representation quality are decoupled, and no single method consistently performs best across the tasks considered here. By grounding our evaluation in independently measured physical quantities, we hope this study serves to highlight the potential of scientific data as a basis for constructing interpretable benchmarks for foundation models.
1. From pixels to physics
This section frames tokenization as an underexamined determinant of transformer representations in astronomy and motivates evaluating it against independently measured physical quantities. It introduces a controlled comparison of four tokenization strategies to characterize their effects on scientific inference.
- Astronomical surveys are producing imaging data at unprecedented scale, while transformer-based foundation models have emerged as a leading paradigm for extracting scientific insight.
- Tokenization converts observations into transformer-processable sequences, yet its influence on learned representations has received comparatively little attention.
- Astronomy enables physically grounded representation evaluation through independently measured quantities such as redshift, stellar mass, and star formation rate.
- The study systematically compares Affine, AIM-like, Jetformer-like, and VQ-VAE-based tokenization within a controlled framework to characterize effects on learned representations and scientific inference.The strategies are respectively linear projection, MLP-based, flow-based, and VQ-VAE-based.
2. AstroPT backbone, tokenization strategies, and experimental setup
The study compares four diverse tokenization strategies within a shared autoregressive AstroPT decoder-only transformer, using DESI-LS galaxy images and standardized tests of reconstruction and physical-property prediction. The comparison follows each tokenizer’s realistic training procedure rather than a fully uniform ablation.
- AstroPT backbone: AstroPT provides a neutral shared backbone: an autoregressive, decoder-only transformer trained on sequences of astronomical data with minimal inductive bias.This architecture reflects the dominant paradigm in current foundation models.
- Tokenization strategies: The four tokenizers span linear and non-linear mappings, continuous and discrete representations, and deterministic and probabilistic decoding.Affine uses linear projections; AIM uses an MLP with patch-wise z-score normalization; JetFormer uses an invertible flow; VQ-VAE uses a learned codebook.
- Experimental considerations: The comparison is not entirely ablative because each tokenizer requires its own competitive training procedure, so the study follows published methodologies as faithfully as possible.Forcing a fully uniform regime could create unrealistic and unfavorable settings for some methods.
- Data selection and pre-processing: 640 000 galaxies are selected from an 8.6 million-galaxy DESI Legacy Surveys Data Release 8 dataset for training.The survey provides homogeneous, wide-field imaging suited to large-scale representation learning.
- Evaluation: Linear and MLP probes predict physical properties for 167 000 held-out galaxies, while SSIM and PSNR assess reconstruction on 5 000 random test galaxies.Properties cover photometry, morphology, and spectroscopy; each property uses k=10-fold cross-validation.
3. Does tokenization affect latent physical knowledge?
Tokenization changes both the accessibility and allocation of latent physical information, with VQ-VAE generally strongest for probes and JetFormer strongest for reconstruction. Reconstruction fidelity and representation quality are therefore decoupled across tokenizers.
- Physical-property probing: VQ-VAE achieves the strongest linear and MLP probe performance for most physical properties, with linear probes usually exceeding MLP probes.This suggests its discrete bottleneck organizes physical information into a more linearly accessible representation.
- Physical-property probing: JetFormer outperforms AIM and Affine on linear probes for photometric and spectrally related quantities, whereas AIM performs better on morphology-driven parameters.JetFormer’s linear and MLP probe performances are within the error spread of each other, while patch-based methods preserve localized spatial information.
- Physical-property probing: Across all tokenizers and probe types, g−r color is predicted more accurately than r −z.Although g−r is more dust-sensitive, its larger dynamical range and stronger variation over the training set make it an easier regression target.
- Representational capacity: JetFormer dedicates more representational capacity to apparent photometry and structural parameters, reaching R2 = 0.33 ± 0.02 and R2 = 0.29 ± 0.02.VQ-VAE reaches R2 = 0.19 ± 0.01 and R2 = 0.16 ± 0.01 for the corresponding groups.
- Reconstruction fidelity: On 5 000 held-out images, JetFormer achieves mean PSNR of 31.11 dB (σ = 3.54) and mean SSIM of 0.762 (σ = 0.138).VQ-VAE attains mean PSNR of 23.57 dB (σ = 0.93) and SSIM of 0.544 (σ = 0.054).
- Reconstruction versus representation: JetFormer reconstructions retain nearly all original-image information, whereas VQ-VAE reconstructions systematically lose fine morphology, specific star formation rate, and absolute photometric scale.Thus, VQ-VAE makes physical content accessible in embedding geometry but discards pixel detail, while JetFormer preserves pixels but scatters information in its embedding.
4. Trade-offs, limitations, and the case for physics-grounded benchmarks
The comparison shows that tokenization shapes learned representations, with reconstruction fidelity decoupled from accessible physical information. These trade-offs motivate downstream-task-based tokenizer selection, while physics-grounded benchmarks and methodological improvements remain important.
- Trade-offs: JetFormer achieves the best reconstruction, whereas VQ-VAE produces the most informative embeddings for physical properties.This demonstrates a decoupling between reconstruction fidelity and the physical information accessible in representations.
- Trade-offs: AIM and Affine occupy a middle ground, performing well on morphology-driven features while the transformer backbone largely drives representation learning.
- Tokenizer selection: Tokenizer choice should follow the eventual downstream task rather than reconstruction metrics alone.VQ-VAE may benefit photometric or spectroscopic inference, JetFormer image generation or reconstruction, and Affine compute-constrained settings.
- Limitations: Linear and MLP probes may understate representation information, and tokenizer training differences make the comparison not purely ablative.Probing reconstructions supports JetFormer’s information accessibility, but a fully probe-independent comparison across all four tokenizers remains open.
- Physics-grounded benchmarks: Astronomy enables objective tokenizer benchmarks through known physics and independently measured quantities, avoiding reliance solely on human-annotated tasks or toy datasets.The framework is intended to inform similar evaluations in domains with available physical ground truth.
A. Embedding Structure via PCA and UMAP
PCA and UMAP reveal coherent physical-property structure in all galaxy embeddings, including distinct colour–distance and morphology–luminosity axes. Tokenization methods differ in geometry, with Affine and AIM fragmenting under UMAP while JetFormer and VQ-VAE retain connected manifolds.
- Shared embedding structure: PCA and UMAP projections reveal coherent gradients in photometric redshift, g −r colour, r-band magnitude, and smoothness fraction.The strongest relationship is between g −r colour and photometric redshift.
- Shared embedding structure: The latent space separates colour and distance from morphology and luminosity, with smoother early-type galaxies tending toward brighter r-band magnitudes.These two semantically distinct axes recover known scaling relations directly from imaging.
- Method-specific geometry: Under PCA, Affine and AIM form open horseshoe arcs, JetFormer an elongated crescent, and VQ-VAE a compact filled lobe.These patterns indicate distinct projected embedding geometries across tokenization methods.
- Method-specific geometry: Under UMAP, Affine and AIM fragment into disconnected islands, whereas JetFormer and VQ-VAE retain a single connected manifold with persistent physical gradients.The gradients remain coherent even across the fragmented UMAP clusters.
- Affine–AIM similarity: Affine and AIM have visually similar embedding geometries, and their apparent transposition reflects PCA sign and rotation ambiguity rather than physical meaning.Their matching UMAP island structures and probe scores support the hypothesis that the MLP tokeniser head adds little beyond a linear projection.
B. Computational Cost and Carbon Footprint
Training the tokenization strategies incurred differing computational costs, with Affine and AIM requiring 3 hours each on one NVIDIA H100 GPU, while JetFormer required 8 hours on two NVIDIA H100 GPUs. The study also estimates energy consumption and commits to releasing trained models and code to reduce redundant retraining.
- Training Cost and Carbon Footprint: Affine and AIM each trained for 3 hours on a single NVIDIA H100 80 GB GPU, consuming an estimated 2.1 kWh.The authors plan to release fully trained models and code upon deanonymization to avoid redundant retraining emissions.
- Training Cost and Carbon Footprint: Energy consumption was estimated from reported GPU TDP values across hardware configurations, providing a common comparison beyond wall-clock runtime.Affine and AIM used one NVIDIA H100 80 GB GPU, JetFormer used two NVIDIA H100 80 GB GPUs, and VQ-VAE used one NVIDIA Quadro GV100 32 GB GPU.