🤖 machine learning

Evolution-Aware MSA Reasoning for Subsampling via Factor Graphs

This paper introduces AP-REASONER, a factor-graph-based optimization framework that treats MSA subsampling as a controllable problem to explicitly balance evolutionary signals like query identity and diversity, thereby outperforming heuristic methods in protein structure prediction and conformational ensemble recovery.

Original authors: Zhangzhi Xiong, Minzhang Li, Haotian Yu, Sixian Shen, Kexin Zhang, Mingrui Li, Jie Zheng, Kewei Tu, Jingyi Yu

Published 2026-07-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhangzhi Xiong, Minzhang Li, Haotian Yu, Sixian Shen, Kexin Zhang, Mingrui Li, Jie Zheng, Kewei Tu, Jingyi Yu

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

Imagine the language of life written in a code called proteins. These aren't just static blocks; they are tiny, folding machines that do everything from building your muscles to fighting off viruses. To understand how a protein works, scientists often look at its "family tree"—a massive collection of related sequences called a Multiple Sequence Alignment (MSA). Think of this family tree like a giant, noisy library of books, where every book is a slightly different version of the same story, written by ancestors who lived millions of years ago. By reading these books together, computers can figure out the protein's 3D shape and how it moves.

However, there's a problem: the library is too big. Modern super-smart computer programs (called protein language models) are like brilliant students who can only read a few hundred pages at a time before their brains get full. If they try to read the whole library, they crash. So, scientists have to pick a small, representative sample of books to study. For a long time, they've done this by guessing—picking random books, or trying to find the most different ones, or filtering out duplicates. It's a bit like trying to pick the best 100 songs for a party by just grabbing a handful from the top of the pile or picking only the ones you've heard before. It works okay, but you might miss the hidden gems that explain how the protein actually dances.

This paper introduces a new way to pick those songs, called AP-REASONER. Instead of just guessing or using simple filters, the authors treat the selection process like a complex puzzle that needs to be solved with math. They built a system that acts like a smart committee meeting. Imagine a room full of protein sequences, and the goal is to pick a specific number of them to represent the whole group. The system uses a "factor graph," which is like a giant web of connections, to reason through the choices. It asks questions like, "If I pick this sequence, does it cover enough ground?" and "Does this one look too much like the one we already picked?"

The magic of this new method is that it has two "control knobs" that scientists can turn to change the outcome. One knob, called $\alpha$ , controls how diverse the group should be. Turn it up, and the system picks sequences that are very different from each other, ensuring a wide variety of evolutionary stories are told. The other knob, $\beta$ , controls how close the selected sequences should be to the original "query" protein. Turn this one, and the system focuses on finding sequences that are very similar to the target, which helps in finding specific shapes or states the protein might take.

The researchers tested this new "reasoning" system against the old guessing methods. They found that when they used AP-REASONER, the computer models were better at predicting how proteins fold and, more impressively, how they change shape. In some tests, the old methods completely missed certain shapes that the new method found easily. For example, with a protein called KaiB, the new method could find a rare, hidden shape using far fewer attempts than the old methods, which had to guess wildly and waste a lot of computer power. The paper suggests that by treating the selection of protein sequences as a careful, controllable optimization problem rather than a random guess, we can get a much clearer picture of how life's building blocks move and function. It's not just about picking a few books; it's about curating the perfect story to understand the character.

Technical Summary: Evolution-Aware MSA Reasoning for Subsampling via Factor Graphs

Problem Statement

Multiple Sequence Alignments (MSAs) provide protein language models (PLMs) with explicit evolutionary context, which is crucial for tasks like contact prediction and conformational analysis. However, the depth of MSAs retrieved from genomic databases often reaches tens of thousands of sequences, far exceeding the token budgets of attention-based PLMs (e.g., MSA Transformer, AlphaFold2), which typically process only a few hundred sequences. Consequently, subsampling is a mandatory step.

Existing subsampling strategies rely on heuristics such as random selection, identity-based filtering (e.g., HHfilter), or diversity-driven greedy sampling (e.g., divMax, divMin). While effective at removing near-duplicates or enforcing specific biases, these methods operate on local similarity relations and lack explicit control over the trade-off between two complementary evolutionary measures: query identity (how similar sequences are to the target) and sequence diversity (the breadth of evolutionary variation). This limitation prevents the explicit balancing of these factors, potentially leading to suboptimal evolutionary signals being retained for downstream tasks.

Methodology: AP-REASONER

The authors propose AP-REASONER, a method that recasts MSA subsampling from a heuristic filtering procedure into an explicit optimization problem solved via Affinity Propagation (AP) on a factor graph.

Factor Graph Formulation

The problem is modeled as a Maximum A Posteriori (MAP) estimation on a factor graph where:

Variables: For each row $i$ in the MSA, a discrete variable $s_i$ indicates which row $k$ represents it. A row $k$ is an "exemplar" (selected for the subset) if and only if $s_k = k$ .
Unary Factors ( $L_i$ ): These measure the plausibility of row $i$ $i$ being represented by exemplar $k$ $k$ . The similarity matrix $S$ $S$ is constructed using a normalized, gap-aware Hamming distance ( $d_{H}^{gap}$ $d_{H}^{g a p}$ ) to quantify sequence redundancy and divergence.
- Off-diagonal entries ( $S_{ik}, i \neq k$ ): Defined as $-\left(\frac{d_{H}^{gap}(i, k)}{L}\right)^\alpha$ . The exponent $\alpha$ acts as a control knob for global diversity. A larger $\alpha$ flattens the penalty for small distances, encouraging the selection of more globally distributed exemplars.
- Diagonal entries ( $S_{kk}$ ): Defined as $p + \beta \frac{d_{H}^{gap}(k, 0)}{L}$ , where $0$ is the query sequence. The coefficient $\beta$ acts as a control knob for query identity. A positive $\beta$ incentivizes selecting exemplars distant from the query, while a negative $\beta$ biases selection toward similar sequences. The query row is anchored with a large fixed preference to ensure it is always retained.
Exemplar-Consistency Factors ( $f_k$ ): These enforce the constraint that if any row selects $k$ as an exemplar, $k$ must also select itself ( $s_k = k$ ).

Optimization and Control

The system performs message-passing reasoning (log-domain max-product) to infer the optimal set of exemplars. To strictly adhere to a fixed token budget $N_B(L)$ , the authors employ a Fixed-Cardinality Wrapper. This involves a bisection search over the global preference offset $p$ to tune the number of inferred exemplars to the target budget, followed by top- $N$ truncation or padding based on diagonal beliefs.

Key Contributions

Optimization Reformulation: The work is the first to recast MSA subsampling as an explicit optimization problem rather than a heuristic sampling procedure, elevating biological measures (diversity and query identity) to explicit objectives.
AP-REASONER: Introduction of an Affinity-Propagation-based factor-graph approach that encodes subsampling objectives into factors and solves them via message passing. It features two interpretable control knobs ( $\alpha$ and $\beta$ ) to explicitly modulate sequence diversity and query identity, respectively.
Empirical Validation: Comprehensive evaluation demonstrating that AP-REASONER outperforms baseline subsamplers (random, divMin, divMax, HHfilter, K-medians) on structure-sensitive downstream tasks and enables the controllable recovery of alternative protein conformations.

Experimental Results

The authors evaluated AP-REASONER on protein language model pretraining, long-range contact prediction (LRCP), and protein multiple conformation prediction.

Pretraining (MLM): On masked language modeling tasks, AP-REASONER performed competitively with random subsampling and K-medians, outperforming divMin and divMax. However, the authors note that MLM performance does not always correlate with downstream structural performance.
Long-Range Contact Prediction (LRCP): On CASP15 targets, models trained with AP-REASONER consistently achieved higher Top-L contact precision than those trained with baseline samplers. The performance gap widened on an "informative subset" of targets. Furthermore, varying the diversity knob $\alpha$ during inference showed a monotonic improvement in LRCP performance as $\alpha$ increased, confirming the knob's utility in promoting diverse subsets.
Conformational Ensemble Prediction: On fold-switching proteins (KaiB, RfaH, MAD2), AP-REASONER demonstrated superior efficiency and precision compared to brute-force random sampling and AF-Cluster.
- It successfully recovered rare functional states (e.g., the Fold-switch state of KaiB) that were missed by random sampling.
- It achieved comparable structural accuracy (RMSD) with significantly fewer samples (e.g., 56 samples vs. 1008 for AF-Cluster on KaiB).
- The $\beta$ knob allowed for the controlled steering of the predicted ensemble between distinct conformational basins (e.g., Ground State vs. Fold-switch State), enabling the targeted recovery of specific functional states.

Significance and Claims

The paper claims that modeling MSA subsampling as a controllable optimization problem solved through factor-graph reasoning offers a principled alternative to heuristic selection. By explicitly balancing query identity and diversity, AP-REASONER not only improves downstream performance on structure prediction but also provides a resource-efficient mechanism for exploring protein conformational landscapes. The ability to steer sampling via $\alpha$ and $\beta$ allows for the predictable recovery of alternative functional states, which is critical for understanding protein dynamics and allosteric mechanisms without the prohibitive computational cost of brute-force sampling.

The authors acknowledge limitations, noting that AP-REASONER is approximately 2–3 $\times$ slower than lightweight heuristic samplers due to message passing and the bisection wrapper. Additionally, the pretraining experiments were conducted on a 260K-scale dataset rather than the full 26M-scale corpus, leaving the scalability of the method in larger regimes as an open question for future exploration.