🧬 biology

UMA-Inverse: Ligand-Conditioned Protein Inverse Folding with a Distogram-Supervised Dense Pair Encoder

This paper introduces UMA-Inverse, a compact, ligand-conditioned inverse folding model that utilizes a dense pair-representation encoder to propagate ligand information beyond the binding interface, achieving competitive performance against LigandMPNN while providing new insights into how all-pair encoders distribute ligand signals.

Original authors: William Sobolewski

Published 2026-07-10✓ Author reviewed ⓘ

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: William Sobolewski

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you're a master chef trying to recreate a secret recipe for a dish that only exists when a specific, rare spice (the ligand) is added to the pot. Your goal is to write down the exact list of ingredients (the protein sequence) that will make the dish taste right and hold its shape.

For a long time, the best chefs used a method called LigandMPNN. Think of this like a neighborhood gossip chain: to figure out how the spice affects the whole pot, the chef asks the ingredient right next to the spice, who asks the neighbor next to them, and so on. The message travels step-by-step. It works well, but if the spice is at the bottom of a huge pot, the ingredients at the very top might never hear about it.

Enter UMA-Inverse, a new, compact kitchen assistant created by William Sobolewski. Instead of a gossip chain, UMA-Inverse uses a dense pair-representation encoder. Imagine this as a magical, all-seeing radar screen that connects every ingredient to every other ingredient and the spice simultaneously. It doesn't just pass a message down the line; it updates the entire map of the pot at once.

The Magic Radar and the "Triangle" Trick

UMA-Inverse is built on a clever shortcut. It uses a six-block system called a PairMixer. Think of this as a team of six chefs who constantly check the distances between every pair of ingredients using a special "triangle multiplication" rule. They look at how three points (two ingredients and the spice, or three ingredients) relate to each other to understand the shape of the dish.

Crucially, this team skips the heavy lifting of "self-attention" (where every ingredient tries to talk to every other ingredient individually in a complex way) to save time. Instead, they focus purely on the geometry—the shapes and distances. To make sure they aren't just guessing, they have a distogram supervisor. This is like a strict head chef who constantly checks the radar screen and asks, "Are you actually seeing the right distances between these ingredients?" This keeps the model honest and ensures it understands the 3D structure, not just the list of words.

The Results: Good, but Not the Best (Yet)

When the authors tested this new assistant against the old neighborhood gossip method (LigandMPNN), the results were a mix of impressive efficiency and honest limitations.

The Score: On a test of small-molecule spices, UMA-Inverse got 56.1% of the ingredients right. The old method, when tested under the exact same conditions by the authors, got 59.8%. For metal spices, UMA-Inverse scored 55.1% versus 64.4% for the old method. For nucleotide (DNA/RNA) spices, it scored 35.3% compared to 53.3%.
The Verdict: The paper explicitly states that UMA-Inverse does not surpass LigandMPNN in accuracy. It trails behind. However, it is much smaller and faster, with only about 3.3 million parameters (compared to the larger models out there).

The Superpower: Hearing the Spice from Afar

Here is where UMA-Inverse shines. While the old gossip chain (LigandMPNN) stops hearing the spice signal once you get about 10 Ångströms away (a tiny distance, but far in molecular terms), UMA-Inverse keeps the signal alive.

The paper measured this by looking at how much the model's predictions changed when the spice was present versus when it was hidden. At distances beyond 10 Å, the old method's signal dropped to near zero. UMA-Inverse, however, retained a strong signal all the way out to 25 Å and beyond. In fact, at the farthest points, UMA-Inverse's signal was up to 30 times stronger than the old method's.

The Catch: Just because the model hears the spice from far away doesn't mean it uses that information to cook a better dish. The paper found that this extra hearing didn't translate into better scores for the final recipe. The model is great at knowing "the spice is over there," but it still struggles to use that knowledge to pick the perfect ingredients as well as the older, gossip-chain method does.

The "DNA" Glitch

There was one specific failure the paper highlights. When the "spice" was a large DNA or RNA molecule, UMA-Inverse failed to improve on a basic model that didn't even look at the spice. The authors explain this is because the radar screen has a limit: it can only track the 50 nearest atoms to the center of the protein. For a tiny spice, that's the whole thing. But for a giant DNA strand with hundreds of atoms, the radar only sees a tiny, useless slice of it. The old gossip method, by contrast, sends a local scout to every part of the DNA, so it doesn't miss anything.

The Bottom Line

UMA-Inverse is a compact, efficient baseline that proves a dense, all-seeing radar screen is a viable way to design proteins. It successfully folds into stable shapes and binds to ligands, but it currently lags behind the established LigandMPNN in getting the exact sequence right.

The authors suggest that the future isn't about choosing one or the other, but perhaps building a hybrid kitchen: using the dense radar for the overall shape and long-distance signals, but keeping the local scouts (like the old gossip chain) to handle the massive, complex details of DNA and RNA. Until then, UMA-Inverse stands as a fascinating, smaller, and surprisingly "aware" tool that shows us exactly how far a ligand's influence can travel in a protein, even if it hasn't quite mastered the art of cooking the perfect dish yet.

Technical Summary: UMA-Inverse

Problem Statement

Protein inverse folding aims to design amino acid sequences that adopt a specific, fixed backbone structure. While autoregressive Graph Neural Network (GNN) models like ProteinMPNN and LigandMPNN have become the standard for this task, they rely on sparse $k$ -nearest-neighbor ( $k$ -NN) graphs. This architecture inherently constrains long-range information flow, as distant residues communicate only through chains of local message-passing steps. This limitation potentially hinders the model's ability to encode allosteric relationships or shell-by-shell geometric dependencies between a ligand and residues far from the binding interface. Furthermore, while LigandMPNN successfully conditions on ligands, it remains the accuracy frontier only for specific splits, with larger, more complex models recently reporting higher recovery rates. The central question addressed is whether a dense pair-representation architecture, inspired by the geometric reasoning capabilities of AlphaFold's Evoformer, can improve ligand-conditioned inverse folding while maintaining computational efficiency.

Methodology

UMA-Inverse is a ligand- and nucleic-acid-conditioned inverse-folding model that replaces the sparse GNN encoder of LigandMPNN with a dense, all-pairs pair-tensor encoder.

Architecture

The model consists of four main stages:

Featurization: Residues are encoded via backbone dihedral angles ( $\phi, \psi, \omega$ ). Ligand and nucleic acid atoms are encoded using atomic features (type, group, period). Pairwise features include 25 backbone distances for residue-residue pairs and 5 backbone-to-ligand distances plus relative frame angles for residue-ligand pairs.
Pair-Tensor Initialization: A dense $N \times N \times d$ tensor ( $d=128$ ) is initialized using node features, radial basis function (RBF) embeddings of inter-node distances, and learned relative position embeddings.
PairMixer Encoder: The core innovation is a six-block encoder based on the PairMixer design. Unlike the full AlphaFold2 Evoformer, it omits triangle self-attention and the sequence track to reduce computational cost. It retains triangle multiplication (updating $z_{ij}$ using $z_{ik}$ and $z_{kj}$ ) and transition MLPs. This allows the model to capture three-body geometric relationships across all residue triplets and residue-ligand pairs.
Autoregressive Decoder: Sequences are decoded in random order. A key distinction from standard mean-pooling baselines is the learned ligand-attention decoder. Instead of uniformly averaging residue-ligand pair information, each residue $i$ computes a position-specific attention weight over the ligand atoms ( $\alpha_{i\ell}$ ) to read a context vector $s_i$ . This allows interface residues to attend sharply to coordinating atoms while distal residues integrate the ligand field more diffusely.

Training and Auxiliary Objectives

Curriculum Learning: Training proceeds in three stages with increasing structure sizes (64 $\to$ 128 $\to$ 384 nodes) to build structural context progressively.
Distogram Supervision: An auxiliary distogram head is added during training to predict binned inter-residue distances from the pair tensor. This forces the encoder to maintain an explicit, structure-predictive geometric representation. The total loss is $L = L_{CE} + 0.2 L_{distogram}$ .
Nucleic Acid Routing: DNA/RNA chains are routed into the ligand atom pool (capped at 50 atoms nearest the protein centroid) to enable nucleotide conditioning, though the paper notes this specific routing strategy has limitations (see Results).

Key Contributions

Dense Pair Architecture for Inverse Folding: The introduction of a stripped-down Pairformer (Triangle Multiplication only) as an encoder for ligand-conditioned inverse folding, demonstrating that explicit all-pair geometric reasoning is viable for sequence design.
Distogram-Supervised Geometry: The use of an auxiliary distogram objective to ensure the pair tensor remains explicitly structure-predictive, preventing the encoder from collapsing into a sequence-only shortcut.
Learned Ligand Attention: A decoder mechanism that conditions residue predictions on position-specific ligand contexts rather than uniform averages, addressing the divergence of distal residue predictions in mean-pool baselines.
Compact Baseline: A model with only ~3.3 million parameters (comparable to LigandMPNN's ~2.5M) that serves as a compact, single-pass alternative to larger generative models.
Representational Characterization: An empirical analysis showing that the dense encoder propagates ligand identity to residues far beyond the binding interface (up to ~30 $\times$ the signal of LigandMPNN at distal shells), a property not observed in sparse graph models.

Results

Sequence Recovery

Evaluated on the standard LigandMPNN test splits (Small Molecule, Metal, Nucleotide) using a matched protocol (10 autoregressive samples, $T=0.1$ , 5Å interface cutoff):

Small Molecule: UMA-Inverse achieves 56.1% recovery vs. LigandMPNN's 59.8% (re-run by authors; published: 63.3%).
Metal: UMA-Inverse achieves 55.1% vs. LigandMPNN's 64.4% (re-run; published: 77.5%).
Nucleotide: UMA-Inverse achieves 35.3%, matching the ligand-agnostic ProteinMPNN (34.0%) but trailing LigandMPNN (53.3%).
Teacher-Forced Recovery: 66.1% on the validation set.

Structural Validation (Pocket-Fixed Design)

In a setting where the binding pocket is fixed and the rest of the protein is redesigned:

Foldability: Designs are confidently folded and ligand-binding-competent under Boltz-2 cofolding (ligand ipTM $\approx$ 0.96).
Pose Accuracy: UMA-Inverse designs yield a median ligand-pose RMSD of 1.53 Å, which is close to the native sequence floor (0.90 Å) but slightly worse than LigandMPNN (1.27 Å).
Distal Recovery: LigandMPNN recovers distal (non-pocket) positions better (53.7% vs. 48.3%).

Distal Signal Analysis

A KL-divergence analysis of ligand-conditioning signal vs. distance reveals a stark difference:

LigandMPNN: Signal collapses to near zero beyond 10 Å.
UMA-Inverse: Retains substantial signal (0.09–0.16 KL) past 25 Å, propagating ligand identity to distal residues several-fold stronger than the sparse graph baseline.

Significance and Claims

The paper positions UMA-Inverse not as a state-of-the-art design tool that surpasses LigandMPNN in accuracy, but as a compact baseline and a vehicle for specific architectural insights.

Viability of Dense Encoders: It demonstrates that a dense, all-pairs encoder supervised by a distogram is a viable architecture for ligand-conditioned inverse folding, achieving competitive results on small-molecule and metal splits despite being a single-pass model.
Representational Difference: The primary contribution is the characterization of how dense encoders distribute information. Unlike sparse graphs that confine ligand signals to the local neighborhood, UMA-Inverse propagates ligand identity globally across the fold.
Limitations and Future Directions: The authors explicitly acknowledge that the current dense architecture fails on nucleotide partners due to the "global atom budget" constraint (capping ligand atoms at 50 to manage the $O(L^3)$ complexity of triangle multiplication). They argue that a hybrid architecture—keeping the dense residue-residue tensor for geometry while using a separate, per-residue atom-context module for ligands—would be the necessary path forward to fix nucleotide design and decouple ligand coverage from cubic computational costs.
Reproducibility: The paper highlights discrepancies between their re-run of LigandMPNN and published numbers (particularly for metals), suggesting that controlled protocol comparisons are essential for accurate benchmarking. It also notes issues with the stability of PDB accession codes in existing splits and the presence of non-physiological metal ions in test sets.

In summary, UMA-Inverse offers a proof-of-concept that dense geometric reasoning can be applied to inverse folding with high efficiency, providing a new perspective on how ligand information propagates through protein structures, even if it currently trails the specialized sparse-graph baselines in raw sequence recovery accuracy.