viguera10/esa-graph-encoder

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}
}