[논문 리뷰] Augmented Memory: Capitalizing on Experience Replay to Accelerate De Novo Molecular Design

Jeff J. Guo, Philippe Schwaller|arXiv (Cornell University)|2023. 05. 10.

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

이 논문은 SMILES 데이터 증강과 경험 재생을 조합하여 샘플 효율성을 향상시키는 강화학습 알고리즘인 Augmented Memory를 제안한다. 고성능 오라클 결과를 여러 학습 업데이트에 걸쳐 재사용하고, 저다양성 분자의 선택적 제거를 통해, 10,000회 오라클 호출 예산 하에 PMO 벤치마크에서 이전 방법들보다 23개 작업 중 19개에서 최고 성능을 기록한다.

ABSTRACT

Sample efficiency is a fundamental challenge in de novo molecular design. Ideally, molecular generative models should learn to satisfy a desired objective under minimal oracle evaluations (computational prediction or wet-lab experiment). This problem becomes more apparent when using oracles that can provide increased predictive accuracy but impose a significant cost. Consequently, these oracles cannot be directly optimized under a practical budget. Molecular generative models have shown remarkable sample efficiency when coupled with reinforcement learning, as demonstrated in the Practical Molecular Optimization (PMO) benchmark. Here, we propose a novel algorithm called Augmented Memory that combines data augmentation with experience replay. We show that scores obtained from oracle calls can be reused to update the model multiple times. We compare Augmented Memory to previously proposed algorithms and show significantly enhanced sample efficiency in an exploitation task and a drug discovery case study requiring both exploration and exploitation. Our method achieves a new state-of-the-art in the PMO benchmark which enforces a computational budget, outperforming the previous best performing method on 19/23 tasks.

연구 동기 및 목표

• 비용이 많이 들거나 제한된 평가 횟수로 인해 평가 횟수가 제한되는 새로운 분자 설계에서의 샘플 효율성 문제를 해결하기 위해.
• 경험 재생과 SMILES 증강을 활용하여 REINVENT와 같은 정책 기반 강화학습 모델의 성능을 향상시키기 위해.
• 새로운 선택적 메모리 정리 메커니즘을 통해 샘플 효율성을 향상시키면서도 분자의 다양성을 유지하기 위해.
• Augmented Memory와 BAR의 새로운 구현을 포함하여 Practical Molecular Optimization (PMO) 벤치마크를 확장하여 철저한 비교를 가능하게 하기 위해.
• 모험적 작업과 실제 약물 발견 시나리오 모두에서 뛰어난 최적화 효율성을 입증하기 위해.

제안 방법

• 이 방법은 이전 세대에서 얻은 고성능 분자를 저장하는 재생 버퍼를 사용하여, 오라클 호출 하나당 여러 번의 정책 업데이트를 가능하게 한다.
• 동일한 분자를 다른 문자열 표현 방식으로 생성하기 위해 SMILES 증강을 적용하고, 동일한 오라클 점수를 여러 학습 업데이트에 재사용한다.
• 기존 모델의 로그 확률과 스케일된 오라클 점수를 조합하여 보조 보상이 되는 증강 확률을 계산한다.
• 알고리즘은 각 에포크 동안 두 번의 학습 업데이트를 수행한다: 원본 샘플과 증강된 SMILES에 대해 각각 동일한 보상 신호를 사용한다.
• 재생 버퍼는 반복적으로 업데이트되며, 불필요한 스케일드를 제거하기 위한 다양성 필터를 사용해 선택적으로 정리될 수 있다.
• 최종 손실은 원본 및 증강된 SMILES 시퀀스에서의 업데이트를 모두 통합하여 정책 학습 효율성을 향상시킨다.

실험 결과

연구 질문

• RQ1 엄격한 계산 예산 하에서 경험 재생이 정책 기반 분자 생성에서 샘플 효율성에 크게 기여할 수 있는가?
• RQ2 경험 재생과 결합된 SMILES 데이터 증강이 학습 효율성에 어느 정도 기여하는가?
• RQ3 선택적 메모리 정리 메커니즘이 샘플 효율성 향상을 유지하면서도 분자의 다양성을 유지할 수 있는가?
• RQ4 Augmented Memory는 REINVENT나 BAR와 같은 최고 수준의 방법들과 비교해 복잡한 분자 설계 작업 전반에서 최적화 성능 측면에서 어떻게 비교되는가?
• RQ5 제안된 방법이 모험적 작업과 복잡하고 보상이 흐린 실제 약물 발견 시나리오 모두에 일반화되는가?

주요 결과

• Augmented Memory는 PMO 벤치마크에서 새로운 최고 성능을 달성하여, 이전 최고 성능을 기록한 REINVENT보다 23개 작업 중 19개에서 승리한다.
• SMILES 증강과 경험 재생을 통해 오라클 점수를 여러 학습 업데이트에 재사용함으로써 샘플 효율성이 크게 향상된다.
• 선택적 메모리 정리의 포함으로 다양성이 향상되었으며, 최적화 성능 손실 없이 일찍이 저다양성 솔루션으로 수렴하는 것을 방지한다.
• 알고리즘은 강력한 일반화 능력을 보이며, 통제된 모험적 작업과 복잡한 실제 약물 발견 사례 연구 모두에서 뛰어난 성능을 발휘한다.
• 실험 결과에 따르면, 데이터 증강과 경험 재생을 조합함으로써 10,000회 오라클 예산 내에서 수렴 속도가 빨라지고 보상 누적이 높아진다.
• 광범위한 추상화 및 비교 분석을 통해 경험 재생과 증강의 핵심적인 역할이 샘플 효율적인 분자 설계에서 입증되었다.