Efficient Edge-aware Attention Network for Graph Generation

Vincent Jung; Alba Carballo-Castro; Yiming Qin; Lonneke van der Plas; Pascal Frossard

cs.LGOct 20, 2026

Efficient Edge-aware Attention Network for Graph Generation

Authors: Vincent Jung, Alba Carballo-Castro, Yiming Qin, Lonneke van der Plas, Pascal Frossard

Organizations: LTS4, EPFL, Lausanne, Switzerland · Università della Svizzera italiana, Lugano, Switzerland

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.

Figures & tables

Panel of Figure 1 — Figure 1: Computational efficiency across datasets. Percentages denote the relative change compared to the original implementation.

Planar
Model	V.U.N. $\uparrow$	Ratio $\downarrow$
Train set	$100$	$1.00$
GT (Original)	$99.5\pm 1.0$	$2.47\pm 0.33$
FlexAttentionGT (Ours)	$100.0\pm 0.0$	$2.21\pm 0.95$
Tree
Train set	$100$	$1.00$

Table 1: Generation performance on Planar, Tree, and SBM.

Panel of Figure 2 — Figure 2: V.U.N. curves over training steps on the Planar and SBM datasets, comparing FlexAttentionGT against the original GT and its lightweight variants. We evaluate 1000 and 50 generation steps, respectively.

Figure 3: Training loss comparison between FlexAttentionGT (Ours) and GT (Original), smoothed with a window size of 20. Similar optimization dynamics suggest that gains stem from the attention layer’s inductive bias.

Panel of Figure 4 — Figure 4: V.U.N. curves comparing normalized and unnormalized networks over training steps on the Planar dataset, evaluated with 1000 generation steps.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

	Inference (samples/s) $\uparrow$		Training (samples/s) $\uparrow$		Memory (MB) $\downarrow$
Dataset	GT	FlexAttentionGT	GT	FlexAttentionGT	GT	FlexAttentionGT
Planar	$294.2$	$561.5$	$132.6$	$217.9$	$9{,}580$	$3{,}323$
Tree	$310.6$	$549.1$	$147.7$	$217.7$	$9{,}538$	$3{,}281$
SBM	$62$	$133.7$	$31.9$	$39.2$	$62{,}840$	$21{,}991$

Appendix

Table 2: Detailed computational efficiency results for the Planar, Tree, and SBM datasets. Inference and training speeds are measured in samples per second. Memory is reported during training, in MB. Tested on A100 GPU.

	Inference		Training
Model	Samples/s $\uparrow$	Mem (MB) $\downarrow$	Samples/s $\uparrow$	Mem (MB) $\downarrow$
GT (Original)	$6.6$	$10{,}412$	$3.4$	$43{,}429$
FlexAttentionGT (Single-pass)	$13.2$	$2{,}584$	$0.9$	$15{,}320$
FlexAttentionGT (Two-pass)	$13.7$	$2{,}595$	$6.1$	$15{,}316$

Appendix

Table 3: Efficiency benchmark on the SBM-Large dataset (batch size 2), with an average size of 500 nodes. Inference and training speeds are measured in samples per second. Memory is reported in MB. Tested on A100 GPU.

	GT (Original)	FlexAttentionGT (Ours)
Attention logit	$q_{i,h}\odot k_{j,h}\in\mathbb{R}^{d_{f}}$	$q_{i,h}\cdot k_{j,h}\in\mathbb{R}$
Softmax	Per-channel ( $d_{f}$ distributions per $(i,h)$ )	Standard (1 distribution per $(i,h)$ )
Edge FiLM bandwidth	$d_{x}=H\cdot d_{f}$ channels	$H$ scalars
Edge FiLM location	Applied to $Y\in\mathbb{R}^{N\times N\times H\times d_{f}}$	Inside fused kernel via score_mod
Edge update source	Full $d_{x}$ -dim bilinear $Y^{\mathrm{flat}}$	$H$ -dim attention map $\alpha$
Node & edge update	Two separate passes	Single fused kernel call

Appendix

Table 4: Comparison of the original and FlexAttention layer variants.

Figure 5: Schematic of the FlexAttentionGT architecture.

Dataset	Min Nodes	Max Nodes
Planar	64	64
Tree	64	64
SBM	44	187
SBM-Large	401	599

Appendix

Table 5: Sizes for different datasets we use in experiments.

Model	Initial Distribution	Train Distortion	Sample Distortion	$\omega$	$\eta$	GPU Type	Training steps
Planar
GT (Original)	Marginal	Polydec	Polydec	50	0.05	V100	34000
FlexAttentionGT (Ours)	Marginal	Polydec	Polydec	50	0.02	V100	34000
Tree
GT (Original)	Marginal	Polydec	Polydec	0	0	V100	26000
FlexAttentionGT (Ours)	Marginal	Polydec	Polydec	0	0	V100	16500

Appendix

Table 6: Training and sampling settings for Planar, Tree, and SBM.

Graph Generation

Graph generation builds models and algorithms that create graphs with desired structures, properties, or distributions.

See all 67 papers

GraphK: Variable-Size Graph Generation with Efficient Edge Construction
Aug 19, 2026Resul Tugay, Eren Oluğ, Elif Ak +1Graph Generation Generative Modeling
Department of Artificial Intelligence and Data Engineering Ataturk University, Turkey · Department of Artificial Intelligence and Data Engineering ITU AI Research and Application Center Istanbul Technical University, Turkey · Faculty of Engineering and Applied Science Memorial University, Canada
Scaling Novel Graph Generation via Lightweight Structure-Guided Autoregressive Models
Jun 2, 2026Alessio Barboni, Massimiliano Lupo Pasini, Bishal Lakha +1Graph Generation Generative Modeling
Boise State University Boise, Idaho, USA · Oak Ridge National Laboratory Oak Ridge, Tennessee, USA
Local Message-Passing for Discrete Graph Generation
Mar 9, 2026Jay Revolinsky, Harry Shomer, Jiliang TangGraph Neural Networks Graph Generation
Michigan State University East Lansing, Michigan, USA · University of Texas - Arlington Arlington, Texas, USA

Aug 19, 2026J/K move · Enter open · S save

GraphK: Variable-Size Graph Generation with Efficient Edge Construction

Resul Tugay, Eren Oluğ, Elif Ak, Sule Gunduz Oguducu

Department of Artificial Intelligence and Data Engineering Ataturk University, Turkey · Department of Artificial Intelligence and Data Engineering ITU AI Research and Application Center Istanbul Technical University, Turkey · Faculty of Engineering and Applied Science Memorial University, Canada

Graph Generation Generative Modeling Graph Representation Learning

Open paper AskPDF arXiv

Graph generation models have advanced significantly with deep learning, yet they remain limited in scalability, flexibility, and ability to model underlying structures. We present GraphK, a novel encoder-sampler-decoder framework for graph generation that overcomes these challenges through structural flexibility and computational efficiency. Unlike autoregressive approaches constrained by vocabulary size (i.e. number of nodes in graph generation), GraphK allows for both upscaling (generating graphs with more nodes than the input) and downscaling, providing a flexible control over output graph size. By learning permutation-invariant latent representations and sampling new node embeddings via maximum likelihood estimation, GraphK generalizes across graph sizes and structures. For edge generation, we employ edge prediction with a KDTree-based top-k neighbor search in the latent space, reducing computational cost. Based on the manifold smoothness assumption, our method effectively captures graph properties. Experiments on synthetic and real-world datasets show that GraphK outperforms existing methods, accurately learns graph structures, and generates synthetic graphs without explicit definitions.

Abstract

Graph Generation

Explore similar work

GraphK: Variable-Size Graph Generation with Efficient Edge Construction

Scaling Novel Graph Generation via Lightweight Structure-Guided Autoregressive Models

Local Message-Passing for Discrete Graph Generation