Diamond Maps for Protein Binder Design: Inference-Time Scaling Survives Stochastic Flow Map Distillation

Arthur Liang

Keywords: protein design, flow maps, model distillation

TL;DR: Distilling Proteína-Complexa with Posterior Diamond Maps gives a single-step stochastic posterior sampler that keeps SMC value-function estimates unbiased, so inference-time scaling survives distillation

Abstract: Conditional protein generative models rely on inference-time scaling techniques such as Sequential Monte Carlo and reward-guided sampling to achieve competitive binder design performance. These techniques estimate a value function over candidate trajectories, which requires many evaluations of expensive pretrained models. Flow map distillation could relieve this bottleneck, but standard flow maps are deterministic and yield biased value-function estimates, collapsing the inference-time scaling that makes generative methods competitive with hallucination in the first place. We apply Posterior Diamond Maps to Proteina-Complexa, distilling its joint latent and C$\alpha$ flow into a single-step stochastic sampler of the posterior $p_{1|t}(z \mid x_t)$ that admits consistent, low-cost value-function estimation. We pair the distilled model with Sequential Monte Carlo using AlphaFold-Multimer interface confidence as the reward, and use the PD-L1 binder design problem as proof of concept. On a 5K-step student, deterministic higher-order integration of the teacher saturates well short of FK-steering, while a $(P, K)$ search-width sweep with the distilled SMC monotonically improves every tail metric tracked. These results are consistent with stochastic flow map distillation preserving the inference-time scaling behavior of the underlying generative model.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 208