기술 요약: 이산 확산 모델의 온라인 적응을 위한 설계 공간에 관하여

문제 정의

분자 최적화는 일반적으로 유효한 화학 구조에 대한 광범위한 사전 분포(prior)를 포착하는 사전 학습된 생성 모델(예: 이산 확산 모델)에서 시작됩니다. 그러나 실질적인 목표는 테스트 시점에 이 사전 분포로부터 샘플링하는 것이 아니라, 제한된 오라클 예산(oracle budget)을 사용하여 특정 작업의 보상 함수 $r(x)$(예: 결합 친화도 또는 단백질 적합도)를 최대화하는 분자를 식별하는 것입니다. 보상을 평가하는 오라클은 종종 비용이 많이 듭니다(예: 실험실 분석 또는 고비용 도킹 파이프라인).

핵심 과제는 **온라인 모델 적응(online model adaptation)**입니다. 즉, 사전 학습된 생성기 $p_{\theta_0}$를 유효성과 다양성을 유지하면서 고보상 영역으로 유도하는 것입니다. 이는 다음의 결정들이 밀접하게 결합된 루프를 포함합니다:

1. 획득(Acquisition): 주어진 오라클 예산 내에서 어떤 후보를 평가할 것인가?
2. 적응(Adaptation): 관찰된 보상을 어떻게 모델 업데이트로 변환할 것인가?
3. 피드백 재사용(Feedback Reuse): 역사적 데이터를 얼마나 유지할 것인가?
4. 사전 편향 제거(Prior Debiasing): 유효성 붕괴 없이 생성기를 사전 학습된 모드(mode)로부터 얼마나 멀리 이동시킬 것인가?

이러한 구성 요소들은 각각 개별적으로 연구되어 왔으나, 이들이 전체 온라인 적응 루프에 통합되었을 때 서로 보완하는지, 간섭하는지, 혹은 중복되는지는 여전히 불분명합니다.

방법론

저자들은 이산 확산 모델의 온라인 적응을 위한 모듈형 "하네스(harness)"에 대한 통제된 연구를 제안합니다. 이 루프(그림 1)는 후보 샘PLING, 오라클 평가를 위한 배치 선택, 그리고 생성기 미세 조정(fine-tuning)을 교대로 수행합니다. 본 연구는 9개의 태스크(6개의 저분자 결합 친화도 태스크 및 3개의 단백질 적합도 태스크)에 대해 5가지 특정 "플러그인" 설계 선택지(표 1)를 평가합니다.

주요 설계 구성 요소

1. 스티어링(Steering, 획득):

   • 톰슨 샘플링(Thompson Sampling): 대리 모델(surrogate model)의 평균 예측에 기반한 탐욕적(greedy) 선택 대신, 저자들은 톰슨 샘플링을 사용하여 후보를 선택합니다. 이는 $\hat{r}_i = \mu_i + \epsilon_i \sigma_i$를 샘플링함으로써 탐색(exploration)과 활용(exploitation)의 균형을 맞추며, 여기서 $\sigma_i$는 앙상블 불일치(불확실성)를 나타냅니다. 이는 평균 예측이 낮더라도 불확실성이 높은 후보를 평가하도록 유도합니다.

2. 적응(Adaptation, 보상 형성 및 편향 제거):

   • CVaR 보상 형성: 평균적인 보상이 아닌 고보상 꼬리 부분(tail)에 집중하기 위해, 훈련 보상은 조건부 가치 위험(Conditional Value-at-Risk, CVaR)을 사용하여 형성됩니다. 보상은 $\tilde{r}(x) = (r(x) - Q_\tau)_+ / (1-\tau)$로 변환되며, 여기서 $Q_\tau$는 $\tau$-분위수(0.8로 설정)입니다. 이는 상위 성과를 내는 샘플에 그래디언트 업데이트를 집중시킵니다.
   • 밀도 엔트로피 정규화(Density Entropy Regularization, 모델 편향 제거): 모델이 사전 분포의 이미 고밀도인 모드로 붕괴되는 것을 방지하기 위해 다음과 같은 정규화 항이 추가됩니다: $\log \tilde{r}_{debias}(x) = \log r(x) - \gamma \log p_{\theta_t}(x)$. 이산 확산에서 정확한 주변 가능도(marginal likelihood) $\log p_{\theta_t}(x)$를 계산하는 것은 난해하므로, 저자들은 확산 ELBO(단일 패스 근사)를 기반으로 한 계산 효율적인 몬테카를로 추정치를 사용하여 음의 로그 가능도를 근사합니다.

3. 목적 함수 불일치 수정(Objective-Mismatch Corrections):

   • 리플레이 버퍼(Replay Buffer): 샘플링 분포가 변화함에 따라 초기 단계에서 발견된 고보상 발견물들을 모델이 잊어버리는 것을 방지하기 위해, 우선순위 퇴출(가장 낮은 보상의 샘플을 먼저 제거) 기능이 있는 고정 용량 FIFO 버퍼가 사용됩니다. 층화 샘플링(stratified sampling)을 통해 미니 배치를 추출함으로써 고보상 이상치(outliers)가 희석되지 않도록 합니다.
   • 무효 출력 페널티(Invalid-Output Penalty): 유효성(예: 유효한 SMILES 문자열)을 유지하기 위해, 파싱에 실패한(예: RDKit에 의해) 후보에는 고정된 음수 페널티 $r_{inv}$를 부여합니다. 이는 탐색 과정에서 무효한 구조가 생성될 때 생성기를 다시 유효한 분자 매니폴드로 끌어당깁니다.

실험 설정

• 태스크: 6개의 단백질 타겟에 대한 저분자 최적화(GenMol 생성기와 FlashAffinity/Boltz-2 오라클 사용) 및 3개 제품군(GB1, CreiLOV, TrpB)에 대한 단백질 적합도 최적화.
• 베이스라인: 전체 온라인 루프는 다음 항목들과 비교됩니다:
  • 오프라인 미세 조정(정적 데이터셋에서 훈련).
  • 고정된 생성기를 사용하는 추론 시점 탐색(Beam Search, MCTS, Best-of-N 등).
  • 온라인 루프의 절제된 버전(구성 요소 하나를 하나씩 제거).
• 지표: 성능은 최적 보상, 상위-$k$ 평균, 정규화된 보상을 기준으로 평가되며, 누적 오라클 호출 횟수 및 GPU 시간 대비 측정됩니다. 이차 지표로는 유효성, 다양성(Tanimoto), QED, 합성 용이성(SA)이 포함됩니다.

주요 결과

1. 구성 요소의 상호 보완성

연구 결과, 제안된 구성 요소들은 특히 고보상 후보가 사전 학습된 사전 분포에서 멀리 떨어져 있는 저분자 태스크에서 상호 보완적인 이점을 제공함을 발견했습니다.

• 저분자: 획득(톰슨 샘플링), CVaR 형성, 밀도 엔트로피 정규화가 가장 명확한 보상 이득을 제공합니다. 톰슨 샘플링은 오라클 호출 할당을 개선하고, CVaR은 업데이트를 보상 꼬리에 집중시키며, 편향 제거는 탐색을 사전 분포의 과밀 영역으로부터 밀어냅니다.
• 단백질 적합도: 절제 실험(ablation) 효과가 일반적으로 더 약합니다. 단백질의 경우 가족별 사전 분포가 이미 고적합 영역을 커버하고 있어, 큰 분포 변화가 덜 중요하기 때문입니다. 그러나 톰슨 샘플링은 오라클 할당을 개선함으로써 여전히 일관된 이득을 제공합니다.

2. 안정화 메커즘

• 리플레이 버퍼와 유효성 페널티는 주요 보상 증폭기라기보다 안정화 장치로서 작동합니다. 리플레이 버퍼는 그래디언트를 평가된 데이터의 시간 평균 분포와 정렬시켜 초기 고보상 발견물이 손실되는 것을 방지합니다. 유효성 페널티는 탐색 과정에서 화학적으로 유효하지 않은 문자열을 생성하여 보상을 해킹하는 현상(페널티 없이는 빠르게 발생하는 실패 모드)을 방지합니다.

3. 온라인 적응 vs 추론 시점 탐색

전체 온라인 루프는 동일한 오라클 호출 횟수 및 GPU 시간 계산 하에서 추론 시점 탐색 베이스라인(예: Beam Search, MCTS)을 크게 능가합니다.

• 효율성: 저분자 태스크에서 온라인 루프는 최적의 탐색 베이스라인 점수에 도달하는 데 최적의 탐색 베이스라인보다 약 58배 적은 오라클 호출만 필요했습니다. 단백질 태스크에서는 16배의 속도 향상을 보였습니다.
• 이유: 탐색 방법은 고정된 생성기로부터 후보를 생성하고 점수를 매기는 데 예산을 소비합니다. 반면, 온라인 적응은 오라클 피드백을 사용하여 생성기 자체를 업데이트함으로써 각 오라클 호출의 가치를 복리로 증가시킵니다. 또한 탐색 방법은 생성에 상당한 컴퓨팅 자원을 소비하여, 고정된 실제 시간 예산 내에서 미세 조정 그래디언트 단계를 수행할 시간을 거의 남기지 않습니다.

4. 분포 변화 분석

성능 향상은 최적의 솔루션이 사전 학습된 사전 분포로부터 큰 이동을 필요로 할 때 가장 두드러집니다.

• 밀도 엔트로피 정규화는 생성기의 질량을 고밀도 사전 영역에서 저밀도 고보상 영역으로 밀어내는 주요 동력입니다.
• CVaR 및 톰슨 샘플링은 고보상 후보가 사전 분포와 관련하여 어디에 위치하느냐에 따라 조건부로 이러한 이동을 촉진합니다.
• 단백질 태스크는 가족별 사전 분포가 이미 고적합 영역 근처에 집중되어 있기 때문에 분포 변화가 약하게 나타납니다.

의의 및 주장

본 논문은 피드백 효율적인 분자 최적화를 위한 실용적인 레시피를 제공한다고 주장합니다: 획득, 보상 형성, 편향 제거, 리플레이 및 유효성 제어를 포함한 온라인 미세 조정.

• 독창성: 본 연구는 이러한 구성 요소들을 단일 이산 확산 프레임워크 내에서 통합하고 체계적으로 평가하여, 이들이 중복되는 것이 아니라 상호 보완적임을 입증했습니다.
• 효율성: 제안된 방법은 특히 오라클이 비싸고 타겟 분자가 사전 학습된 모델과 분포가 다른(out-of-distribution) 시나리오에서 오프라인 미세 조정 및 추론 시점 탐색 베이스라인보다 우수한 성능을 보입니다.
• 강건성: 이 레시피는 서로 다른 미세 조정 목적 함수(DDPP-LB 및 VIDD)와 서로 다른 도메인(저분자 vs 단백질)에 걸쳐 강건함을 보여주었으나, 이득의 크기는 사전 분포와 타겟 분포 사이의 거리에 따라 달라집니다.

저자들은 구성 요소들이 (사전 분포가 넓고 타겟이 먼) 저분자 설계에서 가장 중요하지만, 이 프레임워크가 제한된 피드백을 가진 과학적 발견 작업에서 이산 생성 모델을 유도하는 일반화 가능한 접근법을 제공한다고 결론짓습니다.