기술 요약: CheMatE – 공동 표현 학습을 위한 이중-의미론적 화학 임베더 (Bi-semantic Chemical Embedder)

문제 정의

트랜스포머 모델은 SMILES 문자열을 통해 자연어 처리(NLP) 아키텍처를 화학 분야로 성공적으로 확장해 왔습니다. 그러나 기존의 도메인 적응형 모델들은 두 가지 결정적인 한계에 직면해 있습니다:

1. 도메인 과특화(Domain Over-specialization): SMILES에 대해 집중적으로 사전 학습된 모델은 기초적인 의미론적 능력의 파괴적 망각(catastrophic forgetting)을 겪으며, 다양한 과학적 자연어 작업에서 일반화하는 데 어려움을 겪습니다. 반대로, 일반 언어 모델로부터 적응된 모델은 도메인 특화 사전 학습 과정에서 화학적 이해도가 저하됩니다.
2. 문맥적 제약(Contextual Constraints): 대부분의 분자 임베딩 모델은 짧고 고립된 SMILES 문자열에 최적화되어 있어, 화학적 엔티티가 더 넓은 문맥 내에 등장하는 긴 형태의 과학 문서들을 처리할 수 있는 아키텍처 역량이 부족합니다.

현재의 접근 방식들은 일반적으로 이러한 트레이드오프를 강요하며, 화학적 구문(syntax)이나 일반 과학 언어 중 하나에는 탁월하지만, 두 가지 모두를 동시에 수행하는 경우는 드뭅니다.

방법론

저자들은 ModernBERT 백본을 기반으로 한 긴 문맥 이중-의미론적 인코더인 CheMatE(Matryoshka 임베딩을 이용한 화학 임베더)를 소개합니다. 이 모델은 새로운 2단계 학습 파이프라인을 통해 단일 표현 공간 내에서 분자 구조와 도메인 특화 자연어를 포착하도록 설계되었습니다.

1. 데이터 구축: 대규모 SMILES 주입 (Large-Scale SMILES Injection)

구조적 표기법과 과학적 문맥을 결합한 학습 코퍼스를 만들기 위해, 저자들은 1,440만 개의 문서(ChemPile 논문의 1,145억 토큰 및 필터링된 FineWeb-Education/Web의 1,041억 토큰)를 처리하는 자동화된 파이프라인을 개발했습니다.

• 필터링: 단어 빈도 분석을 기반으로 한 텍스트 화학 점수(Text Chemistry Score, TCS)를 사용하여 화학 관련 콘텐츠를 유지하도록 텍스트를 필터링했습니다.
• 주석 생성 파이프라인: 화학적 엔티티를 주석 처리하기 위한 4단계 프로세스입니다:
  1. NER: Chemical Data Extractor 2(CDE2)를 사용한 화학적 엔티티 탐지.
  2. 변환: OPSIN을 통한 결정론적 IUPAC-to-SMILES 변환.
  3. 조회: 해결되지 않은 이름에 대한 PubChem API로의 폴백(fallback).
  4. 검증: RDKit을 통한 정규화(canonicalization) 및 원자가(valence) 확인.
• 주입: 유효한 SMILES 문자열을 텍스트 내의 해당 화학적 엔티티 언급 바로 뒤에 삽입하여, SMILES 토큰이 긴 문맥의 과학적 서사 내에 자연스럽게 나타나는 혼합 모달리티 코퍼스를 생성합니다.

2. 학습 전략

학습은 두 가지 순차적인 단계로 진행됩니다:

1단계: 지속적 마스크드 언어 모델링 (Continued Masked Language Modeling, MLM)

• 목표: 모델은 표준 MLM 목적 함수(15% 마스킹 비율)를 사용하여 SMILES가 주입된 코퍼스에서 지속적 사전 학습을 수행합니다.
• 효율성 혁신: 멀티 GPU 설정에서 극도로 다양한 시퀀스 길이(최대 8,192 토큰)를 처리하기 위해, 저자들은 BalancedTokenBatchSampler를 개발했습니다. 단순 패딩(padding)은 부하 불균형을 초래하므로, 이 샘플러는 세 가지 제약 조건—원시 토큰 예산(raw-token budget), 패딩된 토큰 메모리 캡(padded-token memory cap), 그리고 이차 비용 천장($\sum L_i^2$)—을 기준으로 배치를 구성합니다. 이를 통해 GPU 간의 계산 부하를 균형 있게 맞추어 패딩 오버헤드를 최대 50% 줄이고 메모리 부족(OOM) 오류를 제거합니다.

2단계: Matryoshka 임베딩을 이용한 대조 학습 (Contrastive Learning with Matryoshka Embeddings)

• 목표: 이중-의미론적 이해를 정교화하기 위해, 모델은 Matryoshka 래퍼가 적용된 **다중 부정 순위 손실(Multiple Negative Ranking Loss, MNRL)**을 사용하여 미세 조정됩니다.
• 합성 쌍 생성: 알고리즘 파이프라인이 주석 처리된 코퍼스에서 직접 앵커-포지티브-네거티브 삼중항(anchor-positive-negative triples)을 생성합니다. 앵커 SMILES를 선택하고, 앵커에 대한 Morgan 핑거프린트의 Tanimoto 유사도를 기반으로 텍스트 세그먼트를 점수화합니다. 유사도가 높은 세그먼트는 포지티브 세트가 되고, 낮은 세가 넥티브 세트가 됩니다.
• Matryoshka 손실: 모델은 하위 차원(768, 512, 256, 128, 64)으로 축소되었을 때도 유익한 임베딩을 생성하도록 최적화되어, 다운스트림 애플리케이션을 위한 유연성을 높입니다.

주요 기여

1. 대규모 SMILES 주입 파이프라인: 과학적 텍스트에 화학적 엔티티를 통합하여 이중-의미론적 학습에 최적화된 219억 토큰의 코퍼스를 생성하는 다단계 자동화 시스템.
2. 비용-예산 분산 배치 샘플러: 단순히 샘리 수(sample counts)가 아닌 이차적 계산 비용을 균형 있게 맞춤으로써, 멀티 노드 GPU 클러스터에서 이질적이고 긴 문맥의 데이터를 효율적으로 학습할 수 있게 하는 새로운 샘플러(BalancedTokenBatchSampler).
3. 이중-의미론적 순차 학습: SMILES를 별개의 모달리티가 아닌 자연어 내의 통합된 토큰으로 취급한 후 대조 학습으로 정교화하는 파이프라인. 이를 통해 단일 인코더가 구조적 화학과 과학적 NLP 모두에 대한 강력한 표현을 학습할 수 있습니다.
4. 통합 벤치마킹: 분자 특성 예측, 재료 과학, 과학적 언어 이해를 아우르는 48개의 다양한 데이터셋(26개 분류, 22개 회귀)에 대한 평가.

결과

CheMatE는 일반 목적 모델(예: SciBERT, ModernBERT) 및 화학 특화 인코더(예: ChemBERTa, MoLFormer, BERT-SMILES)를 포함한 11개의 베이스라인과 비교 평가되었습니다.

• 이중 도메인 성능: CheMatE는 SMILES 작업(1.9)과 NLP 작업(3.2) 모두에서 가장 낮은 평균 순위를 기록했습니다. 이는 두 모달리티 모두에서 동시에 상위 Nemenyi 동등 그룹(equivalence group)에 속한 유일한 화학 특화 모델입니다.
• 절제 연구 (Ablation Study): 대조 학습 단계를 제거했을 때(CheMatE-MLM), 이중-의미론적 점수가 **86.7%**에서 **55.5%**로 크게 떨어졌습니다. 이는 순차적 대조 단계가 교차 모달 성능을 정교화하는 데 매우 중요하다는 것을 확인시켜 줍니다.
• 베이스라인과의 비교:
  • 화학 특화 모델(예: MoLFormer, ChemBERTa)은 SMILES 성능은 좋았으나 NLP 리더보드에서는 최하위에 머물렀습니다.
  • 일반 언어 모델은 NLP 성능은 준수했으나 SMILES 작업에서는 크게 고전했습니다.
  • CheMatE는 이러한 모달리티 트레이드오프를 성공적으로 피하며, SMILES 작업에서 Morgan 핑거프린트보다 우수한 성능을 보였고, NLP 작업에서는 전용 과학 언어 모델과 대등한 성능을 보였습니다.

의의 및 주장

본 논문은 CheMatE가 강력한 SMILES 이해와 과학적 자연어 이해가 인코더 전용 모델에서 서로 배타적이지 않음을 입증한다고 주장합니다. 구조적 신호와 문맥적 신호를 단일 텍스트 기반 프레임워크 내에 통합함으로써, 모델은 전형적인 전문화 트레이드오프 없이 전문 화학 작업과 범용 언어 벤치마크 모두에서 경쟁력 있는 성능을 달와냅니다.

저자들은 CheMatE를 다음과 같은 용도의 다재다능한 기반(foundation)으로 제시합니다:

• 분자 및 자연어 임베딩: 두 모달리티 모두에 대해 고품질의 표현 제공.
• 검색 애플리관: 혼합 의미론적 환경에서의 검색을 위한 백본 역할.
• 미래 생성 모델: 분자 구조가 풍부한 과학적 문맥에 근거하는 검색 증강 생성(RAG) 파이프라인의 기반으로 활용 가능성.

이 연구는 또한 SMILES 주입 파이프라인과 BalancedTokenBatchSampler가 대규모의 이질적인 과학 코퍼스를 학습하는 데 재사용 가능한 도구임을 강조합니다. 저자들은 기존 NER 및 데이터베이스 도구에 대한 의존성(새로운 화합물에 대해 실패할 수 있음)과 대조 학습 데이터셋의 크기가 다소 작다는 한계를 인정하면서도, 본 접근 방식이 과학적 텍스트에 분자 표기법을 직접 병합하여 범용 화학 표현 학습을 수행하는 경로를 유효하게 검증했다고 주장합니다.