[논문 리뷰] Protein Design with Guided Discrete Diffusion (NIPS 2023)

• 이 글에서 다루는 논문: Protein Design with Guided Discrete Diffusion (2023.05, N Fray et al.) https://arxiv.org/abs/2305.20009
• 비슷한 시기에 나온 유사한 논문: Protein Discovery with Discrete Walk-Jump Sampling (2023.06, N Gruver et al.) https://arxiv.org/abs/2306.12360
• 참고하면 좋을 논문: Plug-and-Play Language Models (PPLM) https://arxiv.org/abs/1912.02164 
• Code: https://github.com/ngruver/NOS

---

Introduction

Protein sequence design 문제는 다음과 같은 challenge에 항상 고통받아 왔다.

• 생물학의 오래된 믿음, "Sequence determines structure. Structure determines function."에 따르자면 구조를 기능에 맞게 디자인하는 것이 가장 확실할텐데, 이를 세포가 유전자 단위로 해독하여 생산하려면 결국 서열로 바꿔주어야 한다 (inverse folding). 한가지 구조는 반드시 하나의 서열로만 만들어질 수 있는 건 아니다. (어떤 sequence가 제일 좋은가? 라는 문제를 남긴다). 또한 구조 데이터는 양이 부족하다.
• 그러면 바로 서열을 설계하자니 넓은 sequence space(L 길이 만큼의 서열을 만들자면 20^L의 경우를 탐색해야 한다.)를 검색하기에 느린 방법론들(MCMC)을 써왔고, graident-based 방법론을 도입하자니 discrete space라서 제약이 많다.
• Binding affinity, expression yield와 같은 physiochemical property(function)가 좋은 서열을 설계하고 싶다. 여러 property를 동시에 최적화하는 문제는 pareto optimization이라고도 불리는 단순하지 않은 문제다.

저자들은 각 문제에 맞는 프레임워크를 설계해서 합쳤다!

• Latent Multi-Objective Bayesian Optimization (LaMBO): 서열에서 변경해야 하는 위치(edit position)를 선별하고 변경에 따른 desired property(binding affinity 등) value가 좋아지도록 서열을 샘플링한다.
• DiffusioN Optimized Sampling (NOS): LaMBO-1에서는 서열에서 변경되는 residue 간의 독립성 가정 때문에 같은 residue가 연속적으로 뽑히는 말이 안되는 서열이 샘플링될 가능성이 있다. 멀리 있는 residue 간의 연관성을 고려하여 denoising step을 반복적으로 거치는 controllable categorical diffusion 기반 모델이다.

Methods

Diffusion model은 간단히 표현하자면 데이터를 pure noise에 이를 때까지 조금씩 noise를 첨가해 망가뜨린 다음에, pure noise에서 기존 데이터와 유사하지만 다양한 샘플을 샘플링하는 방법론이다. Diffusion  model을 학습하는 각 과정은 다음의 확률 분포 $p$를 학습하고 사용한다 (많은 디테일이 생략되었다):

$$\text{forward (noise process)}: p(x_t|x_{t-1})\\ \text{reverse (denoising process)}: p(x_{t-1}|x_t,\hat x_0)$$

Diffusion 기반 NOS 모델은 noise의 종류에 따라 두 가지 버전이 있다: NOS-D(Discrete transition), NOS-C(Continuous transition).

NOS-D

Discrete space에서 noise는 각 토큰 $w^{(i)}$를 MASK 토큰으로 대체하는 것을 의미한다. Forward process의 끝에는 pure noise인 $[\text{MASK}]^L$이 남게 된다. 정답(ground truth) 서열에서 denoising procss의 likelihood는 다음과 같고 이를 최대로 만들어주면 된다:

$$L(\theta) = \mathbb{E}_{w_0, t}[-\log p_{\theta}(w_0|w_t)], w_t ~ p(w_t|w_0)$$

위에서 학습한 denoiser $p_\theta(\hat w_0|w_t,t)$를 이용해 reverse process를 정의하면:

$$p_\theta(w_{t-1}|w_t) = \sum{p(w_{t-1}|w_t,\hat w_0)p_{\theta}(\hat w_0|w_t, t)}$$

$p(w_{t-1}|w_t,w_0)$은 Bayes' rule에서 나온 categorical distribution이다. (주: 아마도 우리가 선택한 MASK 토큰 선택의 과정의 chain이라는 뜻?)

NOS-C

$p(w)$를 학습하기 위해 continuous space를 활용하는 방법은 $w$로 표현되는 토큰들을 continuous variable $x_0$으로 임베딩 하는 방법이다. 여기서 prior는 Gaussian distribution $\pi(x) = \mathcal{N}(0, \mathcal{I})$이다.

Forward process:

$$p(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar \alpha_t}x_0, (1-\bar \alpha_t)\mathcal{I}) $$

Loss function:

$$L(\theta) = \mathbb{E}_{w, t}[-\log p_{\theta}(w|x_t)], x_t ~ p(x_t|x_0 = U_\theta w_0)$$

위에서 학습한 denoiser $p_\theta(\hat w|w_t,t)$를 이용해 reverse process를 정의하면:

$$p_\theta(w_{t-1}|w_t) = \sum{p(x_{t-1}|x_t,\hat x_0 = U_\theta \hat w)p(\hat w_0|x_t, t)}$$

$p(x_{t-1}|x_t, x_0)$는 Gaussain distribution이다.

 

LaMBO-2

LaMBO-1로부터 달라진 점은 masked langauge modeling을 사용한 sequence sampler 대신 NOS를 사용과 edit position의 선택을 Gaussian process kernel 대신 ensemble-based uncertainty measure로 대체로 설명할 수 있다. (LaMBO-1를 읽는 것이 더 유익해보인다.)

Results

Sequence recovery는 주어진 정답(ground truth) 서열을 얼마나 맞췄는지를 정량화한 것이다. 하지만 generation task에서 recovery가 얼마나 큰 의미를 주겠는가? OAS database에 있는 서열을 모두 복원할 수 있다면 창의성(새로운 약물)을 발휘할 수 있는 capacity는 전혀 없다는 것을 의미하는가? 하나 고무적인 것은 구조 정보를 학습에 사용하지 않고도 구조 기반 생성 모델(RFDiffusion)과 구조-서열 co-design 모델(DiffAb) 보다 좋은 recovery 성능을 보였다는 것이다.

이 모델의 강력한 장점 중 하나는, multiple-objectives optimization이 가능하다는 것, 원하는 여러 개의 property를 만족하는 서열을 생성할 수 있다는 점이다. "Beta Sheets의 비율"과 "SASA (solvent accessible 면적, surface로 노출된 residue 면적을 의미한다)"을 충족해야 하는 조건으로 주고 서열을 생성했을 때 DiffAb와 MCMC 방법론 등 보다 낮은 energy를 가지는 서열로 수렴했다는 것을 보여준다. 

Naturalness는 자연에 존재하는 서열일 법한지, 세포에서 잘 발현이 될지 등 가능성에 대한 평가 지표이다. Naturalness와 expression yield 및 binding affinity 예측 값을 같이 놓고 보았을 때, 두 가지 특성을 모두 최적화하는 샘플들(Pareto optimal)은 WJS에서 naturalness에 치우치도록 생성한 반면 NOS는 desired property는 우수하지만 naturalness는 낮도록 샘플을 생성한다. 

HER2 항체 라이브러리 wet lab 실험에서도 AbSci에서 발표했던 sucess rate(10%) 보다 더 우수한 binder prediction 결과(25%)를 보였다! Naturalness로 반영되는 expression yield에서도 68개 중 66개인 97%가 발현에 성공했다. 모델은 NOS-D를 사용했다.

Thoughts

• 구조를 암시적으로 latent space에서 인코딩하는 것이 아니라 명시적으로 사용할 방법은 없을까? 구조와 서열 데이터가 다른 방식으로 얻어지고 겹치지 않는 space를 덮고 있다면, 결국 둘다 사용하는 것이 데이터 효율성을 최대로 하는 길이 아닐까? Sequence-structure joint generation을 하지 않더라도 sequence-structure joint embedding을 배우는 것이 sequence(+structure) 생성 태스크에 도움될 것이다.
• Diffusion이 서열 생성에 가장 좋은 방법일까? 1) Pure noise에서 sampling하여 desired sequence를 생성하는 것은 에너지의 낭비일 수 있다. 실제로는 항체와 같은 어느정도 올 수 있는 residue의 확률분포가 정해져있을 수 있다. 항체와 같은 domain-specific 문제를 풀 때 position과 frequency에 대한 prior knowledge를 최대한 활용하는 것이 필요할 수 있다. 2) Conversational AI에서 큰 성공을 거두고 있는 GPT 기반 모델들은 왜 sequence generation에서 이렇다할 성공을 거두지 못하고 있을까? 언어의 구조를 문법이라 하고 단백질의 구조를 3D structure(physical interaction)이라고 하면, 문법을 배우는 것 처럼 language modeling을 통해 zero-shot generation을 할 수 있을까? Conversational AI에 비해 데이터가 너무 적어서 아직은 때가 아닌 걸까?