Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

61 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

πŸŽ—οΈ Mammo-SAE: Interpreting Breast Cancer Concept Learning with Sparse Autoencoders [Deep Breast Imaging Workshop, MICCAI 2025 πŸ”₯]

Krishna Kanth Nakka*

Munich, Bavaria, Germany

paper Website HuggingFace

If you like our project, please give us a star ⭐ on GitHub for the latest update. Every star keeps me motivated to explore.

πŸ“£ Latest Updates

  • βœ… Sep 10-2025: Code Release along with MammoSAE checkpoint along with inference time intervention scripts.
  • βœ… Oct 11-2025: Upload training code for SAEs

πŸ”₯ Highlights

  • Mammo-SAE: First SAE trained on Mammogram Vision-Language Foundation Model (MammoCLIP) to interpret breast concepts such as mass, calcification.
  • Latent Neuron Intervention: We localise highest activated latent neurons in the presence of breast concepts, and then conduct top-k group intervention to understand the performance on downstream concept prediction
  • Breast Concept Localization: We visualize the top-k latent neurons that are responsible for downstream performance and study their alignment with ground-truth regions to understand the model behaviour.

πŸ”Introduction

Overview of Mammo-SAE. The SAE is first trained on patch-level CLIP features xj ∈ ℝd at any given layer, projecting them into a high-dimensional, interpretable sparse latent space z ∈ ℝh, and decoding them back for reconstruction. Once trained, the SAE is used to analyze which latent neurons are activated and what semantic information they encode. We also perform targeted interventions in the latent neuron space to assess their influence on downstream label prediction. We observe the learned latents capture diverse regions such as nipple regions, masses, and background areas. Red boxes indicate ground-truth mass localization.

πŸ› οΈ Installation

  • To install all required dependencies, run:

    pip install -r requirments.txt
    source envs.sh
  • Basic Dependencies:

    - CUDA Version >= 12.8
    - Python >= 3.8
    - Pytorch >= 2.8.0
    

πŸ“¦ Pretrained Checkpoints

  • Sparse Autoencoder (SAE) Weights: Pretrained SAE checkpoints are uploaded to the Huggingface. SAE is trained on the local features of the last layer of the Mammo-CLIP vision embeddings.

    # minimal example to download the MammoSAE
    ckpt_path = hf_hub_download(
        repo_id=f"KKNakka/MammoSAE",
        filename=f"sparse_autoencoder_final.pt",
        repo_type="model"
    )
    
    print(f"Loading SAE checkpoint from: {ckpt_path}")
    state_dict = torch.load(ckpt_path)
    
    # load the weights into the SAE
    autoencoder.load_state_dict(state_dict)
  • MammoCLIP: Download the pretrained MammoCLIP checkpoints from Hugging Face and place them in the Mammo_CLIP_weights/ directory.


πŸ“ VinDr Dataset

  • We follow similar preprocessing pipeline and scripts from the Mammo-CLIP codebase. We already provide the annotations csv file at datasets/Vindr. The CSV file is taken from Mammo-CLIP codebase.

  • We download the VinDr dataset images from the Kaggle project. Please see here for more details. Please put the downloaded images to datasets/Vindr/images_png.

  • The dataset folder should look like below:

    datasets
    |--- VinDr
    |   |--- images_png
    |   |--- README.md
    |   |--- vindr_detection_v1_folds.csv
    

✨ Mammo-SAE Generalization

  • Our MammoSAE is trained on the local features of the finetuned:suspicious calcification prediction model. To understand its effectivness in reconstructing other models (such as for different concept (eg., mass) or for different training phase like pretraining), we conduct the below experiment.

  • We reconstruct the local features at the last layer of the Mammo-CLIP backbone with SAE for different models and concepts, using

    bash scripts/eval_with_sae_reconstruction.sh

🩺 Launch Intervention on Mammo-SAE Latent Neurons

  • We provide scripts to reproduce Figure 2 using the commands below.

🎯 Intervention: Top-k Latent Class-level Neurons Activated

  • To launch an intervention that activates only the top-k neurons, run:

    bash scripts/topk_activate_latent_neuron_intervention.sh

🎯 Intervention: Top-k Class-level Latent Neurons Deactivated

  • To launch an intervention that deactivates only the top-k neurons, run:
    bash scripts/topk_deactivate_latent_neuron_interventions.sh

πŸ”¬ Visualization of Mammo-SAE Latent Neurons

  • To visualize the top-k latent neurons, please run:

    bash scripts/visualizations.sh
  • The visualizations are stored in the ./results/visualizations folder. For example, given an image, we plot top-10 latent neurons corresponding to presence of suspicious calcification class is shown below:

  • We visualizaed the latent neurons that have highest activated aggregated over all class-level images. Future work could explore better ways to localize neurons positions that are interpretable and important for downstream prediction.

βš™οΈ SAE Training Pipeline

  • We first obtain the embeddings the target model and save them to the disk using the script

    bash scripts/save_base_embeddings.sh
  • We then train the SAE using the script

    bash scripts/train_sae.sh

πŸ” Class-Level Top Latent Neurons

  • First, we compute the latent embeddings in SAE space using the script

    bash scripts/save_sae_concept_strengths_global.sh
  • After that, we compute the top-activated latents aggregated over all images using the script below:

    bash scripts/get_top_class_level_concepts.sh

πŸ“ Citation

If you find this paper or repository useful, please consider staring 🌟 this repo and citing πŸ“‘ our paper:

@article{nakka2025mammo,
  title={Mammo-SAE: Interpreting Breast Cancer Concept Learning with Sparse Autoencoders},
  author={Nakka, Krishna Kanth},
  journal={Deep Breath workshop, MICCAI 2025},
  year={2025}
}

πŸ™ Acknowledgement

  • This project is built upon Mammo-CLIP and VisionSAE codebases. Thank you to both of them for open sourcing their codebases.

❀️ Thanks for visiting ✨ Mammo-SAE!

Views

About

Official code release for the paper: "Mammo-SAE: Interpreting Breast Cancer Concept Learning with Sparse Autonencoders"

Topics

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages