Source-linked AI summary
Condensing Graphs via One-Step Gradient Matching
Wei Jin, Xianfeng Tang, Haoming Jiang, Zheng Li, Danqing Zhang, Jiliang Tang, Bing Yin
TL;DR
Large graph datasets make GNN training expensive, while existing condensation methods do not directly handle discrete graph structures and rely on costly nested optimization. DosCond models graph structure probabilistically and uses one-step gradient matching, achieving substantial dataset reduction while retaining much of the original performance. Its analysis and experiments support an efficient graph-condensation approach, subject to assumptions involving gradient gaps and an SGC-based theoretical study.
Problem
Existing condensation methods are not directly applicable to discrete graph structures and are computationally expensive because they use nested optimization.
Method
DosCond models discrete graph structures probabilistically and matches GNN gradients using a one-step strategy without training network weights.
Results
DosCond reduces dataset size by 90% while approximating up to 98% of the original performance.
Takeaways & Limitations
One-step gradient matching provides an efficient way to condense real graphs into a small number of informative discrete graphs.
Takeaways & Limitations
The theoretical analysis assumes the largest gradient gap occurs at the 0-th epoch and studies SGC as a simpler architecture related to GCN.
Abstract
from arXiv · showhide
As training deep learning models on large dataset takes a lot of time and resources, it is desired to construct a small synthetic dataset with which we can train deep learning models sufficiently. There are recent works that have explored solutions on condensing image datasets through complex bi-level optimization. For instance, dataset condensation (DC) matches network gradients w.r.t. large-real data and small-synthetic data, where the network weights are optimized for multiple steps at each outer iteration. However, existing approaches have their inherent limitations: (1) they are not directly applicable to graphs where the data is discrete; and (2) the condensation process is computationally expensive due to the involved nested optimization. To bridge the gap, we investigate efficient dataset condensation tailored for graph datasets where we model the discrete graph structure as a probabilistic model. We further propose a one-step gradient matching scheme, which performs gradient matching for only one single step without training the network weights. Our theoretical analysis shows this strategy can generate synthetic graphs that lead to lower classification loss on real graphs. Extensive experiments on various graph datasets demonstrate the effectiveness and efficiency of the proposed method. In particular, we are able to reduce the dataset size by 90% while approximating up to 98% of the original performance and our method is significantly faster than multi-step gradient matching (e.g. 15x in CIFAR10 for synthesizing 500 graphs). Code is available at \url{https://github.com/amazon-research/DosCond}.
1 INTRODUCTION
Graph neural networks require large datasets and costly training, motivating dataset condensation for graphs. Existing condensation methods face discrete-structure and nested-optimization challenges, while DosCond addresses both through probabilistic graph modeling and one-step gradient matching.
- GNNs extract useful patterns from graph data but are data-hungry, making graph-model training computationally expensive.
- Dataset condensation constructs small synthetic training sets intended to retain sufficient information for training neural networks.
- Existing image-oriented condensation methods cannot directly produce binary graph structures and typically require expensive bi-level optimization.
- DosCond models graph structure probabilistically for differentiable optimization and matches GNN gradients in one step without repeatedly updating network weights.
2 THE PROPOSED FRAMEWORK
DosCond condenses graph datasets by combining differentiable probabilistic modeling of discrete graph structures with one-step gradient matching. This avoids the nested optimization of vanilla gradient matching while retaining a theoretical link between gradient alignment and lower loss on real graphs.
- Framework objective: Dataset condensation seeks a much smaller synthetic graph set whose trained GNN achieves comparable performance to training on the full dataset.The synthetic set contains fewer graphs than the real dataset and includes graph structures, node features, and labels.
- Gradient matching objective: Vanilla gradient matching aligns GNN gradients computed on synthetic and real graphs across a training trajectory, requiring outer synthetic-data updates and inner parameter optimization.The objective can also sample initial model parameters to generalize matching across a distribution of initializations.
- Learning discrete graph structure: DosCond represents graph structure probabilistically so discrete adjacency values can be optimized differentiably through Bernoulli variables and a binary-concrete reparameterization.The method learns adjacency parameters and synthetic node features; temperature annealing moves the relaxation toward discrete samples while preserving gradients.
- One-step gradient matching: One-step gradient matching matches gradients only at model initializations, directly updating synthetic graphs without following the full training trajectory.This removes the expensive nested optimization and avoids tuning some inner- and outer-loop hyperparameters required by multi-step matching.
- Theoretical justification: Theoretical analysis bounds the loss gap after synthetic training using the one-step matching loss and input-matrix norms, supporting lower real-graph loss when the matching loss is reduced.With mean pooling, the norm term can be neglected; with sum pooling, it is included as a regularizer.
- Efficiency: DosCond can require less training cost than coreset methods while yielding comparable performance to its bi-level counterpart.The paper reports that the method is substantially faster because it avoids the nested optimization.
3 EXPERIMENT
Experiments evaluate DosCond’s classification effectiveness and efficiency across graph datasets, comparing it with discrete-structure baselines and examining scalability, graph learning components, and runtime.
- Experimental setup: DosCond is evaluated on molecular, TU, superpixel, and e-commerce graph datasets against Random, Herding, K-Center, and DCG baselines.The protocol learns synthetic graphs, trains a GCN on them, and tests classification performance on the original data.
- Classification performance: 92% and 96.5% of original performance are approximated on ogbg-molbace and ogbg-molhiv using only 0.2% and 0.3% data, respectively.On ogbg-molbace, two synthetic graphs achieve ROC-AUC 0.657 versus 0.714 on the full training set.
- Classification performance: 89%/93%/98% of original DD performance is approximated with 0.2%/2.1%/10.6% data, while increasing synthetic graph count improves classification.More synthetic graphs provide additional learnable parameters and more diverse patterns for training GNNs.
- Classification performance: DosCond outperforms baselines especially with 1 sample per class on CIFAR10, although performance is less promising when the synthetic graph count is limited.The paper attributes the limitation to CIFAR10’s more complex topology and feature information, and reports improved performance with a larger synthetic set.
- Ablation and representation analysis: Learning both synthetic graph structure and node features is necessary for preserving information, because DCG’s feature-only learning underperforms DosCond in most cases.DCG uses randomly selected graph structures and learns only node features.
- Efficiency comparison: DosCond empirically consumes less training cost than Herding, while its runtime increases with synthetic graph count and graph size according to O(N′L(n^2d+nd^2)).Herding requires full-dataset pretraining, whereas DosCond matches gradients for K1 initializations without fully training on the real dataset; costs are lower for ogbg-molhiv than DD because n=26 versus n=285.
3.3 Further Investigation
Further investigations show that DosCond remains effective as synthetic-set size changes, is substantially faster than bi-level matching, preserves discriminative graph representations, and benefits from sparsity-aware regularization.
- Synthetic Graph Size: DosCond outperforms random selection across condensed dataset sizes, although their performance gap narrows as the synthetic set grows.Both methods improve with more samples; random selection eventually approaches whole-dataset training performance.
- Ablation Study: DosCond reaches comparable performance to its bi-level variant in approximately 5 minutes versus 75 minutes, requiring only 6.7% of the training cost.This ablation isolates the efficiency benefit of one-step gradient matching on CIFAR10 with 50 graphs per class.
- Sparsity Regularization: Increasing sparsity leaves classification performance unchanged while producing graphs that require less storage and lower GNN training cost.The sparsity coefficient controls regularization strength; removing the regularizer slightly worsens ogbg-molhiv performance at 1, 10, and 50 graphs per class.
- Representation Quality: DosCond graphs yield class-separated GCN representations comparable to those learned from the full training dataset, unlike random or DCG graphs.The result indicates that DosCond preserves information needed to recover original performance.
- Scale of the Two Terms: The first theoretical term dominates early training with mean pooling, whereas the second term is non-negligible with sum pooling and supports retaining it as a regularizer.The empirical comparison uses DD with mean pooling and ogbg-molhiv with sum pooling.
- Memory Efficiency: DosCond outputs discrete graph structures that can be stored as binary values, making them more memory-efficient than GCond's weighted graphs.GCond requires storing edge weights in floating-point format.
4 RELATED WORK
Graph neural networks exploit graph structure for representation learning and support tasks including graph classification, link prediction, and node classification.
- Graph Neural Networks: GNNs have advanced graph representation learning by effectively exploiting structural information in graph data.The paper identifies graph classification, link prediction, and node classification as basic graph-related tasks where GNNs perform well.
5 CONCLUSION
The paper concludes that DosCond efficiently condenses real graphs into a small number of informative discrete graphs using one-step gradient matching, achieving strong performance with substantial dataset reduction.
- Conclusion: DosCond uses one-step gradient matching to condense real graphs into a small number of informative graphs with discrete structures.The framework is supported by both theoretical and empirical analysis.
- Conclusion: 90% dataset reduction preserves up to 98% of the original performance.The conclusion reports this as a headline experimental result across graph datasets.
A.1 Algorithm
Algorithm 1 initializes graph parameters and repeatedly samples networks and data, computes gradient-matching objectives, and updates graph structure and features before returning the condensed representation.
- Inputs and Initialization: Algorithm 1 takes training data, condensed labels, a GNN, temperature, sparsity controls, learning rates, and epoch counts as inputs or requirements.The algorithm initializes Ω and X′ before optimization.
- Sampling and Construction: Each iteration samples initial network parameters, a uniform random variable, and real and synthetic class-specific graph batches.The sampled random variable is used in the differentiable construction of A′.
- Optimization and Output: The algorithm updates Ω and X′ by descending the gradient of the gradient-matching objective plus sparsity regularization, then returns Ω, X′, and Y′.The returned parameters define the learned condensed graphs and their labels.
A.2 Dataset Statistics and Code
This section provides dataset-statistics tables for graph- and node-classification datasets and notes that code is supplied for reproducibility.
- Code is provided in the supplementary file for reproducibility.
- Table 4 reports statistics for graph-classification datasets.
- Table 5 reports statistics for node-classification datasets.
B.1 Proof of Theorem 1
The proof of Theorem 1 analyzes graph classification with a linearized K-layer SGC under bounded parameters, establishing convexity for the real loss and Lipschitz continuity for the synthetic loss.
- Theorem 1 assumes a linearized K-layer SGC with parameters bounded by ∥θ∥2 ≤ M2.
- The proof establishes that the real-data cross-entropy loss ℓT(θ) is convex.
- The proof establishes Lipschitz continuity for the synthetic-data loss ℓS(θ).
- The proof combines convexity, Lipschitz continuity, and gradient-descent updates to bound optimization behavior over training steps.
- For mean pooling, the analysis introduces a 1/ni factor relative to sum pooling, where ni is the number of nodes in the graph.
B.2 Theorem for Node Classification Case
The node-classification theorem considers a K-layer SGC on one real and one synthetic graph, assuming bounded parameters and using convexity and Lipschitz continuity of the associated losses.
- Node classification is represented with adjacency matrices A and A′ and feature matrices X and X′ for real and synthetic graphs.
- Theorem 2 assumes a K-layer SGC with fθ(A, X, θ) = A^KXW and bounded parameters ∥θ∥2 ≤ M2.
- The proof establishes convexity of the real-data loss ℓT(θ) through the positive semidefiniteness of its Hessian.
- The synthetic-data loss ℓS(θ) has a Lipschitz constant involving (C−1)/(CN′)∥A′^KX′∥.
- The proof combines the convexity inequality with the Lipschitz bound to derive an optimization bound for node classification.