Leveraging Topological Guidance for Improved Knowledge Distillation

TDA algorithms are applied to the data to extract topological features, which are robust to noises or perturbations and encodes the shape of complex data (Adams et al., 2017; Wang et al., 2021). Persistent homology is a fundamental tool in TDA that helps in understanding the shape and structure of data, which involves constructing a filtration (Edelsbrunner & Harer, 2022), typically based on a distance function and tracking variations of $n$-dimensional holes represented by assortments of points, edges, and triangles through a dynamic thresholding process called filtration. In filtration, the appearance and disappearance of these holes are described in the persistence feature, summarized in a persistence diagram (PD), which records the birth and death times as $x$ and $y$ coordinates of planar scatter points (Adams et al., 2017; Edelsbrunner & Harer, 2022).

Since PDs can have a high dimensionality with complex structures or a large number of points that can vary, using PDs directly in machine learning is challenging. To address this problem, persistence image (PI) has been widely used, which is one of the ways to encode geometric information via the lifespan of homological structures present in the data. This representation can be easily integrated into machine learning (Barnes et al., 2021; Edelsbrunner & Harer, 2022). Specifically, the points within the PD are projected onto a two-dimensional grid $\rho:\mathbb{R}\rightarrow\mathbb{R}^{2}$. The grid points are then assigned values determined by a weighted sum of Gaussian smoothing, which is centered around the scattered points within the PD. The grid is represented as a matrix and can be treated as regular image data called PI, as depicted in Figure 2. This allows for the application of convolutional neural networks (CNNs) and machine learning algorithms, and offers a more manageable format for analysis and visualization.

Previous studies showed that topological features complement features of the raw data to achieve improved performance (Som et al., 2020). However, additional processing for TDA and concatenation in networks increase computational time and resources at inference-time, which poses difficulty in implementing the process on small devices having limited computational resources and power. To alleviate this issue, we propose a framework based on KD to infuse topological features into a small model that uses solely the raw image data at test-time.

Conventional Knowledge Distillation. Knowledge distillation obtains a smaller model by utilizing the learned knowledge of a larger model, which was first introduced by Buciluǎ et al. (Buciluǎ et al., 2006) and explored more by Hinton et al. (Hinton et al., 2015). In KD, soft labels are utilized for knowledge transfer from a teacher to a student, which provide richer supervision signals and reduce overfitting. This also leads to better transferability of learned representations. The loss function of conventional KD for training a student is:

where, $\mathcal{L_{CE}}$ is the standard cross entropy loss, $\mathcal{L_{K}}$ is KD loss, and $\lambda$ is a hyperparameter; $0<\lambda<1$. The difference between the output of the softmax layer for a student network and the ground-truth label is minimized by the cross-entropy loss:

where, $\mathcal{H(\cdot)}$ is a cross entropy loss function, $l_{S}$ is the logits of a student, and $y_{g}$ is a ground truth label. The gap between outputs of student and teacher are minimized by KL-divergence loss:

where, $z_{T}=softmax(l_{T}/\tau)$ and $z_{S}=softmax(l_{S}/\tau)$ are softened outputs of a teacher and student, respectively, and $\tau$ is a hyperparameter; $\tau>1$. To obtain the best performance, we adopt early stopping for KD (ESKD) which improves the efficacy of KD (Cho & Hariharan, 2019).

Feature-based Knowledge Distillation. Features from intermediate layers of a network can be utilized in knowledge transfer (Gou et al., 2021; Zagoruyko & Komodakis, 2017; Tung & Mori, 2019). Zagoruyko et al. (Zagoruyko & Komodakis, 2017) suggested activation-based attention transfer (AT), which is computed by a sum of squared attention mapping function, and calculating statistics across the channel dimension. Tung et al. (Tung & Mori, 2019) introduced similarity-preserving knowledge distillation, matching similarity within a mini-batch of samples between a teacher and a student. Since the size of a similarity map is determined by the size of a mini-batch, the size of the extracted similarity maps from the teacher and student is the same even if they generate different sizes of features. In details, the similarity map $M\in\mathbb{R}^{b\times b}$ is obtained as follows:

where $F$ is reshaped features from an intermediate layer of a model, $b$ is the size of a mini-batch, and $c$, $h$, and $w$ are the number of channels, height, and width of the output, respectively. These feature transfer methods are popularly used; however, these are to match knowledge with similar characteristics in a uni-modal manner.

Utilizing Multiple Teachers. Not only one teacher, but multi-teacher distillation has been widely utilized to provide more diverse knowledge in training process (Reich et al., 2020; Liu et al., 2020; Gou et al., 2021). In some cases, the data utilized for training a student cannot be used during testing. Also, teachers trained with different modalities or representations can be utilized in distillation. Thoker and Gall (Thoker & Gall, 2019) train a student with paired samples from two modalities for action recognition. Jeon et al. (Jeon et al., 2022) explored to train a student model with two teachers trained with different representations for wearable sensor data analysis. With this insight, we develop a framework and explore to utilize topological features involving two teachers for image data analysis.

To compute PIs, Scikit-TDA python library (Saul & Tralie, 2019) and the Ripser package are used for generating PDs, as explained in Som et. al. (Som et al., 2020). Firstly, image data is normalized in range in $[0,1]$. To compute level-set filtration PDs, image data is reshaped to row- and column-wise signals, considering different order of context which can extract different topological features (Barnes et al., 2021). By filtration, PDs summarize the different peak and local minima intensities in the data. Specifically, each channel of data is transformed and used to generate a PI, where the PI implies birth-time vs. lifetime information. We utilize row- and column-wise transforms separately for creating images with channels of ($R_{r}$, $G_{r}$, $B_{r}$) and ($R_{c}$, $G_{c}$, $B_{c}$) to collect diverse knowledge, and all created PIs are concatenated in an image. Then, six channels ($P_{Rr}$, $P_{Gr}$, $P_{Br}$, $P_{Rc}$, $P_{Gc}$, $P_{Bc}$) of PI implying persistence knowledge is created. The total dimension size of one PI is $g\times g\times c$, where $g$ and $c$ are a constant value and the number of channels for a sample. The created PI is utilized to train a teacher model that acts as a pre-trained model to transfer topological features to a student model in KD process.

Knowledge Transfer with Logits. Knowledge of logits from two teachers are transferred individually, thus additional process including concatenation or hidden layers is not required. KD loss for logit knowledge of two teachers is explained as follows:

where, $\alpha$ is a constant to balance the effects of different teachers, and $z_{T_{1}}$ and $z_{T_{2}}$ are softened outputs of teachers trained with the raw image data and PIs, respectively.

Knowledge Transfer with Intermediate Features. To transfer sufficient knowledge from two teachers, we utilize features from intermediate layers additionally. Since PI and the raw image data have different statistical characteristics in semantic information, it is more effective to align and convey the information by using the correlation between samples rather than using spatial information. We utilize similarities, as explained in equation 4, to easily integrate information from two teachers and provide topological features to student in distillation. Figure 3 shows an example of similarity maps obtained from different intermediate layers of two WRN16-3 teachers. Note, Teacher1 and Teacher2 denote models learned with the raw image data and PI, respectively. High values represent high similarities. This implies similar patterns can be created when two samples belong to the same category. Since two models are trained with different representations, their highlighted patterns are different. We merge the maps from two teachers with weighted summation as follows:

where, $M^{(l)}_{T_{m}}\in\mathbb{R}^{b\times b}$ is the generated map from the similarity maps of two teachers $M_{T_{1}}$ and $M_{T_{2}}$ in a layer pair ($l^{T_{1}}$ and $l^{T_{2}}$). By merging the maps, the similarities include topological features which can complement the original features to improve the performance. The loss that encourages the student to mimic teachers is:

where $\widetilde{M^{(l)}_{T_{m}}}$ and $\widetilde{M^{(l^{S})}_{S}}$ are normalized map for a merged teacher and a student, $\|\cdot\|_{F}$ is the Frobenius norm (Tung & Mori, 2019), and $L$ collects the layer pairs ($l$ and $l^{S}$). $l^{T_{1}}$, $l^{T_{2}}$, and $l^{S}$, can be selected with the same depth or the end of the same block of networks. Since a student model uses the raw image data only, the gap between the merged features of teachers and the feature of the student can be generated, which makes degradation. To alleviate this problem, an annealing strategy (Jeon et al., 2022) is used, which initializes the student model with weight values of a model trained from scratch. The final loss function is as follows:

where $\gamma$ is a hyperparameter.

Datasets. The CIFAR-10 (Krizhevsky & Hinton, 2009) dataset consists of 60k images distributed among 10 classes, with each class including 5k and 1k images for training and testing, respectively. Each image is a 32$\times$32 sized RGB image. The experiments on CIFAR-10 allows us to evaluate our model’s efficacy with less time consumption. We extend our experiments on CINIC-10 (Darlow et al., 2018) that augments CIFAR-10 formatting but includes a larger set of 270k images whose scale closer to ImageNet. The images are evenly split into each ‘train’, ‘validate’, and ‘test’ sets, with ten classes of 9k images per class. The size of the images is 32$\times$32 as well.

Experimental Settings. In generating PIs by TDA, by referring to the previous study (Som et al., 2020), we set birth-time range and Gaussian function parameter as $[$0, 0.3$]$ and 0.01. The threshold for life-time is set to 0.02. we set $g$ and $c$ of PI as 50 and 6, respectively.

For experiments, we set the batch size $b$ as 128, the total epochs as 200 using SGD with momentum 0.9, and a weight decay of $1\times 10^{-4}$. The initial learning rate $lr$ is set to 0.1 that is decayed by a factor of 0.2 at epochs 40, 80, 120, and 160. Empirically, we set KD hyperparameters $\lambda$, $\tau$, and $\gamma$ as ($\lambda$ = 0.9, $\tau$ = 4, $\gamma$ = 3000) and ($\lambda$ = 0.6, $\tau$ = 16, $\gamma$ = 2000) for CIFAR-10 and CINIC-10, respectively, referred to previous studies (Cho & Hariharan, 2019; Tung & Mori, 2019; Jeon et al., 2023).

We compare with KD based baselines including conventional KD (Hinton et al., 2015), attention transfer (AT) (Zagoruyko & Komodakis, 2017), relational knowledge distillation (RKD) (Park et al., 2019), variational information distillation (VID) (Ahn et al., 2019), similarity-preserving knowledge distillation (SP) (Tung & Mori, 2019), attentive feature distillation and selection (AFDS) (Wang et al., 2020), and multi-teacher based distillation using topological features in KD (Base) (Jeon et al., 2022). AT and SP are utilized with KD. For all baseline methods, the same hyperparameter settings are used as those specified in their papers, and their author-provided code is used for evaluation. $\alpha$ of Base is 0.9, and TGD is 0.99 as a default setting. All experiments were repeated three times, and the averaged accuracy and the standard deviation of performance are reported. More details are explained in appendix.

In this section, we show analysis on various combinations including different capacity of teachers and architectural styles of teacher-student networks.

In this section, we investigate sensitivity for $\alpha$, robustness on noise, and evaluate with feature visualization and model reliability.

We measure the processing time of various models on CIFAR-10 testing set (10k samples). The total processing time is explained in Table 7. A student (WRN16-1) of TGD takes much less time than teachers on both CPU and GPU. Creating PI (6 channels) takes more than 30k seconds on the CPU, which is not efficient in inference time as well. As described in the prior section, the student by TGD outperforms a model learned from scratch by 1.96$\%$ in classification accuracy. These findings clearly highlight the essential necessity of using a compact model for implementation on small devices with limited computational resources and the effectiveness of TGD.