边感知注意力层让图生成模型峰值显存降低约65%
Efficient Edge-aware Attention Network for Graph Generation
做图生成的人可以看看：用 FlexAttention 实现的边感知注意力层，显存省了 65%，吞吐还更快。
论文提出一种用于图生成的边感知注意力层，通过 head-wise score modulation 注入结构信息，并用轻量残差机制更新边表示。该层可基于 PyTorch 的 FlexAttention 实现，绕开了图专用 Transformer 无法利用硬件级注意力优化的问题。实验显示生成质量与基线相当，峰值 GPU 显存降低约 65%，训练和采样吞吐明显提升。作者认为这是构建可扩展图基础模型的实用一步。
Efficient Edge-aware Attention Network for Graph Generation
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.
