Do Topological Characteristics Help in Knowledge Distillation?

Paper 💥

Do Topological Characteristics Help in Knowledge Distillation?

International Conference on Machine Learning (ICML) 2024
Jungeun Kim*, Junwon You*, Dongjin Lee*, Ha Young Kim, Jae-Hun Jung
Yonsei University & POSTECH

Knowledge distillation (KD) aims to transfer knowledge from larger (teacher) to smaller (student) networks. Previous studies focus on point-to-point or pairwise relationships in embedding features as knowledge and struggle to efficiently transfer relationships of complex latent spaces. To tackle this issue, we propose a novel KD method called TopKD, which considers the global topology of the latent spaces. We define global topology knowledge using the persistence diagram (PD) that captures comprehensive geometric structures such as shape of distribution, multiscale structure and connectivity, and the topology distillation loss for teaching this knowledge. To make the PD transferable within reasonable computational time, we employ approximated persistence images of PDs. Through experiments, we support the benefits of using global topology as knowledge and demonstrate the potential of TopKD.

conda env create -f environment.yml
conda activate topkd

mkdir data

*** Please modify the shell script files to match your experimental environment. ***

First, train teacher network to extract persistence diagram (PD) of teacher embedding features and distill the knowledge.

sh train_teacher.sh

Next, generate PI from PCD to integrate PD to DNN. Here, PCD is embedding features of teacher network and PI is an vectorization of PD.

python generate_PCD_PI.py

During training student network, accurately calculating PDs and PIs for each batch requires heavy computational demands. Thus, we utilize RipsNet to approximate PI for rapid calculation.

To train RipsNet with generated PCD and PIs, run the following code:

python train_ripsnet.py

Finally, to mimic PIs of the teacher, train student network with topology distillation loss, KD loss, and CE loss.

To train student network,

sh train_student.sh

We included the weights and PI of some teacher networks, WRN-40-2 and ResNet56, in “save_t_models” and “ripsnet” folders, respectively. The remaining teacher networks will be uploaded later.

Homogeneous architectures

Heterogeneous architectures

Homogeneous architectures

Heterogeneous architectures

Teacher: VGG13, student: MobileNetV2

CIFAR-100:

ImageNet-1K

@InProceedings{pmlr-v235-kim24aj,
  title = 	 {Do Topological Characteristics Help in Knowledge Distillation?},
  author =       {Kim, Jungeun and You, Junwon and Lee, Dongjin and Kim, Ha Young and Jung, Jae-Hun},
  booktitle = 	 {Proceedings of the 41st International Conference on Machine Learning},
  pages = 	 {24674--24693},
  year = 	 {2024}

This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (2023R1A2C200337911, RS-2023-00220762, 2021R1A2C3009648, NRF2021R1A6A1A1004294412, and RS-2023-00219980), and partly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.RS-2020-II201361, Artificial Intelligence Graduate School Program (Yonsei University) and No.2019-0-01906, Artificial Intelligence Graduate School Program (POSTECH)).