FragmentNet: Adaptive Graph Fragmentation for Graph-to-Sequence Molecular Representation Learning

이 논문은 분자를 조절 가능한 세분성의 화학적으로 유효한 조각으로 분해하는 새로운 적응형 토크나이저를 사용하는 그래프-시퀀스 모델인 FragmentNet 을 소개하며, 이러한 조각 수준에서의 사전 학습이 전통적인 원자 수준 또는 경직된 규칙 기반 접근법에 비해 하류 특성 예측 성능을 크게 향상시킨다는 것을 입증합니다.

핵심

컴퓨터에게 화학을 이해하도록 가르친다고 상상해 보세요. 전통적으로 과학자들은 컴퓨터에게 분자를 두 가지 주요 방식으로 보게 해 왔는데, 둘 다 결점이 있습니다:

1. "원자 단위" 접근법: 이는 한 글자씩 읽어가며 소설을 이해하려는 것과 같습니다. "t"를 보고, "h"를 보고, "e"를 보지만, 단어 "the"를 완전히 놓쳐버리는 것입니다. 화학에서 이는 컴퓨터가 개별 원자는 보지만, 기능적 부분 (자동차의 엔진이나 도어 핸들처럼) 을 형성하기 위해 어떻게 그룹화되는지 이해하는 데 어려움을 겪음을 의미합니다.
2. "경직된 규칙" 접근법: 이는 사전에 정의된 변경 불가능한 단어만 있는 사전을 사용하는 것과 같습니다. 새로운 유형의 단어가 나타나면 사전이 이를 처리할 수 없습니다. 화학에서는 분자를 조각내기 위해 고정된 규칙을 사용하는 것을 의미합니다. 이는 어느 정도 작동하지만 유연성이 부족하며 자연에서 발견되는 다양한 화학 구조에 적응하지 못합니다.

FragmentNet 등장: "스마트 레고" 접근법

이 논문은 분자에 대해 컴퓨터를 가르치는 새로운 방법인 FragmentNet을 소개합니다. 단일 원자를 보거나 경직된 규칙을 사용하는 대신, FragmentNet은 학습된 적응형 토크나이저를 사용합니다.

분자를 레고 블록으로 지어진 거대하고 복잡한 구조로 생각하세요.

• 기존 방법은 블록의 모든 작은 플라스틱 돌기 (원자) 를 하나씩 보거나, 구조를 몇 가지 미리 정의된 범주로 강제로 맞추려 했습니다.
• FragmentNet은 구조를 보고 스스로 블록을 의미 있는 덩어리로 그룹화하는 법을 학습합니다. 특정 블록 군집이 "바퀴"를 형성하고, 다른 군집이 "좌석"을 형성하며, 또 다른 군집이 "엔진"을 형성한다고 결정할 수 있습니다. 이러한 덩어리가 바로 "분편 (fragments)"입니다.

작동 원리 (세 가지 마법)

1. 그룹화 학습 (적응형 토크나이저):
   모델은 단순히 블록을 그룹화하는 방법을 추측하지 않습니다. 수백만 개의 분자를 연구하여 어떤 원자 그룹이 화학적으로 함께 붙어 있는지 학습합니다. 여기서 "토큰"이 단순히 글자나 원자가 아니라 분자의 화학적으로 유효한 조각 (예: 전체 기능기) 이 되는 맞춤형 사전을 생성합니다. 이는 컴퓨터에게 "c-a-r"을 단순히 보는 것이 아니라, "ing"가 접미사임을 인식하거나 "car"가 어근임을 인식하도록 가르치는 것과 같습니다.

2. 지도 유지 (공간적 위치 인코딩):
   3D 레고 성을 1D 단어 목록 (시퀀스) 으로 변환할 때, 보통 조각들이 서로 상대적으로 어디에 있는지에 대한 정보가 손실됩니다. FragmentNet은 모든 분편에 특별한 "GPS 태그"를 추가하여 이를 해결합니다. 이러한 태그는 컴퓨터에게 "이 엔진 조각은 이 바퀴 조각에 연결되어 있으며, 좌석에서 세 단계 떨어져 있다"고 알려줍니다. 이를 통해 컴퓨터는 분자가 목록으로 평평하게 펼쳐져 있더라도 분자의 모양을 기억할 수 있습니다.

3. "빈칸 채우기" 게임 (마스크된 분편 모델링):
   모델이 매우 똑똑해지기 위해 "매드 립스"나 십자말풀이와 유사한 게임을 합니다.

   • 컴퓨터는 분편으로 구성된 분자를 봅니다.
   • 분편 중 하나를 숨깁니다 (마스크).
   • 주변 문맥을 바탕으로 그 missing 조각이 무엇인지 추측해야 합니다.
   • 단일 원자가 아닌 전체 **덩어리 (분편)**를 추측하기 때문에, 화학의 "문법"을 훨씬 빠르게 학습합니다. "바퀴"와 "좌석"을 보게 되면, missing 조각이 무작위 플라스틱 블록이 아니라 "엔진"일 가능성이 높다는 것을 학습합니다.

논문이 발견한 것

저자들은 이 새로운 방법을 여러 표준 화학 테스트 (약물이 물에 얼마나 잘 녹는지 또는 혈액 - 뇌 장벽을 통과할 수 있는지 예측 등) 에서 기존 "원자 단위" 방법과 비교하여 테스트했습니다.

• 결과: "스마트 레고" 접근법 (FragmentNet) 이 대부분의 경우 승리했습니다.
• 이유: 그것은 문맥을 학습했기 때문입니다. 전체 분편으로 훈련함으로써 컴퓨터는 특정 원자 그룹이 함께 작동한다는 것을 이해하게 되어 더 나은 예측을 이끌어냈습니다.
• 보너스 기능: 논문은 또한 모델이 이러한 덩어리를 이해하기 때문에, 하나의 "레고 덩어리"를 다른 것으로 쉽게 교체하여 새롭고 유효한 분자를 생성할 수 있음을 보여줍니다. 이는 자동차에서 엔진을 제거하고 자동차가 무너지지 않도록 다른 엔진을 끼워 넣는 것과 같습니다.

주의할 점 (한계)

논문은 한계를 솔직하게 인정합니다. 예산 제약으로 인해 이 실험은 단일 노트북 (MacBook Pro) 에서 수행되었습니다. 거대한 AI 모델들이 사용하는 수십억 개에 비해 상대적으로 작은 데이터셋 (200 만 개의 분자) 을 사용했습니다. 또한 "덩어리"의 두 가지 수준 (매우 작은 조각 대 중간 크기 조각) 만 테스트했습니다.

한 줄 요약

FragmentNet은 개별 원자를 응시하는 것이 아니라 의미 있는 "단어 (분편)"를 인식하고 이러한 단어가 어떻게 결합하여 문장을 형성하는지 이해함으로써 컴퓨터에게 화학을 읽는 법을 가르치는 새로운 도구입니다. 이는 컴퓨터를 화학의 훨씬 더 훌륭한 학생으로 만들어 분자의 거동에 대한 더 정확한 예측으로 이어집니다.

기술 요약

기술 요약: FragmentNet

문제 정의

분자 표현 학습은 전통적으로 분자를 개별 원자로 토큰화하거나 경직된 규칙 기반 분해 (예: BRICS) 를 활용하는 데 의존해 왔습니다. 이러한 접근법은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 원자 수준 토큰화는 종종 광범위한 화학적 맥락을 포착하지 못해, 사전 훈련된 모델이 단순한 베이스라인보다 성능이 떨어지는 '부정적 전이 (negative transfer)'를 초래합니다. 개별 원자를 마스킹하면 결합 규칙과 기능기 상호작용의 학습을 방해하는 화학적으로 일관성 없는 환경을 조성할 수 있습니다.
2. 규칙 기반 분해는 유연성이 부족하며 다양한 화학 공간에 걸쳐 일반화하는 데 어려움을 겪습니다.
3. 시퀀스 기반 방법 (예: SMILES 토큰화) 은 분자 그래프에 내재된 중요한 위상 정보를 종종 상실합니다.

그래프에 적용된 기존 마스킹 언어 모델링 (MLM) 전략은 종종 원자를 마스킹하여 화학적 일관성을 깨뜨립니다. 반면, 서브그래프를 마스킹하는 방법 (예: SimSGT) 은 이들 간의 상호작용을 명시적으로 모델링하지 않아 장거리 의존성 포착을 제한합니다.

방법론

저자들은 적응형 학습 토큰화를 통해 그래프 위상과 시퀀스 모델링 간의 간극을 메우도록 설계된 그래프-시퀀스 모델인 FragmentNet을 소개합니다.

1. 적응형 학습 토크나이저

규칙 기반 방법과 달리, FragmentNet 은 분자 그래프를 조절 가능한 세분성을 가진 화학적으로 유효한 분해로 분해하는 데이터 기반 토크나이저를 사용합니다.

• 반복적 쌍별 병합: 토크나이저는 개별 원자로 시작하여 훈련 코퍼스에서 유도된 학습된 병합 이력에 기반하여 연결된 쌍을 반복적으로 병합합니다.
• 세분성 제어: 병합 반복 횟수 ($T$) 가 토큰 크기를 제어합니다. 분자는 재훈련 없이 첫 $t$번의 병합 ($t \le T$) 을 사용하여 토큰화될 수 있으며, 이를 통해 작업별 세분성 최적화가 가능합니다.
• 미결합 결합 처리: 끊어진 결합은 '더미 원자' (원자 번호 0) 로 표현됩니다. 분해는 끊어진 결합의 수와 유형 (예: 하나의 끊어진 단일 결합을 가진 탄소 대 두 개의 끊어진 단일 결합을 가진 탄소) 으로 구분됩니다.
• 고유성: 입체이성질체와 타우토머를 구분하기 위해 저자들은 Weisfeiler-Lehman (WL) 그래프 해싱 알고리즘을 사용하여 비동형 그래프가 고유한 해시를 받도록 보장합니다.

2. 계층적 인코더 (VQVAE + GCN)

이 모델은 하이브리드 인코더를 사용하여 원자 수준 및 분해 수준의 특징을 통합합니다:

• VQ-VAE: 이산적인 원자 수준 특징을 양자화된 잠재 공간으로 인코딩합니다.
• GCN: 이산적 분해 내의 인접 노드에서 특징을 집계하여 구조적 관계를 포착합니다.
• 통합: 원자 임베딩을 평균화하여 분해 표현을 형성한 후, 이를 GCN 출력과 결합하여 압축된 분해 수준 특징 임베딩을 생성합니다.

3. 화학적 인식 공간 위치 인코딩 (SPEs)

그래프를 시퀀스로 직렬화할 때 분자 위상을 보존하기 위해 FragmentNet 은 세 가지 유형의 위치 인코딩을 사용합니다:

• Hop 기반 인코딩: 최단 경로 거리를 통해 상대적 연결성을 포착합니다.
• WL 절대 위치 인코딩: 이성질체를 구분하기 위해 그래프 구조에 기반한 고유한 역할 ID 를 할당합니다.
• 쿨롱 행렬 인코딩: 역제곱 법칙 거리와 원자 전하에 기반한 상호작용을 모델링합니다.
  이러한 것들은 Transformer 에 포괄적인 공간적 맥락을 제공하기 위해 집계됩니다.

4. 마스킹 분해 모델링 (MFM)

사전 훈련 목적은 개별 원자가 아닌 전체 화학적으로 유효한 분해를 마스킹하는 것입니다.

• 과정: 분해가 [MASK] 토큰으로 대체되고, 모델은 마스킹되지 않은 분해의 맥락을 사용하여 원래 분해를 예측합니다.
• 장점: 이는 NLP 에서 다단어 구를 재구성하는 것과 유사하게 화학적으로 의미 있는 맥락을 보존하여 결합 규칙과 기능적 관계의 학습을 용이하게 합니다.
• 구성: 저자들은 맥락을 보존하기 위해 시퀀스당 단일 토큰으로 마스킹을 제한하며, 200 만 개의 분자로 훈련되었습니다.

5. 아키텍처

RDKit 기술자에서 유래된 분자 기술자 CLS 토큰으로 풍부해진 직렬화된 분해 임베딩은 Transformer 인코더에 의해 처리됩니다. 속성 예측 헤드는 하류 작업을 위해 시퀀스 전체에 대한 최대 풀링을 사용합니다.

주요 기여

1. 새로운 학습 적응형 토크나이저: 구조적 연결성을 보존하면서 분자 그래프를 화학적으로 유효한 분해로 분해하는 방법으로, 조절 가능한 세분성을 허용합니다.
2. 공간 위치 인코딩: Hop, WL, Coulomb 인코딩 세트로 분자 그래프 위상을 시퀀스 호환 형식으로 포착하여 효과적인 그래프-시퀀스 모델링을 가능하게 합니다.
3. 세분성에 대한 실증적 연구: 토큰화 세분성이 중요한 설계 선택임을 보여주는 증명입니다. 논문은 분해 수준 토큰화가 MFM 사전 훈련과 결합될 때 대부분의 속성 예측 작업에서 원자 수준 토큰화보다 우월함을 보여줍니다.

결과

모델은 MoleculeNet 및 말라리아 벤치마크에서 스캐폴드 분할 (80-10-10) 을 사용하여 평가되었습니다.

• 사전 훈련 영향: MFM 으로 사전 훈련된 FragmentNet 은 일관되게 사전 훈련되지 않은 모델보다 우월했습니다.
• 분해 대 원자: MFM 사전 훈련을 통해 분해 수준 변형 (100 회 병합 반복) 은 7 개 데이터셋 중 5 개 (BBBP, Tox21, ToxCast, BACE, ESOL, Lipo, Malaria) 에서 원자 수준 변형 (0 회 병합 반복) 보다 우월했습니다. 사전 훈련 없이 원자 수준 토큰화가 종종 더 좋은 성능을 보였으며, 이는 더 거친 토큰화의 이점이 특히 사전 훈련을 통해 해제됨을 시사합니다.
• 해석 가능성: 어텐션 맵은 용해도 (ESOL) 에 대해 하이드록실 그룹에 집중하거나 항말라리아 활성에 대해 퀴나졸린 코어에 집중하는 등 화학적으로 직관적인 패턴을 드러냈으며, 이는 알려진 약리단과 일치합니다.
• 분해 교체: 학습된 토크나이저는 부분 구조 매칭 없이 화학적으로 유효한 유사체 (예: 이부프로펜 수정) 를 생성할 수 있는 분해 교체 모듈을 가능하게 하여 분자 편집에서의 유용성을 입증했습니다.

중요성과 주장

이 논문은 토큰화 세분성이 분자 표현을 개선하기 위한 핵심 레버라고 주장합니다. 원자 수준에서 분해 수준 모델링으로 전환함으로써 FragmentNet 은 원자 수준 마스킹에서 흔히 발생하는 부정적 전이 문제를 해결하고 더 높은 수준의 구조적 모티프를 포착합니다.

저자들은 그들의 접근 방식이 "화학적 정보에 기반"하여 표준 Transformer 모델에 비해 시퀀스 길이를 단축하고 계산 비용을 낮춘다고 강조합니다. 200 만 개의 분자와 작은 어휘로 단일 노트북에서 훈련된 modest 한 설정임에도 불구하고, 사전 훈련된 분해 모델은 사전 훈련되지 않은 변형에 비해 상당한 향상을 보여주었습니다.

이 연구는 적응형 학습 토큰화와 마스킹 분해 모델링의 결합이 분자 표현 학습을 위한 실행 가능하고 효과적인 전략임을 확립하며, 향상된 하류 성능과 향상된 화학적 해석 가능성을 제공합니다. 저자들은 실험 규모 (단일 노트북, 작은 데이터셋) 와 관련된 한계를 인정하며, 향후 작업은 특정 작업에 대한 최적의 세분성을 탐구하고 더 큰 모델과 데이터셋으로 확장해야 한다고 제안합니다.