Source-linked AI summary
What AstroPT knows about galaxies, and what that can teach us about LLMs
UniverseTBD, :, Kshitij Duraphe, Aman Kumar, Michael J. Smith, Shashwat Sourav
TL;DR
Interpretability claims about concept emergence and linear probes are difficult to validate in language models because language lacks known ground-truth orderings and relationships. The paper uses AstroPT, an LLM-like transformer trained on galaxy images, to test these methods against astronomical structure; galaxy properties emerge in a difficulty-related sequence that is stable across tested objectives and model capacities, while probe directions recover known physical relationships.
Problem
Language models lack ground-truth orderings of concepts and relationships, making claims about emergence and probe-recovered structure difficult to validate.
Method
The study probes frozen AstroPT representations across checkpoints, layers, model sizes, and GPT-like or BERT-like patch-reconstruction objectives using linear probes of galaxy properties.
Results
Galaxy properties emerge in a fixed order matching their known difficulty across training and depth, invariant to tested objectives and sequence-preserving across model sizes, while probe directions recover known physical sign relationships.
Takeaways & Limitations
Astronomy provides a controlled testbed for calibrating developmental and probe-based interpretability methods applied to LLMs.
Takeaways & Limitations
The analysis is correlational, uses catalog-inferred labels, one epoch, a limited parameter sweep, and probes that reveal only linearly accessible structure.
Abstract
from arXiv · showhide
Interpretability research increasingly asks when concepts emerge during training and whether linear probes recover real structure, but in language models these claims are hard to validate because language offers little ground-truth ordering of concepts or relationships among them. We propose the use of astronomical ground truth through AstroPT, a transformer trained on millions of galaxy images, as a calibration testbed. AstroPT is an LLM-like model trained within a domain where the difficulty ordering of concepts and the relations among them are known in advance. Probing frozen representations across checkpoints, layers, model sizes, and objective choices, we find that galaxy properties emerge in a fixed order that tracks their known difficulty---quantities written almost directly into the pixels (band magnitude) become decodable early in training and shallow in the network, while multiband/spectra based and inferred quantities (such as redshift and specific star formation rate) emerge later and deeper. This order is invariant to our tested training objectives, and scales in magnitude but not in sequence with capacity. Our linear probe directions further recover the known physical structure among galaxy properties. Our findings suggest that astronomy offers a controlled sandbox for calibrating mechanistic interpretability methods we otherwise apply to LLMs blind.
1 Astronomy for machine learning
Astronomy provides an LLM-like setting with known difficulty orderings and physical relationships, allowing interpretability methods to be tested against scientific ground truth. AstroPT uses galaxy-image patches and reconstructive training to probe whether representations recover that structure and whether emergence patterns persist across training choices.
- Astronomy supplies known relationships among galaxy properties, enabling tests of whether hidden representations recover structure beyond individual labels.
- Galaxy properties form a difficulty ladder: pixel-close quantities such as apparent magnitude are easier to recover than redshift or specific star formation rate.
- AstroPT splits galaxy images into ordered patches treated as tokens and trains GPT-like or BERT-like transformers by reconstructing dropped image chunks.
- The study uses astronomy as a controlled testbed to examine concept emergence, latent-space geometry, and robustness of probe-based interpretability methods.
- Galaxy properties emerge in a fixed difficulty-related order during pre-training and network depth, with directly recoverable quantities appearing earlier and shallower.
2 Training and probing LLM-like models with astronomical data
AstroPT probes reveal that galaxy properties emerge in a fixed order aligned with their known difficulty, across training progress, network depth, objectives, and model sizes. Probe directions also recover known physical relationships, while the study’s scope is bounded by several methodological limitations.
- Training setup: AstroPT models trained on 8.7M galaxy images with GPT-like and BERT-like objectives across 1M, 21M, and 100M parameters.The models reconstruct dropped image patches as an astronomical analogue of token prediction.
- Emergence order: r-band magnitude emerges earliest and strongest, redshift later and more weakly, while sSFR remains weak after one epoch.The emergence order tracks the properties’ known distance from low-level pixel statistics.
- Emergence order: The same difficulty ordering appears across network depth: photometric features peak shallow, whereas integrated and inferred quantities become accessible deeper.Redshift strengthens in later layers, while sSFR remains weak across depth at the tested capacity.
- Robustness: Changing between autoregressive and masked-autoencoding objectives preserves the emergence order, while greater capacity raises probe performance without reordering properties.This pattern indicates that the sequence reflects the galaxy-learning problem rather than the tested pre-training objective.
- Relational geometry: Probe directions recover known galaxy-relationship signs: luminosity aligns with stellar mass, sSFR anti-aligns, and redshift aligns positively.The mass residual is much less aligned with luminosity than stellar mass is.
- Limitations: The analysis is correlational, uses catalog-inferred labels, covers one epoch and a limited parameter sweep, and measures only linearly accessible structure.These constraints bound how far the findings generalize beyond the tested embedding spaces and configurations.
Contributions
The contributions describe a team effort spanning software, experiment design, pre-training code, figures, equations, inference, writing, and project conceptualization.
- Contributions: KD developed embedding-processing and extraction software, investigated property selection and tokenization, and helped review and edit the manuscript.
- Contributions: AK ran experiments before finalizing the experiment code and reviewed and edited the manuscript.
- Contributions: MJS wrote the AstroPT pre-training code, created Figures 2 and 3, wrote the paper body, and guided project work.
- Contributions: SS selected astronomy equations, contributed to writing and review, made Figure 1, performed inference, and prepared multiple figures.
- Contributions: All authors conceptualized the project.
A Linear probe training and alignment details
The study trains linear ridge probes on frozen galaxy representations and compares normalized probe directions to test whether hidden-space geometry preserves known galaxy relationships.
- Probe training: At each checkpoint and layer, linear ridge probes predict catalog quantities from each galaxy’s representation.Targets include mag r, an absolute r-band luminosity proxy, stellar mass, sSFR, and redshift.
- Probe targets: The luminosity proxy is defined as ℓr ≡ −0.4Mr, with the additive constant absorbed by the linear-fit intercept.This follows from the magnitude–luminosity relation Mr = −2.5 log10 Lr + C.
- Alignment analysis: Normalized probe weights define model-space directions, whose cosine similarity measures whether two properties are represented similarly, oppositely, or orthogonally.Positive, negative, and near-zero cosine values correspond to aligned, anti-aligned, and orthogonal directions.
A.1 Supporting relational-geometry plots
The supporting plots establish label-space galaxy relationships as baselines for evaluating whether probe directions recover the expected physical geometry.
- Label-space baselines: Figure 4 uses three label-space relationships as baselines: sSFR–mass anti-correlation, redshift–mass positive correlation, and luminosity–mass strong correlation.These relationships set the expected signs for the representation-geometry analysis.
- Physical origins: Luminosity and stellar mass are strongly related because more massive galaxies host more stars, providing the expected positive baseline.
- Physical origins: The mildly negative sSFR–mass trend follows from the sub-linear star-forming main sequence and the definition sSFR ≡ SFR/M⋆.The rising quiescent fraction at high stellar mass reinforces this trend, though the effects are weak relative to sSFR’s range.
- Physical origins: The positive redshift–mass relationship primarily reflects Malmquist-style selection, because only intrinsically luminous galaxies exceed a flux limit at higher redshift.
A.2 Training-time stability of relationship geometry
The signs of probe-direction relationships remain stable throughout training under both AR and MAE objectives, including anti-alignment between sSFR and stellar mass and positive alignment between redshift and stellar mass.
- Under AR, sSFR remains anti-aligned with stellar mass while redshift remains positively aligned with stellar mass throughout training.
- The relationship signs are stable across model sizes, suggesting they are not confined to fully pretrained models.
- Under MAE, the same anti-alignment and positive-alignment signs persist throughout training.
A.3 Layer-wise relationship geometry
AstroPT preserves known galaxy-property relationship geometry across network layers, rather than restricting it to the final layer.
- Across layers, luminosity and stellar mass remain strongly aligned, sSFR and stellar mass remain anti-aligned, and redshift remains positively aligned with both.
A.4 Final-checkpoint bootstrap metrics
Final-checkpoint bootstrap metrics support stable physical relationships across objectives and model sizes, with aligned luminosity–mass directions and characteristic signs for sSFR and redshift relative to stellar mass.
- Final-checkpoint metrics are reported across training objectives and model-size choices with 95% bootstrap intervals.
- Luminosity and stellar mass remain strongly aligned across configurations.
- sSFR and stellar mass remain anti-aligned, while redshift and stellar mass remain positively aligned across configurations.
B Residual mass-at-fixed-luminosity and activation patching
The residual mass at fixed luminosity isolates mass variation not explained by the empirical mass–luminosity relation, and activation patching provides evidence that this residual is represented separately from luminosity.
- The mass residual is defined as log M⋆ − (aℓr + b), using a training-split linear mass–luminosity fit.
- The residual measures whether a galaxy is more or less massive than expected from r-band luminosity alone, reflecting variation in mass-to-light structure.
- Across configurations, luminosity and stellar mass decode much better than the residual, which remains weak across objective choices.
- Residual and stellar-mass probe directions are separated from the dominant luminosity–mass axis, especially in larger autoregressive models.
- Patching activations between galaxies with similar luminosity but different mass residual moves residual predictions toward the source more often than control patches.