EGAT-Pool: A Hierarchical Graph Attention Network with Edge-aware SAG Pooling for Robust Environmental Toxicity Prediction

Abstract

Predicting the environmental toxicity of novel chemical compounds is a critical challenge for public health and ecological safety, hindered by the structural complexity of molecules and the severe class imbalance inherent in toxicological datasets. Conventional graph neural networks often struggle to capture the hierarchical nature of toxicophores and exhibit difficulties with out-of-distribution generalization. To address these limitations, a novel hierarchical GNN architecture, EGAT-Pool, is proposed for robust toxicity prediction. EGAT-Pool synergistically combines the dynamic attention mechanism of GATv2 with an edge-aware Self-Attention Graph Pooling (SAG Pooling) strategy, enabling the model to learn multi-scale molecular representations while explicitly incorporating bond semantics to identify critical substructures. Evaluated on the highly imbalanced Tox21 benchmark under a rigorous Murcko scaffold split protocol, EGAT-Pool demonstrates competitive predictive performance with a ROC-AUC of 0.8153 and a Matthews Correlation Coefficient (MCC) of 0.4068, outperforming established GNN baselines (GCN, GAT, and GIN) across all evaluated metrics. Furthermore, it is demonstrated that EGAT-Pool provides mechanistic interpretability through spatial attention heatmaps, which autonomously identify known toxicophores without prior supervision. This “white-box” capability bridges the gap between predictive accuracy and regulatory applicability, establishing EGAT-Pool as a promising tool for high-throughput screening, environmental risk assessment, and the rational design of safer chemicals.