[논문 리뷰] Searching for High-Value Molecules Using Reinforcement Learning and Transformers

Raj Ghugare, Santiago Miret외 3명|arXiv (Cornell University)|2023. 10. 04.

Computational Drug Discovery MethodsComputer Science인용 수 3

한 줄 요약

이 논문은 트랜스포머와 텍스트 기반 분자 표현 방식(SMILES/SELFIES)을 사용하는 강화학습 프레임워크인 ChemRLformer를 소개한다. 이 프레임워크는 정렬된 데이터셋에서의 사전 훈련, 힐클라이밍 리PLAY 버퍼 및 KL 정규화와 같은 타겟된 알고리즘 구성 요소를 결합하여, 복잡한 단백질 도킹 시뮬레이션을 포함한 25개의 분자 설계 과제에서 최신 기술 수준의 성능을 달성한다. 이는 이전 방법들을 단순화한 경험적 분석을 통해 효과적인 설계 선택을 도출한 결과이다.

ABSTRACT

Reinforcement learning (RL) over text representations can be effective for finding high-value policies that can search over graphs. However, RL requires careful structuring of the search space and algorithm design to be effective in this challenge. Through extensive experiments, we explore how different design choices for text grammar and algorithmic choices for training can affect an RL policy's ability to generate molecules with desired properties. We arrive at a new RL-based molecular design algorithm (ChemRLformer) and perform a thorough analysis using 25 molecule design tasks, including computationally complex protein docking simulations. From this analysis, we discover unique insights in this problem space and show that ChemRLformer achieves state-of-the-art performance while being more straightforward than prior work by demystifying which design choices are actually helpful for text-based molecule design.

연구 동기 및 목표

• 텍스트 기반 표현 방식을 사용하는 더 효과적이고 간결한 강화학습 접근법을 개발하기 위해.
• 텍스트 기반 RL에서 성능 향상에 기여하는 아키텍처 및 알고리즘 선택 사항을 규명하기 위해.
• 다양한 데이터셋에서의 사전 훈련과 보상 함수 설계가 분자 최적화 결과에 미치는 영향을 평가하기 위해.
• 트랜스포머가 온라인 RL에서 분자 생성에 대해 안정성과 성능을 어떻게 확보하는지 분석하기 위해.
• 보다 단순하고 분석 기반 설계된 방법이 고가치 분자 발견에서 복잡한 이전 방법들을 능가할 수 있음을 보여주기 위해.

제안 방법

• 텍스트 기반 표현 방식(SMILES 또는 SELFIES 형식)으로 분자를 토큰 단위로 생성하는 트랜스포머 기반 정책 네트워크를 강화학습으로 훈련한다.
• 큰 규모의 정렬된 분자 데이터셋에서의 사전 훈련을 통해 정책를 초기화함으로써 샘플 효율성과 성능을 향상시킨다.
• 훈련 중에 높은 보상을 가진 다양성이 높은 분자를 우선순위로 배치하는 힐클라이밍 전략을 리PLAY 버퍼에 통합한다.
• 사전 훈련된 정책에서의 큰 변화를 방지하기 위해 KL 정규화를 적용함으로써 안정적인 훈련과 효율적인 탐색을 촉진한다.
• 분자 도킹 점수와 약물 유사성 및 합성 가능성 부족에 대한 보상 외부 요소를 결합한 복합 보상 함수를 사용하여 현실적인 분자 생성을 유도한다.
• 계산적으로 비용이 많이 드는 단백질-리간드 도킹 시뮬레이션을 포함한 25개의 다양한 분자 설계 과제에서 프레임워크를 평가한다.

Figure 1 : Autoregressively generating a benzene molecule. : An autoregressive model for sequence generation can be viewed as an RL policy where the actions $a_{t}$ are the next tokens to append to the sequence and the state is the concatenation of all actions taken up to time $t-1$ . A special end-

실험 결과

연구 질문

• RQ1 텍스트 기반 분자 표현 방식(SMILES 대비 SELFIES) 중 어느 것이 RL 기반 분자 생성에서 더 높은 성능을 낼 수 있는가?
• RQ2 정렬된 데이터셋에서의 사전 훈련이 분자 설계에서 RL 에이전트의 샘플 효율성과 최종 성능에 어떤 영향을 미치는가?
• RQ3 리PLAY 버퍼 설계, 정규화, 가능성 페널티 등 알고리즘 구성 요소 중 어떤 것이 텍스트 기반 분자 생성에서 RL 성능 향상에 가장 기여하는가?
• RQ4 완전히 연결된 신경망이 텍스트 기반 RL에서 보상 해킹 행동을 보이는 이유는 무엇이며, 이를 어떻게 완화할 수 있는가?
• RQ5 신경망 아키텍처 선택(예: 트랜스포머 대비 RNN)이 분자 설계의 온라인 RL에서 안정성과 성능에 어떤 영향을 미치는가?

주요 결과

• ChemRLformer는 단백질 도킹 시뮬레이션을 포함한 25개의 분자 설계 과제에서 최신 기술 수준의 성능을 달성하였으며, 더 단순한 아키텍처를 가짐에도 불구하고 이전 방법들을 능가한다.
• 작고 고급 품질의 데이터셋(예: ZINC 1M)에서의 사전 훈련이 더 큰 규모이지만 낮은 품질의 데이터셋에서의 훈련보다 훨씬 뛰어난 성능을 보이며, 데이터 품질이 양보다 더 중요함을 입증한다.
• 힐클라이밍 리PLAY 버퍼 설계가 성능 향상에 상당한 기여를 하였고, 표준 리PLAY 버퍼는 최소한의 이점을 제공한다.
• 사전 훈련된 정책에 대한 KL 정규화는 성능 향상에 기여하지 않으며, 이 설정에서는 효과적인 탐색을 위해 필수적이지 않다는 것을 시사한다.
• 완전히 연결된 신경망은 보상 해킹 행동을 보이며, 보상 함수를 악용하기 위해 긴 탄소 및 질소 원자 시퀀스를 생성하는 비현실적인 분자를 만들어내는 경향이 있다.
• 이론적으로 유리한 점이 있음에도 불구하고, PPO 기반 알고리즘이 분자 최적화에서 바닐라 정책 그라디언트 방법보다 성능이 열 劣하다. 이는 이 분야에서 더 단순한 RL 알고리즘이 더 안정적이고 효과적일 수 있음을 시사한다.

Figure 3 : Comparison between SELFIES and SMILES: The SELFIES representation makes it relatively difficult for ChemRLformer agents to explore effectively leading to generally lower performance on pytdc and docking while scoring higher on diversity. Scores are reported for the transformer model and a