LatentFlow: Visual Analytics for Latent Space Analysis in Molecular Graph Neural Networks

2026-07-23

AI Summary

Existing methods struggle to reveal the dynamic evolution of latent spaces in molecular graph neural networks across different layers and training stages, as well as their relationship to chemical concepts. To address this gap, this work proposes a visual analytics system that enables, for the first time, cross-layer and cross-training-state tracking of latent representations. By clustering molecular embeddings and visualizing their evolutionary trajectories through an enhanced Sankey diagram—linked interactively with representative molecules, key substructures, and domain knowledge—the system significantly enhances model interpretability. Integrating graph neural network embeddings, clustering analysis, substructure extraction, and interactive visualization techniques, the framework is validated through two case studies, demonstrating its effectiveness in helping domain scientists understand latent space dynamics, identify meaningful molecular patterns, and gain deeper insights into model decisions.

Abstract

Chemists and materials scientists increasingly use machine learning models, such as graph neural networks (GNNs), to predict properties of molecules and the outcomes of their reactions. Beyond predictive performance, understanding how these models organize chemical information internally in their latent spaces, i.e., the embeddings of the molecules, is critical. Analyzing latent spaces helps diagnose model behavior and assess whether the learned embeddings are organized in ways that reflect meaningful chemical relationships. Unfortunately, existing methods provide limited support for analyzing latent spaces across layers and across different model states (e.g., training epochs, model configurations, and input data), making it difficult to understand how these latent spaces evolve throughout a model or relate to chemical concepts. We present LatentFlow, a visual analytics system developed in collaboration with a domain expert for analyzing latent spaces in molecular GNNs. LatentFlow groups embeddings into clusters and supports exploration of latent spaces by tracking how these clusters change across layers and model states using a modified Sankey diagram. To support interpretation, LatentFlow links these clusters to representative molecules and their shared substructures, and it allows scientists to introduce their own domain knowledge and compare it with the patterns found in the latent spaces. We evaluate LatentFlow through two case studies. The results show that LatentFlow helps scientists understand how latent spaces evolve, identify meaningful molecular patterns, and better interpret model behavior.

Problem

• latent space analysis
• molecular graph neural networks
• visual analytics
• model interpretability
• chemical representation

Innovation

• latent space analysis
• visual analytics
• graph neural networks
• molecular representation
• cluster evolution