Paper: Spatio-Temporal Graph Knowledge Distillation (ICLR-2024)

Reviewer comment (Evaluation):
The experiment lacks a comparison of runtime. In addition to the comparison of efficiency, the running time of the different methods should also be compared, which is also a very important indicator.

Author response:
Thank you for your concerns about the efficiency comparison of our STGKD model framework. In Section 4.3 'Model Scalability Study,' we conducted a comprehensive study on the model's efficiency, following the practices used in knowledge distillation work on graph data for efficiency comparison. [1]

[1] Zhang, Shichang, et al. "Graph-less Neural Networks: Teaching Old MLPs New Tricks Via Distillation." *International Conference on Learning Representations*. 2021.

Reviewer comment (Writing):
The paper does not provide a clear explanation of the interpretability and explainability of the proposed framework. Although it mentions that the student MLP selectively inherits task-relevant spatio-temporal knowledge from the teacher GNN framework, it does not explain how this knowledge transfer occurs and how the student model utilizes the transferred knowledge. A more detailed discussion on interpretability and explainability would enhance understanding of the proposed approach and its inner workings.

Author response:
We highly appreciate your insightful considerations regarding the interpretability of our model. For the explainability of STGKD in the knowledge distillation process, please refer to the supplementary material, section A.2.1 'Rationale Analysis of STGKD's Robustness.' The information control process within our IB-based KD framework plays a vital role in determining the inherent robustness of KD.

Reviewer comment (Novelty):
As stated in the related works, 'A significant contribution of this work lies in the novel integration of the spatio-temporal information bottleneck into the KD framework.' However, it should be noted that the incorporation of the information bottleneck into knowledge distillation has been previously explored, e.g., see references [3] [4] below. The paper lacks clarification on how the proposed method differs from existing approaches.

Author response:
Thank you for your concerns regarding the novelty of our work. In reference to work [3], our proposed method for information bottleneck distillation simultaneously imposes information bottleneck objectives on both the student model and the knowledge distillation process. The joint optimization objective is formulated as:
$$
\min_{\mathbb{P}(\mathbf{Z}|\mathbf{X})}  (-I(\mathbf{Y},\mathbf{Z})+ \beta_1 I(\mathbf{X}, \mathbf{Z})) + (-I(\mathbf{Y}^{T},\mathbf{Z})+ \beta_2 I(\mathbf{X}, \mathbf{Z}))
$$
In contrast, [3] only addresses the control of IB during the knowledge distillation process, with the optimization objective:
$$
\min_{s} \{ I(X; F_s) - \beta I(Y; F_s) - \gamma I(F_t; F_s) \}
$$
Another significant distinction is that these models are designed for knowledge distillation frameworks targeting classification tasks, wherein a temperature coefficient can be employed to prevent the student from learning erroneous behaviors from the teacher model. However, for regression tasks such as spatio-temporal forecasting, determining how to avoid the student model from receiving incorrect knowledge from the teacher model is also one of the challenges our STGKD method seeks to address.

Reviewer comment (Theory):
As an ICLR submission, the paper lacks theoretical guarantees. For instance, it is better to quantify the model's robustness against noise after adopting such information bottleneck regularizer. Additionally, as reducing complexity of the model is an important idea of knowledge distillation, providing a generalization bound related to the model's complexity would strengthen the method's support.

Author response:
Thank you for your attention to the theoretical guarantees aspect of our work. Regarding the gains in robustness from the IB you mentioned, we have actually conducted a detailed theoretical analysis in the supplementary material, section A.2.1 'Rationale Analysis of STGKD's Robustness.' The information control process within our IB-based KD framework is crucial in determining the inherent robustness of KD. As for your mention of theoretical analysis on the time and space complexity, we have provided a detailed theoretical analysis in the supplementary material, section A.2.2 'Model Complexity Analysis,' demonstrating our model's advantages in scalability.

Reviewer comment (Experiments):
A comparison with existing prompt learning methods is missing for the spatio-temporal prompt learning module, and the similarities and differences should be clarified to prevent confusion.

Author response:
We appreciate your attention to the spatio-temporal prompt learning aspect of our work. Firstly, to date, there are only a few works on prompt learning in spatio-temporal forecasting scenarios, such as PromptST [1] (a paper accepted at CIKM 2023, accepted in the proceedings following the ICLR submission deadline), which introduces prompts to the Spatio-temporal transformer for prompt tuning. However, the prompts in STGKD are designed to introduce spatio-temporal context information to the student MLP, not for prompt tuning paradigms, and serve as an end-to-end encoding module. Secondly, while partly inspired by works [1, 2], our STGKD prompts are more comprehensive, encompassing both static and dynamic spatio-temporal prompts, fully capturing static and dynamic spatio-temporal correlations in forecasting scenarios, significantly differing from existing works.

[1] Zijian Zhang, Xiangyu Zhao, Qidong Liu, Chunxu Zhang, Qian Ma, Wanyu Wang, Hongwei Zhao, Yiqi Wang, and Zitao Liu. 2023. PromptST: Prompt-Enhanced Spatio-Temporal Multi-Attribute Prediction. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management (CIKM '23). Association for Computing Machinery, New York, NY, USA, 3195–3205.

Reviewer comment (Novelty):
Utilizing three types of prompts as input and a learnable embedding method has been done in previous STGNN works [1,2], limiting the novelty of the proposed method.

Author response:
We appreciate your attention to the spatio-temporal prompt learning aspect of our work. Firstly, to date, there are only a few works on prompt learning in spatio-temporal forecasting scenarios, such as PromptST [1] (a paper accepted at CIKM 2023, accepted in the proceedings following the ICLR submission deadline), which introduces prompts to the Spatio-temporal transformer for prompt tuning. However, the prompts in STGKD are designed to introduce spatio-temporal context information to the student MLP, not for prompt tuning paradigms, and serve as an end-to-end encoding module. Secondly, while partly inspired by works [1, 2], our STGKD prompts are more comprehensive, encompassing both static and dynamic spatio-temporal prompts, fully capturing static and dynamic spatio-temporal correlations in forecasting scenarios, significantly differing from existing works.

[1] Zijian Zhang, Xiangyu Zhao, Qidong Liu, Chunxu Zhang, Qian Ma, Wanyu Wang, Hongwei Zhao, Yiqi Wang, and Zitao Liu. 2023. PromptST: Prompt-Enhanced Spatio-Temporal Multi-Attribute Prediction. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management (CIKM '23). Association for Computing Machinery, New York, NY, USA, 3195–3205.

Reviewer comment (Evaluation):
The overall performance improvements of the proposed model on all datasets are not significant.

Author response:
Regarding your questions (a) and (b), we still believe in the significant superiority of our STGKD framework. Notably, our model outperforms the state-of-the-art (SOTA) models on three different datasets (traffic, crime, weather), demonstrating its superior performance and generalizability, given the significant differences among these datasets in aspects like data sparsity and periodicity (for example, crime data is very sparse, while traffic data shows more obvious periodicity). For question (c), we do not perceive the use of models like STGCN as a teacher model to be an unfair comparison. In fact, the performance of the student MLP, when our STGKD framework is applied to such traditional models, surpasses SOTA models, further highlighting the superiority of the STGKD distillation framework.

Reviewer comment (Evaluation):
The STID model, which is similar to the proposed model but without the knowledge distillation module, outperforms most complex STGNNs, raising the question of why transferring knowledge from weaker STGNNs to MLPs can lead to improvement. However, the paper does not offer clear explanations for this phenomenon.

Author response:
Regarding your questions (a) and (b), we still believe in the significant superiority of our STGKD framework. Notably, our model outperforms the state-of-the-art (SOTA) models on three different datasets (traffic, crime, weather), demonstrating its superior performance and generalizability, given the significant differences among these datasets in aspects like data sparsity and periodicity (for example, crime data is very sparse, while traffic data shows more obvious periodicity). For question (c), we do not perceive the use of models like STGCN as a teacher model to be an unfair comparison. In fact, the performance of the student MLP, when our STGKD framework is applied to such traditional models, surpasses SOTA models, further highlighting the superiority of the STGKD distillation framework.

Reviewer comment (Experiments):
The authors have conducted generalizability testing on PEMS data with synthesized data missing. However, this type of distribution is only one specific example of covariate shift, and there are various other types of distribution shifts that need to be considered, e.g., the distribution shifts of traffic patterns during rush hours or seasonal traffic patterns shifts.

Author response:
Regarding your questions (a) and (b), we still believe in the significant superiority of our STGKD framework. Notably, our model outperforms the state-of-the-art (SOTA) models on three different datasets (traffic, crime, weather), demonstrating its superior performance and generalizability, given the significant differences among these datasets in aspects like data sparsity and periodicity (for example, crime data is very sparse, while traffic data shows more obvious periodicity). For question (c), we do not perceive the use of models like STGCN as a teacher model to be an unfair comparison. In fact, the performance of the student MLP, when our STGKD framework is applied to such traditional models, surpasses SOTA models, further highlighting the superiority of the STGKD distillation framework. Regarding question (d) and your mention of distribution shifts in traffic patterns, please refer to section A.6.2 'Visualization of prediction' in the supplementary material. We have compared our model's predictions in different time segments with SOTA models, finding that our model consistently predicts well, indicating stronger generalizability and ability to handle distribution shifts across various time segments.

Reviewer comment (Novelty):
Novelty: My primary concern pertains to the paper's novelty. While the authors posit that the integration of the spatio-temporal information bottleneck into the Knowledge Distillation (KD) framework is a significant contribution, I'd like to highlight that the concept of the information bottleneck has already been explored in the context of knowledge distillation[1]. Moreover, the idea of employing knowledge distillation on dynamic graphs isn't novel either[2,3,4]. It appears the authors are leveraging well-established ideas to tackle specific challenges in spatial-temporal graph forecasting.

Author response:
Thank you for your concerns regarding the novelty of our work. In reference to work [1], our proposed method for information bottleneck distillation simultaneously imposes information bottleneck objectives on both the student model and the knowledge distillation process. The joint optimization objective is formulated as:
$$
\min_{\mathbb{P}(\mathbf{Z}|\mathbf{X})}  (-I(\mathbf{Y},\mathbf{Z})+ \beta_1 I(\mathbf{X}, \mathbf{Z})) + (-I(\mathbf{Y}^{T},\mathbf{Z})+ \beta_2 I(\mathbf{X}, \mathbf{Z}))
$$
In contrast, [1] only addresses the control of IB during the knowledge distillation process, with the optimization objective:
$$
\min_{s} \{ I(X; F_s) - \beta I(Y; F\_s) - \gamma I(F\_t; F\_s) \}
$$
Regarding works [2] and [3] you mentioned (and it seems that there is a missing citation for work [4]), neither includes models of knowledge distillation. The self-supervised signal distillation discussed in [2] is entirely different from the knowledge distillation considered in our work. Our motivation is to address the challenges of spatio-temporal forecasting, especially with large-scale temporal and spatial nodes, where the complexity of STGNNs increases sharply, impeding deployment. Hence, our proposed STGKD method effectively distills the capability of STGNNs to model complex spatio-temporal correlations into an MLP framework.

Reviewer comment (Evaluation):
Evaluation: Another area of improvement is in the choice of baseline models for evaluation. Notably, the absence of other graph knowledge distillation models as baselines seems to be an oversight. Including them would make the comparison more comprehensive and equitable.

Author response:
Thank you for your suggestion to improve upon the baselines. However, we would like to point out that, to the best of our knowledge, our method represents a pioneering work in the realm of knowledge distillation for spatio-temporal forecasting. Directly applying graph knowledge distillation techniques to spatio-temporal forecasting poses adaptability issues. Namely, the design frameworks (i.e., objective functions) of most graph knowledge distillation tasks are tailored for classification tasks, and cannot be directly transferred to regression tasks like spatio-temporal forecasting. Therefore, we have introduced our STGKD method to address this challenge.

Reviewer comment (Experiments):
For a comprehensive evaluation, why were other graph knowledge distillation models not considered as baseline models? Would incorporating them not offer a more balanced and insightful comparison in the context of your study?

Author response:
Thank you for your suggestion to improve upon the baselines. However, we would like to point out that, to the best of our knowledge, our method represents a pioneering work in the realm of knowledge distillation for spatio-temporal forecasting. Directly applying graph knowledge distillation techniques to spatio-temporal forecasting poses adaptability issues. Namely, the design frameworks (i.e., objective functions) of most graph knowledge distillation tasks are tailored for classification tasks, and cannot be directly transferred to regression tasks like spatio-temporal forecasting. Therefore, we have introduced our STGKD method to address this challenge.

Reviewer comment (Novelty):
In terms of novelty and motivation, the manuscript does not make a strong case. While it appears that the authors might be introducing knowledge distillation to the GNN domain for the first time (though this is not explicitly claimed in the paper), this alone does not constitute a substantial contribution. The paper needs to delineate the differences between traditional knowledge distillation approaches applied to CNNs and STGNNs, and the proposed method, explaining why it is particularly effective in the GNN context.

Author response:
Thank you for bringing this to our attention. As stated in the introduction, Spatio-Temporal Graph Neural Networks (STGNNs) have recently been extensively applied in spatio-temporal forecasting, achieving state-of-the-art (SOTA) results compared to CNN-based methods (e.g., ST-ResNet [1] and DeepST [2]). The fundamental reason lies in the STGNNs' ability to capture complex spatial neighbor relationships in spatio-temporal data. However, for large-scale spatio-temporal sequence forecasting, where the number of spatial and temporal nodes increases significantly, STGNNs face substantial challenges in terms of time and space complexity. Therefore, we propose distilling rich spatio-temporal knowledge into Multilayer Perceptrons (MLPs) using knowledge distillation. A critical aspect of our work involves determining how to effectively distill the complex spatio-temporal knowledge from STGNNs into MLPs without compromising accuracy. This approach distinguishes our work from CNN knowledge distillation. Specifically, we employ spatio-temporal prompts to enhance the spatio-temporal context in the student MLP, thereby maximally preserving the spatio-temporal modeling capabilities of the STGNNs.

[1] Zhang, J., Zheng, Y. and Qi, D. 2017. Deep Spatio-Temporal Residual Networks for Citywide Crowd Flows Prediction. *Proceedings of the AAAI Conference on Artificial Intelligence*. 31, 1 (Feb. 2017).

[2]  Junbo Zhang, Yu Zheng, Dekang Qi, Ruiyuan Li, and Xiuwen Yi. 2016. DNN-based prediction model for spatio-temporal data. In Proceedings of the 24th ACM SIGSPATIAL International Conference on Advances in Geographic Information Systems (SIGSPACIAL '16). Association for Computing Machinery, New York, NY, USA, Article 92, 1–4.

Reviewer comment (Theory):
What are the significant differences between applying the knowledge distillation to CNN and STGNN?

Author response:
Thank you for bringing this to our attention. As stated in the introduction, Spatio-Temporal Graph Neural Networks (STGNNs) have recently been extensively applied in spatio-temporal forecasting, achieving state-of-the-art (SOTA) results compared to CNN-based methods (e.g., ST-ResNet [1] and DeepST [2]). The fundamental reason lies in the STGNNs' ability to capture complex spatial neighbor relationships in spatio-temporal data. However, for large-scale spatio-temporal sequence forecasting, where the number of spatial and temporal nodes increases significantly, STGNNs face substantial challenges in terms of time and space complexity. Therefore, we propose distilling rich spatio-temporal knowledge into Multilayer Perceptrons (MLPs) using knowledge distillation. A critical aspect of our work involves determining how to effectively distill the complex spatio-temporal knowledge from STGNNs into MLPs without compromising accuracy. This approach distinguishes our work from CNN knowledge distillation. Specifically, we employ spatio-temporal prompts to enhance the spatio-temporal context in the student MLP, thereby maximally preserving the spatio-temporal modeling capabilities of the STGNNs.

[1] Zhang, J., Zheng, Y. and Qi, D. 2017. Deep Spatio-Temporal Residual Networks for Citywide Crowd Flows Prediction. *Proceedings of the AAAI Conference on Artificial Intelligence*. 31, 1 (Feb. 2017).

[2]  Junbo Zhang, Yu Zheng, Dekang Qi, Ruiyuan Li, and Xiuwen Yi. 2016. DNN-based prediction model for spatio-temporal data. In Proceedings of the 24th ACM SIGSPATIAL International Conference on Advances in Geographic Information Systems (SIGSPACIAL '16). Association for Computing Machinery, New York, NY, USA, Article 92, 1–4.

Reviewer comment (Writing):
For SPATIO-TEMPORAL IB INSTANTIATING, what is the spatio-temporal part here?

Author response:
Thanks for this constructive comment. Our proposed Information Bottleneck (IB) principle is applied in two aspects: one is the control of information within the student model itself, and the other is the control of information transfer from the teacher model to the student model. In both path, the IB-controlled information flow encompasses both temporal and spatial features. The hidden states of the student MLP are rich in temporal and spatial information, while the distilled information from the teacher model is abundant in the temporal and spatial characteristics inherently modeled by the STGNNs.