🤖 machine learning

Modeling Atomic Conformational Ensembles of Proteins via Test-Time Supervision of Boltz-2 on Cryo-EM Density Maps

This paper introduces CryoSampler, a novel method that fine-tunes pre-trained static structure prediction models like Boltz-2 directly on raw cryo-EM density maps via test-time supervision to generate accurate atomic conformational ensembles, thereby bypassing traditional two-stage model building and demonstrating potential for in-domain generalization.

Original authors: Jay Shenoy, Miro Astore, Axel Levy, Frédéric Poitevin, Sonya M. Hanson, Gordon Wetzstein

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jay Shenoy, Miro Astore, Axel Levy, Frédéric Poitevin, Sonya M. Hanson, Gordon Wetzstein

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Seeing Proteins in Motion

Imagine a protein not as a rigid statue, but as a piece of playdough that constantly changes shape. To understand how a protein works (like a key opening a lock), scientists need to see all the different shapes it can take, known as its "conformational ensemble."

For a long time, scientists have had a hard time seeing these shapes.

The Old Way (Two-Stage Process): First, they take a blurry 3D photo of the protein (called a Cryo-EM map) and manually or automatically build a 3D model of it. Then, they use that finished model to teach a computer how to predict shapes.
The Problem: Building that 3D model from a blurry photo is incredibly difficult. It's like trying to build a perfect Lego castle based on a single, slightly out-of-focus photograph. Often, the builders get stuck, leave out pieces, or build the castle with the wrong colors. Because this step is so hard, many interesting "snapshots" of protein movement are lost before they can ever be used to teach computers.

The New Solution: CryoSampler

The authors introduce a new method called CryoSampler. Instead of building the Lego castle first and then teaching the computer, they teach the computer to look at the blurry photo and the Lego instructions simultaneously.

Think of it like this:

The Teacher (Boltz-2): They start with a very smart AI (Boltz-2) that is already an expert at building static Lego castles (predicting a single, still protein shape) from a recipe (the protein's genetic sequence).
The Student (CryoSampler): They take this expert teacher and give it a special "gym workout." They show the teacher a whole gallery of blurry photos (an ensemble of Cryo-EM maps) and say, "Don't just build one castle. Build a whole collection of castles that look exactly like these blurry photos."

How It Works: The Two-Step Training

The paper describes a two-stage training process, which can be visualized as a sculptor and a dreamer:

Stage 1: The Sculptor (The VAE)
The AI starts with a perfect, static statue (the output from Boltz-2). It then learns to gently stretch, twist, and bend that statue to match the blurry photos. It's like a sculptor who takes a clay figure and molds it until it fits perfectly inside a glass case (the Cryo-EM map). The AI learns the specific "offsets" or movements needed to make the clay fit the glass.
- Key Innovation: It does this directly on the raw photos, skipping the messy step of manually building a perfect model first.
Stage 2: The Dreamer (The Latent Diffusion Model)
Once the sculptor has learned how to bend the clay, the AI learns to "dream" up new shapes. It learns the pattern of movement. If the protein usually bends its arm left, then right, the AI learns that rhythm. Now, even if you give it a new protein from the same family (a different recipe), it can "dream" up a collection of shapes that mimic those movements, without needing a blurry photo of that specific new protein.

What They Discovered (The Results)

The paper tested this on real scientific data and found two major things:

Better Model Building: When asked to turn blurry photos into 3D models, CryoSampler did a better job than previous methods. It built more complete structures that fit the photos more accurately, with fewer "broken" parts (like missing atoms or impossible angles).
- Analogy: If other methods built a Lego castle that was missing the roof and had a wobbly tower, CryoSampler built a castle that fit the photo perfectly and stood up straight.
Family Resemblance (In-Domain Generalization): They trained the AI on one type of protein channel (TRPV3) and then asked it to predict the shapes of a sibling protein (TRPV5) that it had never seen before.
- The Result: The AI successfully predicted the sibling's movements. It learned the "family style" of bending and twisting from the first protein and applied it to the second.
- Limitation: The paper notes this only worked for proteins in the same family. They did not claim it could predict the shape of a completely different type of protein (like a muscle protein) just by looking at a channel protein.

Summary

CryoSampler is a new way to teach AI how to see proteins in motion. Instead of forcing scientists to manually fix blurry 3D photos before training the AI, this method lets the AI learn directly from the photos. It acts like a master sculptor who can not only fix a single statue but also learn the "dance moves" of a protein family, allowing it to predict how similar proteins will move in the future.

The paper claims this leads to more accurate 3D models and shows promise for predicting how related proteins move, but it stops short of claiming it works for all proteins or has immediate medical applications.

Technical Summary: Modeling Atomic Conformational Ensembles via Test-Time Supervision of Boltz-2 on Cryo-EM Density Maps

Problem Statement

The prediction of protein atomic conformational ensembles from sequence is a critical frontier in structural biology, yet it is hindered by a scarcity of high-quality experimental ensemble data. While cryo-electron microscopy (cryo-EM) allows proteins to explore a broader range of conformations than X-ray crystallography, integrating this data into training pipelines remains difficult. Currently, the field relies on a two-stage process: experimental cryo-EM density maps are manually or semi-automatically converted into atomic structural ensembles (model building), and these resulting structures are subsequently used to train sequence-to-ensemble predictors. This pipeline is limited by the difficulty of atomic model building into heterogeneous, lower-resolution cryo-EM volumes, often resulting in incomplete structures or inaccurate stereochemistry. Consequently, public repositories like the PDB lack the rich dynamic data present in raw cryo-EM experiments, leaving ensemble prediction models blind to these variations.

Methodology: CryoSampler

The authors propose CryoSampler, a unified framework that bypasses the traditional two-stage process by fine-tuning a pre-trained static structure predictor (Boltz-2) directly on raw cryo-EM density maps. The method is a latent diffusion model (LDM) trained in two distinct stages:

Stage 1: 3D Variational Autoencoder (VAE) for Model Building
- Architecture: A 3D spatial VAE is trained to learn a distribution of atomic deformations. The encoder compresses an input ensemble of cryo-EM maps into a latent representation. The decoder generates a 3D feature grid representing a hybrid neural field of atomic offsets.
- Process: The model takes a static reference structure ( $X_{ref}$ ) generated by Boltz-2 and applies predicted per-atom offsets to deform it into an estimated atomic ensemble ( $\hat{X}$ ).
- Differentiable Rendering: The estimated atomic coordinates are rendered back into an estimated map ensemble ( $\hat{V}$ ) using a differentiable physics-based cryo-EM forward model ( $\mathcal{F}_{EM}$ ).
- Loss Function: The VAE is optimized end-to-end using:
  - Reconstruction Loss ( $\mathcal{L}_{recon}$ ): Inverse normalized cross-correlation (NCC) between the input and rendered maps to ensure map-model fit.
  - Structural Loss ( $\mathcal{L}_{structural}$ ): Includes differentiable MolProbity metrics (Ramachandran outliers, rotamer outliers, clashes) and a novel soft rigid body backbone loss to preserve secondary structure geometry.
  - KL Divergence: Regularizes the latent space.
Stage 2: Latent Diffusion Model
- Training: Once the VAE is trained, its weights are frozen. A flow-matching diffusion model is trained on the latent embeddings ( $\mu$ ) extracted from the VAE for the training set of cryo-EM maps.
- Objective: The model learns a vector field that interpolates between a standard normal prior ( $z_0$ ) and the target latent distribution ( $z_1$ ) using a continuous-time flow matching objective.
- Inference: At inference time, the model can sample new latent codes from Gaussian noise, which are decoded by the frozen VAE to produce atomic conformations. This allows for the generation of ensembles for unseen proteins within the same family without requiring new cryo-EM data (in-domain generalization).

Key Contributions

Direct Supervision on Raw Maps: The paper introduces a principle for fine-tuning foundation models (specifically Boltz-2) directly on raw cryo-EM density maps, eliminating the need for an intermediate, error-prone atomic model building step before training.
Unified Framework: CryoSampler combines model building and ensemble generation into a single procedure, achieving state-of-the-art map-model fit and stereochemical accuracy.
In-Domain Generalization: The authors demonstrate that the method can transfer learned conformational features from one protein family member (e.g., TRPV3) to another (e.g., TRPV5), generating accurate atomic ensembles for the target protein without experimental density maps.

Results

The method was evaluated on four biological systems: TRPV3, integrin $\alpha_V\beta_8$ , neurokinin-1 GPCR, and human P-glycoprotein.

Model Building Accuracy: CryoSampler achieved superior map-model fit (measured by $CC_{volume}$ , $CC_{mask}$ , $CC_{peaks}$ , and $CC_{box}$ ) compared to existing methods like ModelAngelo, E2GMM, and native Boltz-2. For instance, on TRPV3, CryoSampler achieved a $CC_{volume}$ of $0.84 \pm 0.02$ , outperforming ModelAngelo ( $0.79 \pm 0.08$ ) and Boltz-2 ( $0.31 \pm 0.03$ ).
Stereochemistry: The generated models maintained competitive MolProbity scores, ensuring physically valid geometries.
Ensemble Prediction (TRPV5): In a zero-shot transfer experiment, the model trained on TRPV3 maps successfully predicted the conformational ensemble of TRPV5. It achieved a lower Wasserstein ( $W_2$ ) distance ($0.44$) compared to Boltz-2 ($0.61$) and Normal Mode Analysis ($0.58$), indicating higher fidelity to the experimental TRPV5 map ensemble.

Significance and Claims

The authors claim that CryoSampler demonstrates the potential to train next-generation atomic ensemble prediction models directly on raw cryo-EM measurements. By supervising pre-trained models on the raw data distribution rather than processed atomic structures, the method overcomes the limitations of current two-stage pipelines.

The paper modestly frames its generalization claims: while it shows strong in-domain generalization (transferring features within the TRP channel family), it does not claim out-of-domain generalization across diverse protein families in this work. The authors attribute this to the sparsity of public cryo-EM ensemble data. They posit that their approach lays the building blocks for future foundation models that can be trained directly on raw experimental data, potentially extending to other modalities like crystallographic structure factors or NMR spectra. The primary validation remains the ability to perform high-fidelity atomic model building on map ensembles and generate accurate ensembles for related proteins.