README.md

---
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](https://www.nature.com/articles/s41467-025-56067-5) (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}
}
```