Spatio-Temporal Graph Knowledge Distillation

We are grateful for your detailed feedback and suggestions! Please find our updates and modifications in response to your feedback below. If our responses resolve your concerns, we genuinely hope that you can increase your rating score.

1) Data description

Thank you for your concerns about the data description. In fact, you can refer to Section A.4 'Details of the Experimental Datasets' in the supplementary material, where the data processing procedures and characteristics of different datasets are elaborately described. Additionally, we appreciate your pointing out the typos.

2) Limitation discussion

We discussion the limitation of our STGKD as below: Firstly, although our teacher-bounded regression loss could effectively prevent the student model from being misguided by deterministic yet erroneous regression results produced by the teacher model, the threshold $\delta$ is a hyperparameter and can not adaptively adjust. So how to adaptively learn a $\delta$ to control the KD process is a future opportunity. Moreover, our proposed KD paradigm only concern the knowledge transferring between the embedding level, other perspective of KD, e.g., structure-based KD, are not considered in our work. Therefore, how to distillate teacher's structural knowledge to student is also a future direction.

3) Comparison of efficiency

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.

4) Interpretability of our STGKD

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.

4) Weakness of spatio-temporal prompt learning

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.

5) Clarification for the evaluation

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.

We are grateful for your detailed feedback and suggestions! Please find our updates and modifications in response to your feedback below. If our responses resolve your concerns, we genuinely hope that you can increase your rating score.

1) Further discussison about challenges

Thank you for your attention to how we articulated the challenges in the introduction. Indeed, the discussion of scalability and generalization issues in the introduction revolves around spatio-temporal forecasting. Concerning scalability, we highlight that spatio-temporal forecasting often involves large-scale datasets (i.e., substantial temporal and spatial nodes), where the computational time and space complexity of state-of-the-art STGNN models increase dramatically, hindering deployment. As for generalization, when the temporal span of spatio-temporal data widens, inherent distribution shifts occur, affecting the generalizability of the original STGNNs.

2) Differences from previous work

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.

3) Clarification for the theoretical guarantees

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.

We are grateful for your detailed feedback and suggestions! Please find our updates and modifications in response to your feedback below. If our responses resolve your concerns, we genuinely hope that you can increase your rating score.

1) Clarification about the novelty of our work

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.

2) Concerns about the selection of baselines

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.

Thank you for taking the time to review my paper and for providing such detailed feedback and questions. Your comments are very helpful in improving our work! Please find our updates and modifications in response to your feedback below. If our responses resolve your concerns, we genuinely hope that you can increase your rating score.

1) Differences between applying the knowledge distillation to CNN and STGNN

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.

2) Clarification for Spatio-Temporal IB Instantiating

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.

3) Clarification for the spatio-temporal prompts

Thank you for your insightful question. The spatio-temporal prompts we propose are learnable, meaning that for different temporal spans of spatio-temporal sequences as inputs, our model can assign both static and dynamic spatio-temporal prompts simultaneously. This infuses a substantial amount of spatio-temporal context into the MLP framework. These prompts are optimized together with the overall model, rather than being transformation from the spatio-temporal sequence inputs. Therefore, they should be distinguished from the spatio-temporal features themselves.