Table of Contents
Fetching ...

Uncovering Capabilities of Model Pruning in Graph Contrastive Learning

Junran Wu, Xueyuan Chen, Shangzhe Li

TL;DR

This work tackles the reliance of graph contrastive learning on data augmentations that may distort semantics by reframing pre-training around model pruning instead of view generation. The proposed framework, LAMP, uses a dense original graph encoder and a pruned perturbation encoder to produce contrasting embeddings, complemented by a local node-level loss to address hard negatives. The authors provide theoretical results showing pruning can preserve or even improve mutual information with downstream labels compared to augmentation-based views, and they demonstrate strong empirical gains in unsupervised and transfer learning across diverse benchmarks. Overall, LAMP offers a general, domain-agnostic approach to graph representation learning that reduces reliance on potentially destructive augmentations while delivering state-of-the-art performance.

Abstract

Graph contrastive learning has achieved great success in pre-training graph neural networks without ground-truth labels. Leading graph contrastive learning follows the classical scheme of contrastive learning, forcing model to identify the essential information from augmented views. However, general augmented views are produced via random corruption or learning, which inevitably leads to semantics alteration. Although domain knowledge guided augmentations alleviate this issue, the generated views are domain specific and undermine the generalization. In this work, motivated by the firm representation ability of sparse model from pruning, we reformulate the problem of graph contrastive learning via contrasting different model versions rather than augmented views. We first theoretically reveal the superiority of model pruning in contrast to data augmentations. In practice, we take original graph as input and dynamically generate a perturbed graph encoder to contrast with the original encoder by pruning its transformation weights. Furthermore, considering the integrity of node embedding in our method, we are capable of developing a local contrastive loss to tackle the hard negative samples that disturb the model training. We extensively validate our method on various benchmarks regarding graph classification via unsupervised and transfer learning. Compared to the state-of-the-art (SOTA) works, better performance can always be obtained by the proposed method.

Uncovering Capabilities of Model Pruning in Graph Contrastive Learning

TL;DR

This work tackles the reliance of graph contrastive learning on data augmentations that may distort semantics by reframing pre-training around model pruning instead of view generation. The proposed framework, LAMP, uses a dense original graph encoder and a pruned perturbation encoder to produce contrasting embeddings, complemented by a local node-level loss to address hard negatives. The authors provide theoretical results showing pruning can preserve or even improve mutual information with downstream labels compared to augmentation-based views, and they demonstrate strong empirical gains in unsupervised and transfer learning across diverse benchmarks. Overall, LAMP offers a general, domain-agnostic approach to graph representation learning that reduces reliance on potentially destructive augmentations while delivering state-of-the-art performance.

Abstract

Graph contrastive learning has achieved great success in pre-training graph neural networks without ground-truth labels. Leading graph contrastive learning follows the classical scheme of contrastive learning, forcing model to identify the essential information from augmented views. However, general augmented views are produced via random corruption or learning, which inevitably leads to semantics alteration. Although domain knowledge guided augmentations alleviate this issue, the generated views are domain specific and undermine the generalization. In this work, motivated by the firm representation ability of sparse model from pruning, we reformulate the problem of graph contrastive learning via contrasting different model versions rather than augmented views. We first theoretically reveal the superiority of model pruning in contrast to data augmentations. In practice, we take original graph as input and dynamically generate a perturbed graph encoder to contrast with the original encoder by pruning its transformation weights. Furthermore, considering the integrity of node embedding in our method, we are capable of developing a local contrastive loss to tackle the hard negative samples that disturb the model training. We extensively validate our method on various benchmarks regarding graph classification via unsupervised and transfer learning. Compared to the state-of-the-art (SOTA) works, better performance can always be obtained by the proposed method.

Paper Structure

This paper contains 31 sections, 1 theorem, 10 equations, 10 figures, 5 tables, 1 algorithm.

Key Result

theorem 1

Suppose the graph encoder $f$ is implemented by a GNN with at least 2 layers and $f^*$ is the optimal version. Given a general data augmentation function $t$, the optimal pruned encoder $f_p^*$ satisfies,

Figures (10)

  • Figure 1: Framework of LAMP. One branch takes the original graph as input instead of the augmented view. The other branch is pruned from the former online and also embeds the original graph. Besides the ubiquitous NT-Xent loss, the graph encoder is jointly optimized with a local contrastive loss to optimize the hypersphere of contrastive learning.
  • Figure 2: Quantification of structural damage from data augmentation. Percent change in structural entropy of MUTAG and REDDIT-BINARY after data augmentation (i.e., Edge perturbation, Node dropping, and Subgraph with 20% strength from GraphCL).
  • Figure 3: Illustration of hard negative samples. Via contrasting the graph embeddings, the pre-trained model is hard to distinguish this two kinds of graphs.
  • Figure 4: Sensitivity w.r.t. pruning ratio $\gamma$.
  • Figure 5: Sensitivity w.r.t. loss balance $\alpha$.
  • ...and 5 more figures

Theorems & Definitions (1)

  • theorem 1