Reimplement SAE training eval code

docs/sae_experiment_tracking.md

@@ -0,0 +1,84 @@
+# SAE Experiment Logging and Tracking
+
+An objective we have when training an SAE is to have a small number of non-zero elements in the hidden layer (i.e. a small L0 'norm'). The L1 norm is used as a proxy for this.
+
+The loss term of the SAE is:
+
+$$loss = mse + l1\_coeff * sparsity + [ \: ghost \: gadients \: ]$$
+
+The two loss term in the SAE loss function have conflicting goals:
+* The **reconstruction term** (MSE) works to make the autoencoder good at reconstructing the input.
+* The **sparsity** term works to reduce the magnitudes in the hidden layer. 
+
+When training an SAE most metrics are related to measuring reconstruction loss vs sparsity, in different ways. The goal is to monitor these and find pareto improvements.
+
+Weights and Biases is used to track experiments (toggled with `VisionModelSAERunnerConfig.log_to_wandb`)
+
+# Validation
+
+Whilst training an SAE validation metrics are logged every `VisionModelSAERunnerConfig.wandb_log_frequency` steps, whilst the sparsity is logged every `VisionModelSAERunnerConfig.feature_sampling_window` steps. Note that each step is `VisionModelSAERunnerConfig.train_batch_size` number of tokens, for a total of (num_images x num_epochs x tokens_per_image / training_batch) steps.
+
+During validation the following matrics are logged:
+
+#### plots
+
+* `log_feature_density_histogram` - every step since the last log the number of feature activations that are greater than 0 divided by the total amount. See [here](https://arena3-chapter1-transformer-interp.streamlit.app/[1.3.2]_Interpretability_with_SAEs) for a discussion on how to interpret this. 
+
+#### details
+
+* `details/current_learning_rate` - The learning rate at each step of training.
+* `details/n_training_tokens` - The number of training tokens seen so far (Images x tokens per image)
+* `details/n_training_images` - The number of training images seen so far
+
+#### losses
+
+* `losses/mse_loss` - The normalised MSE loss between the SAE input and output, i.e. the reconstruction loss. Normalised such that it is less dependent on the size of the input/output. 
+* `losses/l1_loss`
+* `losses/ghost_grad_loss` - this describes the method of adding an additional term to the loss, which essentially gives dead latents a gradient signal that pushes them in the direction of explaining more of the autoencoder's residual.
+`losses/overall_loss`
+
+#### metrics
+
+
+* `metrics/explained_variance` - explained variance is the ratio of the variance in the reconstructed data to the variance in the original input data. A higher explained variance indicates that the SAE has learned features that capture more of the meaningful structure in the data. This is the mean over the batch.
+* `metrics/explained_variance_std` - This is the std over the batch.
+* `metrics/l0` - the feature activations for that step greater than 0 summed together, and the mean over the batch.
+* `metrics/mean_log10_feature_sparsity`
+
+#### sparsity
+
+* `sparsity/mean_passes_since_fired` - For each activation, Tracks the number of training tokens (images x tokens per image) since activation was greater than 0
+* `sparsity/dead_features`
+* `sparsity/below_1e-5` - feature_sparsity < 1e-5
+* `sparsity/below_1e-6` - feature_sparsity < 1e-6
+
+# Evaluation
+
+There are two types of evaluation, namely that which is called during training and after training. This is because it is a costly procedure and the `TrainingEvalConfig` defines an evaluation run that is performed multiple times during training to give an idea of SAE performance and `PostTrainingEvalConfig` which aims to give a full breakdown of model performance. 
+
+## Training Evaluation
+
+Whilst training an SAE the training evaluation will run every `EvalConfig.log_frequency` and a post-training evaluation will run at the end. A training run can be kicked off with: 
+
+TODO EdS: Document which evals are available
+
+```python
+cfg = VisionModelSAERunnerConfig()
+trainer = VisionSAETrainer(cfg)
+sae = trainer.run()
+```
+
+## Post-Training Evaluation
+
+To manually kick off a post-training evaluation on a saved SAE:
+
+```python
+model = load_model(cfg)
+sae = load_sae(cfg)
+train_data, val_data, val_data_visualize = load_dataset(cfg)
+wandb.init(project=cfg.wandb_project, name=cfg.run_name)
+evaluator = Evaluator(model, val_data, cfg, visualize_data=val_data_visualize)
+evaluator.evaluate(sae, context="post-training")
+```
+
+TODO EdS: CLI implementation for evaluation

src/vit_prisma/sae/config.py

@@ -1,9 +1,15 @@
+import json
+import math
 from abc import ABC
 from dataclasses import fields, field, asdict, dataclass
-import json
-from typing import Any, Optional, cast, Literal
+from enum import Enum
+from typing import Any, Optional, Literal, Dict
+from typing import List
 
 import torch
+from transformers import ViTConfig
+
+from vit_prisma.utils.constants import Evaluation
 
 
 @dataclass
@@ -62,6 +68,8 @@ class RunnerConfig(ABC):
     def __post_init__(self):
         self.hook_point = f"blocks.{self.hook_point_layer}.{self.layer_subtype}"  # change hookpoint name here
 
+        self.num_patch = int(math.sqrt(self.context_size - 1))
+
         # Autofill cached_activations_path unless the user overrode it
         if self.cached_activations_path is None:
             self.cached_activations_path = f"activations/{self.dataset_path.replace('/', '_')}/{self.model_name.replace('/', '_')}/{self.hook_point}"
@@ -85,10 +93,59 @@ def pretty_print(self):
             print(f"  {field.name}: {value}")
 
 
+
+@dataclass
+class EvalConfig(ABC):
+    evaluation_functions: List[Evaluation]
+    eval_frequency: Optional[int]
+    batch_size: int
+    max_evaluation_images: int
+    samples_per_bin: int  # Number of features to sample per pre-specified interval
+    max_images_per_feature: int  # Number of max images to collect per feature
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+@dataclass
+class TrainingEvalConfig(EvalConfig):
+    """Default values to be used when evaluating an SAE during training. These are less
+    compute and time costly as they get run every step.
+    """
+
+    evaluation_functions: List[Evaluation] = field(
+        default_factory=lambda: []
+    )
+    eval_frequency: Optional[int] = 250
+    batch_size: int = 28
+    max_evaluation_images: int = 100
+    samples_per_bin: int = 10
+    max_images_per_feature: int = 20
+
+
+@dataclass
+class PostTrainingEvalConfig(EvalConfig):
+    """Default values to be used when evaluating an SAE post-training. These evaluations
+    are more extensive as they are run to get a complete picture of the quality of the
+    SAE.
+    """
+
+    evaluation_functions: List[Evaluation] = field(
+        default_factory=lambda: [
+            Evaluation.FEATURE_BASIS_EVAL,
+            # Evaluation.NEURON_BASIS_EVAL,
+        ]
+    )
+    eval_frequency: Optional[int] = 1
+    batch_size: int = 28
+    max_evaluation_images: int = 100
+    samples_per_bin: int = 10
+    max_images_per_feature: int = 20
+
+
 @dataclass
 class VisionModelSAERunnerConfig(RunnerConfig):
     """
-    Configuration for training a sparse autoencoder on a language model.
+    Configuration for training a sparse autoencoder on a vision model.
     """
 
     architecture: Literal["standard", "gated", "jumprelu"] = "standard"
@@ -114,7 +171,7 @@ class VisionModelSAERunnerConfig(RunnerConfig):
     train_batch_size: int = 1024 * 4
 
     # Imagenet1k
-    dataset_name: str = "imagenet1k"
+    dataset_name: str = 'imagenet1k'  # imagenet1k | cifar10
     dataset_path: str = "/network/scratch/s/sonia.joseph/datasets/kaggle_datasets"
     dataset_train_path: str = (
         "/network/scratch/s/sonia.joseph/datasets/kaggle_datasets/ILSVRC/Data/CLS-LOC/train"
@@ -143,6 +200,15 @@ class VisionModelSAERunnerConfig(RunnerConfig):
         "/network/scratch/s/sonia.joseph/sae_checkpoints/tinyclip_40M_mlp_out"
     )
 
+    # Evaluation
+    sae_path: str = '/network/scratch/s/sonia.joseph/sae_checkpoints/tinyclip_40M_mlp_out/1f89d99e-wkcn-TinyCLIP-ViT-40M-32-Text-19M-LAION400M-expansion-16/n_images_520028.pt'
+    training_eval: EvalConfig = field(
+        default_factory=TrainingEvalConfig
+    )
+    post_training_eval: EvalConfig = field(
+        default_factory=PostTrainingEvalConfig
+    )
+
     def __post_init__(self):
         super().__post_init__()
         if not isinstance(self.expansion_factor, list):
@@ -234,6 +300,8 @@ def make_serializable(obj):
             elif isinstance(obj, dict):
                 # Recursively process dictionaries
                 return {key: make_serializable(value) for key, value in obj.items()}
+            elif isinstance(obj, Enum):
+                return obj.value
             else:
                 return obj  # Other types are left as-is