Protein Counterfactuals via Diffusion-Guided Latent Optimization

TL;DR

MCCOP is a novel framework that generates minimal, plausible protein mutations to flip predictive models' outputs, aiding interpretability and targeted protein engineering by balancing validity, proximity, and plausibility in a diffusion-guided latent space.

Contribution

Introduces MCCOP, a diffusion-guided latent optimization method for generating biologically plausible counterfactuals in protein models, advancing interpretability and design capabilities.

Findings

MCCOP produces sparser, more plausible counterfactuals than baselines.

Mutations identified by MCCOP align with known biophysical mechanisms.

Demonstrated effectiveness on GFP, stability, and ligase activity tasks.

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.