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

Constraint-Aware Optimization for Robust Protein Stability Prediction
Jun 6, 2026A Shivram, Aneesh S. Chivukula, Manik Gupta +1OOD Generalization Protein Language Models
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
EvoFlows: Evolutionary Edit-Based Flow-Matching for Protein Engineering
Mar 12, 2026Nicolas Deutschmann, Constance Ferragu, Jonathan D. Ziegler +2Protein Design Sequence Modeling
Cradle
Ligand-Conditioned Discrete Diffusion for Protein Sequence-Structure Co-Design
May 15, 2026Chen Wei, Fanding Xu, Minghao Sun +5Protein Language Models Diffusion Language Model Decoding
Xi’an University of Posts & Telecommunications · National University of Singapore · Xi’an Jiaotong University +1

Jun 6, 2026J/K move · Enter open · S save

Constraint-Aware Optimization for Robust Protein Stability Prediction

A Shivram, Aneesh S. Chivukula, Manik Gupta, Sourav Chowdhury

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Constraint-Aware Optimization for Robust Protein Stability Prediction

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

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