[논문 리뷰] Enforcing Predictive Invariance across Structured Biomedical Domains

Wengong Jin, Regina Barzilay외 1명|arXiv (Cornell University)|2020. 06. 06.

Machine Learning in Materials Science참고 문헌 33인용 수 8

한 줄 요약

이 논문은 분자의 골격과 단백질 가족과 같은 복잡하고 구조화된 환경에서 예측의 불변성을 강제함으로써 생물의학 AI 모델의 일반화 성능을 향상시키기 위해 회귀 최소화(RGM) 및 그 구조적 확장형(SRGM)을 제안한다. 일반화를 보유한 환경에 대해 후행적으로 접근 가능한 오라클에 대한 회귀를 최소화하는 방식으로 정의하고, 미리 보지 못한 도메인의 변동을 시뮬레이션하기 위해 기울기 기반의 변형을 사용함으로써, RGM/SRGM는 분자의 성질 예측(41.7 vs. 52.3 MAE, QM9 기준) 및 단백질 안정성 예측(0.793 vs. 0.73 스피어만의 ρ)에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Many biochemical applications such as molecular property prediction require models to generalize beyond their training domains (environments). Moreover, natural environments in these tasks are structured, defined by complex descriptors such as molecular scaffolds or protein families. Therefore, most environments are either never seen during training, or contain only a single training example. To address these challenges, we propose a new regret minimization (RGM) algorithm and its extension for structured environments. RGM builds from invariant risk minimization (IRM) by recasting simultaneous optimality condition in terms of predictive regret, finding a representation that enables the predictor to compete against an oracle with hindsight access to held-out environments. The structured extension adaptively highlights variation due to complex environments via specialized domain perturbations. We evaluate our method on multiple applications: molecular property prediction, protein homology and stability prediction and show that RGM significantly outperforms previous state-of-the-art baselines.

연구 동기 및 목표

• 학습 데이터 분포를 초월한 생물의학 모델의 일반화 문제를 해결함. 특히 분자의 골격, 단백질 슈퍼패밀리와 같이 구조화되어 있고 희소하게 표현된 환경에서의 일반화에 초점한다.
• 기존의 불변성 기반 방법(예: IRM)의 한계를 극복함. 단일 예제나 매우 구조화된 환경에서는 표현 학습에 대한 제약이 약해 성능이 저하되기 때문이다.
• 미리 보지 못한 환경에 대해 후행적으로 접근 가능한 오라클에 대한 성능과의 차이를 정량화함으로써 일반화 성능을 측정하는 회귀 최소화 프레임워크를 개발함.
• 기울기 기반의 표현 변형을 사용해 구조화된 도메인에 대한 회귀 프레임워크를 확장함. 이는 모든 환경를 명시적으로 정의하지 않고도 도메인의 변동을 시뮬레이션할 수 있도록 한다.
• 실제 생물의학 벤치마크에서 제안된 방법의 슈퍼리어리티를 입증함. 분자의 성질 예측, 단백질 안정성 및 동일성 분류 작업을 포함한다.

제안 방법

• 불변 위험 최소화(IRM)를 회귀 최소화 문제로 재정의함: 모델은 보류된 환경에 접근할 수 있는지 여부에 따라 학습된 예측기 간의 차이(회귀)를 최소화한다.
• 학습 데이터의 부분 집합을 샘플링하여 보류된 환경 $E_e$를 정의하고, $E_e$가 포함된 예측기와 포함되지 않은 예측기 간의 손실 차이로 회귀를 계산한다.
• 구조화된 환경(예: 분자의 골격)의 경우, 표현 $\phi$에 대해 별도의 분류기 $g$를 사용해 도메인(예: 골격 또는 구조)을 예측하는 방식으로 기울기 기반 변형을 적용함.
• 변형 $\nabla_\phi \mathcal{L}_{\text{cls}}(g, \phi)$는 골격별 특이한 변동을 강조하며, 변형된 예제에서의 회귀는 새로운 골격 유형으로의 일반화를 측정한다.
• 다양한 보류된 환경 샘플과 변형 방법에 대해 전체적으로 이 회귀를 최소화하도록 메인 모델을 엔드 투 엔드로 최적화한다.
• 희귀하거나 희소한 환경(예: 10개 이하의 예제를 가진 단백질 슈퍼패밀리)을 군집화하여 두 개의 근사 도메인 $E_0, E_1$로 묶어 표준 기준 모델의 훈련을 안정화시키며, SRGM는 기울기 기반 변형을 통해 직접적으로 구조화된 도메인에서 작동한다.

실험 결과

연구 질문

• RQ1 희소하거나 구조화된 환경에서 IRM보다 더 강력한 불변성 유도 성향을 제공할 수 있는가?
• RQ2 모든 환경를 명시적으로 모델링하지 않고도 표현의 기울기 기반 변형이 구조화된 생물의학 도메인에서 일반화를 어떻게 향상시키는가?
• RQ3 분자의 골격, 단백질의 구조 등 구조적 도메인 정보를 활용할 경우 분자의 모델링 및 단백질 모델링 작업에서 일반화 성능가 얼마나 향상되는가?
• RQ4 추론 및 골격 분할 평가 프로토콜에서 기존의 도메인 일반화 기준(예: IRM, MLDG, CrossGrad)을 초월하는가?
• RQ5 학습 및 테스트 분포가 크게 다를 때, 다양한 수준의 도메인 이동에 대해 이 방법은 어떻게 스케일링되는가?

주요 결과

• 골격 분할 조건에서 QM9 데이터셋에서 RGM은 12개의 분자 성질에 대해 평균 MAE를 기존 최고 기준(52.3)에서 41.7로 감소시켜 뚜렷한 일반화 향상을 보였다.
• SRGM은 단백질 안정성 예측 작업에서 0.793의 스피어만 상관계수를 달성하여 이전 최고 기준(0.73)과 ERM 기준(0.736)을 모두 초월했다.
• 동일성 예측 작업에서 SRGM는 23.8%의 정확도를 기록하여 기존 최고 기준(ERM 기준 22.3%)을 뛰어넘었으며, 큰 진화적 격차에 대해서도 강건성을 보였다.
• 절단 실험 결과 SRGM는 도메인 이동 수준이 증가함에 따라 RGM 및 CrossGrad를 일관되게 능가함을 확인했으며, 심각한 분포 이동 조건에서도 효과적임을 입증했다.
• HIV 데이터셋에서 SRGM는 정확도를 0.735로 향상시켜 다음으로 높은 기준인 CrossGrad(0.708)를 크게 앞서며, 구조화된 도메인 모델링의 이점을 입증했다.
• 분자의 성질 예측, 단백질 안정성, 동일성 탐지 등 다양한 생물의학 작업에서 일관된 성능 향상을 보였으며, 이는 낮은 데이터 환경에서의 구조화된 도메인에 대한 일반화 가능성의 타당성을 검증한다.