Technical Summary: Learning Implicit Bias in Generative Spaces for Accelerating Protein Dynamics Emulation

1. Problem Formulation

Generative emulators of protein dynamics offer a computationally efficient alternative to Molecular Dynamics (MD) for reproducing short trajectories. However, these models inherit the limitations of their training distributions. Under long-horizon extrapolation, they tend to revisit known states (local minima) rather than exploring rare, biologically significant conformations separated by high free-energy barriers.

Classical enhanced sampling methods (e.g., metadynamics, umbrella sampling) address this by biasing MD simulations toward underexplored regions, but they incur substantial per-system simulation costs. While recent generative models (e.g., score-based diffusion with SE(3)-equivariant architectures) have enabled trajectory-level emulation, existing approaches to encourage exploration often require injecting guidance via collective variables (CVs) or modifying the training objective, which can compromise the model's fidelity or require retraining.

The core problem addressed is how to accelerate long-horizon exploration in a pretrained protein dynamics emulator without retraining the emulator itself, while maintaining structural validity and avoiding drift off the data manifold.

2. Methodology

The proposed framework introduces a history-dependent bias directly into the generative space of a frozen pretrained emulator. The system consists of three main components:

A. History-Aware Score Estimator (Implicit Bias)

Instead of modifying the frozen emulator, a separate, trainable history-aware score estimator is introduced.

• Mechanism: The estimator takes the current noisy state $x_t$, diffusion time $t$, sequence features $c$, and a history bank $H$ of previously generated structures.
• Architecture: It employs $H$ parallel branches, each trained to denoise the current state toward a specific history structure $h_j \in H$.
• Bias Calculation: The estimator computes a bias score $s_{bias}$ as a weighted sum of the differences between the emulator's score ($s_{emu}$) and the branch-specific scores ($s_{his,j}$).
  $$s_{bias}(x_t, t, H, c) = \sum_{j=0}^{H-1} w_j \left( s_{emu}(x_t, t) - s_{his,j}(x_t, t, H, c) \right)$$
• Weighting: Weights $w_j$ are determined by a softmax over the Kabsch-aligned C$\alpha$ RMSD between the emulator's denoised prediction and the branch's target, prioritizing branches closest to the emulator's current trajectory.
• Effect: The bias points away from directions that would denoise toward previously visited conformations, effectively steering the reverse-time sampling toward unexplored regions.

B. Training with Environment-Support Regularization

To ensure the biased score remains valid and does not drift the sampling process off the emulator's data manifold, the training objective includes two terms:

1. Per-branch Diffusion Loss ($L_{diff}$): Supervises each branch to accurately denoise toward its specific history target.
2. Environment-Support Regularizer ($L_{support}$): Minimizes the Kullback-Leibler (KL) divergence between the biased reverse transition kernel and the original frozen emulator's kernel. This ensures that the biased score proposes next states that the frozen emulator still considers reliable, preventing out-of-distribution degradation.

C. Score-Based Refinement for Long Horizons

Long-horizon rollouts with strong bias can accumulate drift, degrading structural validity. To counter this without fine-tuning the emulator, a refinement step is applied at inference:

• Generated frames are forward-diffused to a small noise level ($t_{ref}$).
• They are then re-denoised using the frozen emulator's score (without bias) for a short number of steps.
• This re-projects drifted samples back onto the data manifold, restoring structural validity while preserving the exploration gains achieved by the bias.

3. Key Contributions

1. History-Dependent Generative Bias: A novel method to inject exploration into a frozen generative model by learning a bias conditioned on rollout history, analogous to metadynamics but operating in the generative space.
2. Manifold-Preserving Training: The introduction of an environment-support regularizer that keeps the learned bias within the support of the pretrained emulator during training.
3. Inference-Time Refinement: A score-based refinement step that corrects structural drift in long-horizon rollouts, enabling the use of stronger biases without sacrificing validity.
4. Empirical Validation: Demonstration of the method's efficacy across short- and long-timescale regimes on both in-distribution (DynamicPDB-80) and zero-shot (Fast-Folding) benchmarks.

4. Experimental Results

Short-Timescale (DynamicPDB-80)

• Diversity: The method increased ensemble diversity (mDiv-mean) by 35% compared to the unbiased frozen emulator.
• Validity: Structural validity (C$\alpha$ and peptide bond distances) was largely preserved, with only a minor drop (approx. 1.5 percentage points) compared to the baseline, which remained significantly higher than other trajectory generation baselines like MDGen.
• Transferability: The bias mechanism successfully transferred to a different pretrained backbone (ConfRover), confirming it is not specific to the emulator architecture.

Long-Timescale (Fast-Folding Proteins, Zero-Shot)

Evaluated on 12 proteins held out from training, simulating microsecond horizons:

• Speedup: The learned bias alone reached the coverage of the unbiased emulator up to 15$\times$ faster.
• Bias + Refinement: Pairing the bias with the refinement step achieved coverage up to 37$\times$ faster than the unbiased baseline.
• State Coverage: The combined approach covered approximately 3$\times$ as many low-energy states (free energy < 4 kcal/mol) compared to the unbiased rollout.
• Validity Recovery: Without refinement, high bias strengths ($\eta=0.4$) caused validity to collapse. With refinement, validity was restored to >90% even at high bias strengths, enabling the 37$\times$ speedup.

Ablation Studies

• Refinement Noise ($t_{ref}$): Validity recovery peaked at $t_{ref}=0.2$, balancing the repair of broken geometries with the preservation of diversity.
• Regularizer ($L_{support}$): Removing the environment-support regularizer caused a sharp drop in structural validity at high bias strengths.
• History Length ($H$): Increasing $H$ improved diversity but reduced validity; $H=3$ was selected as the optimal trade-off.
• Weighting: Distance-aware weighting (Eq. 3) significantly outperformed uniform aggregation.

5. Significance and Claims

The paper claims that this approach successfully decouples exploration acceleration from model retraining. By operating in the generative space of a frozen emulator, the method leverages the high-fidelity dynamics of the pretrained model while overcoming its tendency to revisit known states.

The authors emphasize that:

• The method acts as a "plug-in" accelerator for existing emulators.
• It preserves the structural validity required for biological analysis, which is often lost in aggressive exploration methods.
• It enables the exploration of rare conformational states (metastable basins) that are inaccessible to standard rollouts within reasonable computational budgets.

Limitations noted by the authors:

• The bias is currently unconditional; it cannot target a specific endpoint conformation.
• The reachable conformational space is bounded by the frozen emulator's training distribution; the method redistributes density within this support but cannot generate structures entirely outside it.
• The base emulator does not target the equilibrium distribution, complicating direct free-energy reconstruction via methods like MBAR without further post-hoc processing.

The work positions itself as a step toward controllable, efficient sampling of protein conformational landscapes, with potential applications in drug discovery (e.g., cryptic pocket detection) and structural biology.