/privacy-aware-image-classification-with-kNN

Code Repository for the ISBI 2024 paper: "Integrating kNN with Foundation Models for Adaptable and Privacy-Aware Image Classification"

Primary LanguageJupyter NotebookApache License 2.0Apache-2.0

Integrating kNN with Foundation Models for Adaptable and Privacy-Aware Image Classification

Official code repository for the paper: "Integrating kNN with Foundation Models for Adaptable and Privacy-Aware Image Classification" (ISBI 2024)

News ๐ŸŽ‰

  • Paper is accepted for IEEE ISBI 2024 ๐ŸŽ‰

  • Source code has been released ๐Ÿ“‚

  • Experiments are available in Colab Open In Colab

Overview ๐Ÿง 

Traditional deep learning models encode knowledge within their parameters, limiting transparency and adaptability to data changes. This poses challenges for addressing user data privacy concerns. To overcome this limitation, we propose to store embeddings of the training data independently of the model weights. This enables dynamic data modifications without retraining. Our approach integrates the k-Nearest Neighbor (k-NN) classifier with a vision-based foundation model pre-trained on natural images in a self-supervised manner. This integration enhances interpretability and adaptability while addressing privacy concerns.

system_anim

Figure 1: During pretraining 1), the image encoder is trained to extract representative features. The knowledge-storing phase 2) utilizes the pre-trained (now frozen) encoder to extract and store task-relevant knowledge from the training data. During inference 3), that knowledge allows the classification of query images through majority voting on the top-k similar embeddings.

Key Features ๐Ÿ”‘

  • Open-source implementation including a previously unpublished baseline method and performance-improving contributions
  • Integration of k-NN classifier with recent vision-based foundation models
  • Flexible data storage system for dynamic data modifications without retraining
  • Evaluation of method's performance across established benchmark datasets and medical image classification tasks
  • Assessment of method's robustness in continual learning and data removal scenarios

Results ๐Ÿ“Š

1. Improved classification accuracy demonstrated across established benchmark datasets

Accuracy [%] CIFAR-10 CIFAR-100 STL-10
ResNet-101 87.3 63.6 98.1
CLIP ViT-B/16 92.4 68.0 98.5
CLIP ViT-L/14 95.5 74.2 99.4
DINOv2 ViT-B/14 98.0 87.2 99.4
DINOv2 ViT-L/14 98.5 88.3 99.5

Table 1: Classification accuracy of our k-NN approach for different backbone choices.

2. Applicability of the method to distinct medical image classification tasks confirmed

Accuracy [%] Pneumonia Melanoma
CovXNetโ€  (Mahmud et al. 2020) 98.1 ---
EfficientNetB0โ€  (Cassidy et al. 2022) --- 62.1
Ours (DINOv2 ViT-B/14) 88.1 68.5
Ours (DINOv2 ViT-L/14) 89.9 69.8

Table 2: Comparison of our approachโ€™s strong transfer learning ability for medical image analysis (โ€  refers to fully supervised, end-to-end models).

3. Robustness in continual learning and data removal scenarios demonstrated

Figure 2: Visualization of the methodโ€™s ability for diverse continual learning tasks (left: class incremental learning, right: sample incremental learning).

Figure 3: Illustration of our methodโ€™s classification consistency despite the continuous diminishing of the support set (left: Pneumonia, right: Melanoma).

Getting Started ๐Ÿš€ Open In Colab

To get started with using our method, you can execute the research experiments using Google Colab.

If you have access to a capable GPU and want to run the code locally:

  1. Clone this repository to your local machine.
  2. Find the notebooks for the experiments presented in the paper under the "experiments" directory.
  3. Place the extracted ISIC2018 Task 3 and Pneumonia (ChestXRay2017.zip) datasets inside the /assets/datasets/ folder.
  4. Each notebook includes comprehensive documentation to explain its usage and application.

Acknowledgements ๐Ÿ‘

  • Chroma (The AI-native open-source embedding database")
  • PyTorch
  • timm (Ross Wightman - PyTorch Image Models)
  • pandas
  • scikit-learn
  • Nakata, Kengo, et al. "Revisiting a knn-based image classification system with high-capacity storage." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022. doi

Citation ๐Ÿ“–

If you find this work useful in your research, please consider citing our paper:

@InProceedings{doerricharchut2024kNNIntegration,
  author={Doerrich, Sebastian and Archut, Tobias and Salvo, Francesco Di and Ledig, Christian},
  booktitle={2024 IEEE International Symposium on Biomedical Imaging (ISBI)}, 
  title={Integrating kNN with Foundation Models for Adaptable and Privacy-Aware Image Classification}, 
  year={2024},
  pages={1-5},
  keywords={Adaptation models;Data privacy;Source coding;Training data;Nearest neighbor methods;Data models;Robustness;k-NN classifier;continual learning;transfer learning;few-shot classification;explainability},
  doi={10.1109/ISBI56570.2024.10635560}
}