AstroPT Recovers Galaxy Properties in Known Difficulty Order, a Calibration Test for LLM Probes

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

UniverseTBD, :, Kshitij Duraphe, Aman Kumar, Michael J. Smith, Shashwat Sourav

COLM 2026

cs.LG, astro-ph.IM

2026-08-24

AstroPT treats galaxy images as patch tokens. Linear probes decode magnitude early, redshift later, sSFR weakly; the order holds across objectives and 1M–100M scale.

What problem this solves

Two interpretability claims are hard to falsify in language models. When do concepts arise during training, and do linear probes recover real structure? Language does not come with a ground-truth difficulty order or a signed map of how concepts relate. Astronomy does. Apparent magnitude is close to integrated flux on the pixels. Redshift needs multiple bands or spectra. Specific star-formation rate is more indirect still. Mass tracks luminosity, sSFR anti-correlates mildly with mass on the star-forming main sequence, and flux-limited samples correlate redshift with mass. Those signs are known before any probe is fit.

AstroPT is built to look like an LLM. A galaxy image is split into ordered patches, treated as tokens, and trained by self-supervised reconstruction, either GPT-style autoregression or BERT-style masking. The aim is not to discover new galaxy laws. It is to ask whether the probe tools we run on language models blind can recover known structure when an answer key exists.

Method

The data are 8.7 million DESI Legacy Survey postage stamps from Smith42/galaxies. Model sizes are 1M, 21M, and 100M, in the spirit of a Pythia-style checkpoint sweep. Each size is trained with both AR and MAE for one epoch, with checkpoints saved along the way. Frozen hidden states at every layer get a ridge linear probe, reported as held-out R². Linear decodability does not prove the model used the property during pre-training. It shows the information is linearly available.

Probe targets follow an a priori ladder: r-band magnitude, redshift z, and log sSFR. A geometry analysis adds a luminosity proxy ℓr=−0.4 Mr, stellar mass, and the mass residual at fixed luminosity. Directions are cosine similarities of normalized probe weights. The residual asks whether mass is only a luminosity shortcut; if it were, the residual direction would still align with luminosity.

Results

Properties become decodable in the known difficulty order, both through training and along depth. mag r arrives first, strongest, and shallow. Redshift follows and peaks deeper. sSFR stays weak through one epoch at these sizes. Switching AR for MAE does not reorder the sequence. Scaling from 1M to 100M raises the scores without reshuffling the order. The sequence is a property of the learning problem, not of one training recipe.

Final-checkpoint, best-layer R² with 95% bootstrap intervals:

ConfigR²(ℓr)R²(mass)R²(log sSFR)R²(z)R²(residual)
AR 1M0.4220.5350.2580.4100.101
AR 21M0.4990.6100.2980.5030.109
AR 100M0.5300.6390.3090.5730.110
MAE 100M0.3860.4880.2460.3490.080

AR beats MAE across the board. The mass residual stalls at 0.07 to 0.11, so what is linearly readable is mostly the luminosity-mass axis.

Probe directions recover the known signs. At AR 100M, cosine(luminosity, mass) is 0.915, cosine(sSFR, mass) is −0.420, cosine(redshift, mass) is 0.563. Those signs are stable through training and across layers. The residual aligns with luminosity less than mass does, and in the larger AR models the residual cosine turns negative (−0.180 at AR 21M). Mass is not stored as a luminosity alias. An appendix activation-patching check on galaxy pairs matched in luminosity but not residual moves the residual readout toward the source more often than same-residual and shuffled-probe controls.

Why it matters

Astronomy is used here as a calibration sandbox for LLM probes, not as a proposal to pretrain language models on galaxies. When an answer key exists, linear probes recover the right difficulty order and the right signed geometry. That kind of check is hard to copy in language. For mechanistic interpretability the value is calibration: the same checkpoint sweep, layer-wise probes, and direction cosines can be graded against physics. For scientific foundation models it suggests that emergence ordered by physical complexity is a testable training dynamic rather than a story told after the fact.

The sweep stops at 100M and one epoch. Do not read the R² numbers as a galaxy-foundation-model leaderboard.

Limitations

The authors are direct: the analysis is correlational; labels are catalog-inferred, not pixel ground truth; pre-training is one epoch at up to 100M parameters; linear probes see only what is linearly available. Image patches are not a discrete vocabulary, and the paper does not claim astronomical data helps LLM pre-training. Residual R² is low, so orthogonality is relative to the main axis, not evidence that the residual is cleanly encoded. Activation patching is exploratory and lives in the appendix. The same probes are not run on a real LLM checkpoint, so the calibration loop is missing its last step.

Terms

Source

Related papers

All paper explainers