Mirage: Model-agnostic Graph Distillation for Graph Classification

Q1. Would MIRAGE still be as effective if used on a very large dataset?

Response: We do not expect graph size to have any effect on the performance of MIRAGE since the underlying principles are not dependent on graph dataset size. This is already evident from Table 2. Specifically, we compute the pearson's correlation coefficient between dataset size and the generalization error (AUC in full dataset-AUC of MIRAGE). The correlation is 0.153 with $p$-value of 0.743 and thus rejecting any bias in AUC from dataset size.

Q2. Why is MIRAGE, which only uses the CPU, so much faster than DOSCOND and KIDD, which can use the GPU?

Response: Existing algorithms, namely DosCOND and KiDD, seek to replicate the gradient trajectory of model parameters on the full dataset. Hence, they need to train on the full dataset for distillation (and as we comment in Sec 1, this compromizes the basic premise of data distillation). In contrast, MIRAGE does not require training on the full dataset. Rather, MIRAGE emulates the input data processed by message-passing GNNs. By shifting the computation task to the pre-learning phase, MIRAGE is not only faster, but also independent of the GNN architecture.

Q3. If a datasets is large, but its global computation tree is very similar, will such a datasets affect the effect of MIRAGE?

Response: As we discussed in our response to Q1, the size of the dataset does not affect the performance of MIRAGE. The performance depends primarily on the frequency distribution of computation trees in a dataset. If the distribution is skewed, MIRAGE is expected to perform well since the top-$k$ most frequent computation trees would effectively capture a substantial portion of the distribution mass. On the other hand, if the distribution is not skewed, MIRAGE is not the most suitable distillation technique for such scenarios. This aspect is explicitly acknowledged in our conclusion (Sec 5). The relevant text is reproduced verbatim below.

Finally, MIRAGE relies on the assumption that the distribution of computation trees is skewed. Although we provide compelling evidence of its prevalence across a diverse range of datasets, this assumption may not hold universally, especially in the case of heterophilous datasets. The development of a model-agnostic distillation algorithm remains an open challenge in such scenarios.

Q4. How will MIRAGE perform in multi-classification tasks?

Response: We have added a new dataset, namely IMDB-M, to evaluate MIRAGE's performance on multi-class classification. As we see below, MIRAGE produces the best performance. These results, and the details of the IMDB-M dataset has been added to Table 2 and Table 1 respectively in the revised manuscript.

The size of the dataset (in bytes) is as follows:

---

1. Unlike DOSCOND and KIDD, MIRAGE seems infeasible to be generalized to node classification tasks.

Response: While the primary focus of MIRAGE is on graph classification (as noted in the title itself), MIRAGE does generalize to node classification. We now empirically demonstrate this ability in App. E. (Note: KIDD also does not discuss node classification. Only, DosCOND does).

Generalization to node classification is simple since learning the node embeddings is a necessary step before aggregating them to form the graph embedding (Recollect Eqs. 4 and 5). Specifically, the pipeline follows in an identical manner except the omission of (Eq. 5) and the loss function being applied to node embeddings instead of graph embeddings.

We use ogbg-molbace to analyze performance in node classification. Here, each node is labeled as aromatic or non-aromatic depending on whether it is part of an aromatic ring substructure. The results are tabulated in Table H. Consistent with previous results, MIRAGE outperforms DosCOND in both AUC-ROC and compression (smaller size of the distilled dataset). DosCOND produces a dataset that is $\approx 4$ times the size of that produced by MIRAGE, yet performs more than $7\%$ lower in AUC-ROC. This coupled with model-agnostic-ness further solidifies the superiority of MIRAGE

2. There is no theoretical analysis nor experimental proof to investigate the error when approximating the true graph representation by aggregating the root representations of only the highly frequent trees.

Response: The above question is indeed important. We have now included a new experiment to evaluate the sufficiency of the frequent computation tree patterns.

Sufficiency of frequent patterns: In order to establish sufficiency of frequent tree patterns in capturing the dataset characteristic, we conduct the following experiment. We train the model on the full dataset and store its weights at each epoch. Then, we freeze the model at the weights after each epoch's training and pass both the distilled dataset consisting of just the frequent tree patterns, and the full dataset. We then compute the differences between the losses as shown in Figure I (GCN, GAT, and GIN). The rationale behind this is that the weights of the full model recognise the patterns that are important towards minimizing the loss. Now, if the same weights continue to be effective on the distilled train set, it indicates that the distilled dataset has retained the important information. In the figure, we can see that the difference quickly approaches $0$ for all the models for all the datasets, and only starts at a high value at the random initialization where the weights are not yet trained to recognize the important patterns. Furthermore, gradient descent will run more iterations on trees that it sees more often and hence infrequent trees have limited impact on the gradients. These results empirically establish the sufficiency of frequent tree patterns in capturing the majority of the dataset characteristic.

Additional information relevant to this question: In Fig. 1, we show that the distribution of computational trees is long-tailed and hence by considering only the top-$k$ most frequent trees, majority of the input information processed by the GNN would be retained. Furthermore, in Fig. 5 (b), we show that the training loss against epochs is similar for the full dataset and distilled dataset.

1a(i). Problem motivation: What is the comparative performance of existing distillation methods (e.g., gradient matching) when applied to both randomly initialized and well-trained models? Is there a substantial performance gap between the two scenarios?

Response: Below we present the AUCROC numbers of DosCOND on randomly initialized DosCOND model, warm start (100 epochs) and training DosCOND until convergence (typically 1000 epochs). As visible, there is a noticeable gap in the AUCROC numbers indicating training until convergence is necessary. This discussion has been added App. D.

1a(ii). Regarding computational costs in the distillation process, it's worth examining the resource requirements for calculating the computation tree frequency distribution, particularly in the context of deep GNNs. How does this computational overhead compare to training models on the full dataset?

Response: The time taken to distill each dataset is provided in Fig. 4(a) in the form of a bar plot and the raw numbers are in Table E. MIRAGE is $\approx$ 500 times faster on average than KIDD and $\approx 150$ times faster than DOSCOND. To further expand on this aspect, we have added the following new experiments:

• Comparison to full dataset training time: We now report the time taken to train on the full dataset in Table I in Appendix. On average, the distillation time of MIRAGE is more than $30$ times faster than training on the full dataset. In contrast DosCOND and KiDD are slower than full dataset training due to reasons already highlighted in Sec 1 of our manuscript.
• Impact of hops (layers): In Fig. H (discussed in App F and referred to from Sec 4.4 in main paper), we report the growth in distillation time against the number of hops. We show that even at a larger values of hops, MIRAGE is 2 to 25 times faster, on average, than training on the full dataset.
• Impact of frequency threshold on distillation: we subject the system to variations in threshold parameters, graphically showing the resulting time in Figure 6. Notably, the distillation process exceeds the time of full dataset training solely under extreme threshold values. This divergence occurs when distilled dataset reaches equality with the full dataset in size post-distillation. Conversely, for pragmatic threshold values, the dataset distillation procedure consistently manifests as a significantly faster option to full dataset training.

1b. Technical Motivation: While the concept of the computation tree is inherent to GNNs, we should address the sufficiency of the frequency distribution for the computation tree. Key points to explore include: Why is maintaining a high computation tree frequency distribution sufficient for effective learning? What is the distribution pattern of the computation tree, and does it exhibit a long-tail distribution?

Response: We note that the frequency distribution of computation trees is already reported in our manuscript in Fig. 1 and it is indeed a long-tail distribution. To further establish the sufficiency of frequent tree patterns, we have now dedicated an explicit section (Sec 4.4).

Sufficiency of frequent patterns: In order to establish sufficiency of frequent tree patterns in capturing the dataset characteristic, we conduct the following experiment. We train the model on the full dataset and store its weights at each epoch. Then, we freeze the model at the weights after each epoch's training and pass both the distilled dataset consisting of just the frequent tree patterns, and the full dataset. We then compute the differences between the losses as shown in Figure 5(a) (GCN) and Figure I (GCN, GAT, and GIN). The rationale behind this is that the weights of the full model recognise the patterns that are important towards minimizing the loss. Now, if the same weights continue to be effective on the distilled train set, it indicates that the distilled dataset has retained the important information. In the figure, we can see that the difference quickly approaches $0$ for all the models for all the datasets, and only starts at a high value at the random initialization where the weights are not yet trained to recognize the important patterns. Furthermore, gradient descent will run more iterations on trees that it sees more often and hence infrequent trees have limited impact on the gradients. These results empirically establish the sufficiency of frequent tree patterns in capturing the majority of the dataset characteristics.

2(i). Objectives: Ideally, the distillate data should mirror the distribution of the full dataset, implying that the gradients of the distillate data should align with those of the full dataset for any model (random or well-trained). It's essential to clarify that while distillation can reduce computation costs, the primary objective remains data distillation.

Response: As mentioned in our response to 1b, the frequency distribution of computation trees is long tailed. With the newly added sufficiency experiments, we now also show that the frequent trees are the patterns that are being learned by the model parameters as well (Response to 1b.). Finally in Fig. 5b and J, we also show that the train losses on full dataset and distilled datasets, with their own respective training weights, are also aligned.

2(ii). Objectives: Additionally, considering the architecture transferability for existing graph distillation methods, the advantage of model-agnostic techniques is not significant.

Response: Our manuscript includes experiments to evaluate transferability of model architectures in graph distillation. The results, reported in Table F of Appendix C, reveal there is a noticeable drop in AUCROC when the model architecture for distillation is different than the architecture used for training on the distilled dataset. The same observation has been reported in another recent work as well (See Table 3 in [1]). In KiDD [2], they theoretically establish why the distillation GNN and the training GNN needs to be same (Page 3 in KiDD). These results provide verification of why model-agnostic distillation is a problem of importance.

To better highlight these results and motivate the proposed problem, we now point to this information from the introduction itself.

• [1] Beining Yang, Kai Wang, Qingyun Sun, Cheng Ji, Xingcheng Fu, Hao Tang, Yang You, and Jianxin Li. Does graph distillation see like vision dataset counterpart? NeurIPS, 2023*.(to appear)
• [2] Zhe Xu, Yuzhong Chen, Menghai Pan, Huiyuan Chen, Mahashweta Das, Hao Yang, and Hanghang Tong. Kernel ridge regression-based graph dataset distillation. ACM SIGKDD, 2023

Experiments: 3a. Comprehensive Model Comparison: In evaluating baselines, it would be beneficial to perform a comprehensive comparison across various model architectures. Specifically, assessing performance across different architecture pairs would enhance the evaluation's robustness and provide insights into transferability.

Response: Our evaluation already includes performance assessment across three different GNN architectures namely- GCN, GAT, and GIN, in Table 2. The results establish the proposed technique to be superior on average and consistently ranked among the top-2 distillation algorithms across all dataset-architecture combinations.

In addition, cross architecture transferability, i.e., the situation where the architecture used for distillation is different from the one used for training and inference is also included in Appendix C. As discussed in our response to 2(ii), there is a noticeable drop in AUCROC, which is consistent with existing literature [1]. This result motivates the need for model-agnostic distillation algorithms.

3b. Baselines and Model Initialization: Investigate the significance of GNN model training prior to distillation for baselines. Explore performance outcomes for well-trained, randomly initialized, and warm-up trained models. Understanding the effectiveness of random initialization can help determine the necessity of training on the full data for baseline methods.

Response: Below we present the AUCROC numbers of DosCOND on randomly initialized DosCOND model, warm start (100 epochs) and training DosCOND until convergence (typically around 1000 epochs). As visible, there is a noticeable gap in the AUCROC numbers indicating full training is necessary. This discussion has been added App. D.

Q1. In the results presented in Table 2, why there is a discrepancy compared to those from the original DosCOND paper? Could the authors clarify why the authors deviated from the settings employed in the DosCOND paper?

Response: The results in DosCOND do not exactly align with ours due to the following reasons:

• Input data: The graph datasets used (both in ours and DosCOND) are node and edge labeled. In DosCOND, only the node labels are used for predictive modeling, whereas we use both node and edge labels. We believe this is a better empirical design since generalization error between full and distilled dataset should be measured after incorporating the full input information.
• Metrics: A direct comparison between the AUCROC numbers of DosCOND and ours is possible only in the three datasets of ogbg-molbace, ogbg-mohiv, and ogbg-molbbbp, since DosCOND does not use AUCROC as the evaluation metric in the other datasets. We use AUCROC across all datasets.
  • Performance of DosCOND: In all three datasets, as we show in the table below (columns 2 and 3), DosCOND's AUCROC improves in our evaluation.
  • AUCROC on full dataset training: We observe significantly reduced AUCROC only in ogbg-molhiv (columns 4 and 5). We conducted hyper-parameter search over a larger space than in our initial setup for ogbg-molhiv (full dataset training) to revisit this issue and identified that reducing the weight of the L2 regularizer in the loss function from 0.01 to 0.00001 improves AUCROC to $75.93$ in GCN. A similar improvement is observed in GAT and GIN as well. This update has now been incorporated in the revised draft (highlighted in blue font in Table 2).

• Other things to note: DosCond reports numbers only for GCN, whereas we include GCN, GAT and GIN. The budget used in our experiments is set to 1 graph per class since even with this budget, MIRAGE is smaller than DosCond (details in Table 3).

Q2. It would be beneficial if the authors could delve deeper into the sensitivity of their approach with respect to distillation parameters $\theta_1$ and $\theta_2$. Undertaking further experiments to report the performance across varied distillation parameters might offer greater insights, given that these parameters significantly influence the size of the distilled dataset.

Response: We appreciate this constructive feedback. We have significantly expanded the studies on parameter sensitivity. Below we summarize the experiments conducted and the key insights derived from them.

Impact of Frequency Threshold on distillation time, compression and accuracy (AUCROC): In Section 4.5 (Fig. 6 presents the results), we measure the impact of frequency threshold on the distillation time. MIRAGE is significantly faster than training on full datasets unless the frequency thresholds are set to extreme values where the distilled dataset reaches equality with the full dataset.

In Fig. L in Appendix, we present the relationship between compression and AUCROC. The compression ratio is controlled by increasing the frequency threshold, i.e., the higher the threshold, the more the compression is. As noted, there is marginal improvement in AUCROC with lower thresholds (higher distillation size). This trend can be explained from the long-tailed distribution of computation trees. Even if we incorporate more computation trees, they have negligible impact on capturing more of the distribution mass and hence limited effect on AUCROC.

Impact of number of hops (layers): The details are furnished in Appendix F and H of the appendix. We note the main observations here. First, even under high hop counts from the GNN perspective, the distillation process is more time-efficient than complete dataset training. Moreover, the distilled dataset's performance converges closely with that of the full dataset, as evident in Table 2. Second, the distillation process exceeds the time of full dataset training solely under extreme threshold values. This divergence occurs when distilled dataset reaches equality with the full dataset in size post-distillation.

In Fig. K in Appendix, we present the impact of hops on the AUC. We see mild deterioration in AUC at higher hops. This is consistent with the literature since GNNs are known to suffer from oversmoothing and oversquashing at higher layers[1].

In summary, for pragmatic threshold values, the dataset distillation procedure consistently manifests as a significantly faster option to full dataset training.

[1] Hamed Shirzad, Ameya Velingker, Balaji Venkatachalam, Danica J. Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. In ICML, 2023

To explain the high compression achieved by MIRAGE, we first would like to point to Figure 1 in the paper and the discussion associated with it in Section 3.2. We observe that frequency distribution of computation trees is a long-tail distribution. The plot shows that a compact set of top-k frequent trees(non-isometric) captures a substantial portion of the distribution mass of the dataset. As the distillation process of Mirage builds the condensed dataset using only these frequent trees, hence, the storage is low. Furthermore, these trees are only n layers deep, where n is the number of layers in GNN. Typically n is within 3.

Finally, we also point to the sufficiency experiments in Sec 4.4. The experiments indicate that the loss is similar for both the distilled dataset and the full dataset when using a model with weights trained on the full dataset. Thus, showcasing the GNN is attending to these frequent patterns primarily. This behavior by the GNN is not surprising. Specifically, gradient descent will run more iterations on trees that it sees more often and hence infrequent trees have limited impact on the gradients.