You can deploy this model and start earning money today!

viguera10/
esa-graph-encoder

Name: esa-graph-encoder
Rating: 2.5 (10 reviews)

esa

graph-neural-network

attention

molecular-property-prediction

chemistry

custom_code

mit

Views

No views yet

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

python

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

bibtex

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