EnsembleEGNN: Set-Based Graph Learning for Thermodynamic Ensembles of Cyclic Peptides

This paper introduces EnsembleEGNN, a molecular ensemble foundation model designed to address the fundamental representational mismatch in molecular property prediction: while many molecules in solution exist as Boltzmann-weighted ensembles of conformers rather than single static structures, most structure-based methods process only a single representative conformation. This issue is particularly detrimental for cyclic peptides, which are inherently flexible scaffolds whose thermodynamic ensembles dictate critical properties like membrane permeability, protease resistance, and target binding.

The model encodes a molecular 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. Each peptide is represented as a set of atom tokens with multiple conformations reusing the same nodes with different position coordinates. From the CREMP dataset, the authors extract all available conformers, rank them by Boltzmann weights, retain the top C=5 states, and renormalize weights to sum to one.

The architecture processes each conformer through L EGNN layers with message passing restricted to k-nearest-neighbor graphs within each conformer, ensuring strict equivariance to rotations and translations. Conformer fusion uses attention weights driven by two signals: a learned score from the mean ensemble embedding and a thermodynamic prior from the Boltzmann weight. This produces two ensemble embeddings: a per-atom embedding (attention-weighted average of each atom across conformers) and a global graph-level embedding via Set Attention Block with inducing points, reducing fusion complexity from O(N2) to O(nm).

The multi-task self-supervised pretraining objective combines three losses: masked token recovery (cross-entropy at masked valid nodes with BERT-style corruption: 80% mask, 10% random, 10% unchanged), noisy-coordinate reconstruction (MSE between refined and clean conformer coordinates), and pairwise distance reconstruction (comparing distances from projected per-atom embeddings with distances from Boltzmann-weighted mean structure). The composite loss is L = 0.3(Ltok) + 0.5(Lcoord) + 0.2(Ldist).

On the CREMP-CycPeptMPDB dataset (n=2979 after boundary filtering), the authors compared four variants: BERT-only sequence baseline (PeptideCLM-2 encoder), EnsembleEGNN with random initialization, pretrained EnsembleEGNN, and a Hybrid model co-training both encoders with concatenated embeddings.

Key findings:

• Training EnsembleEGNN from scratch fails entirely (R2 = 0.005)

• Pretrained EnsembleEGNN reaches R2 = 0.477 and Pearson r = 0.699, outperforming the sequence-only BERT baseline (R2 = 0.439, Pearson r = 0.667)

• The Hybrid model improves further to R2 = 0.538 and Pearson r = 0.737, with correspondingly lower MAE (0.304) and RMSE (0.409)

• The Hybrid model consistently achieved the strongest mean R2 performance across all 5 individual cross-validation folds

The authors note that the current architecture at 11.3M parameters outperforms a BERT-style architecture containing 114M parameters. The randomly initialized EnsembleEGNN performing near random demonstrates that pretraining is necessary for the shared EGNN backbone to establish meaningful physical priors before transfer to downstream regression tasks.

The authors position ensemble representations as a general graph foundation model primitive, noting that protein loop flexibility, RNA secondary-structure ensembles, and small-molecule tautomers all involve distributions over molecular structures. The approach's three primary limitations are: (1) EGNN layers perform message passing only within each conformer, with cross-conformer interactions introduced later through pooling rather than explicit cross-conformer message passing; (2) the study is bounded by the coverage and accuracy of precomputed CREMP ensembles; (3) the CREMP-CycPeptMPDB dataset remains a low-data regime (n=2979), warranting caution in cross-paper comparisons.

The paper concludes that directly learning from conformer ensembles can improve cyclic-peptide representation learning, with the hybrid model improving over the sequence-only baseline by +0.099 in R2 and +0.070 in Pearson r. The authors release the model architecture, pretrained checkpoint, and training code under MIT license.

Improvements to AI systems:

1. Ensemble-aware molecular property predictors – Replace single-conformer encoders with the EnsembleEGNN architecture to directly process Boltzmann-weighted conformational ensembles. This enables more accurate prediction of properties that depend on molecular flexibility (e.g., membrane permeability, binding affinity) for cyclic peptides and other flexible molecules.

2. Hybrid sequence-structure-ensemble models – Combine a sequence-based transformer (like PeptideCLM) with the pretrained EnsembleEGNN via concatenated embeddings. This leverages both evolutionary/sequence patterns and thermodynamic structural distributions, improving regression accuracy (R2 +0.099 over sequence-only) and robustness across cross-validation folds.

3. Self-supervised pretraining with multi-objective geometric losses – Use the composite loss (masked token recovery + noisy-coordinate reconstruction + pairwise distance reconstruction) to pretrain graph neural networks on molecular ensembles. This establishes physical priors (e.g., coordinate consistency, distance geometry) that prevent failure when fine-tuning on small downstream datasets (e.g., n≈3000), where random initialization yields near-zero performance.

4. Efficient conformer fusion via Set Attention with inducing points – Adopt the O(nm) fusion mechanism (instead of O(N2)) to aggregate representations across multiple conformers. This reduces computational cost while preserving attention-driven weighting that combines learned embeddings with thermodynamic Boltzmann priors, enabling scaling to larger ensembles or molecules.

5. Transferable backbone for flexible biomolecules – Apply the shared EGNN backbone pretrained on cyclic peptides to other structure-distribution problems (protein loops, RNA ensembles, tautomers). The model’s equivariance to rotations/translations and its ability to handle multiple conformers as node sets with varying coordinates make it a general primitive for any system with conformational heterogeneity.

What the improved AI system can do:

• Predict cyclic peptide properties (permeability, resistance, binding) with higher accuracy than sequence-only or single-structure models, especially in low-data regimes.

• Learn meaningful molecular representations from ensembles without requiring explicit cross-conformer message passing, reducing architectural complexity.

• Fine-tune effectively on small datasets after pretraining, avoiding the failure mode of random initialization (R2 from 0.005 to 0.477).

• Process multiple conformers per molecule in a single forward pass, with attention weights that automatically balance learned features and Boltzmann probabilities.

• Extend to other biomolecular systems where conformational ensembles are critical, providing a drop-in replacement for single-structure encoders in property prediction or generative tasks.

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.