A Vision Transformer library for mechanistic interpretability.
Project description
Vision Transformer (ViT) Prisma Library
ViT Prisma is an open-source mechanistic interpretability library for Vision Transformers (ViTs). This library was created by Sonia Joseph.
Contributors: Praneet Suresh, Yash Vadi, Rob Graham [and more coming soon]
We welcome new contributors. Check out our contributing guidelines here.
Installing Repo
To install as an editable repo:
git clone https://github.com/soniajoseph/ViT-Prisma
cd ViT-Prisma
pip install -e .
How do I use this repo?
Check out our guide.
What's in the repo?
Part One: Mechanistic interpretability tooling, including activation caching, path-patching, and attention head visualization. In progress.
Part Two: Open source mini transformers (ViT "mice"). In progress.
Part Three: Code to train your own vision transformers.
Part One: Mechanistic Interpretability Tooling
Coming soon.
Attention head visualization code
Part Two: Open Source Mini-Transformers (ViT Mice 🐭)
ViT Mice are the mini-versions of the standard Vision Transformers. Just as mice are used in scientific experiments for their small size, ViT Mice serve a similar purpose to illuminate their larger counterparts. By training these mice on both toy datasets and in-the-wild data, we aim to observe their behaviors in various environments.
Categories of ViT Mice
- Toy Data Mice: Trained on controlled, synthetic datasets to understand specific behaviors or to isolate certain aspects of the learning process.
- In-the-Wild Mice: Trained on naturalistic, real-world data reminiscent of models in-production.
ImageNet-1k classification training checkpoints
The detailed training logs and metrics can be found here. These models were trained by Yash Vadi.
Table of Results
Accuracy [ <Acc> | <Top5 Acc> ]
Size | NumLayers | Attention+MLP | AttentionOnly | Model Link |
---|---|---|---|---|
tiny | 1 | 0.16 | 0.33 | 0.11 | 0.25 | AttentionOnly, Attention+MLP |
base | 2 | 0.23 | 0.44 | 0.16 | 0.34 | AttentionOnly, Attention+MLP |
small | 3 | 0.28 | 0.51 | 0.17 | 0.35 | AttentionOnly, Attention+MLP |
medium | 4 | 0.33 | 0.56 | 0.17 | 0.36 | AttentionOnly, Attention+MLP |
dSprites Shape Classification training checkpoints
Original dataset is here.
Full results and training setup are here. These models were trained by Yash Vadi.
Table of Results
Size | NumLayers | Attention+MLP | AttentionOnly | Model Link |
---|---|---|---|---|
tiny | 1 | 0.535 | 0.459 | AttentionOnly, Attention+MLP |
base | 2 | 0.996 | 0.685 | AttentionOnly, Attention+MLP |
small | 3 | 1.000 | 0.774 | AttentionOnly, Attention+MLP |
medium | 4 | 1.000 | 0.991 | AttentionOnly, Attention+MLP |
To do: to train
1-4 layer, attention-only, and full-attention versions of each.
- ImageNet-1k Mice (reconstruction loss)
- Attention-only
- 1-layer, 2-layer, 3-layer, 4-layer, 5-layer
- Full-attention
- 1-layer, 2-layer, 3-layer, 4-layer, 5-layer
- Attention-only
- dSprites
- Smallest possible model that can recognize size.
- Smallest possible model that can recognize position.
Guidelines for training + uploading models
Upload your trained models to Huggingface. Follow the Huggingface guidelines and also create a model card. Document as much of the training process as possible including links to loss and accuracy curves on weights and biases, dataset (and order of training data), hyperparameters, optimizer, learning rate schedule, hardware, and other details that may be relevant.
Include frequent checkpoints throughout training, which will help other researchers understand training dynamics.
Part Three: ViT Training Code 🚀
This repo includes training code to easily train ViTs from scratch or finetune existing models. See our Usage Guide for more information.
Configurations
The Prisma config object houses certain hyperparameters of the models, along with attributes of the training procedure, which can come in handy to track, and design new experiments.
There are several config objects available in the framework.
Below are the configs available as part of the framework:
- InductionConfig
- CircleConfig
- MNISTConfig
- DSpritesConfig
from vit_prisma.configs import <CONFIG_OF_CHOICE>
By inspecting one of the configs available in the framework, you can create your own one.
Base ViT
The base ViT part of the Prisma Project allows instantiating models in a flexible way, where the architecture is defined by the Prisma config object discussed in the previous section.
Instantiating the Base ViT
from vit_prisma.models import base_vit
from vit_prisma.configs import InductionConfig
# The InductionConfig is available in the framework itself, which is setup for the induction dataset
config = InductionConfig.GlobalConfig()
model = base_vit.BaseViT(config)
Using pretrained models
from vit_prisma.models.pretrained_model import PretrainedModel
from vit_prisma.configs import InductionConfig
config = InductionConfig.GlobalConfig()
hf_model = PretrainedModel('google/vit-base-patch16-224-in21k', config)
timm_model = PretrainedModel('vit_base_patch32_224', config, is_timm=True)
Datasets
The framework also has a few datasets that can be synthetically generated, and cached.
Below are the datasets available as part of the framework:
- Circle
- dsprites
- Induction
You will be able to use any other dataset with the framework, for ease of use, you can instantiate it as a Pytorch dataset object.
from vit_prisma.dataloaders.induction import InductionDataset
train = InductionDataset('train')
# If you don't provide a test split to the trainer, the trainer will perform the train/test split for you.
Training and tracking experiments
The trainer has built in support for wandb experiment tracking. Make sure to set up wandb on your localhost, and configure the tracking attributes in the Prisma config object.
from vit_prisma.models.base_vit import BaseViT
from vit_prisma.configs import InductionConfig
from vit_prisma.training import trainer
from vit_prisma.dataloaders.induction import InductionDataset
train_dataset = InductionDataset('train')
config = InductionConfig.GlobalConfig()
model_function = BaseViT
trainer.train(model_function, config, train_dataset)
Errors 💰
If you point out a conceptual error in this code (e.g. incorrect implementation of a transformer, not a minor function import), I will send you $5-20 per bug depending on the bug's severity.
Citation
Please cite this repository when used in papers or research projects.
@misc{joseph2023vit,
author = {Sonia Joseph},
title = {ViT Prisma: A Mechanistic Interpretability Library for Vision Transformers},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/soniajoseph/vit-prisma}}
}
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for vit_prisma-0.1.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 71b54f5e0370ab95ee786b8d3c17895b7941e0c9ae1246872c989c87ed9c0fa7 |
|
MD5 | 6011ebb9dad42814c68e7af53378f852 |
|
BLAKE2b-256 | e9d0ef2b191c718f999b71ab3b9a2ba8594436c541f132b25839226e110732f2 |