[논문 리뷰] Grammars and reinforcement learning for molecule optimization

Egor Kraev|arXiv (Cornell University)|2018. 11. 27.

Machine Learning in Materials Science참고 문헌 8인용 수 5

한 줄 요약

이 논문은 SMILES를 위한 맞춤형 문맥 자유 문법(CFG)과 Transformer 기반 강화학습을 조합하여 화학적으로 타당한 분자를 최적의 성질을 갖도록 생성하는 새로운 방법을 제안한다. 규칙 마스킹을 통해價수 규칙과 사이클 제약 조건을 강제함으로써 타당성을 보장하면서도 효율적이고 고품질의 생성을 가능하게 하여 원자당 모델 단계 수를 크게 줄이며 분자 성질 최적화 분야에서 최신 기준 성능을 달성한다.

ABSTRACT

We seek to automate the design of molecules based on specific chemical properties. Our primary contributions are a simpler method for generating SMILES strings guaranteed to be chemically valid, using a combination of a new context-free grammar for SMILES and additional masking logic; and casting the molecular property optimization as a reinforcement learning problem, specifically best-of-batch policy gradient applied to a Transformer model architecture. This approach uses substantially fewer model steps per atom than earlier approaches, thus enabling generation of larger molecules, and beats previous state-of-the art baselines by a significant margin. Applying reinforcement learning to a combination of a custom context-free grammar with additional masking to enforce non-local constraints is applicable to any optimization of a graph structure under a mixture of local and nonlocal constraints.

연구 동기 및 목표

• 딥러닝을 활용하여 원하는 화학적 성질을 갖는 분자 설계를 자동화한다.
• 복잡한 아키텍처나 큰 사전 정의된 어휘집에 의존하지 않고도 화학적으로 타당한 SMILES 문자열만 생성하는 문제를 해결한다.
• 원자당 모델 단계 수를 줄여 샘플 효율성을 향상시킨다.
• 국소적(價수) 및 비국소적(사이클 닫힘, 길이) 제약 조건 하에서 복잡한 분자 성질을 최적화할 수 있도록 한다.
• 더 단순한 문법 기반 강화학습 접근법이 분자 최적화 분야에서 최신 기준 성능을 뛰어넘을 수 있음을 보여준다.

제안 방법

• 화학적으로 타당한 원자價수를 보장하고 복잡한 사이클 구조를 내재적으로 지원하기 위해 SMILES를 위한 새로운 문맥 자유 문법(CFG)을 설계한다.
• 비국소적 제약 조건을 강제하기 위해 추가적인 마스킹 논리를 도입한다: 종료 거리 마스킹은 모델의 단계 수 한계 내에서 종료를 보장하고, 사이클 ID 전파를 통해 사이클 닫힘과 길이 제약 조건을 확보한다.
• 모든 단계에서 전체 맥락에 대한 어텐션을 활용하여 정보 기반 결정을 내리는 Transformer 기반 정책 네트워크를 사용하여 생성 과정을 수행한다.
• 분자 성질 최적화 문제를 보상 함수에 성질 점수, SA 점수, 바람직하지 않은 특성에 대한 페널티를 포함시킨 best-of-batch 정책 기반 강화학습으로 설정한다.
• 표준 CFG 파서(NLTK 등)를 사용하여 실제 분자를 생성 규칙 시퀀스로 분해하여 학습 및 검증에 활용한다.
• 유효성 보장과 성질 최적화를 분리함으로써 대규모 분자에 대한 효율적이고 확장 가능한 생성을 가능하게 한다.

실험 결과

연구 질문

• RQ1 특수한 문맥 자유 문법(CFG)을 설계하여 화학적으로 타당한 분자를 보장하면서도 가지치기와 허유 고리와 같은 복잡한 구조적 특징을 지원할 수 있는가?
• RQ2 국소적 어텐션 메커니즘만을 사용할 때 비국소적 제약 조건(예: 사이클 닫힘, 길이 제한)을 어떻게 효과적으로 생성 과정 중에 강제할 수 있는가?
• RQ3 best-of-batch 정책 기반 강화학습을 사용하는 Transformer 기반 정책 네트워크가 이전의 RNN 또는 그래프 기반 모델보다 원자당 더 적은 모델 단계 수로 분자 성질 최적화에서 뛰어난 성능을 낼 수 있는가?
• RQ4 규칙 마스킹을 통한 문법 기반 생성이 큰 사전 정의된 어휘집이나 복잡한 메시지 전달 아키텍처의 필요성을 얼마나 줄일 수 있는가?
• RQ5 다양한 보상 형태 설계 전략(예: 앵커링, SA 페널티, 허유 고리 페널티)이 생성된 분자의 다양성과 약물 유사성에 어떤 영향을 미치는가?

주요 결과

• 제안된 방법은 ZINC 벤치마크에서 8.46의 분자 점수를 기록하여 Jin 등(2018)이 보고한 이전 최신 기준 점수인 7.16을 크게 초월했다.
• 강력한 앵커링과 SA 페널티를 적용한 결과, 최상위 분자들은 여섯 개의 고리 원자로 이루어진 허유 고리, 황과 질소를 포함한 오목한 다섯 원자 고리, 그리고 단 한 개의 염소 원자만을 포함한 구조적 다양성을 보였으며, 실제 약물 유사 분자와 유사했다.
• 이전 방법들보다 원자당 모델 단계 수가 크게 줄어들어 더 큰 분자의 효율적 생성이 가능했다.
• 명시적인 사이클 페널티 없이도 강력한 앵커링을 적용할 경우, 허유 고리 수가 적고 원자 다양성이 더 균형 잡힌 분자를 생성했다.
• 더 단순한 아키텍처와 문법 기반 설계에도 불구하고, Kusner 등(2017)의 베이스라인과 Jin 등(2018)의 복잡한 그래프 기반 모델을 모두 능가하는 성능을 보였다.
• 사이클 ID를 활용한 규칙 마스킹을 통해 복잡한 제약 조건(예: 최소/최대 사이클 길이)을 강제하면서도 문법적 타당성과 구문 분석 가능성을 유지할 수 있었다.