Efficient Edge-aware Attention Network for Graph Generation

TL;DR

The paper introduces a new attention mechanism for graph generation that integrates structural edge information directly into attention scores through head-wise modulation. It uses a lightweight residual update for edge representations, enabling dynamic edge updates without excessive memory overhead. The design is compatible with PyTorch's FlexAttention, leading to significant memory reduction and faster training and inference.

Why It Matters

Graph generative models often face scalability issues due to high memory usage during attention computation. This work enables more efficient training and deployment of graph models on standard hardware, paving the way for larger, more practical graph foundation models. The approach maintains generation quality while drastically cutting GPU memory demands.

Key Findings

• The proposed attention layer modulates attention scores using edge-specific information in a head-wise manner.
• Edge updates are performed via a lightweight residual mechanism that preserves dependency on attention scores.
• Implementation using FlexAttention reduces peak GPU memory by approximately 65% and improves training and sampling throughput.