From Neurons to Neutrons: A Case Study in Interpretability (2405.17425v1)

Abstract: Mechanistic Interpretability (MI) promises a path toward fully understanding how neural networks make their predictions. Prior work demonstrates that even when trained to perform simple arithmetic, models can implement a variety of algorithms (sometimes concurrently) depending on initialization and hyperparameters. Does this mean neuron-level interpretability techniques have limited applicability? We argue that high-dimensional neural networks can learn low-dimensional representations of their training data that are useful beyond simply making good predictions. Such representations can be understood through the mechanistic interpretability lens and provide insights that are surprisingly faithful to human-derived domain knowledge. This indicates that such approaches to interpretability can be useful for deriving a new understanding of a problem from models trained to solve it. As a case study, we extract nuclear physics concepts by studying models trained to reproduce nuclear data.

• The paper demonstrates that neural networks can extract low-dimensional, physics-relevant embeddings from nuclear data via mechanistic interpretability.
• Using PCA, the study uncovers ordered proton and neutron embeddings that mirror established nuclear models, including parity splits and helicity patterns.
• Multi-task training enhances model performance while revealing hidden layer features analogous to the semi-empirical mass formula and nuclear shell model.

Interpretability in Neural Networks: Insights from Nuclear Physics

In the paper "From Neurons to Neutrons: A Case Study in Interpretability" (2405.17425), the authors explore the potential of mechanistic interpretability (MI) to uncover scientific insights from neural networks trained on nuclear physics data. They investigate whether neural networks can learn representations that align with established human knowledge and provide new scientific understanding.

Mechanistic Interpretability Approach

Mechanistic Interpretability is an approach that aims to understand how neural networks make predictions by analyzing their internal representations and computational mechanisms. The paper explores the ability of high-dimensional neural networks to learn low-dimensional embeddings that can reveal insights into the training data, particularly in the domain of nuclear physics.

The authors use Principal Component Analysis (PCA) to analyze the embeddings learned by neural networks. Through PCA, they project the proton and neutron embeddings onto lower-dimensional spaces to identify structural patterns that reflect the known nuclear physics concepts such as pairing effects and the Semi-Empirical Mass Formula (SEMF).

Figure 1: Projections of neutron number embeddings onto their first three principal components (PCs). Models were trained on nuclear data (left) or a human-derived nuclear theory (right).

Nuclear Physics Case Study

The case paper focuses on nuclear physics, where the objective is to predict nuclear properties based on the number of protons ($Z$) and neutrons ($N$). The authors argue that the structured representations learned by neural networks can mirror established physics concepts and assist in scientific discovery.

Figure 2: Binding energy per nucleon as given by the SEMF formula (left) and observed in measurements (right).

Embedding Structures

The paper highlights several features of the learned embeddings that correlate with generalization performance and known physics concepts:

1. Helicity: A spiral pattern observed in the embeddings, particularly when trained on binding energy predictions, suggests a geometric structure that parallels the SEMF's theoretical geometry.
2. Orderedness: The ordering of proton and neutron numbers in the PCA space correlates with model generalization, implying inductive biases encoded within the embeddings.
3. Parity Split: The distinction between even and odd proton and neutron number embeddings, reflecting pairing effects known in nuclear physics.

   Figure 3: Projection of proton number (Z) embeddings onto the first two principal components (PCs), superimposed on the neural network's binding energy predictions.

Principal Components Analysis

Principal components analysis played a crucial role in identifying meaningful patterns in the neural network's representations. The paper provides evidence supporting the efficacy of using PCA to capture essential features, illustrating that a few principal components can recover significant portions of the model's performance.

Figure 4: Fitting a helix to the PC-projected embeddings.

Experiments and Observations

1. Multi-Task Learning: By training the model on multiple related tasks simultaneously, the authors observed improved performance and richer learned representations. This multi-task setup enhances the model's ability to generalize and captures the shared structure underlying various nuclear observables.
2. Hidden Layer Features: Analysis of penultimate layer activations revealed features that strongly resemble terms in established nuclear models, such as volume and pairing terms from the SEMF and contributions from the nuclear shell model.

   Figure 5: Test performance over different observables for models trained on a single task versus multiple tasks jointly.

Conclusion

The research demonstrates that neural networks trained on nuclear physics data can learn representations that provide insights into both the model's prediction process and the underlying physical phenomena. Through mechanistic interpretability, the authors effectively reinterpret the learned embeddings and latent features in terms of established theoretical models.

The paper signifies a step forward in using intrinsic model representations to gain scientific understanding, potentially paving the way for discovering new phenomena in various domains. Future research may expand these interpretability techniques to more complex scientific issues and unexplored datasets.