A Mechanistic Interpretability Study of an Astronomy Foundation Model
Abstract
Foundation models are increasingly used across the sciences, but whether their internal computations correspond to known scientific operations or proceed through different mechanisms remains unclear. Understanding these mechanisms is essential for assessing reliability, identifying failure modes, and learning new science from the models. We present, to our knowledge, the first causal mechanistic interpretability study of an astronomy foundation model, AION-1, focusing on its multimodal advantage in redshift prediction, where combining imaging with photometry outperforms photometry alone. We find that this advantage is largest for spatially extended and dust-attenuated galaxies, and that most of it is localized to the first of the encoder's twelve blocks. Within that block, how well a property is linearly decodable does not predict whether it can be localized. Spatial extent is causally used, and one attention head dominates linear probing for it, yet ablating that head has negligible effect on redshift prediction: the information is distributed across heads and is decodable again one layer after we erase it. Dust attenuation shows the opposite pattern: it is only weakly decodable anywhere in the encoder, yet a single head is necessary for most of the multimodal advantage, particularly for the most dust-attenuated galaxies, and keeping that head as the only active head in its block preserves nearly all of the advantage. None of the astrophysically motivated operations we tested explains what this head computes. Whether this reflects a limitation of current interpretability methods or a computation that does not correspond to any known astrophysical operation remains open, and resolving it is a prerequisite for using interpretability as a tool for scientific discovery.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper presents the first causal mechanistic interpretability study of an astronomy foundation model, revealing that linear decodability does not predict causal importance and localizing the multimodal advantage for dust-attenuated galaxies to a single, unidentified attention head. This work significantly advances the field by demonstrating the necessity of causal interventions over observational probes in scientific AI and providing a rigorous framework for interpreting the internal computations of large-scale scientific models.
Methodology
The paper employs a rigorous suite of causal mechanistic interpretability techniques, including linear probing, activation patching, head ablation, and least-squares concept erasure (LEACE). The methodology is sound, correctly distinguishing between decodability (observational) and causal utility (interventional). The decision to anchor the study to a specific downstream task (photometric redshift) rather than generic pretraining objectives is a strong methodological choice that provides ground truth for validating interpretability findings.
Results & evaluation
The experiments are extensive and well-controlled. The use of the AION-1 foundation model (300M parameters) on the DESI Early Data Release provides a realistic and scientifically relevant testbed. The findings regarding the dissociation between linear decodability and causal importance (specifically for spatial extent vs. dust attenuation) are surprising and well-supported by ablation studies. The identification of a single head (Head 3) as necessary for the dust-graded multimodal advantage is a significant empirical result.
Reproducibility
The paper provides high reproducibility standards, detailing the exact architecture of the readout, training hyperparameters, data splits, and specific intervention protocols. The use of standard libraries (sklearn, scipy) for statistical tests and the clear description of the AION-1 encoder structure allow for independent verification.
Limitations
The study is limited to a single model (AION-1 base) and a single task (redshift prediction). The "unidentified computation" of Head 3 remains a black box, and the paper acknowledges that current methods may not be sufficient to characterize it. The generalizability of these specific mechanistic findings to other scientific foundation models or other tasks is not established.
Broader impact
This paper serves as a crucial proof-of-concept for applying mechanistic interpretability to scientific foundation models. It highlights the limitations of current interpretability tools (probing vs. ablation) in complex, multimodal scientific settings. It sets a template for how interpretability can be used to validate or challenge astrophysical hypotheses derived from AI models, potentially influencing the development of more trustworthy scientific AI. The paper presents the first causal mechanistic interpretability study of an astronomy foundation model, revealing that linear decodability does not predict causal importance and localizing the multimodal advantage for dust-attenuated galaxies to a single, unidentified attention head. This work significantly advances the field by demonstrating the necessity of causal interventions over observational probes in scientific AI and providing a rigorous framework for interpreting the internal computations of large-scale scientific models.
Author of this paper? Send a correction.