Protein Counterfactuals via Diffusion-Guided Latent Optimization

Weronika Kłos; Sidney Bender; Lukas Kades

cs.LGMar 11, 2026

Protein Counterfactuals via Diffusion-Guided Latent Optimization

Authors: Weronika Kłos, Sidney Bender, Lukas Kades

Organizations: Machine Learning Group, Technische Universität Berlin, Berlin, Germany · Berlin Institute for the Foundations of Learning and Data (BIFOLD) · BASF Digital Solutions GmbH, Ludwigshafen am Rhein, Germany

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.

Figures & tables

Figure 1: Overview of MCCOP. A non-fluorescent GFP variant is mapped to a continuous joint sequence–structure latent space via a pretrained autoencoder. After smoothing the classification boundary, the counterfactual embedding is optimized by alternating between (1) a sparse gradient step maximizing target class probability and (2) manifold projection using a pretrained diffusion model (DiMA). The counterfactual shown is a rediscovered sample from the held-out test set (red sticks denote the mutated residue).

Dataset	AUROC ( $\uparrow$ )		Avg. $L_{2}$ norm ( $\downarrow$ )
	Before	After	Before	After
Fluorescence	$0.99\pm 0.00$	$0.99\pm 0.00$	$2.21\pm 0.19$	$1.10\pm 0.10$
Stability	$0.94\pm 0.00$	$0.98\pm 0.01$	$1.38\pm 0.13$	$0.36\pm 0.08$
Activity	$0.82\pm 0.00$	$0.93\pm 0.01$	$0.36\pm 0.06$	$0.33\pm 0.04$

Table 1: Predictor AUROC and average

L_{2}

gradient norm before and after smoothing (mean

\pm

std, 3 seeds).

Dataset	Method	Success Rate ( $\uparrow$ )	Adv. Rate ( $\downarrow$ )	Edit Dist. ( $\downarrow$ )
Stability	Genetic Algorithm	$0.55\pm 0.01$	$0.00\pm 0.00$	$7.76\pm 0.06$
	Gradient Descent †	$1.00\pm 0.00$	$1.00\pm 0.00$	$-$
	Stochastic Hill Climb	$0.23\pm 0.00$	$0.00\pm 0.00$	$9.46\pm 0.18$
	MCCOP (ours)	$\resizebox{1165088}{422343}{{1.00}}\pm 0.00$	$0.03\pm 0.00$	$\resizebox{1165088}{422343}{{2.32}}\pm 0.01$
Fluorescence	Genetic Algorithm	$\resizebox{1165088}{422343}{{0.36}}\pm 0.30$	$0.00\pm 0.00$	$5.37\pm 3.09$
	Gradient Descent †	$1.00\pm 0.00$	$1.00\pm 0.00$	$-$

Table 2: Success rate, adversarial rate, and edit distance (mean

\pm

std, 3 seeds). Edit distance computed on successful counterfactuals only (confidence

\geq 0.95

, edit distance

\geq 1

). Discrete methods cannot produce adversarial examples by construction, so no values are bolded in this column. † Gradient Descent achieves 100% adversarial rate; edit distance is undefined.

Figure 2: Physicochemical plausibility across benchmarks (columns: pLDDT, GRAVY, instability index,

R_{g}

; rows: fluorescence, activity, stability). MCCOP (orange) closely matches the original distribution (gray); discrete baselines show broader shifts. Statistical comparisons via Kruskal-Wallis/Dunn’s tests with Benjamini-Hochberg correction: MCCOP achieves significantly higher pLDDT than both baselines across all tasks (adjusted

p<0.02

Figure 3: Per-residue mutation frequency for fluorescence ( A ) and Ube4b activity ( B ). MCCOP (blue) concentrates mutations in functionally relevant regions – chromophore-proximal residues for GFP, E2-binding interface for Ube4b – while baselines distribute mutations nearly uniformly. Shaded regions: known functional motifs ( Sarkisyan et al., 2016 ; Starita et al., 2013 ) .

Figure 4: Structural alignments between original (gray) and counterfactual (colored) proteins for rediscovered ground-truth examples. ( A ) GFP: mutations near the chromophore. ( B ) Stability: core-facing mutations. ( C ) Ube4b: E2-binding interface mutations. Structures predicted by ESM3.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Average wall-clock time per sample across the complete test set. The genetic algorithm is the most expensive method. MCCOP and stochastic hill climbing have comparable execution times, though their computational profiles differ: discrete baselines spend

>

94% of time on re-encoding candidate sequences, while MCCOP spends 99% on diffusion-based manifold projection.

Figure 6: Physicochemical property distributions for counterfactuals with exactly 3 mutations. MCCOP closely tracks the original test set distribution, while the genetic algorithm and stochastic hill climbing show greater deviation, particularly for GRAVY and instability index.

Symbol	Description	Fluorescence	Stability	Activity
$k$	Top- $k$ masked positions	5	5	5
$\lambda_{\text{dist}}$	$L_{2}$ distance weight	0.1	0.1	0.1
$m$	Margin in $\mathcal{L}_{\text{margin}}$	2.2	2.2	2.2
$\alpha$	Projection strength	0.3	0.3	0.3
$t_{\text{diff}}$	Diffusion noise level	100	100	100
$\eta$	Learning rate	$0.5$	$0.5$	$0.5$

Appendix

Table 3: MCCOP hyperparameters and their values across benchmarks.

Hyperparameter	Symbol	Value
Learning rate	$\eta$	$1\times 10^{-2}$
Gradient steps	$T$	$50$
Confidence threshold	$\tau$	$0.95$
Optimizer	–	Adam
Loss function	$\mathcal{L}$	BCEWithLogitsLoss

Appendix

Table 4: Hyperparameters for the gradient descent baseline.

Hyperparameter	Symbol	Value
Maximum steps	$T$	$50$
Confidence threshold	$\tau$	$0.95$
Amino acid alphabet	$\mathcal{A}$	Standard 20
Mutations per step	–	$1$

Appendix

Table 5: Hyperparameters for the Random Mutation baseline.

Hyperparameter	Symbol	Value
Population size	$N$	$40$
Generations	$G$	$30$
Crossover rate	$p_{c}$	$0.5$
Edit distance penalty	$\lambda$	$0.02$
Confidence threshold	$\tau$	$0.95$
Elite fraction	–	$20\%$

Appendix

Table 6: Hyperparameters for the Genetic Algorithm baseline.

Figure 7: Structural alignment of original (gray) and counterfactual (cyan) stability variants across three topologies. Core-facing mutated residues shown as sticks.

Dataset	Sequences	Positive	Negative	Avg. Length	Binarization
Fluorescence	54,025	30,697	23,328	236.96	Bimodal split
Stability	45,901	22,694	23,207	45.06	Remove middle 33%
Activity	60,692	30,293	30,399	102	Remove middle 33%

Appendix

Table 7: Dataset statistics after preprocessing.

Explore similar work

Jun 6, 2026

Constraint-Aware Optimization for Robust Protein Stability Prediction

Multimodal

ΔΔG

predictors integrating protein language models with inverse-folding representations achieve strong in-distribution accuracy on the Megascale dataset but exhibit limited robustness on out-of-distribution (OOD) proteins, persistent forward-reverse bias on paired-mutation benchmarks, and under-representation of rare stabilizing mutations. Existing approaches address these limitations primarily through additional architectural components, leaving optimization-level intervention comparatively underexplored. We introduce a constraint-aware optimization framework combining Balanced Mean Squared Error, a Siamese anti-symmetric regularizer, and a novel OOD-margin consistency loss on the per-position feature representation, requiring no architectural changes to the SPURS backbone. Across eleven benchmarks and three random seeds, the framework improves Spearman correlation on S669 from 0.486 to 0.540 (

σ=0.002

across seeds), matching the published SPURS baseline (0.50) without architectural modification, and on S461 from 0.653 to 0.711, with consistent smaller gains on five additional OOD datasets. A controlled diagnostic on Ssym reveals that anti-symmetric training does not eliminate systematic forward-reverse bias, indicating that gains arise through implicit regularization rather than exact thermodynamic constraint enforcement.

A Shivram, Aneesh S. Chivukula, Manik Gupta +1

Department of Computer Science and Information Systems, Birla Institute of Technology and Science Pilani, Hyderabad Campus, Hyderabad, 500078, Telangana, India · Department of Biological Sciences, Birla Institute of Technology and Science Pilani, Hyderabad Campus, Hyderabad, 500078, Telangana, India

OOD Generalization Protein Language Models Molecular Property Prediction Consistency Regularization

Ask PDF

Mar 12, 2026

EvoFlows: Evolutionary Edit-Based Flow-Matching for Protein Engineering

We introduce EvoFlows, a variable-length protein sequence-to-sequence modeling approach designed for protein engineering. Existing protein language models are poorly suited for optimization tasks: autoregressive models require full sequence generation, masked language and discrete diffusion models rely on pre-specified mutation locations, and no existing methods naturally support insertions and deletions relative to a template sequence. EvoFlows learns mutational trajectories between evolutionarily related protein sequences via edit flows, allowing it to perform a controllable number of mutations (insertions, deletions, and substitutions) on a template sequence, predicting not only which mutation to perform, but also where it should occur. Through extensive in silico evaluation on diverse protein families from UniRef and OAS, we show that EvoFlows generates variants that remain consistent with natural protein families while exploring farther from template sequences than leading baselines.

Nicolas Deutschmann, Constance Ferragu, Jonathan D. Ziegler +2

Cradle

Protein Design Sequence Modeling Flow-Based Generative Modeling Flow Matching

Ask PDF

May 15, 2026

Ligand-Conditioned Discrete Diffusion for Protein Sequence-Structure Co-Design

Proteins perform their biological functions through three-dimensional structures encoded by amino acid sequences, and ligand-binding protein co-design requires models that generate sequence-structure compatible proteins under explicit ligand constraints. Although continuous diffusion and flow-based models support ligand-aware design in coordinate or latent spaces, existing discrete diffusion protein language models mainly operate over sequence or structure tokens without direct small-molecule conditioning. We introduce \textbf{ProtLiD

^2

}, a Protein Ligand-conditioned Discrete Diffusion model for protein sequence-structure co-design. ProtLiD

^2

jointly generates amino-acid sequence and discrete structure tokens while incorporating ligand chemical and geometric information through geometry-aware cross-attention. Trained on over one million ligand-protein complexes, ProtLiD

^2

extends masked discrete diffusion to ligand-aware functional protein design. We further propose maximum confidence-margin guided ReMask decoding, an inference-time self-correction strategy that retains confident predictions and remasks uncertain tokens. ProtLiD

^2

improves global fold confidence over Complexa in whole-protein design, increasing TM-score from 0.672 to 0.802 and pLDDT from 64.55 to 73.00. In pocket co-design, ProtLiD

^2

reduces active-site BB-RMSD from 3.46/3.40Å for FAIR/PocketGen to 1.97Å, and improves ligand-aware pass rates over PocketGen from 14.86% to 59.73% and from 6.08% to 23.49% under stricter docking thresholds. These results support ligand-conditioned discrete diffusion as an effective token-space framework for functional protein co-design. Code will be available at https://github.com/auroua/ProtLiD.

Chen Wei, Fanding Xu, Minghao Sun +5

Xi’an University of Posts & Telecommunications · National University of Singapore · Xi’an Jiaotong University +1

Protein Language Models Diffusion Language Model Decoding Molecular Generation Diffusion Language Models Protein Design

Ask PDF