论文学术研究图生成FlexAttentionGPU 内存Transformer

Efficient Edge-aware Attention Network for Graph Generation

AI 解读已生成 · AI 生成，数字以原文为准

AI 全文总结值得深入基于论文全文 2.2 万字 · 完整阅读

FlexAttentionGT 用融合 kernel 重写边感知图 Transformer，峰值 GPU 内存降约 65%，质量持平。

针对图生成中密集边感知 Transformer 的二次内存瓶颈，作者提出 FlexAttentionGT，在不牺牲归纳偏置的前提下将 DeFoG 原始骨干改造为硬件友好的 fused kernel 实现。做法是：用标准缩放点积注意力替换逐通道外积，以 head-wise FiLM 在 score_mod 内注入边与全局特征，并将单位矩阵拼接到 value 上，让单次 FlexAttention 同时输出节点聚合和注意力图，再用残差更新边表示。在 Planar、Tree、SBM 上与原始 GT 对照，生成质量（V.U.N.、Ratio）匹配或略优；作者报告训练吞吐平均提升 44.8%、推理提升 94.4%、峰值 GPU 显存降低约 65%。消融显示边更新和注意力分数调制不可省略，使用归一化注意力系数比未归一化收敛更稳。边更新仍为二次复杂度，大图需两次 pass，但为图基础模型扩展提供了实用路径。

1FlexAttentionGT 用 fused kernel 重写边感知注意力，保留原版归纳偏置。
2峰值 GPU 内存降到约 三分之一，边更新仍为二次方。
3去掉 per-edge 调制的 Light Edge 完全无法学习。

为什么重要　对关注推理/训练基础设施与图模型的工程同学，给出用 FlexAttention 保留归纳偏置、同时降低内存的具体实现路径。适合　做图生成、Transformer 融合 kernel 优化或推理训练基础设施的工程与研究人员。

来源: arxiv.org
发布: 2026/10/21
作者: Vincent Jung、Alba Carballo-Castro、Yiming Qin、Lonneke van der Plas 等

中文（AI 翻译）

图生成模型通常使用图专用的 Transformer 架构来联合更新节点和边特征。然而，这些架构并不能直接受益于近期的硬件感知注意力实现，限制了其可扩展性。

这种适配并非易事：边和图级特征必须注入到成对注意力分数中，图掩码必须被遵守，而动态边更新仍是二次方内存占用的主要来源。我们提出了一种用于图生成的边感知注意力层，通过逐头分数调制引入结构信息，同时用轻量残差机制更新边表示，该机制保留依赖注意力分数的边更新。

该层可以有效地使用 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.

降低约 65%

峰值内存

+44.8%

训练加速

+94.4%

推理加速

注释 · 速度范围在 23%–116%，来自图1与第4.2节。出自 §4.2/结论

框架: DeFoG
实现: PyTorch FlexAttention
数据集: Planar/Tree/SBM
模型: FlexAttentionGT
评测指标: V.U.N.↑, Ratio↓

深度总结

1 / 8

65%

融合内核边感知注意力· 判断

FlexAttentionGT 在质量持平前提下把峰值内存削减约 65%。

判断

原文称该层用实现，节点与边更新共享单次 fused kernel 调用，训练/推理吞吐分别平均提高 44.8% 和 94.4%，同时指标与原 GT 持平或略优。

注释 · 结论来自第5节与图1，生成质量对比见表1。出自 §5 / 图1

FlexAttentionGT 用融合 kernel 重写边感知图 Transformer，峰值 GPU 内存降约 65%，质量持平。

1FlexAttentionGT 用 fused kernel 重写边感知注意力，保留原版归纳偏置。

2峰值 GPU 内存降到约 三分之一，边更新仍为二次方。

3去掉 per-edge 调制的 Light Edge 完全无法学习。

来源

arxiv.org

发布

2026/10/21

作者

Vincent Jung、Alba Carballo-Castro、Yiming Qin、Lonneke van der Plas 等

中文（AI 翻译）

原文

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.

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.

Efficient Edge-aware Attention Network for Graph Generation

FlexAttentionGT 用融合 kernel 重写边感知图 Transformer，峰值 GPU 内存降约 65%，质量持平。

摘要

中文（AI 翻译）

原文

效率提升幅度

FlexAttentionGT

关键规格

深度总结

Efficient Edge-aware Attention Network for Graph Generation

FlexAttentionGT 用融合 kernel 重写边感知图 Transformer，峰值 GPU 内存降约 65%，质量持平。

摘要

中文（AI 翻译）

原文

效率提升幅度

FlexAttentionGT

关键规格

深度总结