ConRetroBert: EMA Stabilized Dual Encoders for Template-Based Single-Step Retrosynthesis

ConRetroBert 기술 요약

문제 정의
템플릿 기반 단일 단계 역합성 (retrosynthesis) 은 명시적인 반응 템플릿 (SMARTS 변환으로 인코딩됨) 을 선택하고 적용하여 목표 생성물에 대한 반응물 집합을 예측하는 것을 목표로 합니다. 이 접근법은 합성 계획에 필수적인 고유한 해석 가능성과 추적 가능성을 제공하지만, 템플릿 기반 방법들은 역사적으로 템플릿 프리 모델보다 경쟁력이 낮다고 간주되어 왔습니다. 저자들은 이 한계가 템플릿 사용 자체에서 비롯된 것이 아니라, 반응 규칙의 긴 꼬리 (long-tailed) 라이브러리 전반에 걸친 글로벌 분류라는 지배적인 학습 형식에서 비롯된다고 주장합니다. 이러한 형식은 세 가지 문제를 야기합니다: 빈번한 템플릿이 손실 (loss) 을 지배하고, 소프트맥스 분모는 관찰되지 않았으나 화학적으로 유효한 템플릿을 경쟁적인 부정 레이블로 취급하여 (레이블 불완전성), 출력 레이어는 대규모 템플릿 라이브러리에 따라 잘 확장되지 않습니다.

방법론: ConRetroBert
본 논문은 ConRetroBert를 제시하며, 이는 템플릿 기반 역합성을 글로벌 분류가 아닌 밀집 검색 및 순위 결정 문제로 재정의하는 이중 인코더 프레임워크입니다. 이 아키텍처는 생성물 SMILES 와 템플릿 SMARTS 를 위한 두 개의 독립적인 트랜스포머 인코더를 사용하여 입력을 공유 임베딩 공간에 매핑합니다. 학습 과정은 두 단계로 진행됩니다:

1. 1 단계: 대비적 사전 학습 (Contrastive Pretraining)
   모델은 대칭적 대비적 사전 학습 (InfoNCE 손실) 을 사용하여 결합된 생성물 - 템플릿 표현 공간을 학습합니다. 짝지어진 예제들의 미니배치에 대해, 생성물 인코더는 해당 템플릿과 정렬되고, 배치 내의 다른 예제들은 부정 예제로 작용합니다. 이를 통해 템플릿을 nearest neighbor 검색을 통해 검색할 수 있는 임베딩 공간이 구축됩니다.

2. 2 단계: 다중 긍정 리스트형 순위 결정 (Multi-Positive Listwise Ranking)
   이 단계는 검색 공간을 템플릿 제안 정책으로 정제합니다. 전체 어휘에 대해 분류하는 대신, 모델은 각 생성물에 대해 후보 집합 $C(x)$를 순위 결정합니다. 이 집합은 다음을 포함합니다:

   • 관측된 긍정 (Observed Positives): 학습 데이터에서 생성물과 함께 기록된 모든 템플릿.
   • 하드 부정 (Hard Negatives): FAISS 인덱스에서 검색된 화학적으로 유사한 템플릿.
   • 배치 내 부정 (In-Batch Negatives): 배치 내 다른 생성물과 연관된 템플릿.

   모델은 관측된 긍정을 화학적으로 유사한 하드 부정으로부터 분리하도록 장려하는 다중 긍정 리스트형 부정 로그 가능도 목적 함수를 최적화합니다. 이 형식은 전체 라이브러리 대신 검색된 소수의 후보 중에서 선택해야 하는 추론 시의 의사결정 문제와 더 잘 부합합니다.

EMA 안정화 템플릿 적응
해결된 주요 기술적 과제는 템플릿 인코더가 업데이트될 때 하드 부정 채굴 (mining) 의 불안정성입니다. 템플릿 인코더가 너무 빠르게 변경되면 하드 부정 채굴에 사용되는 검색 인덱스가 구식이 됩니다. ConRetroBert 는 지수 이동 평균 (EMA) 메커니즘을 사용하여 이를 해결합니다:

• 실시간 템플릿 인코더는 순위 결정 손실의 그래디언트를 통해 업데이트됩니다.
• 느리게 이동하는 EMA 쉐도우 인코더 ($\alpha = 0.999$로 업데이트됨) 는 에포크 수준에서 검색 인덱스를 재구축하고 하드 부정을 채굴하는 데 사용됩니다.
  이를 통해 템플릿 표현이 순위 결정 목적에 적응하면서도 후보 구성 과정을 불안정하게 하지 않을 수 있습니다.

주요 기여

1. 검색 기반 형식화: 저자들은 템플릿 기반 역합성을 밀집 생성물 - 템플릿 검색 followed by 후보 집합 리스트형 순위 결정으로 재정의하여, 추론 시 의사결정과 정렬된 목적 함수로 글로벌 분류를 대체했습니다.
2. EMA 안정화 학습: 그들은 EMA 쉐도우에 의해 안정화된 학습 가능한 템플릿 인코더를 도입하여, 하드 부정 채굴을 위한 안정적인 검색 인덱스를 유지하면서 템플릿 표현의 통제된 적응을 가능하게 했습니다.
3. 엄격한 평가 파이프라인: 본 논문은 반응물 집합 순위 결정, 템플릿 적용, 중복 제거, 그리고 USPTO-Full 에서 USPTO-50k 로의 누출 통제 전이를 포함한 포괄적인 평가 프로토콜을 제공합니다.
4. 실증적 검증: 지역 및 긴 꼬리 분석을 통해 저자들은 후보 집합 순위 결정이 명시적 규칙 추적 가능성을 유지하면서 템플릿 기반 역합성을 크게 강화함을 입증했습니다.

결과
실험은 USPTO-50k 벤치마크 (지역 영역) 와 누출 통제된 USPTO-Full 체크포인트로부터의 전이 학습을 통해 수행되었습니다.

• 지역 USPTO-50k 성능:

  • 1 단계 (대비적 검색) 에서 2 단계 (동결된 템플릿 인코더 순위 결정) 로 이동하면 최상위 1 위 반응 정확도가 **50.5% 에서 61.4%**로 향상되었습니다.
  • EMA 안정화 템플릿 적응 (2 단계 EMA Mid) 을 추가하면 최상위 1 위 정확도가 **62.4%**로 추가로 향상되었습니다.
  • 모델은 각각 81.6%, 85.3%, **87.8%**의 최상위 3 위, 5 위, 10 위 정확도를 달성했습니다.
  • 성능 향상은 주로 대비적 정렬만으로는 긍정과 하드 부정을 더 잘 분리하는 2 단계 후보 집합 순위 결정 형식에 기인합니다.

• 전이 학습:

  • 누출 통제된 USPTO-Full 체크포인트에서 미세 조정 (fine-tuning) 한 결과는 USPTO-50k 에서 **75.4%**의 최상위 1 위 정확도에 도달하여, 더 큰 사전 학습 데이터로 접근법의 확장성을 입증했습니다.

• 긴 꼬리 및 신뢰성:

  • 검색 기반 예측은 희귀 템플릿의 긴 꼬리 영역에서 특히 강점을 보였으며, 희귀 클래스에 대한 분류 기반 점수 매김을 능가했습니다.
  • 분석 결과, 올바른 최상위 1 위 반응물 예측 중 **42.8%**는 기록된 정답 레이블과 다른 템플릿으로 생성되었으나 적용 및 중복 제거 후 화학적으로 동등한 "2 차 성공"인 것으로 나타났습니다. 이는 모델이 단순히 특정 기록된 레이블이 아니라 화학적으로 유효한 대안들을 검색하는 법을 학습함을 강조합니다.

의의 및 주장
본 논문은 템플릿 기반 역합성이 해석 가능성을 위해 예측 성능을 희생할 필요는 없다고 주장합니다. 하드 부정 채굴과 EMA 안정화를 통한 검색 기반 순위 결정으로 글로벌 분류를 전환함으로써 ConRetroBert 는 높은 정확도를 위해서는 명시적 반응 템플릿을 포기해야 한다는 가정에 도전합니다. 저자들은 그들의 방법이 최첨단 템플릿 프리 및 세미 - 템플릿 모델과 경쟁력 있는 예측을 생성할 뿐만 아니라, 화학적으로 검사 가능하고 특정 변환 규칙으로 추적 가능하며 하류 합성 계획 제약과 호환된다고 주장합니다. 이 연구는 템플릿 기반 방법의 "약점"이 명시적 규칙 사용의 고유한 한계가 아니라 현대 검색 학습 기술을 통해 해결 가능한 형식화 문제임을 시사합니다.