DOCSEARCHbyAlgolia
Submit Docs

Prisma Toolkit

repository·main·Indexed 18 days ago

https://github.com/prisma-multimodal/vit-prisma

An open-source toolkit for mechanistic interpretability in vision and video models. Prisma provides infrastructure for circuit analysis via activation caching and Sparse Autoencoder (SAE) research, supporting models from Huggingface and OpenCLIP including ViTs, CLIP, DINO, and JEPA. It includes tools for training and evaluating SAEs, attention head visualization, activation patching, and a CLIP benchmark evaluation framework.

Tokens
40.8K
Snippets
86
Records
105
Agent score
63%

What's inside Prisma

  1. Overview of Prisma capabilities

    main

    Prisma is an open-source toolkit for mechanistic interpretability in vision and video. It provides infrastructure for two primary research areas:

    1. Circuit-analysis: Finding causal links between internal model components, primarily using activation caching.
    2. Sparse Autoencoders (SAEs): Examining intermediate activations using fine-grained primitives.

    The toolkit includes support for training, evaluating, and using pretrained SAEs for models like CLIP and DINO.

  2. Explore pretrained Vision SAE weights

    main

    Prisma provides a suite of pretrained Sparse Autoencoder (SAE) weights for mechanistic interpretability research in vision and video.

    Recommended Starting Point

    If you are new to steering CLIP's output, it is highly recommended to start with the Layer 11 resid-post SAE, as it offers high quality for steering tasks.

    Available SAE Types

    • CLIP Vanilla SAEs (All Patches): High-quality SAEs for various CLIP layers and sublayers (e.g., mlp_out, resid_post).
    • DINO (Vanilla, all patches): SAEs trained on DINO models across all layers.
    • CLIP Transcoders: Specialized architectures for mapping between layers.

    For a comprehensive list of all available SAEs, including CLIP top k and specific layer configurations, refer to the SAE Table documentation.

  3. Understand SAE loss terms and training objectives

    main

    The training objective for an SAE involves balancing reconstruction accuracy against sparsity. The loss function is defined as:

    loss = mse + l1_coeff * sparsity + [ghost gradients]

    Key components include:

    • Reconstruction term (MSE): Aims to minimize the error between the input and the reconstructed output.
    • Sparsity term: Aims to reduce the magnitudes in the hidden layer (often using the L1 norm as a proxy for the L0 norm).
    • Ghost gradients: An additional term used to provide dead latents with a gradient signal, pushing them to explain more of the autoencoder's residual.

  4. Understand the Polygenic Induction dataset

    main

    The Polygenic Induction dataset is a synthetic dataset designed to uncover induction heads in Vision Transformers (ViTs). It categorizes images based on two dimensions: Orientation (Horizontal 'H' or Vertical 'V') and a specific Pattern sequence.

    There are 12 total classes (IDs 0-11) defined by the following mapping:

    IDOrientationPattern
    0HAAAA
    1HABAB
    2HABBA
    3HAABB
    4HABBB
    5HAAAB
    6VAAAA
    7VABAB
    8VABBA
    9VAABB
    10VABBB
    11VAAAB

  6. Learn basic mechanistic interpretability techniques

    main

    Prisma supports several core mechanistic interpretability techniques for Vision Transformers (ViT). You can explore these through interactive Colab notebooks:

    1. Main ViT Demo: Covers direct logit attribution, attention head visualization, and activation patching (e.g., switching predictions between classes with minimal ablation).
    2. Emoji Logit Lens: Provides a deeper dive into layer- and patch-level predictions using interactive plots.
    3. Interactive Attention Head Tour: Explores the various types of attention heads in a ViT using interactive JavaScript.

    Key features include:

    • Attention head visualization: Visualizing how specific heads focus on image regions.
    • Activation patching: Testing the causal effect of specific activations on model output.
    • Direct logit attribution: Attributing model decisions to specific features.
    • Emoji logit lens: Visualizing layer-wise predictions through segmentation and emoji mapping.

  8. Recommended Sparse Autoencoders (SAEs) for CLIP-ViT-B-32

    main
    When working with the CLIP-ViT-B-32 model, it is empirically recommended to use the Vanilla SAEs (All Patches). These SAEs are trained on all patches (both CLS and spatial), which provides superior performance compared to alternatives like CLS-only, Top K (k=64), or spatial-patches-only models.

  9. Using VJEPA HF in your code

    main

    To use the Huggingface version of VJEPA, you need to use the VJEPAModel and VJEPAImageProcessor classes. The workflow involves loading a configuration from CONFIGS, processing an image with VJEPAImageProcessor, preparing the tensor dimensions (typically repeating the image for temporal dimension and permuting to B x C x T x H x W), and loading the model via VJEPAModel.from_pretrained(model_path).

    from PIL import Image
import requests
from app.vjepa_hf.modeling_vjepa import VJEPAModel, VJEPAImageProcessor
from app.vjepa_hf.configs import CONFIGS
import yaml

def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
    return image

model_name = "vjepa_v1.5_vit_g_256"
config = CONFIGS["v1.5"]["vit_g_256"]
model_paths = yaml.safe_load(open('paths_cw.yaml'))
model_path = model_paths[model_name]["loc"]

image = prepare_img()
processor = VJEPAImageProcessor(crop_size=config.crop_size)
pixel_values = processor(image, return_tensors="pt").pixel_values
pixel_values = pixel_values.repeat(1, 16, 1, 1, 1) # repeating image 16 times for now
pixel_values = pixel_values.permute(0, 2, 1, 3, 4)  # B x C x T x H x W

model = VJEPAModel.from_pretrained(model_path)
outputs = model(pixel_values)

  10. Convert VJEPA checkpoints to Huggingface format

    main

    Use the app.vjepa_hf.convert_vjepa_to_hf module to convert standard VJEPA checkpoints into the Huggingface format. You must specify the model name, the source checkpoint location, the output destination, and the model type (e.g., v1.5).

    python -m app.vjepa_hf.convert_vjepa_to_hf \
    --model_name vit_g_256 \
    --ckpt_loc <vjepa path> \
    --outp_loc <outp path> \
    --model_type v1.5 \
    --convert

  12. Use custom ImageNet-1k classification checkpoints

    main

    Prisma provides several pre-trained ImageNet-1k classification checkpoints for research and analysis. These are categorized by patch size and model architecture (AttentionOnly vs. Attention+MLP).

    Patch Size 32

    These models have inspectable attention heads (patch size 16 heads are often too large for easy JS rendering).

    • tiny (3 layers): Available as Attention+MLP.

    Patch Size 16

    Models are available in various sizes:

    • tiny (1 layer)
    • base (2 layers)
    • small (3 layers)
    • medium (4 layers)

    Each size supports both AttentionOnly and Attention+MLP architectures. Accuracy is reported as [ <Acc> | <Top5 Acc> ].