계산 화학은 분자 시뮬레이션에서 전통적으로 속도-정확도 상충 관계에 직면해 왔습니다. 머신러닝 기반 힘 장(MLFF)은 정확도를 향상시켰지만, 현대 신약 개발에 필요한 방대한 데이터셋으로 인해 여전히 어려움을 겪고 있었습니다. 이를 해결하기 위해 Schrödinger는 Google Cloud와 협력하여 AI 코딩 에이전트인 AlphaEvolve를 사용하여 MLFF 학습 파이프라인을 최적화했습니다. 그들은 이웃 목록 계산과 Ewald 합산이라는 두 가지 핵심 알고리즘을 성능 병목 현상으로 식별했습니다. 주요 목표는 AI 모델 학습을 가속화하는 것이었으며, 특히 PyTorch 코드에서 계산 집약적인 Ewald 합산을 대상으로 했습니다. AlphaEvolve는 이 알고리즘의 보다 효율적인 구현을 생성하는 임무를 맡았습니다. 시스템은 Ewald 합산을 위해 느린 for-loop를 병렬 배치 행렬 곱셈으로 대체하여 PyTorch 코드를 성공적으로 발전시켰습니다. 엄격한 평가를 통해 발전된 코드의 기능적 정확성과 성능 향상이 확인되었으며, 성공률이 크게 증가했습니다. 이러한 최적화는 MLFF 학습 및 추론에서 4배의 속도 향상을 가져왔고, 신약 개발, 촉매 설계 및 재료 개발을 가속화했습니다. Schrödinger는 사용자 정의 GPU 커널 최적화를 위해 이 진화적 접근 방식을 더 탐구할 계획입니다.