add training png

.gitignore

@@ -0,0 +1 @@
+.DS_Store

README.md

@@ -6,26 +6,29 @@ license: cc-by-4.0
 
 by [Louka Ewington-Pitsos](https://www.linkedin.com/in/louka-ewington-pitsos-2a92b21a0/?originalSubdomain=au) and Ram ____
 
-Heavily inspired by [google/gemma-scope](https://huggingface.co/google/gemma-scope) we are releaseing a suite of sparse 8 autoencoders for [laion/CLIP-ViT-L-14-laion2B-s32B-b82K](https://huggingface.co/laion/CLIP-ViT-L-14-laion2B-s32B-b82K). 
+Heavily inspired by [google/gemma-scope](https://huggingface.co/google/gemma-scope) we are releasing a suite of 8 sparse autoencoders for [laion/CLIP-ViT-L-14-laion2B-s32B-b82K](https://huggingface.co/laion/CLIP-ViT-L-14-laion2B-s32B-b82K). 
 
-| 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.201   |          0.731          |                0         |
-| 20    | 278.06   |        0.706            |           0.0000763              | 
-| 22    |   299.96  |        0.684             |       0                   |
+![](./media/mse.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 may contain some strange artefacts when a spot instance is killed and the training is reumed by another instance. Some of the code is ripped directly from [Hugo Fry](https://github.com/HugoFry/mats_sae_training_for_ViTs).
+| 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                |
+
+
+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).
 
 ### Vital Statistics:
 
 - Number of tokens trained per autoencoder: 1.2 Billion
 - Token type: all 257 image tokens (as opposed to just the cls token)
-- Number of images trained per autoencoder: ~4.6 Million
+- 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
@@ -34,15 +37,23 @@ Training logs are available [via wandb](https://wandb.ai/lewington/ViT-L-14-laio
 
 ## Usage
 
-## Error Forumlas
+## Error 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 Explained Variance as 
 
-We calculate MSE as `((batch - reconstruction) ** 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.
+```python
+delta_variance = (batch - reconstruction).pow(2).sum(dim=-1)
+activation_variance = (batch - batch.mean(dim=-1, keepdim=True)).pow(2).sum(dim=-1)
+explained_variance = (1 - delta_variance / activation_variance).mean()
+```
 
-We calculate Explained Variance as
+We calculate dead feature proportion as the proportion of features which have not activated in the last 10,000,000 samples.
 
-## Subjective Interpretability
+## 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 feaure 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 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.
 
 
 ## Automated Sort EVALs