[Paper Review] Graph Mixture Density Networks

Federico Errica, Davide Bacciu+1 more|arXiv (Cornell University)|Dec 5, 2020

Bioinformatics and Genomic Networks40 references4 citations

TL;DR

This paper introduces Graph Mixture Density Networks (GMDN), a novel deep learning framework that combines graph neural networks with mixture density networks to model multimodal conditional output distributions on arbitrary graph-structured data. GMDN achieves state-of-the-art performance in likelihood-based density estimation on both synthetic epidemic simulation benchmarks and real-world molecular property prediction tasks, significantly outperforming baselines in capturing output uncertainty.

ABSTRACT

We introduce the Graph Mixture Density Networks, a new family of machine learning models that can fit multimodal output distributions conditioned on graphs of arbitrary topology. By combining ideas from mixture models and graph representation learning, we address a broader class of challenging conditional density estimation problems that rely on structured data. In this respect, we evaluate our method on a new benchmark application that leverages random graphs for stochastic epidemic simulations. We show a significant improvement in the likelihood of epidemic outcomes when taking into account both multimodality and structure. The empirical analysis is complemented by two real-world regression tasks showing the effectiveness of our approach in modeling the output prediction uncertainty. Graph Mixture Density Networks open appealing research opportunities in the study of structure-dependent phenomena that exhibit non-trivial conditional output distributions.

Motivation & Objective

• To address the challenge of modeling multimodal conditional output distributions in graph-structured data, where standard regression models fail due to unimodal predictions.
• To extend mixture density networks (MDNs) to operate on graphs, enabling conditional density estimation for complex, structure-dependent phenomena.
• To evaluate the model on a novel benchmark involving stochastic epidemic simulations on random graphs, where both structure and multimodality are critical.
• To demonstrate the model's effectiveness in real-world regression tasks involving molecular graphs, particularly in modeling prediction uncertainty.
• To show that GMDN can transfer knowledge across different graph topologies and sizes, improving generalization on unseen graph structures.

Proposed method

• GMDN uses a graph neural network (GNN) encoder to produce a shared latent representation from an input graph, which is then used to parameterize a mixture of output distributions.
• The GNN output is fed into a mixture density head that predicts the mixing weights, means, and variances of a mixture of Gaussian components for the target variable.
• The model is trained via maximum likelihood estimation, optimizing the log-likelihood of the true target values under the predicted mixture distribution.
• The mixture density head allows the model to output a full probability distribution, capturing multiple plausible outcomes for a given input graph.
• The architecture supports both node-level and graph-level output distributions, enabling flexible modeling of structured data with complex conditional dependencies.
• Transfer learning is evaluated by training on one graph family (e.g., Erdős–Rényi) and testing on another (e.g., Barabási–Albert), assessing generalization across topologies.

Experimental results

Research questions

• RQ1 Can a deep learning model effectively model multimodal output distributions conditioned on arbitrary graph structures?
• RQ2 How does incorporating graph structure improve the modeling of complex, non-Gaussian conditional distributions in real-world applications?
• RQ3 Does the GMDN framework generalize across different graph topologies, such as from Erdős–Rényi to Barabási–Albert networks?
• RQ4 Can GMDN better capture prediction uncertainty in molecular property regression compared to standard GNNs or MDNs?
• RQ5 How does the model’s likelihood performance compare to strong baselines like DGN and MDN on both synthetic and real-world graph benchmarks?

Key findings

• On the epidemic simulation benchmark, GMDN achieved a significantly higher log-likelihood than all baselines, with a RAND score baseline of -27.12 and GMDN reaching -0.57 on the alchemy_full dataset.
• On the ZINC_full dataset, GMDN achieved a log-likelihood of -0.75, outperforming the next-best model (DGN) with -0.90, demonstrating superior density estimation.
• On the alchemy_full dataset, GMDN achieved a log-likelihood of -0.57 with a standard deviation of 1.4, indicating robust and consistent performance across training runs.
• In visualizations of output distributions, GMDN clearly captured two distinct modes for dipole moment predictions, while DGN collapsed to a single Gaussian between the modes, failing to represent uncertainty.
• The model trained on ER-100 graphs generalized better to larger and different graph families (e.g., BA-200, BA-500) than models trained on BA-100, suggesting that harder source tasks improve transferability.
• The MAE scores were comparable across models, but GMDN achieved much better log-likelihood, confirming that it models uncertainty more accurately than baselines despite similar point-estimate performance.