Efficient Edge-aware Attention Network for Graph Generation
作者: Vincent Jung, Alba Carballo-Castro, Yiming Qin, Lonneke van der Plas, Pascal Frossard
分类: cs.LG
发布日期: 2026-10-20
备注: Accepted at the Workshop on Graph Foundation Models and the Workshop on Structured Probabilistic Inference & Generative Modeling, ICML 2026, Seoul, South Korea
💡 一句话要点
提出边缘感知注意力网络以解决图生成中的可扩展性问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 图生成 边缘感知 注意力机制 Transformer 深度学习
📋 核心要点
- 现有的图生成模型在可扩展性上存在不足，无法充分利用硬件感知的注意力实现。
- 论文提出了一种边缘感知注意力层，通过头级分数调制和轻量级残差机制来更新边缘表示，增强了模型的结构感知能力。
- 实验结果表明，该方法在生成质量上与基线相当，同时将峰值GPU内存减少约65%，显著提升了训练和采样的效率。
📝 摘要（中文）
图生成模型通常使用特定于图的Transformer架构来联合更新节点和边特征。然而，这些架构未能直接利用最新的硬件感知注意力实现，限制了其可扩展性。适应这一点并非易事：边缘和图级特征必须注入成对的注意力分数中，图掩码必须得到尊重，动态边缘更新仍然是导致二次内存使用的主要来源。我们提出了一种边缘感知注意力层，用于图生成，通过头级分数调制引入结构信息，同时使用轻量级残差机制更新边缘表示，保留依赖于注意力分数的边缘更新。该层可以有效地使用PyTorch的FlexAttention实现。我们在生成质量上与基线相似，同时将峰值GPU内存减少约65%，显著提高训练和采样吞吐量。这些结果表明，所提出的架构是朝着可扩展图基础模型迈出的实际一步。
🔬 方法详解
问题定义：本论文旨在解决图生成模型在可扩展性方面的挑战，现有方法未能有效利用硬件感知的注意力机制，导致内存使用高且效率低下。
核心思路：提出的边缘感知注意力层通过引入结构信息和轻量级的边缘更新机制，优化了注意力分数的计算方式，从而提高了模型的可扩展性和效率。
技术框架：整体架构包括边缘感知注意力层、头级分数调制和残差机制，主要模块协同工作以实现高效的图生成。
关键创新：最重要的创新在于将边缘和图级特征有效地注入到注意力分数中，同时通过轻量级残差机制实现动态边缘更新，显著降低了内存使用。
关键设计：在设计中，采用了PyTorch的FlexAttention实现，设置了适当的参数以确保边缘表示的更新依赖于注意力分数，同时优化了损失函数以提高生成质量。
🖼️ 关键图片
📊 实验亮点
实验结果显示，所提出的边缘感知注意力网络在生成质量上与基线模型相当，同时将峰值GPU内存减少约65%，训练和采样吞吐量显著提高，这表明该方法在实际应用中的有效性和优势。
🎯 应用场景
该研究的潜在应用领域包括社交网络分析、推荐系统和生物信息学等，能够为图生成模型的实际应用提供更高效的解决方案。未来，随着模型的可扩展性提升，可能会推动更大规模图数据的处理和分析，带来更广泛的影响。
📄 摘要（原文）
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.