Interpreting GFlowNets for Drug Discovery: What probes can and cannot show

기술 요약: 신약 개발을 위한 GFlowNet의 해석

문제 정의
분자 설계(특히 SynFlowNet과 같이 학습된 경우)를 위해 훈련된 생성 흐름 네트워크(Generative Flow Networks, GFlowNets)는 화학적 제약 조건을 준수하면서 화학 공간을 탐색할 수 있는 강력한 프레임워크를 제공한다. 그러나 이들의 내부 결정 정책은 여로 불투명하며, 이는 의약 화학자들이 제안된 구조에 대한 해석 가능한 근거를 필요로 하는 신약 개발 분야에서 채택을 제한하는 요소가 된다. 기존의 해석 가능성 방법들은 연속적인 잠재 공간이나 우도 기반 모델을 위해 설계되었으며, GFlowNet의 이산적이고 순차적인 그래프 기반 정책에는 적합하지 않다. 또한, 정책이 진정으로 학습한 화학적 지식과 그래프 구조 및 원자 피처링(featurization)을 통해 자명하게 얻을 수 있는 정보 사이를 구분해야 할 결정적인 필요성이 존재한다.

방법론
저자들은 특정 체크포인트에서 추출된 동결된 임베딩(frozen embeddings)을 사용하여 SynFlowNet(QED 보상 함수로 학습됨)에 적용되는 다중 척도 해석 가능성 프레임워크를 소개한다. 이 프레임워크는 세 가지 상호 보완적인 접근 방식을 통합한다:

1. 그래디언트 기반 돌출도(Saliency) 및 반사실적 분석(Counterfactuals): 통합 그래디언트(Integrated Gradients, IG)를 "Stop" 액션의 로그 확률에 적용하여 원자 수준의 돌출도 맵을 생성한다. 높은 돌출도를 가진 원자들은 모티프(연결 성분 및 고리 시스템)로 그룹화된다. 이후 RDKit 변환 규칙(예: 할로겐 교환, 아미드 수정)을 사용하여 이러한 모티프에 반사실적 섭동(counterfactual perturbations)을 가하고, 이를 통해 QED 변화량($\Delta$QED)을 측정한다.
2. 희소 오토인코더(Sparse Autoencoders, SAEs):
   • 언더컴플리트(Undercomplete, MLP-SAE): 압축형 오토인코더(256 $\to$ 128)를 사용하여 지배적이고 희소한 잠재 요인을 식별하고, 이들이 물리화학적 특성(예: 극성, 친유성)과 갖는 상관관계를 테스트한다.
   • 오버컴플리트(Overcomplete, BatchTopK SAE): $\ell_1$ 정규화 없이도 중첩(superposition) 상태에 있는 미세한 특징 검출기들을 복구하기 위해 BatchTopK 희소성 메커즘($k=64$)을 사용하는 오버컴플리트 오토인코더(256 $\to$ 1024)를 사용한다.
3. 프로빙(Probing) 및 대조군: 임베딩으로부터 화학적 모티프(작용기, 할로겐 등)와 물리화학적 특성의 존재를 예측하기 위해 얕은 피드포워드 분류기와 엄격한 선형 프로브(Ridge/logistic regression)를 훈련시킨다. 결정적으로, 본 연구는 다음과 같은 엄격한 대조군을 활용한다:
   • 기술자(Descriptor) 베이스라인: 프로브의 성능을 RDKit 기술자 및 ECFP4 핑거프린트와 비교한다.
   • 무작위 레이블 대조군: 프로브가 노이즈를 암기하지 않는지 확인한다.
   • 스캐폴드 불일치 분할(Scaffold-Disjoint Splits): 학습/테스트 분할 시 스캐폴드 누출(leakage)을 방지한다.
   • 미학습 네트워크 대조군: 동일한 아키텍처를 가진 무작위 초기화 네트워크에 동일한 프로브를 실행하여, 훈련의 기여도와 아키텍처의 기여도를 분리한다.
   • 특징 제거(Feature Ablation): 단일 특징을 0으로 설정하여 다운스트림 디코딩에서의 특이성을 테스트한다.

주요 결과

• 돌출도 및 반사실적 분석: 통합 그래디언트는 화학적으로 의미 있는 영역(극성 치환기, 방향족 파편)을 강조한다. 이러한 모티프에 대한 반사실적 편집은 QED의 작지만 측정 가능한 변화(예: +0.015 ~ +0.051)를 유발하지만, 효과 크기는 QED 범위에 비해 작으며, 무작위 편집과의 통제된 비교는 향후 과제로 남겨두었다.
• 잠재 구조: 언더컴플리트 SAE는 크기 및 극성과 강하게 상관관계가 있는 요인들을 식별하지만, 이러한 요인들에 대한 MLP 예측기는 복합 QED($R^2=0.25$)에 대해 개별 특성들에 비해 낮은 성능을 보인다.
• 재구성 및 희소성: 오버컴플리트 BatchTopK SAE는 낮은 NMSE(0.0015)로 임베딩을 재구성하며, 64개의 활성 특징이라는 엄격한 배치 평균 희소성을 유지한다. 이는 유사한 활성 수준에서 매칭된 선형 베이스라인(NMSE 0.0026)보다 우수한 성능을 보인다.
• 선형 디코더빌리티(Linear Decodability): 임베딩은 물리화학적 특성(예: QED $R^2 \approx 0.59$, 크기 $R^2 \approx 0.996$)을 잘 디코딩한다. 그러나 이러한 점수는 무작위 분할 하에서 "누출 없는(leakage-free)" RDKit 기술자 베이스라인이 달성한 점수보다 유의미하게 높지 않다.
• 미학습 네트워크 대조군(결정적 발견): 동일한 아키텍처를 가진 무작위 초기화 네트워크가 훈련된 정책만큼이나 모든 조사된 특성을 잘 디코딩한다. 핵심 속성인 약물 유사성(QED)의 경우, 훈련에 의한 이득은 매우 미미하며(+0.005, 미학습 0.581 vs 훈련 0.586), 이는 무작위 시드(seed)의 분산 범위 내에 있다. 분자 크기의 경우 미학습 네트워크가 약간 더 우수하다.
• 특징 특이성: 오버컴플리트 딕셔너리에서 개별 실행은 화학적으로 풍부한 검출기(예: 요오드, 염소, 브롬, 불소 및 붕소를 구분하는 별도의 특징)를 보여준다. 단일 특징 제거 실험은 특성별 효과를 보여주지만(특이성 비율 8~23배), 특징 인덱스는 시드 간에 재현되지 않으며, 오직 작은 핵심 방향들만이 안정적이다.

의의 및 주장
본 논문의 가장 중대한 결과는 **부정적 대조군(negative control)**이다. 즉, GFlowNet 임베딩으로부터 물리화학적 특성을 선형적으로 디코딩할 수 있다는 사실은 학습된 화학적 지식이 아니라 그래프 아키텍처와 원자 피처링을 반영한다는 것이다.

• 지지되지 않는 사항: 높은 프로브 점수(심지어 완벽에 가까운 AUROC라 할지라도)는 습득된 화학적 이해나 학습된 구조-활성 추론의 증거가 아니다. 특성을 디코딩할 수 있는 능력은 주로 입력 표현과 네트워크 아키텍처의 산물이며, 이는 훈련된 정책이 미학습 정책과 유사한 성능을 보인다는 점을 통해 입증되었다.
• 지지되는 사항:
  1. 오버컴플리트 딕셔너리는 임베딩의 유효하고 정확한 희소 분해를 제공한다.
  2. 개별 실행에는 화학적으로 풍려한 하부 구조 검출기(예: 할로겐별 특징)가 포함되어 있으며, 이는 일부 시드 전반에서 안정적이다.
  3. 이러한 특징들은 프로브 디코딩에 대해 특성별 개입 효과를 가진다.
  4. 약물 유사성(QED)은 밀집된 임베딩보다 희소한 잠재 공간에서 더 선형적으로 접근 가능하다.

결론
저자들은 해석 가능성 방법론이 아키텍처 베이스라인에 대해 엄격히 통제되어야 한다고 결론짓는다. SynFlowNet의 경우, 본 프레임워크는 미세한 특징 검출기를 성공적으로 식별하고 희소 표현의 견고성을 검증하였다. 그러나 본 연구는 높은 디코더빌리티 점수를 학습된 화학의 증거로 해석하는 것에 대해 명시적으로 경고한다. 논문은 딕셔너리 내의 안정적이고 화학적으로 풍부한 특징들의 재현 가능한 핵심 집합의 존재로 주장의 범위를 제한하는 한편, 물리화학적 특성을 예측하는 모델의 능력은 훈련보다는 아키텍처의 기능임을 단언한다. 향후 연구에서는 이러한 검출기들이 미학습 네트워크의 임베딩에 대해 SAE를 훈련했을 때도 지속되는지 탐구하고, 다중 목적 설정으로 프레임워크를 확장해야 한다.