SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Work by Dong et al. [2022] suggests that structural bias can be measured through graph representation bias. Leveraging this, we aim for low coherence in node attributes and representations to fully remove bias from our distilled graph. Specifically, we introduce attribute and structural coherence to decrease attribute and structural bias, respectively, by minimizing the variance in sensitive group membership estimation for node attributes and representations. Given a real graph data  $\mathcal{G} = \{\pmb {A},\pmb {X},\pmb {Y},\pmb {S}\}$  and a distilled graph  $\mathcal{G}' = \{\pmb {A}',\pmb {X}',\pmb {Y}'\}$ , we feed them into a  $L$ -layer GNN, where the  $l$ -th layer latent representation in the GNN is denoted as  $\pmb{Z}_l$ . The latent representation for node attribute after  $l$ -hop propagation contains both attribute bias as well as structural bias. Note the node attribute  $\pmb{X}$  and  $\pmb{X}'$  before propagation as  $\pmb{Z}_0$  and  $\pmb{Z}_0'$ , we get a set of latent presentation  $\{Z_0,Z_1,\dots,Z_L\}$  for  $\mathcal{G}$  and  $\{Z_0',Z_1',\dots,Z_L'\}$  for  $\mathcal{G}'$ . The objective to measure bias of  $\pmb{Z}_l'$  is:
$$
C o h \left(\boldsymbol {Z} _ {l} ^ {\prime}\right) = \widehat {V a r} \left(\boldsymbol {\pi} ^ {j} \left(\boldsymbol {Z} _ {l} ^ {\prime}\right)\right) = \widehat {V a r} \left(\frac {\exp \left(- \lambda \| \boldsymbol {C} ^ {j} \boldsymbol {Z} _ {l} ^ {\prime} \| _ {2}\right)}{\sum_ {j} \exp \left(- \lambda \| \boldsymbol {C} ^ {j} \boldsymbol {Z} _ {l} ^ {\prime} \| _ {2}\right)}\right), \tag {12}
$$
where  $C^j = \gamma_j\left(\boldsymbol {I} + \gamma_j\boldsymbol {Z}_l\Pi^j\boldsymbol {Z}_l^T\right)^{-1}$ .  $\Pi^j$  is introduced in Sec 4.1. Since we consider the binary sensitive attribute,  $j$  is set as 0 without losing generality and is omitted in the notations as  $\pi^j (\cdot)\coloneqq \pi (\cdot)$ . After considering all the latent representations, the coherence loss objective is defined as the summation of all coherence over all layers, i.e.,
$$
\mathcal {L} _ {\text {C o h}} = \sum_ {l = 0} ^ {L} \operatorname {C o h} \left(\boldsymbol {Z} _ {l} ^ {\prime}\right) = \sum_ {l = 0} ^ {L} \widehat {V a r} \left(\boldsymbol {\pi} \left(\boldsymbol {Z} _ {l} ^ {\prime}\right)\right). \tag {13}
$$
Prediction loss for GNN training. The GNNs model is trained on distilled graph  $\mathcal{G}' = \{\pmb{A}',\pmb{X}',\pmb{Y}'\}$  with prediction loss. We adopt  $L$ -layer GNNs model, where  $\theta$  is the parameter of the GNN. We also adopt cross-entropy loss by default:
$$
\mathscr {L} _ {C E} = \mathscr {L} \left(\mathrm {G N N} _ {\theta} \left(\boldsymbol {A} ^ {\prime}, \boldsymbol {X} ^ {\prime}\right), \boldsymbol {Y} ^ {\prime}\right), \tag {14}
$$
4.3 Final Objective and Training Algorithm
Outer loop optimization. In the outer loop, we optimize the fair graph synthesizer with gradient matching loss and coherence loss:
$$
\min  _ {\boldsymbol {X} ^ {\prime}, \boldsymbol {A} ^ {\prime}} \mathcal {L} _ {G M} + \alpha \mathcal {L} _ {C o h}, \tag {15}
$$
where  $\alpha$  is a hyperparameter to regularize the debiasing intensity. The distilled node attribute  $X^{\prime}$  and the distilled node label  $\mathbf{Y}^{\prime}$  are initialized with the nodes uniformly sampling from real graph data  $\mathcal{G}$ .
Inner loop optimization. The GNN parameter  $\theta$  is optimized in the inner loop:
$$
\min  _ {\theta} \mathcal {L} _ {C E} \left(\mathrm {G N N} _ {\theta} \left(\boldsymbol {A} ^ {\prime}, \boldsymbol {X} ^ {\prime}\right), \boldsymbol {Y} ^ {\prime}\right). \tag {16}
$$
Instead of using the real graph data  $\mathcal{G}$  to calculate the loss, we use the distilled graph  $\mathcal{G}'$ . It empirically shows good performance and better efficiency. But the adversarial training baseline uses  $\mathcal{G}$  as it needs the sensitive attribute for discriminator training.
Table 1: Comparison on utility and bias mitigation between GNNs with real graph data (denoted as Real), the distilled small graph without debiasing (denoted as Vanilla), and debiased distilled graph (denoted as FGD) as input.  $\uparrow$  denotes the larger, the better;  $\downarrow$  denotes the opposite. The best ones are in bold. The better performers in Vanilla and FGDare underlined.
GCN	SGC	GraphSAGE
Real	Vanilla	FGD	Real	Vanilla	FGD	Real	Vanilla	FGD
Pokec-z	ACC ↑	70.96±0.4%	66.36±1.0%	66.58±0.7%	70.79±0.1%	68.31±0.7%	68.36±0.3%	70.59±0.3%	67.13±0.4%	67.83±0.6%
AUC ↑	78.19±0.2%	70.30±0.4%	70.48±0.3%	77.16±0.0%	73.51±0.6%	72.92±0.4%	77.91±0.2%	70.47±0.0%	70.26±0.5%
F1 ↑	72.16±0.5%	65.66±0.7%	66.48±0.8%	71.21±0.0%	68.09±0.4%	67.94±0.6%	72.07±0.4%	66.34±0.8%	66.62±0.7%
ΔDP ↓	4.13±1.3%	2.84±1.1%	1.75±1.1%	4.64±0.1%	7.60±1.7%	5.77±0.2%	4.54±1.3%	3.74±0.8%	2.17±1.6%
ΔEO ↓	4.57±1.7%	2.19±1.3%	1.19±1.0%	5.26±0.1%	7.88±1.9%	4.78±0.2%	5.25±1.2%	2.50±1.2%	2.56±1.2%
Pokec-n	ACC ↑	71.97±0.3%	50.10±2.7%	54.80±1.5%	71.16±0.0%	68.06±0.9%	68.19±0.8%	71.91±0.3%	52.80±2.2%	58.40±1.6%
AUC ↑	78.15±0.2%	51.09±2.3%	63.75±0.5%	76.34±0.0%	69.96±0.3%	70.18±0.5%	77.56±0.1%	53.95±2.8%	60.06±2.3%
F1 ↑	69.92±0.4%	44.21±4.6%	49.90±5.1%	67.63±0.0%	63.95±0.1%	64.03±0.1%	70.01±0.3%	58.75±6.1%	62.36±2.0%
ΔDP ↓	0.59±0.4%	4.02±0.6%	0.66±0.5%	4.3±0.1%	5.00±0.5%	4.7±0.5%	0.99±0.4%	2.38±2.8%	2.09±1.6%
ΔEO ↓	1.04±0.6%	5.20±1.0%	1.20±1.2%	2.26±0.1%	4.6±0.9%	4.26±1.2%	1.64±0.6%	2.81±3.0%	2.02±1.2%
German	ACC ↑	74.37±0.4%	72.83±0.8%	70.50±0.1%	72.62±1.6%	70.24±0.2%	70.06±0.1%	74.24±0.2%	72.00±0.5%	71.43±0.9%
AUC ↑	74.31±0.2%	57.75±4.8%	57.89±5.5%	74.94±1.2%	54.82±0.4%	53.92±1.4%	71.37±0.4%	57.32±0.4%	58.76±5.9%
F1 ↑	84.24±0.1%	83.51±0.4%	83.09±0.4%	83.13±0.3%	82.43±0.0%	82.36±0.0%	84.18±0.1%	83.12±0.4%	82.93±0.3%
ΔDP ↓	4.8±3.9%	5.38±3.3%	1.93±0.1%	3.36±2.8%	2.10±2.8%	1.3±0.9%	3.00±0.8%	5.33±2.9%	0.76±0.5%
ΔEO ↓	2.50±2.7%	1.04±1.1%	0.61±0.4%	9.38±9.0%	2.44±0.1%	0.28±0.2%	1.64±0.5%	2.89±1.0%	0.52±0.4%
Credit	ACC ↑	80.54±0.0%	76.92±1.9%	77.91±0.1%	79.66±0.1%	77.41±0.0%	77.36±0.7%	80.52±0.0%	77.91±0.5%	77.86±0.1%
AUC ↑	75.89±0.0%	68.82±2.1%	68.87±0.2%	73.39±0.1%	71.78±0.0%	72.02±0.1%	75.89±0.0%	71.12±0.2%	71.36±0.1%
F1 ↑	88.41±0.0%	85.47±2.0%	87.33±0.4%	88.09±0.0%	85.46±0.0%	85.37±0.7%	88.41±0.0%	87.00±0.9%	86.79±0.5%
ΔDP ↓	5.41±0.6%	12.04±5.4%	4.94±4.8%	2.78±0.3%	9.58±0.4%	7.28±2.8%	6.22±0.6%	8.77±2.0%	4.15±0.8%
ΔEO ↓	3.12±0.6%	9.58±5.5%	3.56±3.4%	1.19±0.3%	7.14±0.5%	5.56±0.1%	3.92±0.5%	6.72±4.0%	3.34±0.9%
Recidivism	ACC ↑	94.45±0.0%	70.89±1.8%	70.09±2.6%	85.10±0.1%	73.32±0.2%	73.10±0.5%	94.48±0.0%	73.66±1.0%	70.63±2.0%
AUC ↑	97.76±0.0%	71.95±2.8%	75.81±4.3%	92.24±0.1%	72.23±0.6%	72.44±0.1%	97.78±0.0%	75.84±1.2%	70.47±5.7%
F1 ↑	92.32±0.0%	55.30±3.9%	52.97±8.3%	76.94±0.3%	60.77±0.8%	60.45±0.6%	92.35±0.1%	60.87±2.4%	56.39±6.1%
ΔDP ↓	6.52±0.1%	2.68±1.0%	0.54±0.3%	7.89±0.0%	2.85±0.9%	1.08±0.7%	6.61±0.1%	0.97±0.8%	0.10±0.1%
ΔEO ↓	3.45±0.4%	1.41±0.5%	0.46±0.2%	8.55±0.2%	2.69±0.6%	1.32±0.8%	3.56±0.3%	0.93±0.8%	0.48±0.4%
We design experiments to validate the effectiveness of the proposed framework FGD and answer the following research questions: RQ.1 How well can FGD mitigate the bias in the distilled graph and alleviate the fairness issue of the GNNs trained on the distilled small graph? RQ.2 How well can FGD balance the trade-off between accuracy and bias mitigation compared with other debiasing baselines? RQ.3 Can FGD further improve the utility or bias mitigation as an add-on module to other bias mitigation methods?
5.1 Experimental setting
Datasets. We use five real-world datasets, including Pokec-z, Pokec-n [Dai and Wang, 2021, Takac and Zabovsky, 2012], German, Credit, and Recidivism [Agarwal et al., 2021]. The detailed setting of the datasets is in Appendix F
GNN Models. We adopt three popular GNN variants in our experiments, including GCN [Kipf and Welling, 2016], SGC [Wu et al., 2019], GraphSAGE [Hamilton et al., 2017].
Baselines. Given the absence of fair graph distillation work, we compare our approach with four baselines. (1) Real uses real graph data to train a GNN model. (2) Vanilla applies a vanilla graph distillation algorithm [Jin et al., 2022] for distilled graph data learning and GNN model training. (3) FairGNN[Dai and Wang, 2021] trains a fair GNN model adversarially on real graph data, aiming to achieve good prediction while fooling a discriminator. (4) EDITS [Dong et al., 2022] is a model-agnostic debiasing method that rewrites graph data for fair GNNs.
Evaluation Metrics. We evaluate model performance from model utility and bias measurements. Good performance represents low bias and high model utility. For model utility metrics, we adopt accuracy (ACC), the area under the receiver operating characteristic curve (AUC), and F1 score to measure prediction performance. For bias measurement, we adopt two commonly used fairness metrics, i.e., demographic parity  $(\Delta_{DP})$  and equal opportunity  $(\Delta_{EO})$  [Beutel et al., 2017, Louizos et al., 2015]. Denote the binary label as  $y \in \{0,1\}$ , and sensitive attribute as  $s \in \{0,1\}$ .  $\hat{y} \in \{0,1\}$  denotes the model prediction. The violation of DP and EO are given by  $\Delta_{DP} = |P(\hat{y} = 1| s = 0) - P(\hat{y} = 1| s = 1)|$ , and  $\Delta_{EO} = |P(\hat{y} = 1| y = 1, s = 0) - P(\hat{y} = 1| y = 1, s = 1)|$ .
Figure 4: Trade-off comparison between FGD and other baselines for five real-world graph datasets.
5.2 Debiasing distilled Graph
In response to RQ.1, we assess FGD's bias mitigation and prediction performance across various GNN architectures and datasets, as shown in Table 1. We compare the  $\Delta_{DP}$  and  $\Delta_{EO}$  values of GNNs trained on real graphs (Real), vanilla distilled graphs (Vanilla), and debiased distilled graphs via FGD (FGD). Key findings include: (1) Models trained on Real graphs consistently outperform Vanilla and FGD in utility, though FGD's utility matches or surpasses Vanilla. (2) FGD consistently yields lower bias than Vanilla, and outperforms Real on 4 out of 5 datasets, excluding Token-n.
We compare the coherence bias of distilled graphs generated by Vanilla and FGD methods across five real-world datasets with the GCN architecture (Table 2). Our analysis reveals that FGD reduces unfairness, reflected in lower coherence bias in the distilled graphs. This consistency confirms the effectiveness of coherence bias as a measure of distilled graph bias.
5.3 Trade-Off Comparison
In response to RQ.2, we compare the trade-off between model utility and bias mitigation against other baselines using the GCN architecture. We utilize the Pareto frontier Ishizaka and Nemery [2013] to evaluate our approach's utility-fairness trade-off, using different hyperparameters. The Pareto frontier graphically represents optimal trade-offs in multi-objective optimization. We use AUC as the utility metric and  $\Delta_{DP}$  and  $\Delta_{EO}$  as fairness metrics. Higher AUC and lower  $\Delta_{DP} / \Delta_{EO}$  are preferred, so models with Pareto frontier curves closer to the bottom right corner (AUC on the horizontal axis and  $\Delta_{DP} / \Delta_{EO}$  on the vertical) have better trade-off performance.
Table 2: Coherence bias comparison between vanilla distilled graph (denoted as Vanilla) and fair distilled graph (denoted as FGD). The lower, the better. The best ones are marked in bold. The architecture model is GCN.
Vanilla	FGD
Pokec-z	0.009468	0.002123(-77.57%)
Pokec-n	0.004464	0.000432(-90.32%)
German	0.012772	0.003489(-72.68%)
Credit	0.011864	0.002866(-75.84%)
Recidivism	0.000098	0.000038(-61.22%)
Figure 4 shows the results for models trained on the real graph, the distilled graph debiased by baseline methods (vanilla graph distillation, FairGNN, and EDITS $^5$ ) and the distilled graph debiased by FGD. We can observe: (1) From a model utility perspective, FGD performs comparably to other baselines, like vanilla graph distillation, FairGNN, and EDITS $^6$ , suggesting it preserves sufficient information for node classification. (2) In terms of bias mitigation, all baselines show effectiveness, with FGD exhibiting the best results. (3) When considering the utility-fairness trade-off, FGD's Pareto front curve lies at the bottom right corner of all baselines, signifying it offers the best balance. Thus, FGD outperforms other baselines in balancing model utility and bias mitigation.
5.4 Add-on Module
<sup>5</sup>EDITS publishes fair graph for German, Credit and Recidivism dataset on Github
$^{6}$ Out of memory (OOM) issue appears when running EDITS on Pokec-z and Pokec-n datasets.
In addition to its superior trade-off performance, our method, FGD, can enhance other debiasing baselines like FairGNN and EDITS by acting as an add-on debias module. This compatibility is due to the fact that these baselines can replace the cross-entropy loss in the GNN training module. To answer RQ.3, we conducted experiments on the Credit dataset comparing FairGNN/EDITS performance with and without FGD. As shown in Figure 5, FairGNN/EDITS coupled with FGD delivers better utility-fairness
trade-off, demonstrating FGD's potential to boost other debias methods.
Figure 5: Trade-off comparison between FairGNN, EDITS and FairGNN+FGD, EDITS+FGD on Credit dataset.
Dataset Distillation & Knowledge Distillation. Dataset Distillation (DD) and Knowledge Distillation (KD) are methods to improve the efficiency of training deep neural networks. DD synthesizes a small dataset encapsulating the knowledge of a larger one, achieving comparable model performance [Wang et al., 2018, Kim et al., 2022, Lee et al., 2022, Zhao et al., 2021a, Yang et al., 2022]. It employs a bi-level optimization approach, with dataset condensation (DC) speeding up the process via gradient matching of model parameters. DD also helps with repeated training or privacy applications like continual learning, neural architecture search, and privacy-preserving scenarios. Meanwhile, graph data condensation methods have been developed for node and graph classification tasks [Jin et al., 2022]. KD, on the other hand, enhances computational efficiency through model compression and acceleration. It trains a compact student model using the knowledge from a larger teacher model Gou et al. [2021]To address the scarcity and high complexity of labeled data in GNNs, knowledge distillation (KD) was introduced to enhance existing GNNs Liu et al. [2023a], Wang et al. [2023], also for fairness problem Dong et al. [2023]. While KD focuses on model compression, DD targets data compression, each improving efficiency from model-centric and data-centric perspectives.
Fair Graph Learning. Fairness in machine learning has attracted many research efforts[Chuang and Mroueh, 2020, Zhang et al., 2018, Du et al., 2021, Jiang et al., 2022b, Han et al., 2023, Jiang et al., 2023]. Many technologies are introduced in graph neural networks to achieve fair graph learning in node classification tasks, including optimization with regularization [Jiang et al., 2022a], rebalancing [Zeng et al., 2021], adversarial learning [Dai and Wang, 2021, Bose and Hamilton, 2019, Fisher et al., 2020] and graph rewiring [Köse and Shen, 2021, Dong et al., 2022]. For link prediction, dyadic fairness and corresponding graph rewiring solutions are also developed in [Li et al., 2021]. Another line of work focuses on solving the individual fairness problem on the graph data Song et al. [2022], Dong et al. [2021], Kang et al. [2020].
Despite the ability of graph distillation to condense valuable graph data, this study finds that the vanilla method can worsen fairness issues. Therefore, we introduce a fair graph distillation process to generate fair distilled graph data. As the distilled graph lacks the nodes' sensitive attributes, conventional fair methods cannot be directly applied. However, we identify a consistent geometric phenomenon in graph distillation to estimate these sensitive attributes. We also introduce a new bias metric, coherence, and propose a bi-level optimization framework, FGD, for fair graph distillation. Experimental results validate FGD's effectiveness in mitigating bias while maintaining model utility across various GNN architectures and datasets. Future work will focus on addressing individual fairness issues and non-binary sensitive attribute conditions, among other aspects, as discussed in Appendix H.
We are deeply grateful to the National Science Foundation for their unwavering support. This research was substantially facilitated by the funding from grants IIS-1939716 and IIS-1900990.
L. Takac and M. Zabovsky. Data analysis in public social networks. In International scientific conference and international workshop present day trends of innovations, volume 1. Present Day Trends of Innovations Lamza Poland, 2012.
A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526-9536. PMLR, 2020.
T. Wang, J.-Y. Zhu, A. Torralba, and A. A. Efros. Dataset distillation. arXiv preprint arXiv:1811.10959, 2018.
Y. Wang, B. Hooi, Y. Liu, and N. Shah. Graph explicit neural networks: Explicitly encoding graphs for efficient and accurate inference. In Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining, pages 348-356, 2023.
F. Wu, A. Souza, T. Zhang, C. Fifty, T. Yu, and K. Weinberger. Simplifying graph convolutional networks. In International conference on machine learning, pages 6861-6871. PMLR, 2019.
S. Yang, Z. Xie, H. Peng, M. Xu, M. Sun, and P. Li. Dataset pruning: Reducing training data by examining generalization influence. arXiv preprint arXiv:2205.09329, 2022.
R. Ying, R. He, K. Chen, P. Eksombatchai, W. L. Hamilton, and J. Leskovec. Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pages 974-983, 2018.
Z. Zeng, R. Islam, K. N. Keya, J. Foulds, Y. Song, and S. Pan. Fair representation learning for heterogeneous information networks. In Proceedings of the International AAAI Conference on Web and Social Media, volume 15, pages 877-887, 2021.
B. H. Zhang, B. Lemoine, and M. Mitchell. Mitigating unwanted biases with adversarial learning. In Proceedings of the 2018 AAAI/ACM Conference on AI, Ethics, and Society, pages 335-340, 2018.
B. Zhao, K. R. Mopuri, and H. Bilen. Dataset condensation with gradient matching. *ICLR*, 1(2):3, 2021a.
B. Zhao, K. R. Mopuri, and H. Bilen. Dataset condensation with gradient matching. *ICLR*, 1(2):3, 2021b.
K. Zhou, Q. Song, X. Huang, and X. Hu. Auto-gnn: Neural architecture search of graph neural networks. arXiv preprint arXiv:1909.03184, 2019.
A Training Algorithm
The training algorithm for fair graph distillation is shown in Algorithm 1.
Algorithm 1 Fair Graph Distillation
Input: Training graph data  $\mathcal{G} = \{\pmb {X},\pmb {A},\pmb {S},\pmb {Y}\}$ , hyperparameters  $\alpha$ , temperature  $\gamma$ , number of alternative optimization step  $T_{alt}$ , distilled label  $\pmb{Y}'$ .
Initialize  $X^{\prime}$  based on real attributes, synthesizer model  $\phi$ , and GNNs model  $\theta$ .
for  $t = 1$  to  $T_{alt}$  do
1. Train GNNs model using distilled graph  $\mathcal{G}_t^\prime$  and Equation (15) to obtain  $GNN_{\theta_t}$ .
2. Given GNNs model  $GNN_{\theta_t}$ , calculate the gradient distance  $D\big(\nabla_{\theta}\mathcal{L}(\mathcal{G}),\nabla_{\theta}\mathcal{L}(\mathcal{G}_t^{\prime})\big)$  over the real graph  $\mathcal{G}$  and distilled graph  $\mathcal{G}_t^\prime$
3. Calculate coherence loss based on GNNs model  $GNN_{\theta_t}$ , real graph  $\mathcal{G}$  and distilled graph  $\mathcal{G}_t'$ .
4. Train synthesizer model using prediction loss as Equation (16).
Output: The fair distilled graph  $\mathcal{G}' = \{\pmb{A}',\pmb{X}',\pmb{Y}'\}$ .
B Proof of Theorem 3.4
We consider GNNs model to learn node presentation  $\mathbf{z}_i$  in the real graph  $\mathcal{G}$  and then followed a linear classifier  $\mathbf{W} = [\mathbf{w}_0, \dots, \mathbf{w}_{C-1}]$  and softmax layer, where  $\mathbf{w}_j$  is the weight vector connected to the  $j$ -tj output neuron. We first focus on the relation between the latent representation and the gradient of the linear classification layer. It is easy to obtain the cross-entropy loss  $J_i(J_i')$  for  $i$ -th node with label  $y_i$  in real graph  $\mathcal{G}$  (distilled graph  $\mathcal{G}'$ ) as follows:
$$
J _ {i} = - \log \frac {\exp \left(\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}, \tag {17}
$$
Then we define gradient over weight vector as  $\pmb{g}_{i,j} = \frac{\partial J_i}{\partial\pmb{w}_j}$  and  $\pmb{g}_{i,j}' = \frac{\partial J_i'}{\partial\pmb{w}_j}$  in the real and distilled graph. If  $j = y_{i}$ , we can obtain
$$
\begin{array}{l} \boldsymbol {g} _ {i, y _ {i}} = - \frac {\sum_ {k} \exp (\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i})}{\exp (\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i})} \\ \cdot \frac {\exp \left(\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i}\right) \sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right) - \exp^ {2} \left(\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\left(\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)\right) ^ {2}} \cdot \boldsymbol {z} _ {i} \\ = - \boldsymbol {z} _ {i} + \frac {\exp \left(\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)} \cdot \boldsymbol {z} _ {i}, \tag {18} \\ \end{array}
$$
Similarly, for  $j \neq y$ , we have
$$
\boldsymbol {g} _ {i, j} = \frac {\exp \left(\boldsymbol {w} _ {y i} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)} \cdot \boldsymbol {z} _ {i}, \tag {19}
$$
In other words, the gradient of the loss for  $i$ -th node with label  $y_{i}$  with respect to the weight vector connected to the  $j$ -th output neuron is given by
$$
\boldsymbol {g} _ {i, j} = \frac {\exp \left(\boldsymbol {w} _ {y _ {i}} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)} \cdot \boldsymbol {z} _ {i} - \mathbb {1} _ {j = y _ {i}} \boldsymbol {z} _ {i}. \tag {20}
$$
Based on Assumption 3.1, each model parameter in the last softmax layer satisfies the same distribution. In other words, the expectation of all predictions are the same, i.e.,
$$
\mathbb {E} _ {\mathcal {P} _ {\theta}} \left[ \frac {\exp \left(\boldsymbol {w} _ {0} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)} \right] = \dots = \mathbb {E} _ {\mathcal {P} _ {\theta}} \left[ \frac {\exp \left(\boldsymbol {w} _ {C - 1} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)}{\sum_ {k} \exp \left(\boldsymbol {w} _ {k} ^ {\top} \cdot \boldsymbol {z} _ {i}\right)} \right]. \tag {21}
$$
Note that the gradient calculation is based on backpropagation, the gradient for the last linear classification layer is quite critical for the gradient of other layers. Hence we consider the gradient of the last linear classification layer in the real graph, shown by
$$
\begin{array}{l} \mathbb {E} _ {\theta \sim \mathcal {P} _ {\theta}} \left[ \nabla_ {\boldsymbol {w} _ {j}} \mathcal {L} (\mathcal {G}) \right] = \mathbb {E} _ {\theta \sim \mathcal {P} _ {\theta}} \left[ \frac {1}{N} \sum_ {i = 1} ^ {N} \boldsymbol {g} _ {i, j} \right] \\ = \frac {1}{N C} \sum_ {i = 1} ^ {N} z _ {i} - \frac {1}{N} \sum_ {\{i: y _ {i} = j \}} z _ {i}, \tag {22} \\ \end{array}
$$
Similarly, we have the gradient of the last linear classification layer in the distilled graph as follows: