Mac-Diff: A Diffusion Model for Diverse Protein Conformations

Why one structure is not enough

The success of structure prediction has been so complete that it is easy to forget what it leaves out. A protein in a cell is in constant motion, jostled by water, bending at hinges, opening and closing pockets. Many proteins do their job precisely by switching between distinct shapes, an open state that grabs a molecule and a closed state that acts on it. Freeze the protein into a single most likely structure and you lose exactly the behaviour that makes it interesting.

The gold standard for capturing that motion is molecular dynamics, a simulation that computes the forces on every atom and steps the whole system forward in tiny increments. It works, and it is brutally expensive, often needing supercomputers running for weeks to watch a small protein fold once. That cost is the wall the field keeps hitting, and it is why so much recent effort has gone into teaching generative models to produce the range of shapes a protein visits without simulating every step in between.

The natural tools for this are the same generative models reshaping images and video. If a diffusion model can turn noise into a photograph, perhaps it can turn noise into a protein shape, guided by the sequence the way an image model is guided by a text prompt. The catch is that a protein is not a picture, and the loose way image models connect words to pixels turns out to be the wrong fit for connecting a sequence to a structure.

Denoising a shape, not a picture

Mac-Diff starts by choosing a clever way to describe a protein. Instead of working with raw atomic coordinates, which change the moment you rotate the molecule, it represents the backbone as a grid of relationships between every pair of residues, capturing the distance between them and the angles that orient them. This pairwise description does not care how the protein sits in space, which spares the model from having to learn the same shape over and over in every orientation.

Onto this grid the model applies diffusion. During training it gradually adds noise until the geometric description dissolves into randomness, then learns to run the process backward, denoising step by step until a clean structure emerges. This is the score based flavour of diffusion, where the network learns the gradient that points from noise toward realistic data. Readers who want the mathematical footing for why this style of model converges will find it in our analysis of convergence guarantees for score based generative models.

Generation is the reverse of that equation, where a learned score function steers the noise back toward a real protein shape.

Once the model has produced a clean grid of pairwise geometry, that grid is still not a set of atoms. A separate folding step, built on the Rosetta software suite, turns the predicted distances and angles into actual three dimensional coordinates by finding the atomic arrangement that best satisfies them. The diffusion model dreams up the geometry, and the folding step makes it physical.

The idea that makes it work

The heart of the paper is how the sequence guides the generation, and this is where the authors break from the image playbook. In a text to image model, every word can attend to every pixel, a loose and global connection that works fine when you are painting a cat on grass. A protein is less forgiving. Which residues sit near each other in the folded shape is dictated by physics and chemistry, not by loose association, and most pairs of residues never touch at all.

So the authors built an attention module they call locality aware modal alignment, which forces each residue to attend only to the small set of others it is likely to interact with. Rather than let the sequence talk to the whole structure at once, it restricts the conversation to the neighbours that matter. The mask that decides who talks to whom is not guessed. It is assembled from three ingredients, and the combination is the clever part.

The three terms each carry a different kind of knowledge. A learned attention score captures patterns the model discovers for itself. A contact map, read out of a protein language model, contributes a probabilistic guess about which distant residues fold back to touch each other, the long range structure that a purely local view would miss. And a simple decaying function favours residues that sit close along the chain, encoding the plain fact that neighbours in sequence are often neighbours in space. Together they tell each residue exactly whom to listen to.

Borrowing meaning from a language model

The contact map and the residue features both come from ESM-2, a protein language model trained on hundreds of millions of sequences. Just as a text model learns grammar and meaning from reading, a protein language model learns the evolutionary, structural, and functional regularities of proteins from reading their sequences. Its embeddings carry a rich, compressed sense of what each residue is and how it tends to behave.

Leaning on a language model this way has a practical payoff. Older approaches often needed a multiple sequence alignment, a painstakingly gathered set of related sequences from other organisms, which is not always available and is slow to build. Mac-Diff needs only the single sequence, because the language model has already absorbed the evolutionary context during its own training. That same shift toward learned sequence representations underpins a lot of modern structural biology, a theme our look at the limits of protein language models for localization examines from a different angle.

What the results show

The proof is in whether the generated ensembles match what expensive simulations produce. The team tested Mac-Diff on a benchmark of twelve fast folding proteins whose full motion had been captured in some of the longest molecular dynamics simulations ever run, and compared the two on several measures of how a protein’s shape is distributed.

Two things stand out. First, Mac-Diff did not buy its accuracy by playing it safe. A generative model can look accurate by producing many near identical structures close to the average, but that defeats the purpose. On a plot that pits diversity against fidelity, Mac-Diff sat closest to the ideal corner, holding both at once, where most competitors traded one away for the other. It generated genuinely varied shapes that still matched the real distribution.

Second, it captured the rare and difficult states, not just the common ones. The small protein known as BPTI is famous for visiting five structurally distinct clusters over a millisecond of simulation, and Mac-Diff recovered all five, matching the hardest three more closely than any other method. Recovering the states a protein rarely visits is exactly what a single structure prediction cannot do, and it is where the value of an ensemble lives.

Mac-Diff successfully recovered conformational distributions of fast folding proteins, captured multiple metastable conformations that were observed only in long MD simulation trajectories, and efficiently predicted alternative conformations for allosteric proteins. Wang and colleagues, Nature Machine Intelligence, 2026

The model also handled proteins that switch between functional shapes. For adenylate kinase, an enzyme that opens and closes like a clam, it recovered both states with high accuracy, and across a broader set of proteins known to adopt two distinct conformations it ranked a close second to the strongest competitor. The pattern throughout is of a model that is competitive with the best while being far cheaper to run, a trade that matters enormously when you want to screen many proteins. This whole line of work, generating three dimensional biomolecular structure with diffusion and flow models, is moving fast, as our coverage of a flow model for RNA structure shows in a neighbouring corner of biology.

Where it falls short

The authors are candid about the limits, and a fair reading has to sit with them.

The most important caveat is about what an ensemble is and is not. Mac-Diff produces a collection of plausible shapes and how often each appears, but it does not produce a trajectory. It cannot tell you the order in which a protein moves between states or how fast, the kinetic story that molecular dynamics delivers and that matters for understanding mechanisms. The authors are explicit that their samples are snapshots without the thread of time connecting them, and they point to combining the model with short simulations as a way to recover that missing motion.

There is also a ceiling set by the training data. The model was fine tuned on a modest collection of simulations covering short timescales, which limits how much genuine flexibility it has seen, especially the slow, large motions that unfold over long times. On one harder benchmark it trailed a competitor that had been trained on a dataset roughly ten times larger, and the authors trace part of the gap to that difference in scale rather than to the method itself. More and longer simulation data would likely close it.

Finally, the pipeline has moving parts whose errors compound. The diffusion model produces geometry, a separate folding step turns that into atoms, and each stage introduces its own approximations. The reliance on a language model’s contact map means the structures inherit whatever that model gets wrong, and the whole system was evaluated on proteins deliberately kept distinct from its training set, which is rigorous but also means real world proteins closer to the training data might behave differently. A generated ensemble is a strong hypothesis about a protein’s flexibility, not a measurement of it.

Why it matters

Step back and the contribution is a shift in what structure prediction aims at. The first wave answered the question of what shape a sequence folds into. Mac-Diff pushes toward the harder and more useful question of what shapes a sequence explores, and it does so cheaply enough to be practical. That reframing, from the single most likely structure to the distribution of functional states, is where a great deal of biology actually happens.

The methodological idea has legs beyond proteins too. The insight that a sequence should align tightly and locally with a structure, rather than loosely and globally the way words align with pixels, is a lesson about matching a generative model’s inductive bias to the physics of its domain. Paired with an open release of the code and models, that makes Mac-Diff both a useful tool for studying protein flexibility and a template for how to adapt diffusion to scientific problems where locality is not optional. As generative models keep moving from pictures into the machinery of life, the ones that respect the structure of their domain, as this one does, are the ones worth watching. For a wider view of diffusion escaping the image world, our piece on latent diffusion for moving 3D shapes traces the same migration.

A reference implementation

The code below is a compact, runnable version of the core ideas. It builds a score based diffusion model over a pairwise geometric tensor, adds the locality aware attention that blends a learned score with a contact prior and a distance decay, corrupts a batch with the variance preserving schedule, and trains the network to predict the noise. The authors full code and models are linked under the block.

# Mac-Diff style conditional diffusion for protein geometry, compact reference
# Score based denoising of a pairwise residue tensor, steered by a locality
# aware attention that blends a learned score with a contact and distance prior.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


def vp_schedule(t, beta_min=0.1, beta_max=20.0):
    # variance preserving noise level, alpha_bar at continuous time t in [0, 1]
    log_mean = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min
    alpha_bar = torch.exp(2.0 * log_mean)
    return alpha_bar


class LamaAttention(nn.Module):
    # each residue attends locally, mask blends a learned score, a contact map,
    # and a distance decay along the chain
    def __init__(self, dim, lam=0.5, alpha=0.5, beta=0.5, h=4.0):
        super().__init__()
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.lam, self.alpha, self.beta, self.h = lam, alpha, beta, h

    def forward(self, seq_feat, contact):
        L, d = seq_feat.shape
        q, k, v = self.q(seq_feat), self.k(seq_feat), self.v(seq_feat)
        learned = torch.softmax(q @ k.t() / d ** 0.5, dim=-1)

        idx = torch.arange(L).float()
        dist = (idx.view(-1, 1) - idx.view(1, -1)) ** 2
        locality = torch.exp(-dist / (2 * self.h ** 2))       # neighbours along the chain

        prior = self.alpha * contact + self.beta * locality
        W = self.lam * learned + (1.0 - self.lam) * prior      # the relational matrix
        W = W / (W.sum(dim=-1, keepdim=True) + 1e-8)
        return W @ v                                          # updated residue features


class MacDiff(nn.Module):
    # predicts the noise on the pairwise geometric tensor, conditioned on sequence
    def __init__(self, channels=5, seq_dim=64):
        super().__init__()
        self.lama = LamaAttention(seq_dim)
        self.to_pair = nn.Linear(2 * seq_dim, 32)
        self.net = nn.Sequential(
            nn.Conv2d(channels + 32 + 1, 64, 3, padding=1), nn.SiLU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.SiLU(),
            nn.Conv2d(64, channels, 3, padding=1),
        )

    def forward(self, x_noisy, seq_feat, contact, t):
        L = seq_feat.shape[0]
        updated = self.lama(seq_feat, contact)                # [L, seq_dim]
        ui = updated.unsqueeze(1).expand(L, L, -1)
        uj = updated.unsqueeze(0).expand(L, L, -1)
        pair = self.to_pair(torch.cat([ui, uj], dim=-1))      # [L, L, 32]
        tmap = torch.full((L, L, 1), float(t))
        cond = torch.cat([pair, tmap], dim=-1).permute(2, 0, 1).unsqueeze(0)
        inp = torch.cat([x_noisy, cond], dim=1)
        return self.net(inp)


def train_step(model, x0, seq_feat, contact, opt):
    t = torch.rand(1).item()
    ab = vp_schedule(torch.tensor(t))
    noise = torch.randn_like(x0)
    x_noisy = ab.sqrt() * x0 + (1 - ab).sqrt() * noise        # forward corruption
    pred = model(x_noisy, seq_feat, contact, t)
    loss = F.mse_loss(pred, noise)                            # score matching on the noise
    opt.zero_grad()
    loss.backward()
    opt.step()
    return loss.item()


if __name__ == "__main__":
    L, seq_dim, channels = 24, 64, 5
    seq_feat = torch.randn(L, seq_dim)                        # stand in for ESM-2 features
    contact = torch.rand(L, L)                                # stand in for a contact map
    contact = (contact + contact.t()) / 2
    x0 = torch.randn(1, channels, L, L)                       # the pairwise geometric tensor

    model = MacDiff(channels, seq_dim)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for step in range(200):
        loss = train_step(model, x0, seq_feat, contact, opt)
        if (step + 1) % 50 == 0:
            print(f"step {step + 1}  loss {loss:.4f}")

Conclusion

The core achievement here is a generative model that treats a protein as the moving thing it actually is. Rather than predict one shape, Mac-Diff learns to sample the range of shapes a sequence can take, and it matches long molecular dynamics simulations closely on how those shapes are distributed while running thousands of times faster. It recovered rare functional states that a single prediction cannot reach, and it did so without needing the multiple sequence alignments that older methods lean on.

The conceptual shift worth remembering is the insistence on locality. Borrowing diffusion from image generation was the obvious move, but the authors saw that the loose, everything talks to everything attention of image models is wrong for proteins, where physics dictates a tight web of local contacts. Building an attention module that respects that structure, fed by a protein language model’s sense of which residues touch, is what let the sequence steer a realistic and varied ensemble rather than a blurry average.

The approach travels because its central lesson is about matching a model to its domain. The same discipline, aligning a generative model’s attention with the real geometry of the problem rather than importing a bias from pictures, applies well beyond proteins. Paired with an open release, that makes Mac-Diff both a practical tool for probing protein flexibility and a worked example of adapting diffusion to a scientific setting where locality carries real meaning.

The honest limits keep it grounded. The model gives a distribution of shapes without the order or timing that connects them, it has seen only a limited slice of protein motion in its training data, and its multi stage pipeline compounds the errors of each part. A generated ensemble is a strong, testable hypothesis about how a protein moves, not a substitute for measuring it, and the authors are careful to say so.

Future directions follow naturally. Fold in far more simulation data to capture the slow motions the model still misses, marry the fast sampling to short simulations so the kinetic thread returns, and extend the idea to capture how flexibility differs across related proteins in different species. If those steps hold up, the field gains a way to see not just the shape a protein settles into, but the full repertoire of shapes through which it does its work.

Frequently asked questions