Robust Graph Representation Learning via Neural Sparsification

In real-world scenarios, graphs often contain abundant neighbor information, with many features and connections, many of which are irrelevant to downstream tasks, leading to suboptimal performance. To address this, the authors propose NeuralSparse, a supervised graph sparsification technique that removes redundant, task‑irrelevant edges to enhance model generalization.

1. Introduction

Connections between two nodes may be unrelated to the target downstream task (e.g., noisy or accidental links). Directly aggregating all original neighbors can therefore introduce task‑irrelevant information and degrade the performance of graph models. The authors illustrate this with synthetic examples where random neighbor sampling fails to align with label information, and they aim for a sparsification method that links same‑label nodes more tightly, making class boundaries clearer.

Previous unsupervised sparsification approaches cannot guarantee optimal downstream performance, while predefined sampling distributions lack adaptability to specific tasks. Existing supervised methods are often computationally intensive. NeuralSparse addresses these issues by learning task‑guided edge selection through a sparsification network jointly optimized with a downstream GNN.

2. Method

2.1 Overview

The overall framework consists of two networks: a Sparsification Network that selects a subgraph and a GNN that performs classification on the sparsified graph. The loss function jointly optimizes both components.

Given a graph with node features, the goal is to learn a sparse adjacency matrix that retains edges most relevant to the downstream prediction task. The authors model the probability of each edge being kept and approximate the distribution using an estimator, introducing two sets of parameters for the sparsification and GNN components.

2.2 Model Structure

The sparsification network is a multi‑layer MLP that outputs a subgraph \(g\). The GNN can be any standard graph neural network (e.g., GCN, GraphSAGE) that consumes the sparsified graph. To construct subgraphs, the method first samples a \(k\)-nearest‑neighbor subgraph, balancing the trade‑off between model capacity and irrelevant information.

Edge filtering is performed on the sampled subgraph. A softmax layer computes the probability of each candidate edge being retained, and a Gumbel‑Softmax trick is used to sample discrete edge selections during training, with a temperature parameter annealed over time.

The computational complexity of the sparsification algorithm is linear in the number of original edges and vertices, i.e., \(O(|E|+|V|)\).

3. Experiments

The experimental setup evaluates NeuralSparse on several backbone GNNs (GCN, GraphSAGE, GIN, GAT) and compares against baselines such as no sparsification, spectral sparsifier, Rank Degree (SS/RD), DropEdge, and LDS. Results show that NeuralSparse consistently improves performance across all backbones and remains stable under different hyper‑parameter settings.

Additional visualizations on transaction data demonstrate that NeuralSparse effectively reduces irrelevant edges, leading to clearer graph structures aligned with the target task.

Overall, NeuralSparse provides a task‑aware graph sparsification solution that enhances GNN performance while maintaining computational efficiency.