[논문 리뷰] Protein Discovery with Discrete Walk-Jump Sampling

Nathan C. Frey, Daniel Berenberg외 11명|arXiv (Cornell University)|2023. 06. 08.

Single-cell and spatial transcriptomics인용 수 4

한 줄 요약

이 논문은 매끄러운 에너지 기반 모델과 스코어 기반 노이즈 제거를 조합하여 효율적이고 고성능의 단백질 서열 생성을 가능하게 하는 새로운 비자기적 생성 모델링 프레임워크인 이산 워크점프 샘플링(dWJS)을 소개한다. 단일 노이즈 수준에서 훈련하고, 매끄러운 데이터 다양체 위에서 랑주뱅 MCMC를 사용함으로써 dWJS는 빠르고 다양한 샘플링을 실현하며, 한 번의 실험에서 97–100%의 발현성과 정제도를 보이며, 35%는 결합 친화도가 향상된 항체를 생성한다.

ABSTRACT

We resolve difficulties in training and sampling from a discrete generative model by learning a smoothed energy function, sampling from the smoothed data manifold with Langevin Markov chain Monte Carlo (MCMC), and projecting back to the true data manifold with one-step denoising. Our Discrete Walk-Jump Sampling formalism combines the contrastive divergence training of an energy-based model and improved sample quality of a score-based model, while simplifying training and sampling by requiring only a single noise level. We evaluate the robustness of our approach on generative modeling of antibody proteins and introduce the distributional conformity score to benchmark protein generative models. By optimizing and sampling from our models for the proposed distributional conformity score, 97-100% of generated samples are successfully expressed and purified and 70% of functional designs show equal or improved binding affinity compared to known functional antibodies on the first attempt in a single round of laboratory experiments. We also report the first demonstration of long-run fast-mixing MCMC chains where diverse antibody protein classes are visited in a single MCMC chain.

연구 동기 및 목표

• 단백질 서열 생성에서 이산 에너지 기반 모델(EBM)의 훈련 및 샘플링 과제를 해결한다.
• 항체 설계에 있어 기존 자기적 및 확산 기반 모델의 취약성, 불안정성 및 비효율성을 극복한다.
• 재현성 있는 훈련 기법(예: 리플레이 버퍼 또는 냉각 기법)이 필요로 하지 않는 단일 노이즈 수준만을 요구하는 단순화되고 일반화 가능한 프레임워크를 개발한다.
• 높은 생물물리학적 타당성과 실험 성공률을 갖춘 새로운 기능성 항체를 제자리에서 발견할 수 있도록 한다.
• 단백질 설계에서 생성 모델을 평가하기 위한 새로운 평가 지표인 분포 적합도 스코어(DCS)를 도입한다.

제안 방법

• 이산 데이터 다양체의 매끄러운 형태 위에서 에너지 함수를 학습하는 새로운 패러다임인 매끄러운 이산 샘플링(SDS)을 제안한다.
• 매끄러운 에너지 기반 모델에서 샘플링을 위해 랑주뱅 MCMC를 사용하여 데이터 분포의 효율적 탐색을 가능하게 한다.
• 일단 단계의 노이즈 제거를 적용하여 샘플을 진정한 이산 데이터 다양체로 다시 투영함으로써 서열 유효성을 유지한다.
• 훈련을 별도의 스코어 기반 모델과 에너지 기반 모델로 분리한다: 노이즈가 첨가된 데이터에서 스코어 모델을 훈련하고, 매끄러운 데이터에서 EBM을 훈련한다.
• 신경적 경험 베이즈(NEB) 형식을 활용하여 이산 데이터에 대한 커널 밀도 추정과 경험 베이즈를 통합한다.
• 기존 EBM 정규화 기법(예: 드롭아웃, L2 정규화 등)이 필요 없이도 훈련 안정성과 샘플 품질을 향상시키는 새로운 이산 EBM 아키텍처를 설계한다.

Figure 1: Selected samples from a single Markov chain Monte Carlo sampling run. Protein color corresponds to different antibody germlines (classes). Discrete walk-jump sampling exhibits fast mixing and explores diverse modes of the distribution in a single chain.

실험 결과

연구 질문

• RQ1 단일 노이즈 수준, 비자기적 생성 모델이 자기적 및 확산 기반 베이스라인을 능가하여 기능성 항체 서열을 생성할 수 있는가?
• RQ2 워크점프 샘플링을 활용한 매끄러운 이산 샘플링이 한 개의 MCMC 체인 내에서 다양한 항체 게놈 라인 클래스를 신속하게 혼합하고 탐색할 수 있는가?
• RQ3 제안된 방법이 EBM 훈련에서 리플레이 버퍼, L2 페널티, 시뮬레이션 냉각과 같은 복잡한 훈련 기법이 필요 없음을 제거할 수 있는가?
• RQ4 분포 적합도 스코어(DCS)가 단백질 생성 모델 평가를 위한 신뢰할 수 있는 스칼라 기준이 될 수 있는가?
• RQ5 dWJS는 한 번의 실험 라운드 내에서 얼마나 높은 수준의 새로운 생물물리학적으로 타당하고 기능적으로 향상된 항체를 생성할 수 있는가?

주요 결과

• 생성된 항체 서열 중 97–100%가 in vitro 실험에서 한 번의 라운드 내에 성공적으로 발현되고 정제되었다.
• 생성된 기능성 설계 중 35%가 알려진 기능성 항체와 동일하거나 향상된 결합 친화도를 보였다.
• 장기적으로 빠른 혼합 속도를 가진 MCMC 체인이 실현되어 한 개의 마르코프 체인 내에서 다양한 항체 단백질 클래스를 탐색하였다.
• 기존 EBM 훈련 안정화 기법인 리플레이 버퍼, L2 페널티, 기각 샘플링이 필요 없어졌다.
• dWJS는 샘플 품질과 실험 성공률 측면에서 자기적 대규모 언어 모델, 확산 모델, 스코어 기반 베이스라인을 모두 능가했다.
• 분포 적합도 스코어(DCS)는 생성 모델을 효과적으로 평가했으며 실험적 기능성과 상관관계를 보였다.

Figure 2: Discrete walk-jump sampling. a The noising and denoising process is applied to antibody proteins. b Discrete inputs $x$ are smoothed with isotropic Gaussian noise, $\varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d})$ , to noisy inputs, $y=x+\varepsilon$ . A discrete energy-based model (dEBM) pa