Understanding how language models make decisions is important to ensure that their use can be trusted. Mechanistic interpretability offers one lens through which to understand how transformer architecture models [1] perform the computations required to get an output. An oft-used tool in mechanistic interpretability is to attribute individual network parts to specific capabilities by ablating those parts and observing capability degradation.
However, choosing how to ablate neurons in language models is still an unsolved problem. The traditional closed-form methods are zero ablation and mean ablation [2,3], as well as an additional, more randomised method of activation resampling in the case of causal scrubbing [3,4], but little empirical analysis has been done to optimise these methods [4].
Understanding exactly how neuron activations deviate, and what baseline they deviate from, is a broadly applicable question that is underexplored, and has the potential to improve techniques for model pruning and analysis into model sparsity.
In this paper, we 1) describe a simplistic working model of neuron activations, 2) suggest an improved, closed-form method of neuron ablation using modal activation, called 'peak ablation', and 3) run experimental analysis on various ablation methods to compare the degree to which they harm model performance.
Mechanistic interpretability is a field of research focusing on understanding how neural network models achieve their outputs. [5,3,6,7,8,9]. A common method used in mechanistic interpretability, is 'ablate and measure' [3]. We investigate more precisely how different ablation methods affect performance, and propose 'peak ablation' as another possible method.
Most relevantly, recent research [4] investigates hyperparameter selection to optimise activation patching for causal scrubbing. Our research differs; instead of interpolating activations between similar inputs, we set neurons' values for all inputs, and do not limit only to resampling.
Pruning: Model pruning [10] is a common practice wherein reduced neural network parameter counts lessen memory and performance costs. In particular, structured pruning of large features [11] is interested in the removal on the scale of neurons and attention heads, and can often achieve a large reduction in parameter count [12]. Our work seeks to question the assumption of using masks that set neuron values to zero.
Modularity: Research into activation sparsity [13], modularity [14], mixture of experts [15,14,16], and unlearning by pruning [17,18] all investigate how different subsets of activations are responsible for different tasks. These implicitly set activations to zero.
We work with two causal text models, Mistral 7B [19] and Meta's OPT 1.3B [20], a masked text model, RoBERTa Large [21], and a vision transformer, ViT Base Patch16 224 [22].
To get a general sense of performance, the above models were evaluated by looking at top1 prediction accuracy1 , as well as cross-entropy loss on various datasets. For text models, we assess on EleutherAI's 'The Pile' [23]. For image models, we assess on Imagenet-1k [24], an image dataset with 1000 different classes. We evaluate on deterministic subsets of 100,000 text tokens and 1000 images respectively
The objects of study are attention pre-out neurons, sometimes called 'z'-hook activations. We define attention pre-out neuron activations π¦ π = π (π₯ π ) = preout(π₯ π ) as π¦ π = βοΈ π π΄ π,π π π π₯ π , where π΄ π,π = softmax((π πΎ π₯ π ) β’ (π π π₯ π )), where π π , π πΎ , π π are the attention query, key, and value matrices respectively. We focus on attention neurons rather than MLP neurons, as these do not have an activation function that privileges positive activations, making analysis more difficult. To ablate a neuron, we replace π¦ π = π (π₯ π ) with some constant.
In Figure 1, we showcase some plots of neuron probability distributions. We see an example of many attention pre-out neuron activation distributions within the same layer. We note that most neurons follow a roughly Gaussian or double-exponential distribution about zero, but note that there is a minority of neurons that are not distributed at zero. As most neurons are zero-centred and symmetric, it makes sense that zero and mean ablation work quite well. We see in (left) an average of distributions of all neurons in a layer, (centre) a bi-modal neuron with both peaks not at zero, and (right) another example of a neuron with an atypical distribution. X-axis shows neuron value, and Y-axis shows probability of a neuron taking that value.
Our hypothesis, based on activation profiles such as those seen in Figure 1 is that neurons have a 'baseline' or 'default-mode' activation (typically at zero), when the input contains no relevant features, which is then deviated from as neurons fire in proportion to various features they are tuned to pick up. In residual stream models [25], information is limited to the width of the residual stream [26], and as the residual stream typically grows exponentially in size [27], noise can become amplified. This is supported by the common redundancy of many circuits [28], even in transformer models trained without the use of dropout [29].
In particular, we expect that ablating neurons should have two contributors to reduced performance. These are 1) removing the relevant contextual computed information that the neuron is providing, and 2) taking the model activation out of distribution, by adding 'noise'. Ablating neurons to a constant value should cause some constant increase in loss for term 1, and different constant should contribute to different values of term 2. As we increase the distance from the 'default-mode' value, the neuron would further degrade the performance by taking the residual stream further out of distribution, thus in some sense, 'adding noise'.
We choose four main methods of ablating neurons, see Table 1 for a summary. These are:
The most common form of ablation, which involves replacing a neuronal activation of any with a zeroed out activation. That is, setting βπ₯ π : π π (π₯ π ) = 0.0 Mean Ablation: A still relatively-common method of ablation, which involves first collecting activations of various neurons on a distribution of inputs, and averaging the activations to find a mean activation. That is, for some dataset π·, for π₯ π β π·, let π π (π₯ π ) = 1
π π π (π₯ π ) Activation Resampling: Inspired by [3,4], we also try general neuron resampling, by setting activations to those found by giving a randomised input. 2 For text model, we take activations by a) sampling random generated characters, b) sampling random tokens, and c) using OPT to generate a random text. For ViT, we use randomly generated pixel values.
Naive Peak Ablation: Observing that neuronal activations frequently exhibit a prominent peak, we propose an ablation method targeting their modal activation. For bin size π, the neuron π activations π π (π₯ π ) for each π₯ π β π· are sorted into bins π π [π] such that π¦ π β€ π π (π₯ π ) < π¦ π + π. The bin π π [π πππ₯ ] with the highest occurrence is selected, and π π (π₯ π ) is set to π¦ ππππ₯ + π 2 .
Under the working model described in Section 3.3, we expect that ablating neurons to different values should have different impacts to performance, with there being a value which leads to some minimal drop in performance due to minimal noise being added to the residual stream. We randomly select attention neurons in increments of 10% and ablate them until the model is fully pruned, and at each step, assess performance by evaluating the Top1 accuracy and Cross-Entropy Loss in the chosen dataset with each ablation method, described in Table 1. The neurons are selected deterministically across three separate seeds, summarised in Table 2 4. Results
In Figure 2, we see the results for random pruning of OPT 1.3b and Mistral 7b with the different methods of ablation. We can see that Peak ablation has the most consistent pattern, causing the lowest amount of degradation, with mean ablation and zero ablation coming a close second and third, and Random resampling causes by far the most degradation. Of the three resampling methods, choosing random tokens causes the lowest degradation.
In Figure 3, we see that for ViT, zero, mean, and peak ablation have statistically insignificant differences in performance, while resampling causes some small additional degradation. We can see that almost all of the performance loss is based on the specific neurons being selected rather than the ablation method being chosen, even between zero, mean, and peak ablation.
In RoBERTa, we see that in the first 75% of pruning, the three main methods of peak, mean and zero ablation are very close, with Peak edging slightly better performance. Beyond 75%, the three methods become more noisy; resampling of IDs ends up having the best performance overall in both Top1 and Cross-Entropy Loss at the task of token unmasking and de-randomization.
In Table 2, we see that in different models and in different regimes, the different methods have different merits in reducing performance, with Peak ablation working overall best in the most cases. Surprisingly, although Random resampling seems to add a lot of noise to the activations, random token ID resampling can sometimes work well, such as in RoBERTa.
The analysis presented seems to suggest that when evaluating and understanding neurons in the attention layers of language models, the ideal centring method seems to depend significantly on the model. In decoder models, a good method is to find the largest peak, with a close second being zero ablation. This similarity is expected, as most neurons are centred at zero. This has downstream effects on improving the way we can look at one of the most crucial aspects of how neural networks work -their activations.
We have seen that neurons can have activations that are: non-Gaussian, non-symmetric, multi-modal, non-zero-centred. We hypothesise that taking into consideration this fact has the potential to make interpretability analysis into more fruitful, and centring activations by their peak seems a potential natural method.
Future work could: 1) investigate other potential better methods for neuron recentring, 2) more thoroughly investigate the differences between 'well-behaved' symmetric zero-centred distributions, and those that deviate from this norm, 3) find more efficient ways of computing the peak activations for larger models.