Source-linked AI summary
Differentiable Graph Module (DGM) for Graph Convolutional Networks
Anees Kazi, Luca Cosmo, Seyed-Ahmad Ahmadi, Nassir Navab, Michael Bronstein
TL;DR
Many GNNs require a known, fixed graph, although practical data may have noisy or unknown structure and inductive settings introduce unseen nodes. DGM jointly learns probabilistic or discrete latent graphs and downstream representations, achieving strong results across healthcare, imaging, graphics, and vision tasks. Its main computational limitation is quadratic complexity in the number of input nodes.
Problem
Many GNNs assume a known, fixed graph, while relevant applications may involve unknown structure and require learning graphs for representation and convolution.
Method
DGM predicts a probabilistic graph, can sample a discrete graph, and integrates graph learning with graph convolution and downstream optimization.
Results
DGM supports applications spanning healthcare, brain imaging, computer graphics, and computer vision, including multimodal and inductive settings.
Takeaways & Limitations
The generic DGM pipeline can be adapted to graph-convolution methods across diverse downstream tasks and data domains.
Takeaways & Limitations
The method still has quadratic complexity in the number of input nodes because it computes all pairwise distances.
Abstract
from arXiv · showhide
Graph deep learning has recently emerged as a powerful ML concept allowing to generalize successful deep neural architectures to non-Euclidean structured data. Such methods have shown promising results on a broad spectrum of applications ranging from social science, biomedicine, and particle physics to computer vision, graphics, and chemistry. One of the limitations of the majority of current graph neural network architectures is that they are often restricted to the transductive setting and rely on the assumption that the underlying graph is {\em known} and {\em fixed}. Often, this assumption is not true since the graph may be noisy, or partially and even completely unknown. In such cases, it would be helpful to infer the graph directly from the data, especially in inductive settings where some nodes were not present in the graph at training time. Furthermore, learning a graph may become an end in itself, as the inferred structure may provide complementary insights next to the downstream task. In this paper, we introduce Differentiable Graph Module (DGM), a learnable function that predicts edge probabilities in the graph which are optimal for the downstream task. DGM can be combined with convolutional graph neural network layers and trained in an end-to-end fashion. We provide an extensive evaluation of applications from the domains of healthcare (disease prediction), brain imaging (age prediction), computer graphics (3D point cloud segmentation), and computer vision (zero-shot learning). We show that our model provides a significant improvement over baselines both in transductive and inductive settings and achieves state-of-the-art results.
1 INTRODUCTION
Graph neural networks commonly assume a known, fixed graph, but many applications involve noisy, partial, or unknown structure. DGM learns latent graphs end-to-end, with continuous and discrete sampling strategies, and is evaluated across multiple domains.
- Most GNN architectures assume the underlying graph is given and fixed, although real data may have noisy, partial, or completely unknown structure.
- DGM learns an optimal latent graph jointly with the downstream task through an end-to-end pipeline.
- DGM provides continuous and discrete sampling strategies, with the discrete version designed to learn larger sparse graphs more efficiently.
- The evaluation covers healthcare, brain imaging, computer graphics, and computer vision applications, including disease prediction, age prediction, point-cloud segmentation, and zero-shot learning.
- The model shows significant improvement over baselines and achieves state-of-the-art results across the reported applications.
2 STATE OF THE ART
Prior graph-learning methods include attention-based, dynamic, spectral, and spatial approaches, but often retain fixed-graph, transductive, dense, or specialized limitations. DGM is positioned as an end-to-end inductive method for latent graph discovery with topology control.
- Graph-learning research includes attention-based, dynamic, spectral, and spatial methods for representing or constructing latent graph structure.
- Dynamic graph methods construct graphs during training, but the graph-building latent space is not optimized for recovering the underlying latent graph.
- Spectral approaches require an initial graph and remain transductive because validation or test nodes must be embedded into the training graph.
- DGM learns latent graphs end-to-end, supports inductive inference, scales to larger node sets, and allows designers to control graph topology.
3 BACKGROUND
Graph neural networks use graph structure for message passing, but standard methods assume that structure is known and fixed. DGM targets settings where the graph is unknown by learning it for downstream representation and prediction.
- Graph neural networks: GNNs produce node embeddings by applying shared message-passing functions over local graph neighborhoods.Edge convolution supports general learnable neighbor interactions, while graph convolution is a node-wise linear special case.
- Latent graphs: Learning an unknown graph both represents data structure and supplies the support for graph-based convolutions.
- Latent graphs: Most GNN architectures keep the given graph topology fixed, limiting applications where the underlying structure is unavailable or not explicit.Dynamic graph methods construct graphs during training but do not necessarily learn the latent graph with a dedicated loss.
4 METHOD
DGM learns graph structure from layer features and combines graph construction with diffusion in an end-to-end architecture. Continuous sampling yields differentiable weighted graphs, while discrete sampling produces sparse graphs with lower diffusion complexity.
- Architecture: DGM and the Diffusion Module form a multilayer architecture that jointly learns graph structure from layer outputs and network parameters.Each layer contains both modules, and the resulting features can feed a final MLP for node prediction.
- Differentiable Graph Module: DGM transforms input features into lower-dimensional auxiliary features, then uses distances between them to assign edge probabilities.The transformation may use an initial graph when available, but an initial graph is optional.
- Continuous sampling: Continuous DGM converts edge probabilities into a differentiable weighted adjacency matrix parameterized by network parameters, temperature t, and threshold T.This strategy generalizes the earlier latent-graph method.
- Discrete sampling: Discrete DGM samples k edges per node with the Gumbel-Top-k trick, producing a sparse unweighted adjacency matrix.Sparsity lowers the computational and memory complexity of diffusion, although stochastic sampling makes inference nondeterministic.
- Loss and training: For discrete sampling, a compound loss rewards edges associated with correct classifications and penalizes edges associated with errors.The graph-loss gradient approximates the gradient of expected downstream performance over sampled graphs.
5 EXPERIMENTS AND RESULTS
The experiments include ablations and comparisons across healthcare, brain imaging, 3D graphics, and computer vision applications.
- Experiments: Experiments cover benchmark ablations, healthcare prediction, brain-imaging age prediction, 3D point-cloud segmentation, and zero-shot learning.
5.1 Benchmark graph based datasets
On citation-graph benchmarks, DGM learns sparse task-adapted graphs and improves graph convolution beyond using the input graph alone. Ablations show that performance depends on sparsity, embedding geometry and dimension, layer count, and access to input connectivity.
- 5.1.2 Graph embedding space geometry: Euclidean embeddings outperform hyperbolic embeddings consistently across all three datasets and show lower standard deviation.The results indicate no advantage from representing node relations hierarchically for these datasets.
- 5.1.3 Graph embedding space dimension: The best embedding dimension is 4 for Cora and CiteSeer but 8 for PubMed, so the suitable latent-space size depends on dataset complexity.Higher dimensions represent more complex relations, whereas smaller spaces provide a regularization effect against overfitting.
- 5.1.4 Number of DGM Layers: A single DGM layer improves classification over conventional input-graph convolution on all datasets, while only PubMed benefits from a second layer.The result suggests that PubMed contains more complex node relations than the other benchmark datasets.
- 5.1.5 Graph embedding function f: Removing input connectivity causes a 10–20% accuracy drop, especially on Cora and CiteSeer, showing that their graph structure contains important class information.Using the input graph in the DGM embedding function is essential for reasonable performance on these datasets.
- 5.1.6 Diffusion function g: Up to 3.4% better accuracy with DGM than without it, while using the same diffusion and classification parameters and at most 150 extra graph-prediction parameters.The comparison holds across diffusion functions in most cases, showing that DGM can learn a graph more suitable for the task even when an input graph is available.
5.2 Application to Healthcare and Brain imaging
The healthcare and brain-imaging experiments evaluate DGM on disease and age prediction with multimodal patient data, including inductive settings where graphs are not fixed in advance. The model learns graphs directly from patient features and significantly outperforms state-of-the-art methods while remaining computationally competitive.
- Experimental setting: Healthcare experiments use multimodal patient data for disease classification and brain-age prediction, motivated partly by brain age’s relevance as a neurological-disorder biomarker.The Tadpole dataset contains 564 patients with imaging, demographic, and genotype features; UK Biobank contains 14,503 individuals.
- Experimental setting: The experiments distinguish transductive learning, where test nodes are present without labels, from inductive learning, where test nodes are absent during training and introduced only at testing.The medical architecture uses DGM blocks with diffusion layers and an MLP graph-embedding function because no input graph is provided.
- Method: Unlike prior GNN approaches using handcrafted patient-population graphs, DGM learns the graph directly from input patient features.Compared methods include linear, handcrafted-graph GNNs, DGCNN, and LDS.
- Results: The model significantly outperforms state-of-the-art methods on all healthcare and brain-imaging tasks, while its runtime is on par with DGCNN and about three orders of magnitude faster than LDS.Results use 10-fold cross validation; LDS does not support inductive learning.
5.3 Application to computer vision and Computer graphics
DGM is applied to 3D point-cloud segmentation and zero-shot learning by replacing fixed or handcrafted graph construction with learned graph structure. It improves ShapeNet segmentation broadly, while zero-shot results are competitive but remain below a knowledge-graph-based method.
- 3D point-cloud segmentation: For point-cloud segmentation, DGM replaces DGCNN’s feature-space kNN sampling while retaining the remainder of the network architecture, including k = 20.The DGM uses feature depth 16.
- 3D point-cloud segmentation: dDGM improves mIoU over the original kNN sampling scheme on almost all ShapeNet part classes, with significantly better overall performance.The evaluation averages IoUs across all testing shapes; the dataset contains 16,881 point clouds sampled at 2,048 points each.
- Zero-shot learning: In zero-shot learning, the model replaces the input knowledge graph in SGCN-style graph convolution with a graph learned by DGM from semantic class features.The model uses two graph-convolution layers paired with two DGM layers.
- Zero-shot learning: 74.7% Top-1 accuracy makes dDGM outperform GCNZ at 70.5% but remain below DGP at 77.3% on AWA2 zero-shot learning.Unlike GCNZ and DGP, dDGM does not use the knowledge graph.
- Zero-shot learning: The predicted zero-shot graph resembles the knowledge graph locally but fails to capture its hierarchical structure.The authors suggest imposing additional constraints, such as a more tree-like structure, as future work.
6 DISCUSSION AND CONCLUSION
The paper presents DGM as a generic, end-to-end module that predicts probabilistic graphs, samples discrete graphs, and supports graph-convolutional downstream tasks. Its broad evaluation covers multimodal and inductive applications, while computational complexity and graph heterogeneity remain open issues.
- Contribution: DGM predicts a probabilistic graph, samples a discrete graph from it, and supplies that graph to graph-convolutional operators through a weighted edge-probability loss.The module is designed for end-to-end optimization with the downstream task.
- Scope: The method is demonstrated across healthcare, brain imaging, computer graphics, and computer vision, including multimodal datasets and inductive inference setups.Applications include disease prediction, age and gender prediction, 3D point-cloud segmentation, and zero-shot learning.
- Limitations and future work: Computing all pairwise node distances gives the method quadratic complexity, although neighborhood restriction and tree-based algorithms could reduce it to O(n log n).The authors also identify fixed-k sampling as insensitive to heterogeneous node degree distributions.