METAL LAB

What AstroPT knows about galaxies, and what that can teach us about LLMs

arXiv:2608.226142026-08-22

Researchers used a galaxy-image AI (AstroPT) as an answer key to test whether interpretability tools that probe language models actually find real concepts

The team trained AstroPT, a GPT/BERT-style transformer, on millions of galaxy images and proposed it as a calibration testbed for interpretability methods normally applied to language models. Because astronomy already has known ground-truth difficulty and relationships among galaxy properties, they could check whether a model's learned representations emerge in the correct order and encode the correct relationships. They found that galaxy properties become decodable in the expected difficulty order, and that the probe directions reproduce the known signs of physical relationships between properties.

METAL LAB explanatory visual

Using AstroPT as an answer key to calibrate interpretability methods

Evidence statusMeasured results reported

  1. Input: galaxy image patchesMillions of galaxy images split into ordered patch sequences, trained self-supervised with GPT-like (AR) or BERT-like (MAE) objectives
  2. Known difficulty ladderr-band magnitude (near-direct from pixels), redshift (needs multiple bands), and sSFR (inferred) form a difficulty order known in advance from astrophysics
  3. Probing for emergence orderLinear probes measure when each property becomes decodable across training progress and network depth, checking it matches the known difficulty order
  4. Probing for relationship geometryCosine similarity between probe directions checks whether known sign relationships (luminosity-mass, sSFR-mass, redshift-mass) appear inside the model
  5. Causal check on mass residualTests whether the part of mass not explained by luminosity is separately represented, and uses activation patching to probe whether that information is causally used
An explanatory diagram made by METAL LAB, not a figure supplied by the paper's authors.

What they did

  1. Quantities nearly readable straight from pixels (band magnitude) became decodable early in training and in shallow layers, while quantities requiring multi-band or inferred information (redshift, specific star formation rate) emerged later in training and deeper in the network.
  2. This emergence order held across both GPT-like autoregressive (AR) and BERT-like masked-autoencoding (MAE) training objectives, and stayed fixed across model sizes (1M/21M/100M parameters), with larger models only raising scores, not reordering them.
  3. Comparing the direction vectors learned by linear probes (simple linear regressions that read a specific property out of the model's hidden representations), luminosity and stellar-mass directions were strongly aligned, sSFR and mass were anti-aligned, and redshift and mass were positively aligned, matching the known sign structure from galaxy physics.
  4. The 'residual' part of stellar mass that luminosity alone cannot predict was represented in a direction nearly orthogonal to the luminosity direction, suggesting the model does not treat mass merely as a stand-in for luminosity.
Figure 1: Overview of AstroPT. Galaxy images are split into ordered patch sequences and used to train a GPT (pictured here) or BERT-style transformer via a patch reconstruction objective. The learned embeddings are then used for downstream probes of galaxy properties.
Figure 1: Overview of AstroPT. Galaxy images are split into ordered patch sequences and used to train a GPT (pictured here) or BERT-style transformer via a patch reconstruction objective. The learned embeddings are then used for downstream probes of galaxy properties.
Figure 2: Concepts emerge in a fixed, difficulty-ordered sequence. Held-out mag​r, redshift, and sSFR R2 as a function of pre-training progress and model depth. Directly recoverable photometric information (r-band magnitude) is decoded earliest and most strongly; redshift follows more weakly; specific star formation rate remains weakly retrievable after one epoch at our model parameter counts. Likewise, our directly recoverable quantities peak in shallow layers, while integrated and inferred quantities peak deeper. The order is fixed across model size, with capacity raising the attained score rather than reordering the properties.
Figure 2: Concepts emerge in a fixed, difficulty-ordered sequence. Held-out mag​r, redshift, and sSFR R2 as a function of pre-training progress and model depth. Directly recoverable photometric information (r-band magnitude) is decoded earliest and most strongly; redshift follows more weakly; specific star formation rate remains weakly retrievable after one epoch at our model parameter counts. Likewise, our directly recoverable quantities peak in shallow layers, while integrated and inferred quantities peak deeper. The order is fixed across model size, with capacity raising the attained score rather than reordering the properties.
Table 1: Final-checkpoint bootstrap metrics across objective and model-size choices. Here ℓr=−0.4​Mr is the absolute r-band log-luminosity proxy, M⋆ is stellar mass, and ϵM|ℓr is the residual mass at fixed luminosity. Intervals are 95% bootstrap intervals.
ConfigR2​(ℓr)R2​(M⋆)R2​(ϵM|ℓr)R2​(log⁡sSFR)R2​(z)
AR 1M0.422 [0.420, 0.424]0.535 [0.533, 0.536]0.101 [0.099, 0.104]0.258 [0.256, 0.260]0.410 [0.408, 0.412]
AR 21M0.499 [0.496, 0.501]0.610 [0.608, 0.612]0.109 [0.104, 0.114]0.298 [0.295, 0.302]0.503 [0.499, 0.506]
AR 100M0.530 [0.526, 0.533]0.639 [0.636, 0.641]0.110 [0.103, 0.116]0.309 [0.304, 0.313]0.573 [0.568, 0.576]
MAE 1M0.356 [0.354, 0.358]0.454 [0.452, 0.456]0.085 [0.082, 0.087]0.248 [0.245, 0.250]0.328 [0.326, 0.330]
MAE 21M0.350 [0.347, 0.354]0.446 [0.443, 0.449]0.070 [0.065, 0.075]0.228 [0.224, 0.232]0.323 [0.319, 0.327]
MAE 100M0.386 [0.381, 0.390]0.488 [0.484, 0.491]0.080 [0.075, 0.085]0.246 [0.243, 0.250]0.349 [0.343, 0.354]
Figure 3: Probe geometry recovers the signs of known galaxy relationships. Across model sizes, luminosity and stellar-mass directions are strongly aligned, sSFR and stellar-mass directions are anti-aligned, redshift and stellar-mass directions are positively aligned, and the mass residual is much less aligned with luminosity than stellar mass is, matching the known sign structure of these relationships.
Figure 3: Probe geometry recovers the signs of known galaxy relationships. Across model sizes, luminosity and stellar-mass directions are strongly aligned, sSFR and stellar-mass directions are anti-aligned, redshift and stellar-mass directions are positively aligned, and the mass residual is much less aligned with luminosity than stellar mass is, matching the known sign structure of these relationships.
Figure 4: Label-space galaxy relationships used as baselines for the representation-geometry analysis. Left: specific star formation rate is anti-correlated with stellar mass. Middle: redshift is positively correlated with stellar mass. Right: absolute r-band luminosity is strongly correlated with stellar mass. These label-space relationships set the expected signs for the probe-direction analysis in the main text.
Figure 4: Label-space galaxy relationships used as baselines for the representation-geometry analysis. Left: specific star formation rate is anti-correlated with stellar mass. Middle: redshift is positively correlated with stellar mass. Right: absolute r-band luminosity is strongly correlated with stellar mass. These label-space relationships set the expected signs for the probe-direction analysis in the main text.

Findings

  • r-band magnitude (a directly observable quantity) was decoded earliest and most strongly, in both training progress and shallow layers; redshift followed more weakly; sSFR remained only weakly decodable within one training epoch (Fig. 2).
  • This difficulty ordering was invariant across both AR and MAE training objectives and across all tested model sizes (1M/21M/100M parameters), with capacity raising average probe performance rather than reordering properties.
  • Probe direction comparisons showed luminosity and stellar-mass directions strongly aligned, sSFR and stellar-mass directions anti-aligned, and redshift and stellar-mass directions positively aligned, matching the known sign structure of galaxy relationships (Fig. 3, compared against label-space baselines in Fig. 4).
  • The mass residual not explained by luminosity was much harder to decode (lower R2) than luminosity or stellar mass, and its probe direction was nearly orthogonal to the luminosity direction, far less aligned than the stellar-mass direction is (Fig. 10, Fig. 11, Table 1).
  • In an activation-patching experiment swapping activations between galaxy pairs with similar luminosity but different mass residuals, the residual prediction moved toward the source galaxy more often than in same-residual or shuffled-probe control conditions (Fig. 12).
Figure 5: Training-time relationship geometry under the AR objective. Left: the sSFR direction is anti-aligned with the stellar-mass direction through training. Right: the redshift direction is positively aligned with the stellar-mass direction through training. The signs are also stable across model sizes, suggesting that the relationships are not confined to our fully pre-trained models.
Figure 5: Training-time relationship geometry under the AR objective. Left: the sSFR direction is anti-aligned with the stellar-mass direction through training. Right: the redshift direction is positively aligned with the stellar-mass direction through training. The signs are also stable across model sizes, suggesting that the relationships are not confined to our fully pre-trained models.
Figure 6: Training-time relationship geometry under the MAE objective. Left: the sSFR direction is anti-aligned with the stellar-mass direction through training. Right: the redshift direction is positively aligned with the stellar-mass direction through training. The signs are also stable across model sizes, suggesting that the relationships are not confined to our fully pre-trained models.
Figure 6: Training-time relationship geometry under the MAE objective. Left: the sSFR direction is anti-aligned with the stellar-mass direction through training. Right: the redshift direction is positively aligned with the stellar-mass direction through training. The signs are also stable across model sizes, suggesting that the relationships are not confined to our fully pre-trained models.

Where it can be used

  • Testing methods that ask when concepts emerge during training and at what network depth, first in a domain with known ground truth (astronomy) before applying them to language models.
  • Validating whether linear probe directions actually capture real relationships between concepts, a methodology that could extend to other scientific foundation models.
  • Using galaxy datasets as a benchmark for causal-tracing techniques such as activation patching.
Figure 7: Layer-wise relationship geometry under the AR objective. The luminosity and stellar-mass directions remain strongly aligned across layers, the sSFR and stellar-mass directions remain anti-aligned, and the redshift direction is positively aligned with both mass and luminosity. This suggests the relationship geometry is distributed across layers rather than confined to the final layer.
Figure 7: Layer-wise relationship geometry under the AR objective. The luminosity and stellar-mass directions remain strongly aligned across layers, the sSFR and stellar-mass directions remain anti-aligned, and the redshift direction is positively aligned with both mass and luminosity. This suggests the relationship geometry is distributed across layers rather than confined to the final layer.
Figure 8: Layer-wise relationship geometry under the MAE objective. Luminosity and stellar mass remain strongly aligned across layers, sSFR and stellar mass remain anti-aligned, and redshift remains positively aligned with both mass and luminosity. The same sign structure appears under MAE and AIM, supporting the claim that the recovered geometry is not tied to a single training objective.
Figure 8: Layer-wise relationship geometry under the MAE objective. Luminosity and stellar mass remain strongly aligned across layers, sSFR and stellar mass remain anti-aligned, and redshift remains positively aligned with both mass and luminosity. The same sign structure appears under MAE and AIM, supporting the claim that the recovered geometry is not tied to a single training objective.

Limits and open work

  • Galaxies are not language and image patches are not discrete tokens; the paper does not claim astronomical data itself is useful for pre-training LLMs.
  • The analysis is correlational, and the physical labels used as ground truth are themselves catalog-inferred values that may carry their own error.
  • Pre-training was limited to one epoch and a relatively small sweep of model sizes ({1M, 21M, 100M} parameters).
  • Linear probes only reveal information that is linearly accessible in the tested embedding spaces, potentially missing nonlinearly encoded structure.
  • The activation-patching experiment on the mass residual is described as exploratory and would need further validation.
Figure 9: Empirical mass–luminosity relation used to define the residual mass at fixed luminosity. The horizontal axis is a log-luminosity proxy from absolute r-band magnitude, ℓr=−0.4​Mr. The dashed line is the linear fit on the training split, and color shows the residual ϵM|ℓr.
Figure 9: Empirical mass–luminosity relation used to define the residual mass at fixed luminosity. The horizontal axis is a log-luminosity proxy from absolute r-band magnitude, ℓr=−0.4​Mr. The dashed line is the linear fit on the training split, and color shows the residual ϵM|ℓr.
Figure 10: Residual mass-at-fixed-luminosity remains weak across objective choices. Each panel reports final-checkpoint, best-layer R2 for the luminosity proxy ℓr, stellar mass M⋆, and the residual ϵM|ℓr, with error bars showing 95% bootstrap confidence intervals (300 resamples). Luminosity and stellar mass are consistently easier to decode, while the residual remains weak across AR/MAE objectives.
Figure 10: Residual mass-at-fixed-luminosity remains weak across objective choices. Each panel reports final-checkpoint, best-layer R2 for the luminosity proxy ℓr, stellar mass M⋆, and the residual ϵM|ℓr, with error bars showing 95% bootstrap confidence intervals (300 resamples). Luminosity and stellar mass are consistently easier to decode, while the residual remains weak across AR/MAE objectives.

Why it matters

Interpretability research on language models asks when concepts emerge and whether probes recover real structure, but language offers no ground truth to check these claims against. Astronomy provides a domain where the difficulty and relationships among concepts are already known, giving researchers an answer key to calibrate the same interpretability tools before applying them blindly to LLMs.

Figure 11: Residual probe directions are separated from the dominant luminosity-mass axis, with error bars showing 95% bootstrap confidence intervals (300 resamples). Across objective choices, luminosity and stellar-mass directions remain strongly aligned. The residual direction is much less aligned with luminosity than the stellar-mass direction is, indicating that the residual probe is not reading out the same luminosity direction. The separation is strongest in the larger AR models and weaker in the MAE setting.
Figure 11: Residual probe directions are separated from the dominant luminosity-mass axis, with error bars showing 95% bootstrap confidence intervals (300 resamples). Across objective choices, luminosity and stellar-mass directions remain strongly aligned. The residual direction is much less aligned with luminosity than the stellar-mass direction is, indicating that the residual probe is not reading out the same luminosity direction. The separation is strongest in the larger AR models and weaker in the MAE setting.
Figure 12: Causal tracing of the mass-luminosity residual. Patching activations between galaxy pairs with similar luminosity but different mass residual moves the residual prediction toward the source more often than same-residual and shuffled-probe controls.
Figure 12: Causal tracing of the mass-luminosity residual. Patching activations between galaxy pairs with similar luminosity but different mass residual moves the residual prediction toward the source more often than same-residual and shuffled-probe controls.

Terms in this paper

  • AstroPT · A transformer trained on galaxy images split into patch sequences, using GPT-like or BERT-like self-supervised objectives
  • linear probe · A simple linear regression fit on a model's frozen hidden representations to test whether a target property can be read out
  • redshift · A measure of how much a galaxy's light has stretched due to cosmic expansion, requiring multiple photometric bands or spectra to estimate
  • specific star formation rate (sSFR) · Star formation rate divided by stellar mass, an inferred and indirect physical quantity
  • cosine similarity · A measure of how aligned two direction vectors are; positive means similar direction, negative means opposite

Original abstract (English)

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.

Authors · UniverseTBD

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: UniverseTBD et al., arXiv:2608.22614, CC BY-SA 4.0