Finding and using interpretable latents in a neutrino foundation model with sparse autoencoders

Raphaël Bonnet-Guerrini; Johann Ioannou-Nikolaides; Inar Timiryasov; Vincenzo Piuri

Source-linked AI summary

Finding and using interpretable latents in a neutrino foundation model with sparse autoencoders

Raphaël Bonnet-Guerrini, Johann Ioannou-Nikolaides, Inar Timiryasov, Vincenzo Piuri

arXiv:2608.26090v1Aug 26, 2026astro-ph.HEcs.AIcs.LGhep-ex

TL;DR

The paper asks what physical information a neutrino foundation model encodes and which downstream heads use it. It applies rigorously validated sparse autoencoder dictionaries, finds a rich physical atlas that the direction head barely uses, and trains an uncertainty head that exploits quality and brightness features to improve selected-event angular resolution.

Problem
The paper addresses which physical structures are encoded in a shared foundation-model representation and which downstream heads use them, a question standard post-hoc explanations do not resolve.
Method
The authors apply sparse autoencoders to PolarBERT and validate candidate read-outs with held-out tests, matched nuisance controls, independent dictionary trainings, and interventions.
Results
At 20% selection efficiency, the uncertainty head achieves 3.2° median angular resolution versus 20.2° using the best single detector observable, while the direction head barely uses the validated atlas.
Takeaways & Limitations
Mechanistic interpretability reveals validated latent physics in the representation and identifies physically interpretable information that a downstream uncertainty task causally uses.
Takeaways & Limitations
The final-layer dictionary does not recover a validated geometry read-out, although a layer-2 dictionary recovers depth with validation AUROC 0.986.

Abstract

from arXiv · show

We present a first application of sparse-autoencoder-based mechanistic interpretability to particle physics. Studying a neutrino foundation model pretrained on IceCube data and fine-tuned for direction reconstruction, we identify a validated atlas of physical concepts in the model representation, using a strict validation protocol consisting of held-out tests, matched nuisance controls, and replication across independent dictionary trainings. Causal interventions show that the direction head barely draws on this atlas. Motivated by this underused information, we train an uncertainty head on the same event-level representation to predict the model's angular reconstruction error. Unlike the direction head, it depends causally on quality and brightness features from the atlas. At $20\%$ selection efficiency, this interpretable estimator improves the median angular resolution from $20.2^\circ$ to $3.2^\circ$. These results suggest that mechanistic interpretability can reveal learned latent physics encoded within a model's internal representation and help design downstream tasks that exploit it.

1 Introduction

The paper investigates which physical concepts a neutrino foundation model represents and which downstream heads use. Sparse autoencoders and strict validation reveal a physical atlas, while interventions show the direction head uses little of it and an uncertainty head exploits quality and brightness features.

Modern neutrino reconstruction uses sparse Cherenkov-light measurements in irregular detectors, whose geometry, ice, noise, and trigger effects enter learned representations.
Foundation models share representations across task-specific heads, raising the question of which encoded structures each head actually uses.
Sparse autoencoders decompose superposed activation patterns into a few learned directions intended to separate overlapping concepts.
SAE features require held-out tests, matched controls, dictionary variation, and interventions because sparsity or apparent interpretability does not establish causality.
PolarBERT analysis identifies a validated atlas containing detector quality, auxiliary activity, event brightness, and detector depth.
At 20% selection efficiency, an uncertainty head reaches 3.2° median angular resolution versus 20.2° for selection using the best single detector observable, while the direction head uses little of the atlas.

2 Setup

The study analyzes PolarBERT’s frozen event-level representation with sparse dictionaries and separates training, validation, discovery, and causal testing. Its protocol combines reconstruction fidelity, sparse feature selection, matched controls, and intervention-based causal judgments.

PolarBERT as a controlled testbed: PolarBERT is an eight-block transformer pretrained on IceCube pulse sequences and fine-tuned for neutrino direction reconstruction.
PolarBERT as a controlled testbed: Each event is a variable-length pulse sequence described by arrival time, charge, DOM identifier, and an auxiliary flag linked to incomplete digitization.
PolarBERT as a controlled testbed: Events contain at most L = 127 input pulses; auxiliary observables measure the auxiliary-pulse fraction and its share of total input-window charge.
PolarBERT as a controlled testbed: Every downstream head reads the final CLS representation, which is evaluated against truth direction using angular error Δψ_i = arccos(û_i · u_i).
Sparse dictionaries: BatchTopK enforces an average of k active latents per event while allowing event-level variation; a concept-independent sweep selects k = 16.
Sparse dictionaries: SAE reconstruction preserves the direction head closely: fidrec ≃ 0.989, with 0.65° mean angular deviation versus 56° under mean ablation.
Validation protocol: The protocol labels features as candidates, validated read-outs, or causal features based on read-out quality, nuisance selectivity, and downstream intervention tests.

3 From candidates to a validated atlas

The atlas retains only a few validated sparse read-outs after held-out, nuisance-control, and replication tests. These read-outs organize a continuous quality-brightness axis, while geometry and morphology are not localized as single validated final-layer latents.

Most of the 1024 latents and several concept candidates fail validation, leaving only a few interpretable read-outs.Association alone is weak because look-elsewhere effects can produce convincing candidates across many unlabelled directions.
Validated read-outs: zbc is a validated bright-clean read-out that separates bright-clean from strongly auxiliary, dim events with AUROC 0.911.It also distinguishes bright-clean from dim-clean events with validation AUROC 0.821, showing dependence on total charge within the clean class.
Validated read-outs: zaux captures auxiliary-dominated, low-charge events, while z′bc is a weaker clean-side feature that passes the same controls at lower discrimination.zbs rises with total charge but fails selectivity controls, so it is treated as a brightness-associated support direction rather than a validated brightness read-out.
Structure of the atlas: The atlas represents quality and brightness as a continuous axis distributed across many sparse coordinates, with only a few coordinates sharp enough to validate individually.The same axis appears in both SAE and PCA descriptions, suggesting it reflects the representation geometry rather than one method alone.
Limits of the atlas: No final-layer latent survives matched geometry contrasts for event radius, depth, or sub-detector location, while a layer-2 depth feature reaches validation AUROC 0.986.This absence does not imply geometry is missing from the model, because earlier-layer information can remain distributed across pulse tokens.
Stability and causal tests: Analogous clean features appear across all dictionary seeds, but complete validation replication is weaker for zbc and zaux.A control-surviving zbc analog appears in every tested dictionary, with validation AUROC between 0.84 and 0.99.
Stability and causal tests: Pulse interventions raise zaux and lower zbc, while pulse removal suppresses zbc and zbs and charge rescaling moves zbs with the charge scale.These perturbations support the physical labels of the validated read-outs and separately support zbs's brightness association.

4 Causality is head-relative

Causal interventions reveal that the direction head has a real, replicable sensitivity axis, but it largely ignores the named physical atlas. In contrast, an uncertainty head causally uses interpretable brightness and clean-quality features to rank reconstruction error.

4.1 The atlas is inert for the direction head: 56.4° is the direction-head output change when the final CLS state is replaced by its validation-set mean, establishing the intervention scale.Zeroing zbc shifts angular error by only +0.06°, comparable to +0.05° for matched controls.
4.1 The atlas is inert for the direction head: The clean pair, clean family, and auxiliary family generally produce null direction-head effects across atlas scales and dictionary seeds.A negative seed shows about 10° changes because a family includes a reconstruction-critical support latent that fails concept validation.
4.1 The atlas is inert for the direction head: 15.0° and 10.9° are the increases in direction-head angular error after zeroing zbs on clean-bright and clean-dim events, respectively.Perturbing zbs above or below its natural value damages reconstruction, producing a V-shaped response replicated across four dictionary draws.
4.2 The direction head’s causal axis: The direction head has a stable causal axis, but no verified latent aligns strongly with any named concept.The leading axis differs by 5.6° across dataset splits, while no verified latent exceeds |cos| = 0.15 with it.
4.2 The direction head’s causal axis: The direction-head causal subspace aligns with causal latents from a reconstruction-anchored functional dictionary at |cos| ≃0.8.Perturbing CLS along Jacobian-derived directions changes the predicted direction, with an event-dependent rather than fixed-rotation effect.
4.3 An uncertainty head with causal latents: The uncertainty head reaches AUROC 0.925 for distinguishing the highest and lowest true-error quartiles, matching a full-CLS linear classifier and exceeding 0.787 from detector observables.A simple linear regressor has comparable performance, indicating that most error-relevant information is linearly accessible in the shared representation.
4.3 An uncertainty head with causal latents: Zeroing zbc shifts predicted log error by M1 = +0.31, while the clean pair, family, and dense core produce M1 = +0.32, +0.78, and +0.37.These uncertainty-head effects exceed matched controls across all dictionary seeds, whereas zaux remains inert.

5 Outlook

The study separates physical information represented by PolarBERT from information used by its direction head, then demonstrates that an uncertainty head can exploit the underused atlas. This interpretable estimator provides substantially sharper event selection and may support analyses where Monte Carlo truth is unavailable.

PolarBERT contains validated read-outs of event quality, auxiliary activity, brightness, and detector depth, but the direction head uses them only marginally.Its aggregate sensitivity is instead dominated by a stable causal axis without a simple physical interpretation.
The uncertainty head reads the atlas that the direction head leaves unused, predicting angular reconstruction error from the same event-level representation.Its causal dependence includes bright-clean and related clean-side features.
At 20% selection efficiency, the uncertainty reader reaches 3.2° median angular resolution versus the best single detector observable’s 20° scale.
Interpretable causal selection matters for real data because cut correctness cannot be checked against Monte Carlo truth.The approach contrasts with black-box regression whose inputs would remain opaque.
Sparse dictionaries can reveal which physical information is available in a shared representation and guide downstream tasks or learning objectives that exploit it.The paper frames mechanistic interpretability as both a diagnostic and a design aid for foundation-model representations.
The validation protocol is tested on Monte Carlo, while extension to real detector data must rely on controls and replication because truth labels are incomplete.Future directions include likelihood analyses, more event-level labels, and other particle-physics foundation models.

A Supplementary analysis protocol

This supplementary section establishes the fixed verdict rules and the head-aligned sensitivity basis used to analyze the representation. Later analyses apply these rules to atlas results, robustness checks, and interventions.

The protocol fixes verdict rules before held-out evaluations and retains partial or failed replications.

A.1 Choice of representational site

The final-block CLS state is selected as the common representational site because directional information improves with depth and then plateaus. The protocol combines held-out probes, matched controls, Jacobian sensitivity, and replication.

Ridge probes are fit from CLS residual states at every depth to true neutrino direction using disjoint concept_dev and concept_test splits.The probes use α = 1; concept_dev contains 69,632 events and concept_test contains 59,648.
61.0° final-block mean angular error closely matches the frozen direction head’s 59.0° on the same events.Directional information improves monotonically with depth and plateaus from block 6 onward.
The final-block CLS state becomes the common site for dictionaries, probes, and head interventions after the depth scan supports its use.
Figure 8 reports held-out mean angular error for a linear CLS probe at each layer, with lower values indicating better performance and the frozen direction head as reference.
Verdict rules use matched-control interventions, including the direction-head selectivity statistic Δ+ − Δ−, with control statistics computed identically.
The within-clean brightness control pool additionally excludes 43 charge-tracking latents, while latent 1006 remains a secondary clean-side read-out after surviving nuisance controls.
For a unit activation-space direction v, Jacobian norm ∥J_i v∥ measures local head sensitivity; sensitivity directions are aggregated over 8,000 discovery events and retained only if replicated.The eigenvectors rank directions by average effect on head output, and normalized eigenvalues give their aggregate sensitivity share.

B.1 Isolation of the named read-outs

The isolation analysis identifies validated SAE read-outs through discovery rankings, held-out performance, and matched-control tests. Sparse clean and auxiliary latents are useful discriminants, while a population-level brightness coordinate is not promoted as a validated per-event read-out.

Latent 640 is the only primary clean-trigger candidate, while latent 1006 is retained as a secondary read-out after surviving nuisance controls.Latents 751, 946, and 369 fail the controls; latent 973 is dense and non-selective.
Latent 195 is the primary auxiliary candidate, achieving 0.832 discovery AUROC and 0.840 held-out AUROC.Its minimum evaluable nuisance-control AUROC is 0.854.
Table 6 ranks clean-trigger and auxiliary contrasts using AUROC, top-1% enrichment, firing rate, minimum control AUROC, and post-criteria status.
Reference-dictionary families test whether concepts distributed across multiple coordinates could explain a single-latent null.
On independent concept_test events, sparse clean latent zbc discriminates more sharply than PC10, whereas zaux and PC10 perform similarly for auxiliary-dominated events.
PC10 outperforms zbs for high-charge events because zbs tracks charge only in population averages and is a poor per-event classifier.Therefore zbs is not treated as a validated brightness read-out.

B.4.1 Morphology proxies

Morphology information is evaluated with sparse probes and layer-dependent isolation tests. The results indicate that morphology is distributed across many dictionary coordinates, while depth information is localized to earlier layers.

Morphology probes: k-sparse probes test whether time extent, linefit speed, linearity, charge concentration, and vertical extent concentrate in few SAE coordinates.Candidates are selected on concept_dev and evaluated on concept_test, with dense ridge probes providing reference ceilings.
Morphology probes: Table 9 compares dense CLS and SAE probes with single-latent and multi-latent sparse reconstructions across layers 8 and 3.The k90% statistic reports the smallest selected-latent count reaching 90% of the dense-SAE R2.
Morphology proxies: zbs is brightness-associated rather than a validated brightness read-out because its population means are charge-monotonic but per-event discrimination is near chance.High/low charge contrasts the top and bottom quintiles of total charge.
Morphology probes: Morphology information is mostly distributed across the dictionary rather than represented by a single validated latent.The greedy selection path yields upper bounds rather than exact circuit sizes, so absent single-latent validation does not imply absent morphology information.
Depth isolation: The final-layer dictionary has no validated depth coordinate, whereas layer-2 depth reaches AUROC 0.986 under matched controls.Neither early-layer coordinate passes the direction-head causal criteria, so the geometry read-out is localized without becoming a direction-head bottleneck.

B.5 Robustness across dictionaries

The validated physical roles are more reproducible than individual SAE coordinates across independent dictionary trainings. Clean-primary and bright-clean analogues are comparatively robust, while auxiliary read-outs vary more by seed.

Cross-seed robustness: Across four independently trained dictionaries, the clean-primary role passes every control scheme in three of four draws.The auxiliary role is less stable: only the reference seed passes all evaluable controls, and seed 3 has no isolated auxiliary candidate.
Cross-seed robustness: Table 11 reports functionally analogous read-outs across SAE seeds together with their held-out validation quantities.
Cross-seed robustness: Across the tested expansion and sparsity grid, every dictionary contains a control-surviving bright-clean analogue with validation AUROC between 0.84 and 0.99.The robust object is the physical role and broader response axis, not a particular latent index.

C.1 Family-level removals

Family-level removal tests find no selective direction-head effect for clean or auxiliary feature sets. This null result also holds when removing a complete 194-latent clean family rather than a single coordinate.

Family-level removals: None of the clean- or auxiliary-side family removals produces a selective direction-head effect.Control means exceed medians for several families because a minority of dense control sets removes substantial reconstruction mass.
Family-level removals: Table 12 summarizes family removals on causal_test using selective angular-error shifts and matched-control statistics.
Family-level removals: Removing the clean pair changes positive-class error by only +0.033° and fails both effect and control-normalized criteria.The pair targets the secondary coordinate 1006 and provides a more reliable verdict than testing one sparse coordinate.
Family-level removals: Removing the complete 194-latent clean family is also null, so the direction-head result is not explained by testing only one sparse coordinate.

C.2 Cross-seed single-latent replication

Cross-seed tests distinguish brightness-support features from clean and auxiliary read-outs. Brightness-support coordinates affect direction reconstruction consistently, while identifiable bright-clean latents remain outside the causal set.

Single-latent replication: The primary clean and auxiliary read-outs never pass direction-head causal criteria in any evaluable seed.This makes the direction-head null independent of the reference dictionary.
Cross-seed tests: Table 13 reports cross-seed direction-head shifts, while Table 14 reports recipient and donor shifts for writing tests.
Single-latent replication: A seed-1 secondary clean latent shifts angular error by +0.082° but is isolated rather than replicated across the other three seeds.The corresponding clean pair is also null.
Brightness-support features: Brightness-support coordinates change angular error by 11.6°–14.8° in every seed, with V-shaped dose responses around their natural activations.Both decreasing and increasing activation damages the prediction, so the head does not interpret activation monotonically as brightness.
Writing tests: Writing clean, auxiliary, or brightness-related activations into matched recipient events produces no effect larger than controls.These writing tests provide no evidence that the direction head uses or responds to those read-outs in the tested transfers.
Functionally trained dictionaries: Functionally trained dictionaries contain no identifiable bright-clean latent without reconstruction, but adding reconstruction recovers AUROCs of 0.87 and 0.878 outside the causal set.None of the causal latents validates against the tested physical concept vocabulary.

C.5 Uncertainty-head replication

Cross-seed interventions replicate the uncertainty head’s positive response to clean-side latent removals, while auxiliary-side interventions remain null. The pre-specified single-carrier uncertainty dose replicates in 3/4 seeds, and a joint dose restores the monotone response in the remaining draw.

Removing zbc, the clean pair, clean family, or clean core shifts predicted error upward in every seed.
Auxiliary-side interventions remain null, preserving the controlled contrast.
Table 16 reports cross-seed uncertainty-head removal effects as shifts M1 in predicted log error.
The pre-specified single-carrier uncertainty dose replicates in 3/4 seeds.After removing support-carrier contamination, seed 3’s core effect remains significant, with M1 decreasing from +0.681 to +0.302 (z = 6.8).