Efficient graph generation method cuts memory use and speeds training
Efficient Edge-aware Attention Network for Graph Generation
Machine Learning
Summary
Generating graphs on computers often involves using special attention methods that look at how nodes and edges relate. These methods can require a lot of memory and run slowly because of how they handle edges and masks. The authors propose a new way to do attention on graphs that uses less memory by cleverly adjusting how it scores connections and updates edges. Their approach keeps the quality of generated graphs while making the process much faster and more efficient on common hardware.
What this means in practice
- •For machine learning engineers: Train and sample graph generative models faster and with less GPU memory using the efficient edge-aware attention layer.
- •For graph software developers: Build scalable graph generation tools that handle large graphs efficiently by integrating the proposed edge-aware attention mechanism.
Authors
Vincent Jung, Alba Carballo-Castro, Yiming Qin, Lonneke van der Plas, Pascal Frossard
Abstract
Graph generative models commonly use graph-specific Transformer architectures to jointly update node and edge features. However, these architectures do not directly benefit from recent hardware-aware attention implementations, limiting their scalability. The adaptation is, however, non-trivial: edge and graph-level features must be injected into pairwise attention scores, graph masks must be respected, and dynamic edge updates remain a major source of quadratic memory usage. We propose an edge-aware attention layer for graph generation that incorporates structural information through head-wise score modulation, while updating edge representations with a lightweight residual mechanism that retains attention-score dependent edge updates. This layer can be effectively implemented using PyTorch's FlexAttention. We achieve similar generation quality to our baseline while reducing peak GPU memory by around 65% and substantially improving training and sampling throughput. These results suggest that the proposed architecture is a practical step toward scalable graph foundation models.