기술 요약: 분자 생성과 특성 예측의 결합을 통한 시너지 효과

문제 정의

본 논문은 기계 학습의 오랜 과제인 생성 및 예측 작업에서 동시에 뛰어난 성능을 내는 모델 개발에 초점을 맞춥니다. 약물 설계 분야에서 성공은 미탐사 화학 공간에서 새로운 분자를 생성하는 능력과 이러한 새로운 분자로도 외삽이 가능한 강력한 특성 예측 능력 사이의 균형을 이루는 데 달려 있습니다. 현재 분자 생성과 특성 예측은 주로 분리된 접근 방식으로 다뤄집니다. 트랜스포머 기반 모델은 각 작업 individually(개별적으로) 최첨단 성능을 보이지만, 결합 학습은 상당한 아키텍처 및 최적화 과제를 안고 있습니다. 생성 부분과 예측 부분을 단일 모델로 통합하면 종종 과도한 정규화나 생성 목적과 예측 목적 간의 기울기 간섭이 발생하여, 결합 학습이 가질 수 있는 잠재적 시너지 효과를 모델이 포기하게 만듭니다.

방법론: Hyformer

이러한 과제를 해결하기 위해 저자들은 생성 및 예측 기능을 단일 공유 파라미터 세트 내에서 통합한 트랜스포머 기반 결합 모델인 Hyformer를 제안합니다.

아키텍처

Hyformer는 자기회귀적 트랜스포머 디코더와 양방향 트랜스포머 인코더를 결합합니다. 핵심 혁신은 $[TASK] \in \{[LM], [PRED], [MLM]\}$인 작업 토큰에 의해 제어되는 교번 어텐션 메커니즘에 있습니다:

• **생성 모드 ($[LM]$):** 모델은 인과적 자기 어텐션 마스크 ($M_\rightarrow$) 를 사용하는 디코더로 작동하며, 각 토큰이 왼쪽 문맥에만 주의를 기울이도록 제한합니다.
• **예측/사전 학습 모드 ($[PRED], [MLM]$):** 모델은 전체 문맥 접근을 허용하는 양방향 자기 어텐션 마스크 ($M_\leftrightarrow$) 를 사용하는 인코더로 작동합니다.

이러한 교번 방식은 결합 모델에서 전형적인 기울기 간섭을 완화합니다. 인과적 마스크 하의 소프트맥스 함수의 야코비안은 양방향 마스크 하의 경우와 다르게 어텐션 모듈의 기울기를 변경하여, 업데이트 중 목적 함수들이 충돌하지 않도록 합니다.

학습 전략

Hyformer는 2 단계 학습 프로세스를 사용합니다:

1. 결합 사전 학습: 모델은 unlabeled 데이터를 사용하여 결합 손실 함수로 사전 학습됩니다:
   $$\ell_{Hyformer} = \ell_{LM} + \mu\ell_{MLM} + \eta\ell_{PRED}$$
   여기서 $\ell_{LM}$은 언어 모델링 손실, $\ell_{MLM}$은 마스킹 언어 모델링 손실이며, $\ell_{PRED}$는 분석적으로 계산 가능한 분자 기술자 (예: 분자량, 소수성) 기반의 예측 손실입니다. 작업 간의 균형은 작업 확률 벡터 $p_{[TASK]}$로 제어됩니다.
2. 파인튜닝: 모델은 결합 손실을 사용하여 레이블이 지정된 하류 작업에서 파인튜닝됩니다:
   $$\ell_{Hyformer} = \ell_{LM} + \lambda\ell_{pred}$$
   모델은 파인튜닝 중 작업 확률에 따라 생성 및 예측 작업 사이를 교번하며, 하류 작업에 집중하고 생성 능력을 유지하기 위해 마스킹 언어 모델링 손실을 생략합니다.

샘플링

Hyformer는 무조건부 및 조건부 샘플링을 모두 지원합니다:

• 무조건부: 자기회귀 부분을 사용하여 $x \sim p_\theta(x)$로 분자를 샘플링합니다.
• 조건부: 특성 조건 $Y$를 만족하는 분자를 생성하기 위해, 모델은 $p_\theta(x)$에서 $K$개의 후보를 샘플링하고, 결합 학습된 예측기 $p_\theta(y|x)$가 해당 후보들을 특성 $Y$를 가진 것으로 분류하는 경우에만 이를 수락합니다. 이는 생성과 제어 간의 더 긴밀한 정렬을 위해 통합된 표현 공간을 활용하는 best-of-$K$ 샘플링의 변형으로 간주됩니다.

주요 기여

1. 새로운 결합 모델: 교번 어텐션 메커니즘을 사용하여 단일 파라미터 세트에서 생성 및 예측 작업을 통합한 Hyformer 도입.
2. 시너지 효과 입증: Hyformer가 다음에서 베이스라인을 능가한다는 실증적 증거 제시:
   • 조건부 분자 생성.
   • 분포 외 (OOD) 특성 예측.
   • (프로빙을 통한) 분자 표현 학습.
3. 성능 동등성: Hyformer가 순수 생성 모델 및 순수 예측 모델의 최첨단 성능과 경쟁할 수 있음을 보여, 결합 학습이 개별 작업의 성능을 희생하지 않아도 된다는 것을 입증.
4. 실제 적용 사례: 새로운 항균 펩타이드 (AMP) 발견을 위한 사용 사례를 통해 약물 설계에서 결합 모델링의 유용성 입증.

실험 결과

조건부 분자 생성

GuacaMol 데이터셋에서 평가된 결합 파인튜닝된 Hyformer는 높은 유효성을 유지하면서 목표 특성 (QED, SA, logP) 전반에 걸쳐 가장 낮은 평균 절대 편차 (MAD) 와 표준 편차 (SD) 를 달성했습니다. 이는 비결합 변형 및 MolGPT, GraphGPT 와 같은 베이스라인보다 훨씬 우수한 성능을 보였으며, 결합 파인튜닝이 제어 가능성을 향상시킨다는 것을 확인시켜 주었습니다.

분포 외 (OOD) 특성 예측

Lo-Hi 벤치마크의 Hit Identification 작업에서 Hyformer는 모든 데이터셋 (DRD2, HIV, KDR, Sol) 에서 가장 높은 평균 AUPRC 를 달성하여 지문 기반 베이스라인 (k-NN, GB, SVM, MLP) 과 딥러닝 모델 (Chemformer, Graphformer) 을 능가했습니다. 이는 모델이 학습 중 보지 못한 분자 골격에도 일반화할 수 있음을 보여줍니다.

표현 학습

MoleculeNet 벤치마크의 새로운 프로빙 프로토콜을 사용하여 결합 사전 학습된 Hyformer 에서 추출된 표현은 선형 및 KNN 프로빙 모두에서 가장 예측력이 높았습니다. Hyformer 는 10 개 데이터셋 중 4 개에서 선형 프로빙, 5 개에서 KNN 프로빙에서 최상의 성능을 달성하여 Uni-Mol 과 같은 최첨단 예측 모델 및 기타 결합 모델을 능가했습니다.

무조건부 생성 및 AMP 설계

• 무조건부 생성: GuacaMol 및 MOSES 벤치마크에서 Hyformer 는 분포 충실도 (FCD, KL 발산), 유효성, 그리고 고유성 사이의 균형을 이루며 파레토 프론트에 일관되게 위치하여, 최첨단 SMILES 기반 및 그래프 기반 모델과 경쟁했습니다.
• 항균 펩타이드 설계: 실제 약물 설계 작업에서 Hyformer 는 높은 적합도와 AMP 확률을 가진 펩타이드를 생성하여 PepCVAE, AMPGAN, HydrAMP, AMP-Diffusion 과 같은 베이스라인을 크게 능가하면서도 경쟁력 있는 퍼플렉시티와 다양성을 유지했습니다. 어텐션 분석은 모델이 항균 활성과 관련된 생물학적으로 관련성 있는 아미노산 (아르기닌, 라이신, 트립토판, 프롤린) 을 우선시함을 보여주었습니다.

중요성과 주장

본 논문은 Hyformer 가 SMILES 기반 표현을 위해 분자 생성과 특성 예측을 성공적으로 통합하여 실제 약물 발견 파이프라인에 통합할 수 있는 길을 열었다고 주장합니다. 저자들은 결합 학습의 시너지 효과, 즉 생성 과정에 대한 제어력 향상, OOD 데이터에 대한 강건성, 그리고 고차원 예측 표현의 학습이 개별 작업의 성능을 희생하지 않고 달성 가능함을 강조합니다.

저자들은 공유 파라미터가 작업별 전문화를 제한할 수 있는 내재적 트레이드오프를 인정합니다. 그들은 향후 연구가 용량을 더 유연하게 할당하기 위해 동적이거나 모듈형 어텐션 아키텍처를 탐구할 수 있음을 제안합니다. 또한, 이러한 시너지 효과가 3D 구조나 전사체학과 같은 다른 모달리티에 어느 정도 적용되는지는 여전히 열린 질문이라고 지적합니다. 이 연구는 악의적인 응용을 위한 해결책이 아니라 새로운 잠재적 의약품 설계를 용이하게 하는 도구로 나아가는 한 걸음으로 제시됩니다.