Paper: 2607.21561 Authors: Aaron Feller, Kris Deibler, Maxim Secor Categories: cs.LG, q-bio.BM

The Gap

Molecular property prediction in ML almost universally picks one “representative” 3D structure per molecule and runs a GNN on it. This single-conformer assumption is well-known in computational chemistry to be a simplification, but nobody in the ML community has shipped an architecture that actually encodes an entire conformational ensemble into a single embedding. For cyclic peptides—small ring-shaped molecules of enormous pharmaceutical interest—this gap is especially painful: they’re highly flexible, constantly shifting between conformations, and the biologically relevant property (e.g., membrane permeability) depends on the full thermodynamic distribution, not any one snapshot.

The logical path from problem to conclusion:

    Problem: Single conformer used for prediction
                            |
                            v
    Assumption: One 3D structure is sufficient
                            |
                            v
    Reality: Cyclic peptides have many conformations;
             properties depend on the full ensemble
                            |
                            v
    Gap: No ensemble-aware molecular encoder exists
                            |
                            v
    Method: EnsembleEGNN (EGNN + Set Attn + pretrain)
                            |
                            v
    Evidence: R^2: 0.005 --> 0.477 --> 0.538
              (scratch)  (pretrain) (hybrid)
                            |
                            v
    Conclusion: Ensemble encoding captures thermodynamic
                information that single conformers miss

The Increment

One sentence: Before this paper, every graph neural network for molecular property prediction treated a molecule as a single rigid shape; after, there’s a working architecture that encodes a whole bag of conformations into one embedding and proves it boosts cyclic peptide permeability prediction.

Core Mechanism

EnsembleEGNN has three stages: per-conformer encoding, set-level pooling, and prediction. First, each conformer in the ensemble is represented as a molecular graph with 3D coordinates. These graphs are fed through a shared Equivariant Graph Neural Network (EGNN)—the same weights process every conformer independently, producing a fixed-size vector representation for each one. The equivariance property ensures that rotating or translating a conformer doesn’t change its representation, which is critical since conformers have no canonical orientation.

Second, the N conformer representations are pooled by a Set Attention Block (SAB), a transformer-like module from the Set Transformer family. The SAB treats the set of conformer vectors as a sequence, applies multi-head self-attention across them, and produces a single output vector—the ensemble embedding. Crucially, this operation is permutation-invariant: it doesn’t matter which conformer is first or last. The result is a single embedding that encodes relationships between conformations, not just individual ones.

Third, during pretraining on the CREMP cyclic peptide dataset, the model learns through three self-supervised objectives simultaneously: masked atom recovery (predict hidden atoms), noisy-coordinate reconstruction (denoise perturbed 3D positions), and pairwise distance reconstruction (predict inter-atom distances). After pretraining, the model is fine-tuned on downstream property prediction (membrane permeability in CycPeptMPDB).

    Conf_1    Conf_2    ...    Conf_N
      |         |               |
      v         v               v
    [EGNN]    [EGNN]    ...   [EGNN]     <-- shared weights
      |         |               |
      v         v               v
    rep_1     rep_2     ...    rep_N
      |         |               |
      +----+----+----...---+---+
           |
           v
    [Set Attention Block]
           |
           v
    ensemble embedding
           |
           v
    [prediction head]

Now here’s the version that actually sticks in your head.

Structural metaphor: the casting director.

Imagine you’re a casting director predicting whether an actor will deliver in a blockbuster role. The old way (single-conformer models) is like watching one audition tape and calling it. But actors are like molecules—they perform differently each take, bringing different energies and choices. You’d never hire based on one take.

EnsembleEGNN’s way: you watch N different takes of the same scene. You assign the same expert coach (shared EGNN weights) to analyze each take independently—technique, timing, emotional range. Each coach report is a conformer representation. Then you bring all N reports into a deliberation room (Set Attention Block) where the coaches cross-reference each other’s findings: “Take 3 had the best timing, but Take 7 showed more range; here’s how they relate.” The panel produces a single dossier (the ensemble embedding) that captures something no single take could: the actor’s range, consistency, and underlying talent.

The pretraining on CREMP is like training your coaches on thousands of hours of footage from many actors before they evaluate your candidate. Without that background training, the coaches are lost—R-squared of 0.005, essentially random. With it, they recognize meaningful patterns and deliver real predictive power.

Key Concepts

  • Conformational ensemble: A molecule is not a rigid statue. At room temperature, a cyclic peptide constantly twists, bends, and flexes through many distinct 3D shapes. Each shape is a “conformer.” The full collection of these shapes, weighted by how thermodynamically likely each one is, is the “conformational ensemble.” Think of a rubber bracelet: you can hold it in an oval, a twisted figure-eight, or folded in half—all valid shapes it cycles through naturally. For drug design, different conformations expose different surfaces to a biological target, so the ensemble tells you what the molecule actually does in solution, not just what it looks like frozen in place.

  • Set Attention Block (SAB): You have N items (conformer vectors) and need to combine them into one. A simple average would treat all conformers equally, but they’re not equally likely or informative. The SAB is a transformer-style attention mechanism applied over the set: each conformer representation attends to every other, producing a context-aware summary. It’s like a roundtable discussion where each food critic hears what the others are saying before forming a final opinion. The result is a single vector that encodes inter-conformer relationships—something no element-wise pooling (mean, max) can capture.

  • Multi-task self-supervised pretraining: Before predicting any property labels, the model first solves three “homework problems” on unlabeled molecular data. (1) Masked atom recovery: hide some atoms and predict what’s there—like fill-in-the-blank for molecules. (2) Noisy-coordinate