Graph Learning on Ensembles of Cyclic Peptides: An Investigation of Molecular Ensemble Modeling

AI summary

Overview

Research area: Machine learning for molecular property prediction, specifically geometric deep learning / graph neural networks applied to cyclic peptide conformational ensembles, with a foundation-model pretraining paradigm.

Technical level: Intermediate. The paper is written for readers comfortable with graph neural networks, self-supervised pretraining, and basic statistical evaluation metrics, though its core framing (single structure vs. ensemble) is accessible without specialist chemistry background.

Scope in one sentence: The paper introduces EnsembleEGNN, a model that encodes an entire set of 3D conformers of a cyclic peptide into one embedding, and tests whether that ensemble-level representation improves membrane permeability prediction over a sequence-only baseline.

What This Paper Is About

Most molecular property-prediction models represent a molecule by a single 3D structure, even though flexible molecules in solution exist as distributions over many conformations weighted by their thermodynamic likelihood. This mismatch is especially costly for cyclic peptides, whose flexibility determines properties such as membrane permeability and protease resistance. The paper's goal is to build and evaluate a model that learns from the full conformer ensemble rather than collapsing it to one structure upfront.

Key Contributions

1. An ensemble geometric encoder. EnsembleEGNN encodes each conformer separately using shared Equivariant Graph Neural Network (EGNN) layers, then fuses the conformer representations into a single thermodynamically informed embedding using a Boltzmann-weighted attention mechanism and a Set Attention Block (SAB).

2. A multi-task self-supervised pretraining objective. The model is pretrained on the CREMP cyclic peptide ensemble dataset with a composite loss combining masked token recovery (weight 0.3), noisy-coordinate reconstruction (weight 0.5), and pairwise distance reconstruction against a Boltzmann-weighted mean structure (weight 0.2).

3. A controlled evaluation against sequence and no-pretraining baselines. Five-fold cross-validation on CREMP-CycPeptMPDB (n = 2979 after boundary filtering) compares a random-initialized model, a BERT sequence-only baseline, the pretrained geometric model, and a co-trained hybrid, showing that pretraining is necessary and that conformer-aware geometry beats sequence alone.

4. A released implementation. The authors state they release the model architecture, pretrained checkpoint, training code, and manuscript figure generation under an MIT license at https://github.com/AaronFeller/EnsembleEGNN.

Main Findings

• Training from scratch fails. The randomly initialized EnsembleEGNN reached R² = 0.005 ± 0.004 and Pearson r = 0.113 ± 0.016 on held-out CREMP-CycPeptMPDB data, which the authors describe as near-random performance after 20 epochs of fine-tuning.

• Pretraining is what makes the geometric model work. The pretrained EnsembleEGNN reached R² = 0.477 ± 0.009 and Pearson r = 0.699 ± 0.009, with MAE 0.319 ± 0.002 and RMSE 0.436 ± 0.004.

• The ensemble model beats the sequence-only baseline. BERT-only, using the PeptideCLM-2 encoder, reached R² = 0.439 ± 0.053, Pearson r = 0.667 ± 0.035, MAE 0.334 ± 0.014, and RMSE 0.451 ± 0.022 — below the pretrained EnsembleEGNN on every reported metric.

• Co-training the best of both modalities gives the strongest result. The Hybrid model (EnsembleEGNN co-trained end-to-end with the BERT sequence encoder, concatenating embeddings) reached R² = 0.538 ± 0.029, Pearson r = 0.737 ± 0.016, MAE 0.304 ± 0.008, and RMSE 0.409 ± 0.013. The paper reports this as an improvement of +0.099 in R² and +0.070 in Pearson r over the sequence-only baseline.

• The hybrid wins on every individual fold. Fold-wise R² for the Hybrid ranged from 0.494 to 0.565, versus 0.402 to 0.472 for BERT-only and 0.451 to 0.509 for EnsembleEGNN.

• A smaller model outperforms a larger one. The paper reports that the architecture at 11.3M parameters outperforms a BERT-style architecture containing 114M parameters.

• Conformer fusion is made cheap. Using a SAB with inducing points reduces conformer fusion complexity from O(N²) quadratic attention to O(nm), making it effectively linear in the number of conformers when the number of inducing points is fixed.

Methodology in Plain English

The researchers represent each peptide as a set of atom tokens plus multiple sets of 3D coordinates, one per conformer. From CREMP they take the top C = 5 conformers ranked by Boltzmann weight and renormalize those weights to sum to one.

Each conformer is passed independently through six shared EGNN layers with a hidden dimension of 448, with message passing restricted to each atom's 12 nearest neighbors (k = 12). Restricting to a fixed neighbor count lets the model handle arbitrary molecule sizes, and EGNN layers are equivariant to rotations and translations, so no coordinate alignment or spatial augmentation is needed.

To merge the conformers, the model computes an attention weight for each conformer that combines a learned score (from a 2-layer MLP applied to the conformer's mean atom feature) with the conformer's Boltzmann weight. That produces two things: a per-atom ensemble embedding (an attention-weighted average of each atom across conformers) and a global graph-level embedding produced by a Set Attention Block over all conformer embeddings. The fusion module has 2 set-attention blocks, 4 inducing points, 8 attention heads, and the conformer-prior bias enabled.

Pretraining corrupts the inputs: 15% of valid node positions are selected for masked-token supervision, with 80% replaced by a mask token, 10% by a random token, and 10% left unchanged. Gaussian noise with standard deviation 0.15 Å is added to coordinates, and the model must refine them back. Pretraining ran for 80 epochs at batch size 8 with learning rate 2×10⁻⁴ and weight decay 1×10⁻⁴.

Downstream, all models are fine-tuned for 20 epochs with two regression-head layers, head dropout 0.1, and early stopping patience 3, using the same data splits. The BERT baseline uses the aaronfeller/peptideclm-2-hybrid-base encoder with batch size 16, learning rate 3×10⁻⁴, and weight decay 1×10⁻². The Hybrid starts from the best pretraining checkpoint and uses separate learning rates of 2×10⁻⁴ (head), 1×10⁻⁴ (ensemble), and 2×10⁻⁵ (CLM). The target is log(P_app), evaluated by MAE, RMSE, Pearson r, and R².

Why This Matters

Impact on research. The paper reframes conformer aggregation as a general foundation-model primitive rather than a peptide-specific trick, arguing that the same issue arises for protein loop flexibility, RNA secondary-structure ensembles, and small-molecule tautomers. The striking gap between random initialization (R² = 0.005) and pretrained transfer (R² = 0.477) is a concrete data point that self-supervised geometric pretraining establishes physical priors that supervised fine-tuning on a small dataset cannot.

Real-world applications:

• Cyclic peptide drug design — better membrane permeability prediction supports the development of orally available peptide therapeutics, which are otherwise limited by poor permeability.
• Prioritizing compounds for wet-lab screening — the authors frame the model as a hypothesis-generating tool to guide which candidates go to physical assays, reducing the cost of high-throughput screening.
• Early-stage drug discovery more broadly — any program where molecular flexibility governs binding, resistance, or transport could in principle adopt an ensemble-aware encoder.
• Molecular representation learning research — the ensemble tokenization scheme (shared atoms, multiple coordinate sets) and the attention-based fusion mechanism are reusable components for other conformational data.

Industry relevance. The work was conducted partly at Novo Nordisk's Molecular AI group, indicating direct pharmaceutical interest in cyclic peptide modalities. The compute profile is also relevant: the released design scales linearly rather than quadratically with the number of conformers and outperforms a model roughly ten times larger in parameter count.

Future Directions

1. Richer cross-conformer interaction. The current EGNN layers pass messages only within a conformer, so interactions across conformers enter only through later pooling. The authors identify explicit cross-conformer message passing as an architectural gap to close.

2. Scaling conformer datasets and model capacity. The authors expect that scaling both the ensemble data and model size will compound the representation advantages of self-supervised geometric pretraining.

3. Integrating high-fidelity conformer generation. The study is bounded by the coverage and accuracy of the precomputed CREMP ensembles; the authors propose generating conformers directly rather than relying on existing libraries.

4. Generalization outside the training distribution. The paper raises the question of whether the pretrained model can transfer beyond its training distribution, noting the current work is a low-data regime (n = 2979) that warrants caution in cross-paper comparisons.

Target Audience

This paper is most valuable to machine learning researchers working on molecular graph foundation models and geometric deep learning, particularly those interested in how to represent conformational ensembles rather than single structures. It is also relevant to computational chemists and medicinal chemists working on cyclic peptides and membrane permeability, to practitioners looking for a small, efficient pretrained geometric encoder (11.3M parameters) they can co-train with a sequence model, and to researchers studying transfer learning in low-data molecular regression settings. Readers without background in equivariant graph neural networks or self-supervised pretraining will find the architecture section demanding, though the experimental comparisons and conclusions are stated plainly enough to follow.

Authors’ 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.