Package: vit-prisma
Vision Transformer (ViT) Prisma LibraryViT Prisma is an open-source mechanistic interpretability library for Vision Transformers (ViTs). This library was created bySonia Joseph.Contributors:Praneet Suresh,Yash Vadi,Rob Graham[and more coming soon]We welcome new contributors. Check out our contributing guidelineshere.Installing RepoTo 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 outour 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 ToolingComing soon.Attention head visualization codeAttention Head Demo NotebookPart 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 MiceToy 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 checkpointsThe detailed training logs and metrics can be foundhere. These models were trained by Yash Vadi.Table of ResultsAccuracy[ <Acc> | <Top5 Acc> ]SizeNumLayersAttention+MLPAttentionOnlyModel Linktiny10.16 | 0.330.11 | 0.25AttentionOnly,Attention+MLPbase20.23 | 0.440.16 | 0.34AttentionOnly,Attention+MLPsmall30.28 | 0.510.17 | 0.35AttentionOnly,Attention+MLPmedium40.33 | 0.560.17 | 0.36AttentionOnly,Attention+MLPdSprites Shape Classification training checkpointsOriginal dataset ishere.Full results and training setup arehere. These models were trained by Yash Vadi.Table of ResultsSizeNumLayersAttention+MLPAttentionOnlyModel Linktiny10.5350.459AttentionOnly,Attention+MLPbase20.9960.685AttentionOnly,Attention+MLPsmall31.0000.774AttentionOnly,Attention+MLPmedium41.0000.991AttentionOnly,Attention+MLPTo do: to train1-4 layer, attention-only, and full-attention versions of each.ImageNet-1k Mice (reconstruction loss)Attention-only1-layer, 2-layer, 3-layer, 4-layer, 5-layerFull-attention1-layer, 2-layer, 3-layer, 4-layer, 5-layerdSpritesSmallest possible model that can recognize size.Smallest possible model that can recognize position.Guidelines for training + uploading modelsUpload your trained models to Huggingface. Follow theHuggingface guidelinesand 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 ourUsage Guidefor more information.ConfigurationsThe 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:InductionConfigCircleConfigMNISTConfigDSpritesConfigfromvit_prisma.configsimport<CONFIG_OF_CHOICE>By inspecting one of the configs available in the framework, you can create your own one.Base ViTThe 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 ViTfromvit_prisma.modelsimportbase_vitfromvit_prisma.configsimportInductionConfig# The InductionConfig is available in the framework itself, which is setup for the induction datasetconfig=InductionConfig.GlobalConfig()model=base_vit.BaseViT(config)Using pretrained modelsfromvit_prisma.models.pretrained_modelimportPretrainedModelfromvit_prisma.configsimportInductionConfigconfig=InductionConfig.GlobalConfig()hf_model=PretrainedModel('google/vit-base-patch16-224-in21k',config)timm_model=PretrainedModel('vit_base_patch32_224',config,is_timm=True)DatasetsThe framework also has a few datasets that can be synthetically generated, and cached.Below are the datasets available as part of the framework:CircledspritesInductionYou will be able to use any other dataset with the framework, for ease of use, you can instantiate it as a Pytorch dataset object.fromvit_prisma.dataloaders.inductionimportInductionDatasettrain=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 experimentsThe 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.fromvit_prisma.models.base_vitimportBaseViTfromvit_prisma.configsimportInductionConfigfromvit_prisma.trainingimporttrainerfromvit_prisma.dataloaders.inductionimportInductionDatasettrain_dataset=InductionDataset('train')config=InductionConfig.GlobalConfig()model_function=BaseViTtrainer.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.CitationPlease 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}}
}