Protein Counterfactuals via Diffusion-Guided Latent Optimization

AI summary

Overview

Research area: Machine learning for computational protein science — specifically counterfactual explanations and diffusion-guided latent optimization applied to protein sequence design.

Technical level: Advanced. The paper assumes familiarity with diffusion models, latent-space optimization, protein language model embeddings, and counterfactual explanation literature.

Scope: The paper introduces MCCOP, a gradient-based framework that edits protein sequences in a continuous joint sequence–structure latent space, using a pretrained diffusion model as a manifold prior, to produce minimal, structurally plausible mutations that flip a property predictor's output to a desired target state.

What This Paper Is About

Deep learning predictors can flag a protein as unstable, inactive, or non-fluorescent, but they do not tell a protein engineer which specific mutations would fix the problem. The authors frame this as an algorithmic recourse problem: given a protein predicted to lack a desired property, find the smallest biologically plausible edit that changes the prediction. MCCOP addresses this by searching in a continuous latent space rather than directly over discrete sequences, so that the resulting edits remain foldable and structurally realistic.

Key Contributions

1. Framework. MCCOP combines predictor-guided gradient descent with diffusion-based manifold projection and gradient-sensitivity masking to produce sparse, valid, and plausible protein counterfactuals, without task-specific retraining of the generative model.
2. Quantitative evaluation. On three benchmarks, the authors report near-perfect success rates with 3–5× fewer mutations than discrete baselines and near-zero adversarial rates.
3. Mechanistic interpretability. MCCOP's recovered mutations concentrate in known functional regions (chromophore packing in GFP, the E2-binding interface in Ube4b, core-facing residues for stability), and in several cases the method exactly recovers ground-truth opposite-label sequences from held-out test data.

Main Findings

• Unconstrained gradient optimization is entirely adversarial. Every counterfactual produced by the plain gradient-descent baseline (no smoothing, no manifold projection, no masking) decoded back to the original sequence, with a 100% adversarial rate on all three datasets and undefined edit distance.
• Smoothing improves robustness without sacrificing accuracy. Adding spectral normalization, Jacobian regularization, Softplus activations, and FGSM-based adversarial augmentation reduced average L2 gradient norms by up to 4× while maintaining or improving AUROC. Activity AUROC rose from 0.82 ± 0.00 to 0.93 ± 0.01; stability from 0.94 ± 0.00 to 0.98 ± 0.01; fluorescence stayed at 0.99 ± 0.00. Gradient norms fell from 2.21 ± 0.19 to 1.10 ± 0.10 (fluorescence), 1.38 ± 0.13 to 0.36 ± 0.08 (stability), and 0.36 ± 0.06 to 0.33 ± 0.04 (activity).
• High success with few edits on stability and activity. MCCOP reached 1.00 ± 0.00 success on both, with 2.32 ± 0.01 (stability) and 2.46 ± 0.33 (activity) mutations, versus 7.76 ± 0.06 / 9.46 ± 0.18 and 6.24 ± 3.67 / 10.91 ± 0.02 for the genetic algorithm and stochastic hill climbing. Adversarial rates were 0.03 ± 0.00 and 0.02 ± 0.02 respectively. MCCOP reached early stopping after a median of 2–10 steps, while hill climbing exhausted its budget in over 95% of cases.
• Fluorescence is the harder task. MCCOP's success rate there was 0.19 ± 0.00, which the authors attribute to the need for precise chromophore geometry possibly exceeding the sparsity budget of k = 5. The successful counterfactuals were nevertheless the sparsest of any method, at 1.37 ± 0.01 mutations, with an adversarial rate of 0.01 ± 0.00.
• Discrete baselines are weak on success rate. Stochastic hill climbing scored 0.23 ± 0.00 (stability), 0.13 ± 0.00 (fluorescence), and 0.03 ± 0.00 (activity); the genetic algorithm scored 0.55 ± 0.01, 0.36 ± 0.30, and 0.17 ± 0.15.
• Counterfactuals stay close to the original distribution. Across pLDDT, GRAVY hydrophobicity, instability index, and radius of gyration, MCCOP outputs were nearly indistinguishable from the original test-set distribution, whereas discrete baselines shifted broadly, especially in hydrophobicity and instability index. Using Kruskal-Wallis/Dunn's tests with Benjamini-Hochberg correction, MCCOP achieved significantly higher pLDDT than both baselines across all tasks (adjusted p < 0.02).
• Mutations land where biology expects them. For GFP, mutations concentrated in the chromophore-proximal region (residues 63–69) and β-barrel strands forming the chromophore cavity, with a small number of distal mutations (e.g., residues 181, 216). For Ube4b, mutations clustered at the E2-binding interface (residues 66–71), through which Ube4b recruits UbcH5c for ubiquitin transfer. For stability, spanning diverse topologies, MCCOP frequently targeted core-facing residues, pointing to hydrophobic core consolidation.
• Ground-truth sequences are rediscovered. MCCOP exactly recovered existing opposite-label sequences in 16 (fluorescence), 18 (activity), and 4 (stability) cases, several from the held-out test set.
• Computational profile differs from discrete methods. Discrete baselines spent 97% (hill climbing) and 94% (genetic algorithm) of their time re-encoding candidate mutations through the CHEAP encoder backed by ESMFold; MCCOP spent 99% of its time on the diffusion-based manifold projection. MCCOP and stochastic hill climbing had comparable per-sample wall-clock times; the genetic algorithm was roughly an order of magnitude more expensive.

Methodology in Plain English

The approach has four moving parts.

A continuous representation of proteins. Each protein sequence is mapped into a continuous latent space using CHEAP, an encoder built on ESMFold activations that compresses both evolutionary and structural information. Its decoder returns both a sequence and backbone coordinates, with near-perfect round-trip reconstruction (greater than 99% residue accuracy). Because the decoder is position-wise — each output residue depends only on its own latent row — masking rows in latent space directly enforces sparsity in sequence space. Both encoder and decoder stay frozen.

A smoothed predictor. A three-layer MLP is trained on flattened CHEAP embeddings to classify proteins into two classes per task. The authors argue that a jagged predictor produces high-frequency gradients that steer optimization toward adversarial, unfoldable examples. They therefore smooth it four ways: spectral normalization on all linear layers, Jacobian regularization penalizing the gradient norm, Softplus activations (β = 1), and adversarial augmentation using FGSM perturbations (ε = 0.01) that decode to the original sequence and are kept with the original label — teaching invariance to semantically meaningless perturbations.

A sparsity mask. At each optimization step, the framework computes a per-position sensitivity as the L2 norm of the gradient of the counterfactual loss with respect to that latent row. Only the top-k positions (k = 5) keep their gradients; everything else is hard-reset to the original embedding. The mask can also be user-defined, for example to freeze catalytic residues.

A manifold projection. After each gradient step, the embedding is partially diffused to noise level 100 and denoised using DiMA, a pretrained diffusion model over protein embeddings. The denoised result is blended with the pre-projection embedding: the new point is (1 − α) times the candidate plus α times the projection, with α = 0.3. This acts as an implicit projection onto the manifold of plausible proteins, countering the tendency of pure gradient descent to produce adversarial examples. The optimization stops early once the target-class confidence reaches τ = 0.95 and the decoded sequence differs from the original.

Evaluation. Three datasets are used: TAPE Fluorescence (GFP homologs, binarized into bright/dark via Otsu's method, predicted dark → bright), TAPE Stability (proteolysis measurements with the middle 33% quantile removed, unstable → stable), and Ube4b Activity (roughly 100k mutations in the U-box domain mapped to auto-ubiquitination activity, middle 33% removed, inactive → active). After preprocessing the datasets contain 54,025 (fluorescence), 45,901 (stability), and 60,692 (activity) sequences, with average lengths 236.96, 45.06, and 102, and embedding dimension D = 1024. Evaluation excludes samples the predictor misclassifies, leaving n = 2209 (fluorescence), 2093 (stability), and 2600 (activity). Results are reported as mean ± std over three seeds. Baselines are stochastic hill climbing (greedy random single-site mutations), a genetic algorithm with edit-distance-penalized fitness, and unconstrained gradient descent in latent space.

Why This Matters

Impact on research. The paper is one of the first demonstrations that diffusion-guided counterfactual optimization — a technique developed for images and tabular data — transfers to proteins, where the manifold constraint is enforced not by pixel realism but by folding and epistasis. It also makes a methodological point about model interpretation: a counterfactual is only as trustworthy as the predictor it explains, so a counterfactual that faithfully reproduces a spurious correlation is itself diagnostic. The authors explicitly note that their framework identifies correlational, not causal, relationships, and that establishing causal links would require interventional experiments.

Real-world applications:

• Antibody and enzyme stabilization: Given a candidate flagged as unstable, MCCOP proposes a handful of mutations (2–3 on average in the authors' stability task) that would restore stability while preserving structure.
• Directed-evolution hypothesis generation: Instead of screening thousands of variants, a lab could prioritize MCCOP's sparse suggestions — 2 mutations versus thousands of variants in typical directed evolution, as the authors note.
• Model auditing and debugging: When a predictor is suspected of relying on spurious correlations, counterfactuals reveal exactly which residues drive the prediction.
• Constrained editing: Because the mask can be user-defined, catalytic residues can be frozen while surrounding positions are optimized.

Industry relevance. One co-author is affiliated with BASF Digital Solutions GmbH, and the work was co-financed by the BASLEARN TU Berlin/BASF Joint Laboratory, signaling industrial interest in computational protein engineering. The authors are candid about practical trade-offs: the diffusion projection accounts for 99% of MCCOP's runtime, and they recommend executing it only every n optimization steps for performance-critical applications. They also note a caveat in the timing comparison — discrete baselines only appear slower to re-encode candidates because the shared evaluation protocol forces them through the CHEAP encoder.

Future Directions

• Multi-objective counterfactuals: jointly optimizing properties such as stability and binding affinity by combining predictor gradients.
• Experimental validation: synthesizing top-ranked variants for closed-loop wet-lab validation, since the current plausibility assessment relies on computational proxies (ESM3 pLDDT, radius of gyration, physicochemical indices) rather than experiments.
• Diverse counterfactual sets: generating alternative mutational strategies rather than a single solution, to enrich understanding of the fitness landscape.
• Regression targets and encoder error: extending beyond binary classification to continuous properties requires replacing the margin loss with MSE or quantile losses, and the CHEAP encoder–decoder's reconstruction error may create artifacts for proteins far from ESMFold's training distribution.
• More principled geometry and spectra: the authors question two implicit assumptions — that plausible proteins lie near a low-dimensional smooth manifold, and that the sequence–function mapping varies smoothly enough for Gaussian smoothing to be valid. They note that protein fitness landscapes are rugged, with higher-order epistasis producing abrupt transitions even between sequences differing by a single residue, and suggest signal filtering tuned to a desired frequency as an alternative to smoothing.

Target Audience

This paper is most valuable to machine learning researchers working on counterfactual explanations and diffusion-guided optimization who want to see the paradigm extended beyond images and tabular data; to computational protein engineers and protein designers seeking sparse, mechanistically interpretable mutation suggestions; and to interpretability researchers interested in active recourse as a complement to passive attribution methods such as attention visualization or sparse autoencoders. Readers should already be comfortable with diffusion models, latent-space optimization, and protein language model representations.

Authors’ abstract

Deep learning models can predict protein properties with unprecedented accuracy but rarely offer mechanistic insight or actionable guidance for engineering improved variants. When a model flags an antibody as unstable, the protein engineer is left without recourse: which mutations would rescue stability while preserving function? We introduce Manifold-Constrained Counterfactual Optimization for Proteins (MCCOP), a framework that computes minimal, biologically plausible sequence edits that flip a model's prediction to a desired target state. MCCOP operates in a continuous joint sequence-structure latent space and employs a pretrained diffusion model as a manifold prior, balancing three objectives: validity (achieving the target property), proximity (minimizing mutations), and plausibility (producing foldable proteins). We evaluate MCCOP on three protein engineering tasks - GFP fluorescence rescue, thermodynamic stability enhancement, and E3 ligase activity recovery - and show that it generates sparser, more plausible counterfactuals than both discrete and continuous baselines. The recovered mutations align with known biophysical mechanisms, including chromophore packing and hydrophobic core consolidation, establishing MCCOP as a tool for both model interpretation and hypothesis-driven protein design. Our code is publicly available at github.com/weroks/mccop.