그래프 변환 정책 네트워크를 이용한 화학 반응 제품 예측

초록

본 논문은 반응물과 시약을 하나의 그래프로 결합하고, 제품을 생성하는 과정을 그래프 변환 시퀀스로 모델링한다. 그래프 신경망과 강화학습을 결합한 Graph Transformation Policy Network(GTPN)를 제안하여, 사전 정의된 반응 규칙 없이도 반응 트리플(원자쌍, 새로운 결합) 집합을 직접 학습한다. USPTO 대규모 데이터셋에서 기존 최첨단 모델 대비 Top‑1 정확도를 약 3% 향상시켰다.

상세 분석

GTPN은 화학 반응을 “그래프 변환” 문제로 재구성함으로써 두 가지 핵심 능력을 동시에 확보한다. 첫째, Graph Neural Network(GNN)를 통해 입력된 반응물·시약 그래프의 원자와 결합 정보를 고차원 임베딩으로 변환한다. 메시지 패싱 단계에서 각 원자는 이웃 원자와 결합 유형을 집계해 업데이트되며, 이는 기존 GNN 기반 반응 예측 모델과 동일한 구조이지만, 여기서는 중간 그래프가 매 변환 단계마다 갱신된다는 점이 차별점이다. 둘째, 강화학습 기반 정책 네트워크(Policy Network, PN)가 변환 순서를 스스로 결정한다. 행동(action)은 (ξ, u, v, b) 형태의 4‑tuple로, ξ는 “계속” 혹은 “종료” 신호, (u, v)는 변환 대상 원자쌍, b는 새로운 결합 유형을 의미한다. 이러한 행동을 MDP로 정의하고, Advantage Actor‑Critic(A2C) 손실에 즉시 보상과 최종 보상을 결합해 학습한다.

핵심 설계 요소는 다음과 같다.

1. Node Pair Prediction Network(NPPN): 전체 원자쌍 중 반응에 관여할 가능성이 높은 K개를 선별한다. 로컬·글로벌 두 가지 스코어링 방식을 실험했으며, 전역 어텐션 기반 모델이 더 높은 정확도를 보였다. 시약 원자는 변형되지 않으므로, 시약에 속한 원자쌍은 사전에 마스킹한다.
2. Policy Network의 서브액션: ξ를 먼저 예측함으로써 불필요한 (u, v, b) 예측을 방지하고, 변환 단계가 가변적일 수 있게 한다. ξ가 0이면 이후 모든 서브액션을 무시하고 에피소드를 종료한다.
3. 보상 설계: 각 서브액션(신호, 원자쌍, 결합)마다 즉시 보상을 부여하고, 최종 제품이 정답과 일치하면 추가 보상을 제공한다. 이는 정책이 “작은 실수에 민감한” 화학 반응 공간에서 빠르게 수렴하도록 돕는다.
4. 손실 함수 확장: A2C 손실 외에 (i) value 손실, (ii) 원자쌍 이진 손실, (iii) 길이 패널티, (iv) Top‑K 순위 손실을 가중합한다. 특히 Top‑K 손실은 정답 원자쌍이 후보 집합에 포함되도록 강제한다.

학습 과정에서 “첫 번째 오류” 이후의 행동을 무시하도록 ζ 벡터를 도입함으로써, 잘못된 그래프 업데이트가 연쇄적으로 발생하는 문제를 완화한다. 이는 강화학습에서 흔히 발생하는 “오프‑폴리시” 오류를 억제하고, 안정적인 수렴을 가능하게 한다.

실험에서는 USPTO‑15k와 전체 USPTO 데이터셋을 사용했으며, GTPN은 각각 Top‑1 정확도 82.39%와 83.20%를 달성해 기존 최고 성능(약 79~80%)보다 3%p 상승했다. 오류 분석 결과, 대부분의 실수는 (a) 잘못된 원자쌍 선택, (b) 부적절한 결합 유형 예측, (c) 과도한 변환 단계(불필요한 종료 지연)에서 발생한다. 모델은 복잡한 다중 단계 반응에서도 비교적 일관된 변환 순서를 학습했으며, 생성된 트리플 시퀀스를 통해 인간이 반응 메커니즘을 해석할 수 있다.

전반적으로 GTPN은 (1) 그래프 기반 화학 지식을 자연스럽게 통합, (2) 사전 정의된 템플릿 없이도 다양한 반응을 학습, (3) 변환 순서와 길이를 자동으로 결정하는 유연성을 제공한다는 점에서 기존 GNN‑only 혹은 템플릿‑기반 접근법을 능가한다.