Activation additions in a simple MNIST network

Abstract

Team Shard’s recent activation addition methodology for steering GPT-2-XL provokes many questions about what the structure of internal model computation must be in order for their edits to work. Unfortunately, interpreting the insides of neural networks is known to be very difficult, so the question becomes: what is the minimal set of properties a network must have in order for adding activation additions to work? I make some progress on this question by gathering data on whether or not activation additions work on an old 3 layer, fully connected MNIST network I had laying around. Space is left in case the reader would like to predict what the results were.

I find that when you add an a - b activation vector for all a and b I tested to any hidden layer of the model (before or after nonlinear layers), it makes the model less likely to answer b when prompted with images of b, and more likely to answer a when prompted with images of b, generalization behavior is destroyed, going from ~0.089 perplexity to ~7.4 perplexity when the dataset is cleaned of vector b.

Colab notebook here

Methodology

I used a simple 3 layer fully connected neural network trained to 0.089 log loss on the MNIST test dataset. The following is the model schematic:

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

To make an a - b vector with a and b images of particular numbers, I found the first image in the training dataset which matched that number, and used that

I tried all combinations of a - b vectors, with coefficient 1, and intervened at layer 0, since for the particular addition vector 3 - 1, intervening at layer 0 had the best results, as measured by perplexity when the dataset was cleaned of 1 examples. The following are the layer interventions for the vector 3 - 1 and their corresponding perplexities:

LayerPerplexity
07.80
112.69
212.69
317.01
417.01

Where the layer number corresponds to intervening just after the model layer as represented in the model schematic given earlier in this section.

Finally for every a - b vector, and each number, I took the average logit distribution the a - b patched model assigned to all test images with that number, turned that into a probability distribution, and plotted the results.

Results

Here is a table of perplexity when the activation addition a - b is made on layer 0, along with the unpatched model when scored on the same dataset (this is only dependent on b, since we remove b from the scoring dataset).

a\b0123456789
00.115.012.99.313.98.46.812.58.919.4
19.70.12.58.60.810.31.81.90.92.0
21.62.40.14.74.28.33.36.33.110.3
319.012.715.90.121.61.024.417.118.716.4
410.40.92.214.10.122.24.57.28.90.8
517.123.326.79.525.20.10.97.62.219.5
64.03.66.717.91.82.90.117.76.29.9
76.74.06.710.02.88.413.20.111.64.5
84.31.84.40.72.81.41.13.40.11.0
94.94.51.37.02.014.75.00.61.40.1
Normal0.0940.0960.0920.0840.0900.0920.0850.0910.0780.085

Here’s also a random sampling of graphs, showing the average effect of the patch for each number. These graphs are, in my opinion, really weird. There’s probably a bunch of interesting observations one could make about patterns in them. Hypotheses are welcome! But for now I note it really does seem like generalization was destroyed.



No comments.