Graph Learning on Ensembles of Cyclic Peptides: An Investigation of Molecular Ensemble Modeling
By Aaron Feller, Kris Deibler, Maxim Secor
"Introduces EnsembleEGNN: encodes molecular ensembles via shared EGNN layers and Set Attention pooling; pretrained on CREMP dataset boosts cyclic peptide property prediction from R²=0.005 to 0.477."
Abstract
Molecular property prediction from structure often uses a single representative conformation, even though many molecules exist as conformational ensembles in solution. We introduce EnsembleEGNN, a molecular ensemble foundation model that encodes an ensemble by first encoding each conformer with shared Equivariant Graph Neural Network (EGNN) layers, then pooling the resulting conformer representations with a Set Attention Block. We pretrain the model on CREMP, a cyclic peptide ensemble dataset, using a multi-task self-supervised objective combining masked token recovery, noisy-coordinate reconstruction, and pairwise distance reconstruction. On the CREMP-CycPeptMPDB dataset, training EnsembleEGNN from scratch fails entirely ($R^2=0.005$). However, the pretrained model reaches $R^2=0.477$ and Pearson $r=0.699$, outperforming the sequence-only BERT baseline ($R^2=0.439$, Pearson $r=0.667$). When EnsembleEGNN is co-trained end-to-end with the BERT sequence encoder, the hybrid model improves further to $R^2=0.538$ and Pearson $r=0.737$. These results demonstrate that encoding conformational ensembles into a single thermodynamically informed embedding improves cyclic-peptide property prediction.
Technical Analysis & Implementation
Technical Breakdown§
Core Methodology§
EnsembleEGNN is a molecular ensemble foundation model designed to incorporate thermodynamic ensemble information into property prediction. The key idea is to encode each conformer in an ensemble using a shared Equivariant Graph Neural Network (EGNN) backbone, then aggregate the per-conformer representations into a single ensemble embedding via a Set Attention Block.
Architecture§
1. Conformer Encoding: Each conformer is a 3D graph with node features (atom types) and edge features (interatomic distances). An EGNN with $L$ layers processes each conformer independently, but weights are shared across all conformers. The EGNN layer updates node embeddings $h_i^{(l)}$ and coordinates $x_i^{(l)}$ via: $$ m_{ij} = \phi_e\left(h_i^{(l)}, h_j^{(l)}, \|x_i^{(l)} - x_j^{(l)}\|^2\right) $$ $$ h_i^{(l+1)} = h_i^{(l)} + \phi_h\left(h_i^{(l)}, \sum_{j \neq i} m_{ij}\right) $$ $$ x_i^{(l+1)} = x_i^{(l)} + \phi_x\left(h_i^{(l)}, h_i^{(l+1)}\right) \cdot \frac{x_i^{(l)} - x_j^{(l)}}{\|x_i^{(l)} - x_j^{(l)}\|} $$ After $L$ layers, a graph-level readout (e.g., mean pooling) produces a conformer representation $z_k$ for each conformer $k$.
2. Ensemble Aggregation: The set of conformer representations $\{z_1, z_2, \ldots, z_m\}$ is passed through a Set Attention Block (SAB) from the Set Transformer architecture. The SAB uses multi-head attention to capture interactions between conformers: $$ \text{SAB}(Z) = \text{FF}\left(\text{MAB}(Z, Z)\right) $$ where MAB is a multi-head attention block with skip connections. The output is a single embedding $z_{\text{ens}}$ that summarizes the ensemble.
3. Prediction Head: The ensemble embedding is fed into a small MLP to predict the target property.
Pretraining Objective§
The model is pretrained on the CREMP dataset (cyclic peptide ensembles) with a multi-task self-supervised objective:
- Masked Token Recovery (MTR): Randomly mask atom types and predict them from context.
- Noisy Coordinate Reconstruction (NCR): Add noise to conformer coordinates and reconstruct original positions.
- Pairwise Distance Reconstruction (PDR): Predict pairwise distances between atoms.
These tasks force the model to learn chemically meaningful representations of both local and global structure.
Implementation Details§
- EGNN: 4 layers, hidden dimension 128.
- Set Attention: 2 blocks, 8 attention heads.
- Optimizer: AdamW with learning rate 1e-4, batch size 32.
- Training: Pretrain for 500 epochs, then fine-tune on CREMP-CycPeptMPDB for property prediction.
Code Snippet (PyTorch-style)§
import torch
from egnn_pytorch import EGNN
from set_transformer import SetAttentionBlock
class EnsembleEGNN(torch.nn.Module):
def __init__(self, dim=128, num_layers=4, num_heads=8):
super().__init__()
self.egnn = EGNN(dim=dim, num_layers=num_layers)
self.sab = SetAttentionBlock(dim=dim, num_heads=num_heads)
self.mlp = torch.nn.Sequential(
torch.nn.Linear(dim, 64),
torch.nn.ReLU(),
torch.nn.Linear(64, 1)
)
def forward(self, feats, coords, edges, batch_indices):
# feats: (total_nodes, dim) but here we use atom types as one-hot
# coords: (total_nodes, 3)
# edges: adjacency
# batch_indices: (total_nodes,) indicating conformer id
h = self.egnn(feats, coords, edges) # returns updated node features
# Graph readout per conformer
from torch_scatter import scatter_mean
z_conformer = scatter_mean(h, batch_indices, dim=0)
# Set attention expects (batch, set_size, dim) but here we treat ensemble as set
# Reshape to (batch, num_conformers, dim)
# Assume batch size=1 for simplicity
z_ens = self.sab(z_conformer.unsqueeze(0))
out = self.mlp(z_ens.squeeze(0).mean(dim=0)) # pool over set
return outResults§
- Training from scratch on the target dataset fails (R²=0.005).
- Pretrained EnsembleEGNN achieves R²=0.477, Pearson r=0.699.
- A hybrid model (EnsembleEGNN + sequence BERT) achieves R²=0.538, Pearson r=0.737.
This demonstrates the value of encoding conformational ensembles via learned pooling and pretraining on relevant ensemble data.