기술적 요약: 그룹 상대적 정책 최적화를 통한 분자 최적화의 상각 (AMORTIX)

1. 문제 제기

신약 개발에서 분자 설계는 종종 처음부터 분자를 생성하는 대신 고정된 시작 구조를 수정 (예: 스캐폴드 장식, 링커 설계) 하는 것을 요구합니다. 이러한 작업은 휴리스틱 프록시부터 자유 에너지 섭동과 같은 고정밀 시뮬레이션에 이르기까지 비용이 많이 드는 '오라클 (scoring functions)'을 사용하여 평가됩니다.

현재 최첨단 방법들은 주로 인스턴스 최적화 패러다임을 따릅니다. 유전 알고리즘 (Mol GA) 과 안내된 이산 확산 모델 (GenMol) 과 같은 알고리즘은 모든 새로운 입력 구조에 대해 반복적이고 검색 기반의 최적화를 수행합니다. 이 접근 방식은 설계 작업당 수천 번의 오라클 호출을 필요로 하여, 시작 구조의 수에 비례하여 선형적으로 증가하는 높은 계산 비용을 초래합니다.

반면 상각 최적화 (amortized optimization) 는 추론 시 오라클 호출 없이 단일 순전파로 최적화된 분자를 생성할 수 있는 전이 가능한 정책을 학습하는 이론적 해결책을 제시하지만, 기존 방법들은 실제 적용에서 어려움을 겪습니다. 이들은 다양한 구조적制約에 대한 일반화에 실패하거나 학습 불안정성을 겪습니다. 저자들은 근본 원인을 최적화 난이도의 이질성으로 규명했습니다. 일부 시작 구조는 쉽게 고점수 유도체를 생성하는 반면, 다른 구조들은 화학적 제약으로 인해 수정 여부와 관계없이 modest 점수만 산출됩니다. 전역 기준선 (예: REINFORCE) 이나 학습된 가치 크리틱 (예: PPO) 을 사용하는 표준 강화 학습 (RL) 방법은 이러한 분산 하에서 학습을 안정화하지 못하며, 보상 신호가 서로 다른 시작 구조 간에 비교 불가능해지기 때문입니다.

2. 방법론: AMORTIX

저자들은 엄격한 구조적 제약 하에서 추론 시 오라클 호출 없이 분자 구조를 최적화하도록 설계된 상각 그래프 트랜스포머 모델인 AMORTIX를 제안합니다.

2.1 정책 아키텍처

AMORTIX 는 분자 그래프 표현에 직접 작동하는 디코더 전용 그래프 트랜스포머를 사용합니다. 생성 과정은 계층적 행동 공간을 가진 순차적 마르코프 결정 과정 (MDP) 으로 모델링됩니다:

1. 작업 선택: 새로운 원자 추가, 기존 원자 수정, 또는 종료 중 하나를 선택합니다.
2. 대상 선택: 현재 그래프에서 연결할 특정 원자를 선택합니다.
3. 결합 지정: 결합 차수 (단일, 이중, 삼중 등) 를 결정합니다.

화학적 유효성을 보장하기 위해 각 단계에서 유효성 마스크가 적용되어 원자 원자가 규칙을 위반하는 행동을 0 으로 만듭니다. 이는 생성된 모든 구조가 구성상 화학적으로 유효하도록 보장합니다. 모델은 RL 미세 조정 전에 화학적 타당성을 학습하기 위해 ChEMBL 데이터베이스에서 지도식 교사 강제 (supervised teacher forcing) 를 통해 사전 훈련됩니다.

2.2 그룹 상대적 정책 최적화 (GRPO) 를 통한 학습

핵심 혁신은 시작 구조의 이질성을 해결하기 위해 분자 도메인에 그룹 상대적 정책 최적화 (GRPO) 를 적용하는 데 있습니다.

• 그룹 형성: 각 학습 단계에서 시작 구조의 배치 (batch) 가 샘플링됩니다. 각 시작 구조 $S_i$에 대해 정책은 $G$개의 완성된 궤적 (분자) 그룹을 생성합니다.
• 상대적 이점: 전역 기준선이나 가치 크리틱을 사용하는 대신, 구조 $i$의 궤적 $j$에 대한 이점 $A_{i,j}$는 해당 그룹의 평균 보상 $\mu_i$에 대해 그 궤적의 보상 $R(O_{i,j})$을 중심화하여 계산됩니다:
  $$A_{i,j} = R(O_{i,j}) - \mu_i, \quad \text{where } \mu_i = \frac{1}{G} \sum_{k=1}^{G} R(O_{i,k})$$
• 정규화: 저자들은 그룹 표준 편차로 나누는 것과 궤적 길이 정규화를 생략합니다. 이는 모든 완성이 유사하게 점수 (매우 쉽거나 거의 불가능한) 를 받는 구조들이 불균형한 기울기 가중치를 받는 것을 방지하고, 간결한 완성에 대한 패널티를 피합니다.
• 업데이트: 정책은 배치당 단일 기울기 단계로 업데이트되며, 모든 생성 단계에 걸친 기여를 집계합니다. 이 접근 방식은 학습 신호를 구조 간 난이도 변화로부터 분리합니다.

3. 주요 기여

1. 실패 모드 식별: 이 논문은 이질적인 분자 구조 난이도로 인한 보상 신호의 높은 분산이 구조적 제약이 있는 최적화에서 상각 RL 의 주요 실패 모드임을 규명했습니다. 제거 실험 (ablation studies) 은 전역 기준선을 사용한 REINFORCE 나 가치 크리틱을 사용한 PPO 가 이 설정에서 학습을 효과적으로 안정화하지 못함을 입증했습니다.
2. AMORTIX 프레임워크: 제로 추론 시 오라클 호출로 분포 외 (out-of-distribution) 스캐폴드에 일반화되는 그룹 상대적 RL 로 훈련된 그래프 트랜스포머 기반 정책.
3. 성능 우위: AMORTIX 는 목표 지향적 스캐폴드 장식 작업에서 상각 및 인스턴스 최적화 기준선 모두를 능가하며, 여러 시작 구조로 확장될 때 전체 오라클 호출을 훨씬 적게 요구합니다.

4. 실험 결과

저자들은 AMORTIX 를 세 가지 다른 작업에서 평가했습니다:

4.1 PMO 벤치마크 (샘플 효율성)

PMO 벤치마크 (10,000 회 오라클 호출 예산을 가진 22 개 작업) 에서 AMORTIX 는 인스턴스 최적화기로 실행되었습니다 (단일 작업 de-novo 생성에는 시작 구조 이질성이 없으므로 전역 기준선을 사용한 표준 REINFORCE 사용).

• 결과: AMORTIX 는 누적 상위 10 개 AUC 점수 16.433을 달성하여 상각 정책 중 1 위, 모든 패러다임 중 전체 2 위를 기록했으며, 이산 확산 인스턴스 최적화기 GenMol 에 이어 2 위를 차지했습니다.

4.2 목표 지향적 스캐폴드 장식

이 평가는 단일 표적 (GSK3$\beta$, JNK3) 및 다중 표적 (Kinase MPO) 최적화에 대한 보지 않은 스캐폴드 (분포 외) 에 대한 일반화를 테스트합니다.

• 설정: 스캐폴드는 Murcko 분해 및 클러스터링으로 분할되어 테스트 스캐폴드가 훈련 데이터와 위상적으로 구별되도록 했습니다.
• 결과: AMORTIX 는 세 가지 작업 모두에서 최고 점수를 기록하여 상각 기준선 (LibINVENT, GraphXForm, DrugEx v3) 과 인스턴스 최적화기 (Mol GA, GenMol) 를 능가했습니다.
  • Kinase MPO: AMORTIX 는 0.619점을 기록하여 LibINVENT (0.455) 와 인스턴스 최적화기 Mol GA (0.461) 및 GenMol (0.479) 보다 현저히 높은 성능을 보였습니다.
  • 안정성: 제거 실험은 AMORTIX-REINFORCE 와 AMORTIX-PPO 가 안정적인 AMORTIX-GRPO 에 비해 심각한 불안정성 (이점 추정치의 높은 표준 편차) 을 보임을 확인했습니다.
• 비용 효율성: AMORTIX 의 고정 훈련 예산 (약 50,000 회 오라클 호출) 은 약 100 개의 테스트 스캐폴드를 Mol GA 에 대해, 150 개를 GenMol 에 대해 처리한 후 회수됩니다. 이 손익분기점을 넘어서면 AMORTIX 는 누적 오라클 비용을 수배 낮추는 효과를 제공합니다.

4.3 사례 연구: 전구약물 전달

모델은 CNS 침투를 개선하기 위한 전구약물 설계를 위한 수정 규칙을 학습하기 위해 단 4 개의 예시 약물로 미세 조정되었습니다.

• 결과: 52 개의 보지 않은 부모 약물에 적용된 AMORTIX 는 구조적으로 깨끗한 전구약물을 성공적으로 생성했습니다. 명시적 지시 없이도 세 가지 서로 다른 분해 가능한 모티프 (알킬 에스터, 알킬 카바메이트, 벤질 에스터) 를 복원했습니다.
• 성능: 생성물의 31% 는 "순수" 전구약물 (분해 가능한 접합부만 포함) 이었으며, 63% 는 적어도 하나의 분해 가능한 핸들을 포함했습니다. 평균 예측 혈액 - 뇌 장벽 (BBB) 침투 향상도는 +0.26 이었습니다.

5. 중요성 및 주장

이 논문은 AMORTIX 가 분자 최적화의 계산적 부담을 추론 시 검색에서 훈련 시 학습으로 효과적으로 전환할 수 있음을 보여준다고 주장합니다. 동일한 시작 구조를 공유하는 완결 그룹 내에서 보상을 정규화하기 위해 GRPO 를 적응시킴으로써, 이 방법은 이질적인 난이도가 존재할 때 학습을 안정화합니다.

저자들은 AMORTIX 를 반복적 검색 방법에 대한 확장 가능한 대안으로 위치시켜, 반복적인 오라클 평가의 prohibitive 비용 없이 고처리량 설계를 가능하게 합니다. 그들은 훈련 과정이 여전히 오라클 집약적이지만, 정책의 상각적 성격이 많은 시작 구조가 관련된 시나리오에서 매우 효율적임을 지적합니다.

6. 한계

저자들은 몇 가지 한계를 명시적으로 인정합니다:

• 2D 위상: 이 프레임워크는 3D 입체 역학이나 입체화학을 명시적으로 포착하지 않고 2D 분자 위상에서만 작동합니다.
• 연결성: 현재 행동 마스킹은 단일 연결된 시작 스캐폴드를 가정합니다. 분리된 부분 그래프 (예: 분자 조각 연결) 로 확장하려면 아키텍처 수정이 필요합니다.
• 가산적 행동: 생성 과정은 가산적 행동 (원자/결합 추가) 으로 제한됩니다. 스캐폴드 호핑과 같은 작업에 필요한 원자 삭제 또는 치환은 지원하지 않습니다.
• 훈련 비용: 추론은 저렴하지만 훈련 과정은 수만 번의 오라클 호출을 필요로 하므로, 특정 사용 사례에 대한 상각 트레이드오프를 정확히 정량화하기 위해 추가 검증이 필요합니다.