finish writeup

.gitignore

@@ -1,2 +1,3 @@
 .DS_Store
-pad.ipynb
+pad.ipynb
+cruft

11_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 11
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/1259_grid.png)
+![](./examples/1462_grid.png)
+![](./examples/628_grid.png)

14_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 14
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/476_grid.png)
+![](./examples/70_grid.jpg)
+![](./examples/843_grid.png)

17_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 17
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/157_grid.png)
+![](./examples/568_grid.png)
+![](./examples/606_grid.png)

20_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 20
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/195_grid.png)
+![](./examples/363_grid.jpg)
+![](./examples/392_grid.png)

22_resid/README.md

@@ -2,8 +2,8 @@
 
 ## Examples
 
-3x3 grids cherry picked from the first 100 features which activate on more than 9/500,000 laion images.
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
 
 ![](./examples/308_grid.png)
 ![](./examples/464_grid.png)
-![](./examples/575_grid.png)
+![](./examples/575_grid.png)

2_resid/README.md

@@ -0,0 +1,8 @@
+# Layer 2
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images. Only 2 features from among the first 2048 actually activated in this case.
+
+![](./examples/1162_grid.png)
+![](./examples/1173_grid.png)

5_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 5
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/1952_grid.png)
+![](./examples/2023_grid.png)
+![](./examples/648_grid.png)

8_resid/README.md

@@ -0,0 +1,9 @@
+# Layer 8
+
+## Examples
+
+3x3 grids cherry picked from the first 100 features which activate on the CLS token on more than 9/500,000 laion images.
+
+![](./examples/1597_grid.png)
+![](./examples/1735_grid.png)
+![](./examples/187_grid.png)

README.md

@@ -10,36 +10,40 @@ Heavily inspired by [google/gemma-scope](https://huggingface.co/google/gemma-sco
 
 ![](./media/mse.png)
 
-| Layer |   MSE    | Explained Variance | Dead Feature Proportion |
-|-------|----------|--------------------|-------------------------|
-|   2   |  267.95  |        0.763       |        0.000912         |
-|   5   |  354.46  |        0.665       |        0                |
-|   8   |  357.58  |        0.642       |        0                |
-|  11   |  321.23  |        0.674       |        0                |
-|  14   |  319.64  |        0.689       |        0                |
-|  17   |  261.20  |        0.731       |        0                |
-|  20   |  278.06  |        0.706       |        0.0000763        |
-|  22   |  299.96  |        0.684       |        0                |
 
+| Layer |   MSE   | Explained Variance | Dead Feature Proportion | Active Feature Proportion | Sort Eval Accuracy (CLS token) |
+|-------|---------|--------------------|-------------------------|---------------------------|-------------------------------|
+|   2   | 267.95  |        0.763       |        0.000912         |           0.001            |               -               |
+|   5   | 354.46  |        0.665       |            0            |           0.0034           |               -               |
+|   8   | 357.58  |        0.642       |            0            |           0.01074          |               -               |
+|  11   | 321.23  |        0.674       |            0            |           0.0415           |            0.7334              |
+|  14   | 319.64  |        0.689       |            0            |           0.07866          |            0.7427              |
+|  17   | 261.20  |        0.731       |            0            |           0.1477           |            0.8689              |
+|  20   | 278.06  |        0.706       |        0.0000763        |           0.2036           |            0.9149              |
+|  22   | 299.96  |        0.684       |            0            |           0.1588           |            0.8641              |
 
-Training logs are available [via wandb](https://wandb.ai/lewington/ViT-L-14-laion2B-s32B-b82K/workspace) and training code is available on [github](https://github.com/Lewington-pitsos/vitsae). The training process is heavily reliant on [AWS ECS](https://aws.amazon.com/ecs/) so may contain some strange artifacts when a spot instance is killed and the training is resumed by another instance. Some of the code is ripped directly from [Hugo Fry](https://github.com/HugoFry/mats_sae_training_for_ViTs).
+
+
+![](./media/sort-eval.png)
+
+Training logs are available [via wandb](https://wandb.ai/lewington/ViT-L-14-laion2B-s32B-b82K/workspace) and training code is available on [github](https://github.com/Lewington-pitsos/vitsae). The training process is heavily reliant on [AWS ECS](https://aws.amazon.com/ecs/) so the weights and biases logs may contain some strange artifacts when a spot instance is killed and the training is resumed by another instance. Some of the code is ripped directly from [Hugo Fry](https://github.com/HugoFry/mats_sae_training_for_ViTs).
 
 ### Vital Statistics:
 
 - Number of tokens trained per autoencoder: 1.2 Billion
-- Token type: all 257 image tokens (as opposed to just the cls token)
+- Token type: all 257 image tokens (as opposed to just the CLS token)
 - Number of unique images trained per autoencoder: 4.5 Million
 - Training Dataset: [Laion-2b](https://huggingface.co/datasets/laion/laion2B-multi-joined-translated-to-en)
 - SAE Architecture: topk with k=32
 - Layer Location: always the residual stream
 - Training Checkpoints: every ~25 million tokens
-- Number of features: 65536
+- Number of features per autoencoder: 65536 (expansion factor 16)
 
 ## Usage
 
-First install our pypi package and PIL
+First install our pypi package and PIL (pillow)
 
-`pip install clipscope PIL`
+`pip install clipscope pillow`
 
 Then
 
@@ -73,9 +77,9 @@ print('latent shape', output['latent'].shape) # (1, 65536)
 print('reconstruction shape', output['reconstruction'].shape) # (1, 1024)
 ```
 
-## Error Formulae
+## Formulae
 
-We calculate MSE as `(batch - reconstruction).pow(2).sum(dim=-1).mean()` i.e. The MSE between the batch and the un-normalized reconstruction, summed across features. We use batch norm to bring all activations into a similar range.
+We calculate MSE as `(batch - reconstruction).pow(2).sum(dim=-1).mean()` i.e. The MSE between the batch and the reconstruction, summed across features. We use batch norm to bring all activations into a similar range. We normalize the input before training using batch norm.
 
 We calculate Explained Variance as 
 
@@ -89,9 +93,41 @@ We calculate dead feature proportion as the proportion of features which have no
 
 ## Subjective Interpretability  
 
-To give an intuitive feel for the interpretability of these models we run 500,000 images from laion2b selected at random through the final trained SAE for each layer and record the latent activations for each. We then winnow down to the first 100 features which activate for at least 9 images. We cherry pick 3 of these and display them in a 3x3 grid for each layer. We do this twice, one for the CLS token and once for token 137 (near the middle of the image). Below are the 6 grids for feature 22. Other grids are available for each layer.
+To give an intuitive feel for the interpretability of these models we run 500,000 images from laion2b selected at random through the final trained SAE for each layer and record the latent activations on the CLS token for each. We then winnow down to the first 100 features which activate for at least 9 images. We cherry pick 3 of these and display them in a 3x3 grid for each layer. Below are the 3 grids for feature 22. Other grids are available for each layer in the `README.md` for that layer.
+
+**feature 308**
+
+![](./22_resid/examples/308_grid.png)
+
+**feature 464**
+
+![](./22_resid/examples/464_grid.png)
+
+**feature 575**
 
+![](./22_resid/examples/575_grid.png)
 
 ## Automated Sort EVALs
 
+We performed automated [Sort Evals](https://transformer-circuits.pub/2024/august-update/index.html) following Anthropic except with the *formatted dataset examples* for both features being replaced by a 3x3 grid of the top 9 activating images and the *formatted query example* being replaced by an image activating for only one of those features but not included in the 3x3s. Our methodology was as follows:
+
+1. pass 500,000 laion images from the training dataset through the SAE, record the first 2048 latent activations for each image (only 2048 to conserve space)
+2. ignore all features which do not activate for at least 10 of those images
+3. for the remaining features, if there are more than 100, select the first 100
+4. if there are fewer than 100, the comparison will not be fair so we cannot proceed (this is why we do not have sort eval scores for layers 8, 5 and 2)
+5. randomly select 500 pairs from among the 100 features, and perform a sort eval for each pair using gpt-4o
+4. select 400 samples randomly from among these 500 sort evals 5 times, each time recording the accuracy (n correct / 400) for that subsample of 400
+5. calculate the mean and standard deviation of these 5 accuracies.
+
+The outcomes are plotted below. Active Feature Proportion is the proportion of features which activate for at least 10 images across the 500,000 image dataset for that layer. For the CLS token at layer 2 only 2/2048 features were "active" in this sense.
+
+![](./media/sort-eval.png)
+
+![](./media/active-feature-proportion.png)
+
 ## Token-wise MSE
+
+All layers were trained across all 257 image patches. Below we provide plots demonstrating the reconstruction MSE for each token (other than the CLS token) as training progressed. It seems that throughout training the outer tokens are easier to reconstruct than those in the middle, presumably because these tokens capture more important information (i.e. foreground objects) and are therefore more information rich.
+
+![](./media/layer_22_training_outputs.png)
+![](./media/layer_22_individually_scaled.png)