FRIGID: Scaling Diffusion-Based Molecular Generation from Mass Spectra at Training and Inference Time

Motivation

• Known: Tandem mass spectrometry(MS/MS)를 통한 소분자 구조 규명은 retrieval-based 방법과 fingerprint-to-molecule 두 단계 de novo 생성 방법이 주류이며, autoregressive SMILES/SAFE 모델과 graph diffusion decoder들이 활용되어 왔다.
• Gap: 기존 fingerprint-to-molecule decoder는 drug discovery용 아키텍처를 재활용하여 MS/MS 데이터의 fragment-centric 특성을 충분히 반영하지 못하고, autoregressive 모델은 inference-time에서 자동 오류 수정이 어려우며, graph diffusion 모델은 대규모 학습 및 추론 속도 측면에서 확장성이 부족하다.
• Why: 고처리량 구조 규명은 신약 개발, 환경 화학, 대사체학 등 다양한 과학적 발견 워크플로우의 핵심 병목이므로, 학습과 추론 양쪽에서 확장 가능하고 스스로 오류를 교정할 수 있는 de novo 생성 모델은 실용적 파급력이 크다.
• Approach: SAFE 시퀀스를 생성하는 masked diffusion language model(FRIGID-base)을 MIST fingerprint와 chemical formula로 조건화하여 대규모 비표지 구조 데이터로 사전학습하고, ICEBERG forward fragmentation model을 이용해 생성된 구조의 스펙트럼 불일치 fragment를 식별·재마스킹하여 inference-time에 반복적으로 정제한다.

Achievement

Figure 4. Inference-time scaling on NPLIB1 and MassSpecGym.

1. FRIGID-base 프레임워크 제안: SAFE 시퀀스를 생성하는 최초의 diffusion language model 기반 de novo 구조 규명 모델로, 수억 개의 비표지 구조로 사전학습 가능하며 NGBoost로 시퀀스 길이를 예측해 탐색 공간을 제약한다.
2. Inference-time scaling 기법 도입: ICEBERG와 같은 forward fragmentation model을 활용해 hallucinated fragment를 식별하고 해당 subgraph만 targeted remasking·재디노이징함으로써 무작위 마스킹이나 단순 diffusion step 증가보다 표본 효율적인 성능 향상을 달성한다.
3. 벤치마크 성능 대폭 개선: NPLIB1에서 기존 최고 방법 대비 Top-1 정확도를 3배 향상시키고, MassSpecGym에서 18% 이상의 Top-1 정확도(상대적으로 70% 향상)를 달성했다.
4. 로그-선형 스케일링 법칙 발견: 학습 데이터 규모와 inference-time 연산량 모두에 대해 성능이 log-linear하게 향상되는 경향을 실증적으로 보여, 향후 de novo 구조 규명 개선을 위한 새로운 방향을 제시한다.

How

Figure 2. (a) FRIGID-base includes a masked diffusion language model (MDLM) that generates SAFE (Noutahi et al., 2023) s

• MIST 인코더로 실험 스펙트럼을 fingerprint로 변환하고, 별도 formula encoder로 precursor chemical formula를 인코딩한 뒤 cross-attention을 통해 diffusion decoder에 조건 정보로 주입한다.
• Fingerprint-to-SAFE decoder는 masked diffusion language model(MDLM) 구조로, self-attention block과 multi-head attention block, FFN layer들로 구성되며 [M](mask) 토큰을 점진적으로 디노이징해 SAFE 문자열을 생성한다.
• NGBoost로 주어진 chemical formula에 대한 SAFE 시퀀스 길이 분포를 모델링하여 생성 전 적절한 시퀀스 길이를 샘플링한다.
• 생성된 후보 구조는 먼저 chemical formula 일치 여부로 필터링하고, 중간 fingerprint와의 Tanimoto 유사도로 순위를 매긴다.
• Inference-time에서는 ICEBERG forward fragmentation model로 생성 구조의 시뮬레이션 스펙트럼을 예측하고, 실험 스펙트럼과의 불일치(hallucination)를 통해 문제가 되는 SAFE 토큰(subgraph)을 식별한다.
• 식별된 불일치 토큰만 선택적으로 재마스킹한 뒤 diffusion 모델로 재디노이징하는 과정을 반복하여 spectrum-structure cycle consistency를 강화한다.

Originality

• SAFE라는 fragment 기반 표현을 diffusion language model과 결합하여 MS/MS의 fragment-centric 특성을 활용한 최초의 시도이다.
• 기존의 naive inference-time scaling(더 긴 denoising trajectory, 강한 guidance)과 달리, forward fragmentation model(ICEBERG)을 활용해 스펙트럼-구조 일관성 신호를 targeted remasking에 직접 활용하는 새로운 inference-time 오류 교정 메커니즘을 제안한다.
• 학습 및 추론 양 측면에서의 스케일링 법칙(로그-선형)을 체계적으로 규명하여 diffusion 기반 de novo 구조 규명의 확장 가능성을 실증적으로 보여준 최초의 연구이다.

Limitation & Further Study

• 발췌된 내용만으로는 inference-time scaling의 계산 비용(반복 횟수 대비 실제 wall-clock 비용) 및 ICEBERG의 예측 오류가 remasking 품질에 미치는 영향에 대한 심층 분석이 명확히 드러나지 않는다.
• MIST 기반 fingerprint 예측 오류가 파이프라인 전체 정확도에 미치는 상한/하한 영향에 대한 ablation이 부족해 보이며, fingerprint 예측 성능 자체의 개선이 없다면 FRIGID 성능 향상의 한계가 있을 수 있다.
• NPLIB1과 MassSpecGym이라는 특정 벤치마크에 국한된 평가로, 다양한 화합물 클래스(예: 매우 큰 분자, 비정형 골격)에 대한 일반화 가능성은 추가 검증이 필요하다.
• 후속 연구로는 fingerprint 예측기와 diffusion decoder의 공동 학습(joint training), ICEBERG 이외의 다양한 forward fragmentation model과의 결합, 그리고 실제 실험실 환경에서의 검증이 필요하다.

Evaluation

총평: FRIGID은 diffusion language model과 forward fragmentation model을 결합해 de novo 구조 규명의 학습·추론 확장성 문제를 동시에 해결한 인상적인 연구로, 벤치마크 성능 향상 폭이 크고 스케일링 법칙 발견이 후속 연구에 중요한 방향을 제시한다.