Attributing Learned Concepts in Neural Networks to Training Data
By now there is substantial evidence that deep learning models learn certain
human-interpretable features as part of their internal representations of data.
As having the right (or wrong) concepts is critical to trustworthy machine
learning systems, it is natural to ask which inputs from the model's original
training set were most important for learning a concept at a given layer. To
answer this, we combine data attribution methods with methods for probing the
concepts learned by a model. Training network and probe ensembles for two
concept datasets on a range of network layers, we use the recently developed
TRAK method for large-scale data attribution. We find some evidence for
convergence, where removing the 10,000 top attributing images for a concept and
retraining the model does not change the location of the concept in the network
nor the probing sparsity of the concept. This suggests that rather than being
highly dependent on a few specific examples, the features that inform the
development of a concept are spread in a more diffuse manner across its
exemplars, implying robustness in concept formation.
Nicholas Konz, Charles Godfrey, Madelyn Shapiro, Jonathan Tu, Henry Kvinge, Davis Brown
# Attributing Learned Concepts in Neural Networks to Training Data

###### Abstract

By now there is substantial evidence that deep learning models learn certain human-interpretable features as part of their internal representations of data. As having the right (or wrong) concepts is critical to trustworthy machine learning systems, it is natural to ask which inputs from the model's original training set were most important for learning a concept at a given layer. To answer this, we combine data attribution methods with methods for probing the concepts learned by a model. Training network and probe ensembles for two concept datasets on a range of network layers, we use the recently developed _TRAK_ method for large-scale data attribution. We find some evidence for _convergence_, where removing the 10,000 top attributing images for a concept and retraining the model does not change the location of the concept in the network nor the probing sparsity of the concept. This suggests that rather than being highly dependent on a few specific examples, the features that inform the development of a concept are spread in a more diffuse manner across its exemplars, implying robustness in concept formation.

+
Footnote †: \(\dagger\) Work done at Pacific Northwest National Laboratory.

## 1 Introduction

Given the role that concepts play in understanding and explaining human reasoning, measuring their use in neural networks is important for the goal of developing explainable and trustworthy AI. Driven by this, substantial effort has gone into developing methods that measure the presence of a concept within a neural network. Relatedly, a growing body of empirical work shows that deep neural networks learn to encode features as _directions_ in their intermediate hidden layers (Merullo et al., 2023; Wang et al., 2023). A common approach to finding these directions (or **concept vectors**) is via linear probing (Alain and Bengio, 2016). While probing has well-known shortcomings (Ravichander et al., 2020), it is hard to overstate the impact that concept probing has had on deep neural network interpretability. Prominent examples include probing for syntactic concepts in 'BERTology' (Rogers et al., 2021) and chess concepts in the AlphaZero network (McGrath et al., 2021).

A separate thread in explainability research explores **data attribution**, which, rather than measuring the importance of a concept for a model prediction, quantifies the impact of individual _training datapoints_ on a given model prediction (e.g., which images in the training set were most relevant for a classifier's prediction "zebra"?). Data attribution methods have proven to be effective for identifying brittle predictions, quantifying train-test leakage, and tracing factual knowledge in language models back to training examples (Ilyas et al., 2022; Park et al., 2023; Grosse et al., 2023).

In this work, we explore an interplay between concept vectors and data attribution, with the goal of obtaining a better understanding of how neural networks utilize human-understandable concepts. Namely, we ask the natural question:

_Which examples in a model's training data were important for learning hidden-layer concepts?_
Overall, we find that the process of learning a concept is robust to both removal of examples (no small subset of examples are critical to learning a concept) and stable across independent model training runs. While this stability may not be surprising from a human perspective, given that concepts are, by design, supposed to be relatively unambiguous between observers, it is interesting that a similar phenomenon is seen in models.

## 2 Experimental Methods

In this section we describe the two methods that are central to our study: using a linear probe to detect a concept and then applying data attribution to concept predictions. Let \(f(x)\) be an image classification neural network (the "base model") and assume that \(f(x)\) has been pretrained on a training set \(\mathcal{X}_{\mathrm{tr}}\). Assume that \(\mathcal{X}_{\mathrm{val}}\) is the corresponding validation set. We write \(f_{\leq i}\) for the composition of the first \(i\) layers of \(f\). Finally, for each of the concepts \(c\) that we study, we assume that we have a concept training and test set \(\mathcal{X}_{\mathrm{tr}}^{c}\) and \(\mathcal{X}_{\mathrm{val}}^{c}\) where elements of these sets are labeled by whether or not the elements are examples of \(c\).

### Probing for Concept Learning

The purpose of training a concept probe is to detect whether a specific human-interpretable concept is encoded in a hidden layer of a model. If the linear probe can effectively separate encoded exemplars of the concepts from encoded examples that are unrelated to the concept, then we take this as evidence that the model has learned the concept. More specifically, having chosen the \(h\) hidden layer of \(f\) for investigation and a concept \(c\) captured by concept training and test sets \(\mathcal{X}_{\mathrm{tr}}^{c}\) and \(\mathcal{X}_{\mathrm{val}}^{c}\), we follow the common approach of training an affine linear probe \(g\) on the outputs of \(f_{\leq i}(\mathcal{X}_{\mathrm{tr}}^{c})\)[12, 15]. Because \(g\)'s decision boundary is linear, it is effectively summarized by a normal vector which, when the probe is effective, we take to point in the "direction" of concept \(c\) (up to a sign). This vector is called a _concept activation vector_ (CAV). For short-hand we will use \(g_{\leq i}(x):=g(f_{\leq i}(x))\) to describe the "subnetwork + probe" model.

### Attribution of Concept Predictions

The data attribution question that we seek to answer in this paper is: "which examples in \(f\)'s original training set \(\mathcal{X}_{\mathrm{tr}}\) were most important for it learning a concept \(c\)?" Since we can quantify how well \(f\) learned a concept at a layer \(i\) by the accuracy of a trained probe \(g_{\leq i}(x)\) on \(\mathcal{X}_{\mathrm{val}}^{c}\), it is more convenient to ask "which examples in the network's original training set \(\mathcal{X}_{\mathrm{tr}}\) were most important for the concept predictions of \(g_{\leq i}\) on a set of test images?"

A _data attribution method_\(\tau(x_{\mathrm{val}},x_{\mathrm{tr}};h)\) is a function that assigns a real-valued score to a training point \(x_{\mathrm{tr}}\in\mathcal{X}_{\mathrm{tr}}\) according to its importance to the prediction of a model \(h\) on some test/validation

Figure 1: Schematic of our approach for concept attribution: (1) train \(N\) models with different random seeds (in green) on the training set. (2) We choose a hidden layer \(i\), append a probing classifier \(g\) to its output, freeze the weights of \(f_{\leq i}\), and train (\(g\circ f_{\leq i}\)) on the concept dataset. (3) We calculate attributions with TRAK [12] for \(g\circ f_{\leq i}\) on elements of the test set in terms of the original training data and aggregate across fixed layers and concepts.
point \(x_{\rm val}\)[Park et al., 2023]. We will define the expected importance of a training point \(x_{\rm tr}\) to concept predictions of \(g_{\leq i}\) as

\[\tau_{c}(x_{\rm tr}):=\mathbb{E}_{x_{\rm val}\sim\mathcal{X}_{\rm val}}\tau(x_{ \rm val},x_{\rm tr};g_{\leq i}), \tag{1}\]

as suggested by Park et al. [2023]. For all attribution experiments we use a recently developed data attribution method, TRAK (Tracing with the Randomly-projected After Kernel) [Park et al., 2023]. For details on this method we defer to the original paper.

TRAK requires an ensemble of \(M\) trained models; we use \(M=20\), but found similar results for as few as \(M=5\). Each model in the ensemble is a "subnetwork+ probe" \(g_{\leq i}^{(j)}\), where \(1\leq j\leq M\). \(g_{\leq i}^{(j)}\) is created by training the same base model \(f\) on \(\mathcal{X}_{\rm tr}\) to obtain \(f^{(j)}\), then training a probe for a concept \(c\) on the \(i^{th}\) layer of \(f^{(j)}\) with the concept training set \(\mathcal{X}_{\rm tr}^{c}\), using \(\mathcal{X}_{\rm val}^{c}\) for validation.

The first step of TRAK is to "featurize" (process \(\mathcal{X}_{\rm tr}\)) and score (process \(\mathcal{X}_{\rm val}\)) each \(g_{\leq i}^{(j)}\), which we run in parallel over the ensemble. After this, the attribution scores of all \(M\) networks are aggregated, resulting in \(|\mathcal{X}_{\rm val}\times\mathcal{X}_{\rm tr}|\) final scores total, one for each pair \((x_{\rm val},x_{\rm tr})\). An important note here is that TRAK requires a task/loss function and corresponding target labels to evaluate the predictions of the models \(g_{\leq i}^{(j)}\) -- in our case, concept prediction and binary concept labels, respectively. For consistency, unlike the concept probe training and validation sets \((\mathcal{X}_{\rm tr}^{c},\mathcal{X}_{\rm val}^{c})\) which use manually-defined concept labels, we use one of the trained \(g_{\leq i}^{(j)}\) to assign these labels to \(\mathcal{X}_{\rm tr}\) and \(\mathcal{X}_{\rm val}\) for attribution. We summarize our experimental design in Fig 1.

## 3 Datasets, Base Models and Concepts

We use "ImageNet10p" as the training and validation sets \(\mathcal{X}_{\rm tr}\) and \(\mathcal{X}_{\rm val}\), respectively, to train the base models. ImageNet10p is defined by randomly sampling \(10\%\) of the images of each class from the ImageNet [Deng et al., 2009] training and validation sets. The resulting ResNet-18 models obtained about \(45\%\) top-1 accuracy on \(\mathcal{X}_{\rm val}\) (see Appendix F for training details), from which we assume that the model learned meaningful enough visual representations for concepts to be present. We build two different concept probing datasets,1 and show example images of each concept in Appendix B.

Footnote 1: We initially experimented with using concepts from the Broden dataset [Bau et al., 2017] but we found probes trained on this dataset did not generalize well to arbitrary ImageNet images.

Concept 1: Snakes.We define the "Snakes" concept with the \(17\) ImageNet snake classes \(477\)-\(493\). The probe training set \(\mathcal{X}_{\rm tr}^{c}\) is constructed with (i) all examples of these classes in \(\mathcal{X}_{\rm tr}\) and (ii) the same number of non-snake images randomly sampled from \(\mathcal{X}_{\rm tr}\). The probe validation set \(\mathcal{X}_{\rm val}^{c}\) is constructed in the same manner as \(\mathcal{X}_{\rm val}\); this gives a concept training/validation split of \(4,374\)/\(170\).

Concept 2: High-Low Spatial Frequencies.We define the "High-Low Frequencies" concept with images where a directional transition between a region with high spatial frequency to a region with low spatial frequency is present [Schubert et al., 2021].2 We created this concept dataset by computing the top \(0.001\%\) highest-activating ImageNet images (by \(L_{\infty}\) norm) of the 'high-low frequency neurons' defined in layer mixed3a of InceptionV1 [Schubert et al., 2021]. These images are used with an equal number of random non-concept images to define the concept training and validations sets \(\mathcal{X}_{\rm tr}^{c}\) and \(\mathcal{X}_{\rm val}^{c}\), respectively, resulting in a training/validation split of \(362\)/\(20\). We found the high selection threshold necessary to identify images where the concept is clearly present.

Footnote 2: An analogous variant of this concept was arguably also discovered in biological neurons [Ding et al., 2023].

## 4 Experiments and Results

### Concept Attribution

In Fig. 2 we display the images in the training set \(\mathcal{X}_{\rm tr}\) which received the highest and lowest concept prediction attribution scores \(\tau_{c}(x_{\rm tr})\) (Eq. (1)) for each concept, for various layers of the base model. In Fig. 3 we show how the presence of each concept varies with network layer depth, where the two concepts were most present on average in layer3. Additional highest-attributed images are in Appendix C.1.
How does learned concept attribution vary between network layers?For the snakes concept, full snake images appear to be important for concept learning in deeper network layers, while images that possess textures common for snakes are most important for the earlier layer1. The concept does not appear to be present in very early layers (layer1.0.conv2), which is reasonable given that "snakes" is an abstract concept (see also Fig. 3). These observations are compatible with the conventional wisdom that deeper network layers learn more complex abstract features (such as objects), while earlier layers learn more basic features (such as textures).

The high-low frequency concept is fairly present throughout the network (Fig. 2, right and Fig. 3). Highest-attributed (and certain lowest-attributed) training set images for this concept contain transitions from high to low spatial frequency, such as pomegranate seeds over a flat background (layer1.0.conv2, image 1), baseball threading alongside a flat casing (layer1, image 4), interwoven threads over a smooth background, (layer2, image 1), and fur over a smooth background (layer3, image 2). In comparison to the snakes concept probe which has increasing accuracy with network depth, likely due to its connection to the base models classification task, the high-low frequency concept fades after layer 3 as it is synthesized into higher-level concepts related to the label classes.

Are the concepts that a model learns the result of a few select exemplars?We analyze the importance of images in the base model's training set \(\mathcal{X}_{\text{tr}}\) for concept learning by (1) removing the \(T\) highest-attributed images from \(\mathcal{X}_{\text{tr}}\) to obtain \(\mathcal{X}_{\text{tr}}^{-T}\) (\(T\in\{100,1,000,10,000\}\)), (2) re-training the \(M\) base models on \(\mathcal{X}_{\text{tr}}^{-T}\), and (3) training concept probes on each of them for a given layer.3 If the probe concept detection validation accuracy changes after the training set is pruned of the most important examples of a concept, then we conclude that these examples were primarily responsible for the model learning the concept. If this does not happen, it suggests that a model learns a concept in a more flexible way, from a broad range of examples. For this experiment we measure concepts in the layer where both were most present on average, layer3. Our results are shown in the middle and right plots in Figure 4 (where these first experiments correspond to sparsity equal to \(10^{2}\)), concept validation accuracy did not change on models trained on \(\mathcal{X}_{\text{tr}}^{-T}\) for varying \(T\) compared to the baseline

Figure 3: **Concept presence within different network layers.** Average concept detection validation accuracy of probes trained on different layers, for each concept; confidence bands are std. deviation over 5 base models.

Figure 2: **Training set attributions for concept learning.** The four highest and two lowest attributed training set images (decreasing \(\tau_{c}(x_{\text{tr}})\) from left to right) for concept learning at different network layers. **Left half:** snakes concept. **Right half:** high-low frequencies concept.
of those on \(\mathcal{X}_{\mathrm{tr}}\), for either concept (Figure 3). This provides further evidence that the learning of a concept is diffuse among exemplars and does not depend on a few special examples. We note that this result may not be surprising given that the attribution scores are mostly similar across examples from \(\mathcal{X}_{\mathrm{tr}}^{c}\), Figure 4 left, (e.g., if all examples have the same importance for learning a concept, removing a fraction of them will not have a large effect). In particular, the high-low frequencies concept could be learnable from the majority of images in the training set, especially if it suffices for the probes simply to learn to be boundary detectors; we investigate this with "relative probing" to discriminate between generic object boundaries and the high-low frequency concept in Appendix D.

Finally, given evidence that semantically meaningful representations tend to be sparse in the neuron basis (Gurnee et al., 2023), we also trained sparse probes, thinking that in this regime removing a fraction of exemplars might have a larger effect. In the middle and right plots in Figure 4 and Figure 5 the concepts also remained robust in this setting.

How similar are different probes trained for the same concept/layer?We show how probe accuracy changes with respect to network layer depth in Fig. 3. For a fixed layer, different probes typically converge to similar performance.4

Footnote 4: In Appendix E we discuss why we do not compare probes via CAV similarity.

In the case of the high-low frequency concept, we see that probe accuracy is highest at intermediate layers, and comparatively low at the earliest and latest layers. This is consistent with the original work of Schubert et al. (2021), which discovered "high-low frequency detector neurons" in intermediate layers of InceptionV3 networks (but not in the earliest or latest layers). This stands in contrast to the snakes concept, where probe accuracy increases monotonically with network depth. One possible explanation for this observation is that the snakes concept dataset was obtained from a subset of ImageNet classes. As such, the base models have been trained to correctly classify positively-labelled concept images using their output logits. Here our observations are consistent with Alain and Bengio (2016), which trained multi-class linear probes for ImageNet classification and found monotonically increasing accuracy with depth.

Figure 4: **Left: sorted distribution of concept learning attribution scores for the training set \(\mathcal{X}_{\mathrm{tr}}\), averaged over the validation set \(\mathcal{X}_{\mathrm{val}}\), for both concepts. Right: Effect of re-training the base model on \(\mathcal{X}_{\mathrm{tr}}\) with the top \(T\) concept-attributed training examples removed (\(\mathcal{X}_{\mathrm{tr}}^{-T}\)), for sparse probe concept detection on layer3. High-low frequency probe results averaged over 5 base models, with entirely overlapping std. deviation confidence intervals shown (probes at different \(T\) for the same base model obtained the same performance due to the small size of \(\mathcal{X}_{\mathrm{val}}^{c}\)).**

Figure 5: Dependence of concept detection validation accuracy of sparse concept probes on their sparsity \(k\).
### Sparse Concept Probing and Attribution

How does forcing probes to be sparse affect concept detection?Forcing a probe to be sparse (at most \(k\) of the CAV elements are non-zero) allows for even more interpretable concept directions (Gurnee et al., 2023). To evaluate this for our concepts, we trained probes with a range of sparsities on different layers of \(f\), using an approach similar to iterative hard thresholding (Jin et al., 2016). After the first half of the training epochs, we set all but the \(k\) parameters of highest absolute value to zero, freeze all of the zeroed parameters from updating, then continue training.

In Fig 5 we show how the concept detection validation accuracy changes with probe sparsity \(k\) at multiple layers, for both concepts. Reasonably, probe accuracy typically increases with \(k\), and we see a similar relative accuracy ranking of different layers as in the non-sparse case (Fig. 3). We see that the concepts are both somewhat learnable with sparse probes, but not nearly to the degree of the non-sparse probes (Fig. 3).

Fig. 6 displays the highest- and lowest-attributed images for probes of varying sparsity \(k\), for each concept at layer3 (compare with the attributions of non-sparse probes in Fig. 1). For the snakes concept, the attributions are different than those for the non-sparse probe, and yet very similar among the sparse probes. Interestingly, we see that almost all of the highest- or lowest-attributed images possess a "honeycomb"-like texture which appears similar to snake scales. For the high-low frequencies concept, we see that the \(k=1\) sparse probe obtained the same attributions as the non-sparse probe, yet the probes with more non-zero entries both obtained the same distinct attributions, which also appear to have examples of the concept.

## 5 Limitations

In order to experimentally vary factors including model layer, concept and concept training data, we were forced to restrict other variables. We only experiment with ResNet18 image classifiers trained on ImageNet10p and two concept datasets (snakes and high-low frequency) -- adding additional base models and concepts would increase the breadth of this study. Another interesting future direction would be conducting analogous experiments on a different modality (e.g., natural language) or task. Finally, we only use TRAK for data attribution, and it would be interesting to know the extent to which our experimental results are particular to TRAK.

## Conclusion

In this paper we explored the importance of individual datapoints in concept learning. We found evidence of "convergence" in several senses, including stability under removal of exemplars and across independent training runs. Although more extensive experiments are needed (with better aggregation methods), our results suggest a robustness to the way that concepts are learned and stored in a deep learning model.

Figure 6: **Training set attributions for sparse probe concept learning. The four highest and two lowest attributed training set images (decreasing \(\tau_{c}(x_{\text{tr}})\) from left to right) for concept learning at layer3 for probes of different sparsity \(k\). Left half: snakes concept. Right half: high-low frequencies concept.**
## Acknowledgments

This research was supported by the Mathematics for Artificial Reasoning in Science (MARS) initiative at Pacific Northwest National Laboratory. It was conducted under the Laboratory Directed Research and Development (LDRD) Program at at Pacific Northwest National Laboratory (PNNL), a multiprogram National Laboratory operated by Battelle Memorial Institute for the U.S. Department of Energy under Contract DE-AC05-76RL01830.

The authors would also like to thank Andrew Engel for useful conversations and feedback related to the paper.
