ComFe: An Interpretable Head for Vision Transformers

TL;DR

ComFe introduces a scalable, interpretable classification head for Vision Transformers that maintains competitive accuracy, enhances robustness, and identifies meaningful component features without additional annotations or extensive hyperparameter tuning.

Contribution

It presents ComFe, the first interpretable head for large-scale Vision Transformers, enabling interpretability and robustness without finetuning the backbone or needing extra annotations.

Findings

Achieves competitive performance on ImageNet-1K

Provides improved robustness over previous methods

Identifies consistent component features within images

Abstract

Interpretable computer vision models explain their classifications through comparing the distances between the local embeddings of an image and a set of prototypes that represent the training data. However, these approaches introduce additional hyper-parameters that need to be tuned to apply to new datasets, scale poorly, and are more computationally intensive to train in comparison to black-box approaches. In this work, we introduce Component Features (ComFe), a highly scalable interpretable-by-design image classification head for pretrained Vision Transformers (ViTs) that can obtain competitive performance in comparison to comparable non-interpretable methods. To our knowledge, ComFe is the first interpretable head and unlike other interpretable approaches can be readily applied to large-scale datasets such as ImageNet-1K. Additionally, ComFe provides improved robustness and outperforms previous interpretable approaches on key benchmark datasets while using a consistent set of hyperparameters and without finetuning the pretrained ViT backbone. With only global image labels and no segmentation or part annotations, ComFe can identify consistent component features within an image and determine which of these features are informative in making a prediction. Code is available at github.com/emannix/comfe-component-features.