Skip to content

README.md

Latest commit

 

History

History
149 lines (118 loc) · 4.42 KB

README.md

File metadata and controls

149 lines (118 loc) · 4.42 KB
Raw

PatchSAE: Training & Usage Guide

📋 Table of Contents

🔧 Train SAE

Train a sparse autoencoder on CLIP features.

Prerequisites

  • CLIP checkpoint
  • Training dataset (images only)

Training Command

PYTHONPATH=./ python tasks/train_sae_vit.py

📝 Configuration files will be added soon

Outputs

  • SAE checkpoint (.pt file)

Monitoring

📊 Extract SAE Latent Data

Extract and save SAE latent activations for downstream analysis.

Prerequisites

  • CLIP checkpoint
  • SAE checkpoint (from training step)
  • Dataset (can differ from training dataset)

Run with Original CLIP

PYTHONPATH=./ python tasks/compute_sae_feature_data.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --vit_type base

Run with Adapted CLIP (e.g., MaPLe)

  1. Download MaPLe from the official repo or Google Drive

  2. Run extraction:

PYTHONPATH=./ python tasks/compute_sae_feature_data.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --vit_type maple \
    --model_path /PATH/TO/MAPLE_CKPT \  # e.g., .../model.pth.tar-5
    --config_path /PATH/TO/MAPLE_CFG \  # e.g., .../configs/models/maple/vit_b16_c2_ep5_batch4_2ctx.yaml

Output Files

All files will be saved to: {root_dir}/out/feature_data/{vit_type}/{dataset_name}/

  • max_activating_image_indices.pt
  • max_activating_image_label_indices.pt
  • max_activating_image_values.pt
  • sae_mean_acts.pt
  • sae_sparsity.pt

Analysis

Explore the extracted features with our patchsae/analysis/analysis.ipynb

🧩 Compute Class-Level SAE Latents

Compute class-level SAE activation patterns.

Prerequisites

  • CLIP checkpoint
  • SAE checkpoint
  • SAE feature data (from previous step)
  • Dataset (must be the SAME dataset used in the extraction step)

Run with Original CLIP

PYTHONPATH=./ python tasks/compute_class_wise_sae_activation.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --threshold 0.2 \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --vit_type base

Run with Adapted CLIP (e.g., MaPLe)

PYTHONPATH=./ python tasks/compute_class_wise_sae_activation.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --threshold 0.2 \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --vit_type maple \
    --model_path /PATH/TO/MAPLE_CKPT \  # e.g., .../model.pth.tar-5
    --config_path /PATH/TO/MAPLE_CFG \  # e.g., .../configs/models/maple/vit_b16_c2_ep5_batch4_2ctx.yaml

Output File

  • cls_sae_cnt.npy - Matrix of shape (num_sae_latents, num_classes)

🎯 Steer Classification

Evaluate classification using feature steering with SAE latents.

Prerequisites

  • CLIP checkpoint
  • SAE checkpoint
  • Class-level activation data (cls_sae_cnt.npy from previous step)
  • Dataset (must be the SAME dataset used for class-level activations, though can be a different split)

Run with Original CLIP

PYTHONPATH=./ python tasks/classification_with_top_k_masking.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --cls_wise_sae_activation_path /PATH/TO/cls_sae_cnt.npy

Run with Adapted CLIP (e.g., MaPLe)

PYTHONPATH=./ python tasks/classification_with_top_k_masking.py \
    --root_dir ./ \
    --dataset_name imagenet \
    --sae_path /PATH/TO/SAE_CKPT.pt \
    --cls_wise_sae_activation_path /PATH/TO/cls_sae_cnt.npy \
    --vit_type maple \
    --model_path /PATH/TO/MAPLE_CKPT \  # e.g., .../model.pth.tar-5
    --config_path /PATH/TO/MAPLE_CFG \  # e.g., .../configs/models/maple/vit_b16_c2_ep5_batch4_2ctx.yaml

Output File

Output will be saved to eval_outputs/:

  • metrics.csv - Contains class-wise True Positive Rate (TPR = TP/(TP+FP+TN+FN)) for each masking configuration
    • Results for both "on" and "off" conditions
    • For k values in [1, 2, 5, 10, 50, 100, 500, 1000, 2000, SAE_DIM]