Adaptive Group Robust Ensemble Knowledge Distillation (AGRE-KD)

Paper | Video Presentation

AGRE-KD tackles subgroup robustness issues that appear when distilling an ensemble of teachers into a lighter student. The method reweights teachers whose gradients differ from a biased reference model so that the student inherits signals that help underrepresented groups. This repository builds on top of milkshake; please cite both projects when using this code.

Overview

• Our paper show that distillation without care can amplify spurious correlations even with debiased teachers.
• AGRE-KD adaptively upweights teachers that provide complementary gradients relative to a biased model.
• Experiments across several benchmarks show improved worst-group accuracy and competitive average accuracy.

Setup

1. Clone the repository

   git clone git@github.com:<your-username>/AGRE-KD.git
   cd AGRE-KD

2. Create and activate a virtual environment

   python -m venv .venv
   source .venv/bin/activate   # On Windows: .venv\Scripts\activate

3. Install dependencies

   pip install -r requirements.txt

4. Prepare datasets
   The milkshake datamodules handle downloads when possible. Ensure you have the necessary dataset credentials if required (e.g., for Civil Comments).

Running train_teachers.py or ensemble_KD.py will create experiment logs under training_logs/ and a summary teacher_erm.pkl mapping that caches available checkpoints.

Workflow

1. Train teacher models

Train biased and/or debiased teachers before distillation. Config files inside cfgs/ store default hyperparameters per dataset.

python train_teachers.py -c cfgs/mnist.yaml --seed 1

Repeat the command with different --seed values to populate the teacher pool. Model checkpoints are stored in training_logs/{biased,debiased}_teachers/.

Teacher architecture and checkpoint versions are controlled with the --teacher_model flag together with its corresponding version argument:

• --teacher_model resnet --teacher_resnet_version 50
• --teacher_model vit --teacher_vit_version base
• --teacher_model bert --teacher_bert_version large
• --teacher_model convnextv2 --teacher_convnextv2_version base

2. Distill into a student

Use ensemble_KD.py to combine teachers and train the student with the desired ensembling method.

python ensemble_KD.py \
  -c cfgs/mnist.yaml \ 
  --seed 1 \
  --teacher_size 5 \
  --ensemble_method AGREKD \
  --erm_teachers False

This command loads five debiased teachers, applies AGRE-KD to train the student defined in the config. Results and metrics are stored in logs/<...>/. Boolean flags use explicit True/False values (e.g., --erm_teachers True). Additional CLI help is available via --help.

Available ensemble methods

• AVERAGE_LOSS: Standard distillation; averages KL losses from all teachers.
• AEKD: Adaptive Ensemble KD in gradient space using an SVM-based weighting.
• RANDOM: Samples one teacher per batch for distillation.
• ENS: Majority vote ensemble (no student training, evaluation only).
• AGREKD: Proposed method—weights teachers whose gradients diverge from a biased model.

For example, to distill from Vision Transformer teachers to ResNet18 student you can run:

python ensemble_KD.py \
  -c cfgs/waterbirds.yaml \
  --teacher_model vit \
  --teacher_vit_version base \
  --teacher_size 5 \
  --ensemble_method AVERAGE_LOSS

If you set --teacher_model mix, the script cycles through the dataset-specific list of architectures defined in the code and requires --teacher_size to be large enough to sample from each family; you can still override per-architecture versions in the config file.

Troubleshooting

• Dataset download issues: check the milkshake documentation for dataset-specific prerequisites.
• Missing teacher_erm.pkl: run train_teachers.py at least once; the mapping file is generated automatically.
• CUDA memory errors: reduce --batch_size in the config or lower --teacher_size.

Citation

If you find this repository helpful, please cite:

@article{kenfacj2025adaptive,
  title={Adaptive Group Robust Ensemble Knowledge Distillation},
  author={Kenfack, Patrik Joslin and A{\"i}vodji, Ulrich and Kahou, Samira Ebrahimi},
  journal={Transactions on Machine Learning Research},
  year={2025},
  url={https://openreview.net/forum?id=G2BEBaKd8Y}
}