metadata
license: mit
tags:
- graph-neural-network
- attention
- molecular-property-prediction
- chemistry
ESA — Edge-Set Attention
Graph encoder from the paper An end-to-end attention-based approach for learning on graphs (Buterez et al., Nature Communications 2025).
ESA treats edges (or nodes) as a set and applies masked multi-head attention over them, encoding graph structure as an attention bias mask. No message-passing, no PyG required at inference time.
Usage
from transformers import AutoModel, AutoConfig
model = AutoModel.from_pretrained("viguera10/esa-graph-encoder", trust_remote_code=True)
# Inputs are plain PyTorch tensors
import torch
# Two molecules batched together:
# mol0: 4 nodes, edges 0->1, 1->2, 2->3 (and reverse)
# mol1: 3 nodes, edges 0->1, 1->2 (and reverse)
node_features = torch.randn(7, 84) # total 7 nodes, 84 node features (ChemProp one-hot)
edge_index = torch.tensor([ # global node indices
[0,1,1,2,2,3, 4,5,5,6],
[1,0,2,1,3,2, 5,4,6,5],
])
edge_attr = torch.randn(10, 14) # 10 edges, 14 edge features (ChemProp one-hot)
batch_mapping = torch.tensor([0,0,0,0, 1,1,1]) # node -> graph id
output = model(
node_features=node_features,
edge_index=edge_index,
batch_mapping=batch_mapping,
edge_attr=edge_attr,
)
# output.last_hidden_state: Tensor[2, 256] — one embedding per graph
Configuration
Key parameters in config.json:
| Parameter | Description |
|---|---|
apply_attention_on |
"edge" (ESA) or "node" (NSA) |
hidden_dims |
Feature dimension at each layer |
num_heads |
Attention heads at each layer |
layer_types |
Sequence of "S" (SAB), "M" (masked SAB), "P" (PMA pooling) |
dim_output |
Dimension of the output graph embedding |
xformers_or_torch_attn |
"torch" (default) or "xformers" |
Citation
@Article{Buterez2025,
author={Buterez, David
and Janet, Jon Paul
and Oglic, Dino
and Li{\`o}, Pietro},
title={An end-to-end attention-based approach for learning on graphs},
journal={Nature Communications},
year={2025},
month={Jun},
day={05},
volume={16},
number={1},
pages={5244},
issn={2041-1723},
doi={10.1038/s41467-025-60252-z},
url={https://doi.org/10.1038/s41467-025-60252-z}
}