Condensing Graphs via One-Step Gradient Matching

What came before

DD (bi-level meta-learning) and DC (Zhao et al. 2021, gradient matching with multi-step unrolled inner training) work on continuous images. GCond had extended gradient matching to graphs for node classification but keeps the full multi-step bi-level inner loop and produces a weighted, real-valued adjacency matrix that must be stored in float and thresholded post hoc.

The problem

Two obstacles block a direct port of image gradient matching to graphs: (1) the adjacency matrix is discrete (0/1), and no existing condensation method outputs binary values; (2) the multi-step bi-level inner loop used by DC/GCond has complexity that scales with the product of inner and outer iterations, which is especially costly for graphs because the GNN forward pass itself scales with the number of edges (quadratic potential in condensed node count per graph).

The idea

Model each entry of the condensed adjacency matrix as a Bernoulli random variable (parameterized by a learnable logit, made differentiable via the binary-concrete/Gumbel-softmax reparameterization), and replace multi-step bi-level gradient matching with "one-step" matching that only aligns gradients at the initial random weights $\theta_0$, discarding the training trajectory — theoretically justified by a generalization bound showing the loss gap is controlled by the initial-step gradient distance plus a norm term.

How it works

Optimizes condensed node features $\mathbf{X}'$ and adjacency logits $\Omega$ (with $\mathbf{A}'_{ij}= \sigma((\log\alpha-\log(1-\alpha)+\Omega_{ij})/\tau)$, $\alpha\sim\mathrm{Uniform}(0,1)$, annealed temperature $\tau$) by minimizing, per class, a one-step gradient-matching loss $D(\nabla_\theta \ell(f_{\theta_0} (\mathbf{A}'(\Omega),\mathbf{X}'),\mathcal{Y}'), \nabla_\theta \ell(f_{\theta_0}(\mathcal{T}),\mathcal{Y}))$ averaged over $K_1{=}1000$ sampled random initializations $\theta_0$, plus a sparsity regularizer on $\sigma(\Omega)$ and (for sum-pooling GNNs only) a second regularizer derived from the generalization bound. A 3-layer GCN is the network in the loop; labels $\mathcal{Y}'$ are fixed hard labels. For node classification the same one-step loss is substituted directly into GCond's framework.

Evidence

Table 1 (graph classification, GCN evaluator, hard labels): ogbg-molbace at 2 graphs/0.2% budget reaches 0.657 ROC-AUC vs. 0.714 whole-dataset (92% retained); ogbg-molhiv at 0.3% budget reaches 0.731 vs. 0.757 whole (96.5% retained); DD dataset at 0.2%/2.1%/10.6% budgets retains 89%/93%/98% of whole-dataset accuracy. Table 3 (node classification, ported into GCond's protocol): DosCond matches GCond within ~2 points on Cora/Citeseer/Pubmed/Arxiv/Flickr while running up to 40x faster per 100 epochs (e.g. Cora 3.5s vs. 75.9s). Table 2: on CIFAR10 (50 graphs/class) DosCond takes ~5 minutes to reach the accuracy DosCond-Bi (the bi-level ablation) takes 75 minutes to reach — 6.7% of the training cost. Key ablation: DCG (a DC variant that only learns features on randomly-selected, i.e. non-learned, structure) trails DosCond by a large margin on graph classification, showing that learning discrete structure jointly with features (not just features) is necessary to preserve classification-relevant information.

Limitations

Performance on CIFAR10-as-graph is comparatively weak, attributed to its more complex topology needing more synthetic parameters than the tested budgets provide; datasets are still small-to-modest scale (molecular graphs, TU datasets, one 1,109-graph e-commerce set) with no ImageNet-scale or web-scale graph tested; the one-step approximation is justified by a bound that assumes the largest gradient gap occurs at initialization, an assumption not directly verified beyond the two loss-term-scale plots; and the paper's own future work flags a lack of interpretability for the learned discrete structures.