Tokenization for Molecular Foundation Models

기술 요약: 분자 파운데이션 모델을 위한 토큰화(Tokenization)

문제 정의
에너지 저장부터 제약 발견에 이르기까지 다양한 산업 분야에서 정확한 화학적 성질 예측은 매우 중요합니다. 트랜스포머 아키텍처가 자연어 처리(NLP)를 혁신한 것과 달리, 분자 파운데이션 모델에 이를 적용하는 데에는 근본적인 병목 현상인 '토큰화'가 존재합니다. 현재의 분자 모델들은 주로 고정된 어휘집을 사용하여 SMILES(Simplified Molecular Input Line Entry System) 문자열을 원자 단위 토큰으로 분리하는 "원자 단위(Atom-wise)" 토큰화에 의존하고 있습니다.

이 접근 방식의 주요 한계는 OpenSMILES 사양을 완전히 커버할 수 없다는 점입니다. 원자 단위 토크나이저는 대괄호로 둘러싸인 원자(동위원소, 카이랄 중심, 전하 및 명시적 수소 수를 인코딩함)를 하나의 분리 불가능한 토큰으로 취급합니다. 이러한 특징들의 모든 가능한 조합을 모두 포함하려면 28조 개 이상의 토큰이 필요합니다. 결과적으로, 기존 모델들은 3,000개 미만의 어휘집을 사용하며, 이로 인해 상당한 커버리지 공백이 발생합니다. 새로운 대괄호 원자를 마주했을 때, 이러한 폐쇄형 어휘 토크나이저는 일반적인 미지 토큰인 [UNK]를 사용하게 되며, 이는 카이랄성이나 특정 동위원소 조성과 같은 중요한 화학적 정보를 가릴 수 있습니다. 또한, 기존의 개방형 어휘 대안들(BPE 기반 방식 등)은 서로 다른 화학적 실체(예: 황-탄소 결합과 원소 스칸듐)를 동일한 토큰으로 혼동하는 모호성 문제를 겪기도 합니다.

방법론
저자들은 OpenSMILES 사양을 완전히 충족하면서도 계산 효율성을 유지하도록 설계된 새로운 토큰화 프레임워크인 Smirk와 그 압축 변형인 Smirk-GPE를 제안합니다.

1. Smirk 토큰화: 이 방식은 OpenSMILES 사양에 정의된 글리프(glyph)를 기반으로 SMILES 문자열을 2단계 캐릭터 수준 분해를 수행합니다.

   • 1단계: 원자로 분해 (예: OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • 2단계: 대괄호 원자를 구성 글리프로 분해 (예: [C@@H] $\rightarrow$ [ C @ @ H ]).
   • 이 접근 방식은 대괄호와 내부 기호를 별개의 토큰으로 취급함으로써 모호한 시퀀스(예: 결합으로서의 Sc와 원소로서의 [Sc])를 구분합니다. 결과적으로 생성되는 어휘집은 165개의 토큰으로 고정되며, 학습이 필요 없고, 어떤 OpenSMILES 인코딩된 분자라도 [UNK] 토큰 없이 토큰화할 수 있음을 보장합니다.

2. Smirk-GPE (Glyph Pair Encoding): 대괄호 원자를 완전히 분해함으로써 발생하는 시퀀스 길이(fertility) 증가 문제를 해결하기 위해, 저자들은 Smirk-GPE를 구현했습니다. 이 변형은 글리프 토큰에 특화된 바이트 페어 인코딩(BPE) 방식의 압축 전략을 적용합니다. 일반적인 BPE가 문자열을 병합하는 것과 달리, Smirk-GPE는 토큰 ID 상에서 병합 규칙을 학습하여, 화학적으로 의미 있는 병합(예: 황과 탄소 글리프의 결합)이 원자 기호(예: 스칸듐)와 모호하게 겹치지 않도록 보장합니다.

3. 평가 프레임워크:

   • 내재적 지표(Intrinsic Metrics): 저자들은 fertility(평균 시퀀스 길이), 정규화된 엔트로피(압축 효율), 토큰 불균형, 그리고 [UNK] 토큰의 빈도를 사용하여 토크나이저를 평가했습니다.
   • 저비용 프록시(Low-Cost Proxy): 모든 토크나이저에 대해 전체 트랜스포머 모델을 학습시키는 것이 계산적으로 매우 비용이 많이 든다는 점을 인식하여, 저자들은 n-gram 모델을 프록시로 활용했습니다. 이들은 16억 개의 SMILES 문자열로 n-gram 모델을 학습시켰으며, 다운스트림 성능을 추정하기 위해 교차 엔트로피 손실(cross-entropy loss)과 정보 손실(KL-divergence를 통한)을 측정했습니다.
   • 외재적 검증(Extrinsic Validation): n-gram 프록시를 검증하기 위해, 저자들은 11개의 서로 다른 토크나이저와 3개의 분자 인코딩을 사용하여 18개의 인코더 전용 RoBERTa 모델을 처음부터 사전 학습시켰습니다. 이후 이 모델들을 MoleculeNet과 tmQM의 6개 회귀 및 7개 분류 작업에 대해 미세 조정(fine-tuning)했습니다.

주요 결과

• 커버리지: Smirk와 Smirk-GPE는 OpenSMILES 사양의 100% 커버리지를 달-성하여 [UNK] 토큰의 사용을 제거한 유일한 평가 대상 토크나이저였습니다. 반면, 기존의 화학 특화 토크나이저들(SPE, APE 및 다양한 BPE 변형 포함)은 [UNK] 토큰을 무시할 수 없는 빈도로 생성했습니다(tmQM 데이터셋에서 최대 약 50%).
• 정보 손실: 커버리지가 제한적인 토크나이저는 특히 전이 금속과 입체 화학이 풍부한 데이터셋(예: tmQM)에서 상당한 정보 손실을 보였습니다. 예를 들어, MoLFormer 토크나이저는 [UNK] 토큰으로 인해 tmQM에서 40.3 nats/molecule의 손실을 입는 반면, Smirk 변형들은 이러한 저하를 완화합니다.
• 성능 상관관계: 연구 결과, n-gram 지표(교차 엔트로피 및 정보 손실)와 트랜스포머 기반 모델의 다운스트림 성능 사이에 강한 선형 상관관계가 있음이 밝혀졌습니다. 이는 토크나이저 품질을 평가하기 위한 저비용 프록시로서 n-gram의 활용 가능성을 입증합니다.
• 다운스트림 영향:
  • Smirk는 tmQM 데이터셋에서 사전 학습 품질과 다운스트립 성능에 긍정적인 영향을 미쳤습니다.
  • 소형 유기 분자가 주를 이루는 MoleculeNet 작업에서는 Smirk가 표준 원자 단위 토큰화와 유사한 성능을 보였습니다.
  • 커버리지가 낮은 토크나이저(SPE/APE)는 베이스라인에 비해 사전 학습 및 다운스트림 성능 모두에 부정적인 영향을 미쳤습니다.
  • 분자 인코딩(SMILES vs. SELFIES)의 선택은 토크나이저의 선택에 비해 미미한 영향을 미치는 것으로 나타났습니다.

의의 및 주장
본 논문은 화학을 위한 파운데이션 모델이 정보의 손실을 피하기 위해 화학적 공간의 전체 범위를 인코딩해야 한다고 주장합니다. 저자들은 현재의 토크나이저들이 대괄호 안의 정보(예: Cisplatin의 카이랄성이나 특정 동위원소)를 의도치 않게 가림으로써, 단순히 이론적인 문제가 아니라 임상 및 산업적으로 중요한 분자들에 영향을 미칠 수 있는 상당한 정보 손실을 유발한다고 주장합니다.

이 연구의 의의는 다음과 같습니다:

1. 강건성(Robustness): 개방형 어휘 토크나이저(Smirk/Smirk-GPE)가 화학적 공간에 대한 견고한 커버리지를 제공하여, 미지 토큰과 관련된 정보 손실을 방지함을 입증했습니다.
2. 효율성(Efficiency): n-gram 모델이 토크나이저 성능을 평가하는 신뢰할 수 있는 저비용 프록시 역할을 할 수 있음을 확립하여, 하이퍼파라미터 튜닝 및 모델 선택의 계산 부담을 줄였습니다.
3. 해석 가능성(Interpretability): Smirk를 통해 연구자들이 대괄호 내의 정보가 풍부한 내용을 직접 조작할 수 있게 함으로써, 원자 단위 토큰화의 해석 가능성 이점을 확장하는 동시에 원자 단위 토큰화의 위험 요소인 미지 단어 오류를 제거했습니다.

저자들은 현재의 벤치마크(예: MoleculeNet)가 요소 및 입체 화학의 다양성 부족으로 인해 제한된 커버리지를 가진 토크나이저의 결함을 충분히 드러내지 못할 수 있지만, 전체 화학적 공간을 인코딩할 수 있는 토크나이저로 전환하는 것이 신뢰할 수 있는 분자 파운데이션 모델을 위해 필수적이라고 결론짓습니다. 또한, 커뮤니티가 벤치마크의 범위를 엄격하게 평가하고 더 다양한 화학적 특징을 포함하는 데이터셋을 확장할 것을 권장합니다.