{"id":13737028,"url":"https://github.com/sayakpaul/probing-vits","last_synced_at":"2026-03-16T22:31:15.182Z","repository":{"id":39732751,"uuid":"469015692","full_name":"sayakpaul/probing-vits","owner":"sayakpaul","description":"Probing the representations of Vision Transformers.","archived":false,"fork":false,"pushed_at":"2022-10-05T11:41:11.000Z","size":34945,"stargazers_count":321,"open_issues_count":0,"forks_count":20,"subscribers_count":9,"default_branch":"main","last_synced_at":"2025-03-25T13:46:11.742Z","etag":null,"topics":["attention","explaining-vits","image-recognition","keras","pre-training","self-supervision","tensorflow","transformers","vits"],"latest_commit_sha":null,"homepage":"https://keras.io/examples/vision/probing_vits/","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sayakpaul.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-03-12T07:44:18.000Z","updated_at":"2025-02-21T22:11:20.000Z","dependencies_parsed_at":"2022-07-08T18:01:01.412Z","dependency_job_id":null,"html_url":"https://github.com/sayakpaul/probing-vits","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fprobing-vits","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fprobing-vits/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fprobing-vits/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fprobing-vits/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sayakpaul","download_url":"https://codeload.github.com/sayakpaul/probing-vits/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248455281,"owners_count":21106590,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["attention","explaining-vits","image-recognition","keras","pre-training","self-supervision","tensorflow","transformers","vits"],"created_at":"2024-08-03T03:01:33.832Z","updated_at":"2026-03-16T22:31:15.148Z","avatar_url":"https://github.com/sayakpaul.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook"],"sub_categories":[],"readme":"# Probing ViTs\n\n[![TensorFlow 2.8](https://img.shields.io/badge/TensorFlow-2.8-FF6F00?logo=tensorflow)](https://github.com/tensorflow/tensorflow/releases/tag/v2.8.0)\n[![HugginFace badge](https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-yellow.svg)](https://huggingface.co/spaces)\n\n_By [Aritra Roy Gosthipaty](https://github.com/ariG23498) and [Sayak Paul](https://github.com/sayakpaul) (equal contribution)_\n\nIn this repository, we provide tools to probe into the representations learned by different families of Vision Transformers (supervised pre-training with ImageNet-21k, ImageNet-1k, distillation, self-supervised pre-training):\n\n* Original ViT [1] \n* DeiT [2]\n* DINO [3]\n\nWe hope these tools will prove to be useful for the community. Please follow along with [this post on keras.io](https://keras.io/examples/vision/probing_vits/) for a better navigation through the repository. \n\n**Updates**\n\n* June 3, 2022: The project got the [Google OSS Expert Prize](https://www.kaggle.com/discussions/general/328914).\n* May 10, 2022: The project got a mention from Yannic Kilcher in [ML News](https://youtu.be/pwSnC8jlh50?t=712). Thanks, Yannic!\n* May 4, 2022: We're glad to receive the [#TFCommunitySpotlight award](https://twitter.com/TensorFlow/status/1521558632768409600?s=20\u0026t=hXgrZOfT_26AuTC_RyCZ_g) for this project.\n\n## Self-attention visualization\n\n| Original Image | Attention Maps | Attention Maps Overlayed |\n| :--: | :--: | :--: |\n| ![original image](./assets/bird.png) | ![attention maps](./assets/dino_attention_heads_inferno.png) | ![attention maps overlay](./assets/dino_attention_heads.png) |\n\nhttps://user-images.githubusercontent.com/36856589/162609884-8e51156e-d461-421d-9f8a-4d4e48967bd6.mp4\n\n\u003csmall\u003e\u003ca href=https://www.pexels.com/video/a-computer-generated-walking-dinosaur-4096297/\u003eOriginal Video Source\u003c/a\u003e\u003c/small\u003e\n\nhttps://user-images.githubusercontent.com/36856589/162609907-4e432dc4-a731-40f4-9a20-94e0c8f648bc.mp4\n\n\u003csmall\u003e\u003ca href=https://www.pexels.com/video/a-dog-running-in-a-grass-field-4166343/\u003eOriginal Video Source\u003c/a\u003e\u003c/small\u003e\n\n## Supervised salient representations\n\nIn the [DINO](https://ai.facebook.com/blog/dino-paws-computer-vision-with-self-supervised-transformers-and-10x-more-efficient-training/) blog post, the authors show a video with the following caption:\n\n\u003e The original video is shown on the left. In the middle is a segmentation example generated by a supervised model, and on the right is one generated by DINO. \n\nA screenshot of the video is as follows:\n\n\u003cimg width=\"764\" alt=\"image\" src=\"https://user-images.githubusercontent.com/36856589/162615199-b5133e51-460e-4864-a83e-5b8007339ff7.png\"\u003e\u003cbr\u003e\n\n\nWe obtain the attention maps generated with the supervised pre-trained model and find that they are not that salient w.r.t the DINO model. We observe a similar behaviour in our experiments as well. The figure below shows the attention heatmaps extracted with\na ViT-B16 model pre-trained (supervised) using ImageNet-1k:\n\n| Dinosaur | Dog | \n| :--: | :--: | \n| ![](./assets/supervised-dino.gif) | ![](./assets/supervised-dog.gif) | \n\nWe used this [Colab Notebook](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/vitb16-attention-maps-video.ipynb) to conduct this experiment.\n\n## Hugging Face Spaces\n\nYou can now probe into the ViTs with your own input images.\n\n| Attention Heat Maps | Attention Rollout |\n| :--: | :--: |\n| [![Generic badge](https://img.shields.io/badge/🤗%20Spaces-Attention%20Heat%20Maps-black.svg)](https://huggingface.co/spaces/probing-vits/attention-heat-maps) | [![Generic badge](https://img.shields.io/badge/🤗%20Spaces-Attention%20Rollout-black.svg)](https://huggingface.co/spaces/probing-vits/attention-rollout) |\n\n## Visualizing mean attention distances\n\n\u003cdiv align=\"center\"\u003e\n\u003cimg src=\"./assets/vit_base_i21k_patch16_224.png\" width=450/\u003e\n\u003c/div\u003e\n\n## Methods\n\n**We don't propose any novel methods of probing the representations of neural networks. Instead we take the existing works and implement them in TensorFlow.**\n\n* Mean attention distance [1, 4]\n* Attention Rollout [5]\n* Visualization of the learned projection filters [1]\n* Visualization of the learned positioanl embeddings\n* Attention maps from individual attention heads [3]\n* Generation of attention heatmaps from videos [3]\n\nAnother interesting repository that also visualizes ViTs in PyTorch: https://github.com/jacobgil/vit-explain.\n\n\n## Notes\n\nWe first implemented the above-mentioned architectures in TensorFlow and then we populated the pre-trained parameters into them using the official codebases. In order to validate this, we evaluated the implementations on the ImageNet-1k validation set and ensured that the reported top-1 accuracies matched. \n\nWe value the spirit of open-source. So, if you spot any bugs in the code or see a scope for improvement don't hesitate to open up an issue or contribute a PR. We'd very much appreciate it. \n\n\n## Navigating through the codebase\n\nOur ViT implementations are in `vit`. We provide utility notebooks in the `notebooks` directory which contains the following:\n\n* [`dino-attention-maps-video.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/dino-attention-maps-video.ipynb) shows how to generate attention heatmaps from a video. (Visually,) best results were obtained with DINO.\n* [`dino-attention-maps.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/dino-attention-maps.ipynb) shows how to generate attention maps from individual attention heads from the final transformer block. (Visually,) best results were obtained with DINO.\n* [`load-dino-weights-vitb16.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/load-dino-weights-vitb16.ipynb) shows how to populate the pre-trained DINO parameters into our implementation (only for ViT B-16 but can easily be extended to others). \n* [`load-jax-weights-vitb16.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/load-jax-weights-vitb16.ipynb) shows how to populate the pre-trained ViT parameters into our implementation (only for ViT B-16 but can easily be extended to others).\n* [`mean-attention-distance-1k.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/mean-attention-distance-1k.ipynb) shows how to plot mean attention distances of different transformer blocks of different ViTs computed over 1000 images.\n* [`single-instance-probing.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/single-instance-probing.ipynb) shows how to compute mean attention distance, attention-rollout map for a single prediction instance.\n* [`visualizing-linear-projections.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/visualizing-linear-projections.ipynb) shows visualizations of the linear projection filters learned by ViTs.\n* [`visualizing-positional-embeddings.ipynb`](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/visualizing-positional-embeddings.ipynb) shows visualizations of the similarities of the positional embeddings learned by ViTs.\n\nDeiT-related code has its separate repository: https://github.com/sayakpaul/deit-tf.\n\n## Models\n\nHere are the links to the models where the pre-trained parameters were populated:\n\n* [Original ViT model (pretrained on ImageNet-21k and fine-tuned on ImageNet-1k)](https://huggingface.co/probing-vits/vit_b16_patch16_224_i21k_i1k)\n* [Original ViT model (pretrained on ImageNet-1k)](https://huggingface.co/probing-vits/vit_b16_patch16_224_i1k)\n* [DINO model (pretrained on ImageNet-1k)](https://huggingface.co/probing-vits/vit-dino-base16)\n* [DeiT models (pretrained on ImageNet-1k including distilled and non-distilled ones)](https://tfhub.dev/sayakpaul/collections/deit/1)\n\n\n## Training and visualizing with small datasets\n\nComing soon!\n\n\n## References\n\n[1] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale:  [https://arxiv.org/abs/2010.11929](https://arxiv.org/abs/2010.11929)\n\n[2] DeiT: https://arxiv.org/abs/2012.12877\n\n[3] DINO: https://arxiv.org/abs/2104.14294\n\n[4] Do Vision Transformers See Like Convolutional Neural Networks?:  [https://arxiv.org/abs/2108.08810](https://arxiv.org/abs/2108.08810)\n\n[5] Quantifying Attention Flow in Transformers: https://arxiv.org/abs/2005.00928\n\n## Acknowledgements\n\n- [PyImageSearch](https://pyimagesearch.com)\n- [Jarvislabs.ai](https://jarvislabs.ai/)\n- [GDE Program](https://developers.google.com/programs/experts/)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fprobing-vits","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsayakpaul%2Fprobing-vits","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fprobing-vits/lists"}