Look Before You Leap: Improving Structure-Based Drug Optimization with Attribution-Guided Genetic Operators

Motivation

• Known: Graph 기반 분자 특성 예측 모델(D-MPNN 등)은 높은 정확도로 후보 분자를 스코어링하는 데 널리 쓰이며, GB-GA와 같은 진화 알고리즘은 2D 분자 그래프에 직접 crossover/mutation을 적용해 훨씬 복잡한 생성 모델들과 견줄 만한 성능을 보인다. 또한 SHAP, integrated gradients, GNNExplainer 등 attribution 기법들은 분자 모델의 예측을 해석하는 post-hoc 도구로 널리 활용되어 왔다.
• Gap: 기존 attribution 기법은 예측 결과를 설명하는 해석 도구로만 사용되었을 뿐, 분자 생성/최적화 과정에서 구조적 개입 위치를 능동적으로 안내하는 operational 요소로는 거의 활용되지 않았다. 또한 GB-GA의 crossover/mutation은 어떤 원자가 fitness에 기여하거나 저해하는지에 대한 정보를 전혀 활용하지 않고 균일 무작위로 위치를 선택한다.
• Why: 예측 모델이 왜 특정 분자가 좋은 점수를 받는지에 대한 구조적 정보를 활용해 어디에 어떻게 개입할지 판단하는 능력은, 생성적 분자 설계 및 AI 기반 과학적 발견의 agentic 접근으로 나아가는 데 핵심적인 역량이며, 이 논문은 그 최소 형태의 실증적 증거를 제공한다.
• Approach: D-MPNN 기반 surrogate scorer와 별도의 attribution 모델을 사용해 원자별 attribution score를 계산하고, 이를 GB-GA의 site-selection 확률(bias probability b=0.5)에 반영하여 fitness 기여도가 낮게 추정된 원자(⌊n/10⌋개)를 우선적으로 crossover/mutation 대상으로 선택하는 modular한 변형을 제안한다.

Achievement

Figure 2. Comparison of top-25 scores (mean ± SEM) under the

1. Attribution 기반 site-selection 프레임워크 제안: single-atom ablation, Shapley value attribution, subgraph removal 세 가지 attribution 방법을 GB-GA의 site-selection 메커니즘으로 활용하는 modular하고 attribution 방법에 대해 agnostic한 프레임워크를 설계함.
2. DOCKSTRING 벤치마크에서 일관된 성능 향상 입증: 세 가지 attribution-guided 변형 모두 DOCKSTRING의 세 가지 타겟 전반에서, favorable initialization과 random initialization 조건 모두에서 unmodified GB-GA를 체계적으로 능가함을 보임.
3. 개선 폭과 surrogate 신뢰도의 연관성 규명: attribution-guided 방법의 성능 개선 정도가 surrogate 모델의 신뢰도(reliability)와 상관관계를 가짐을 확인하여, attribution guidance의 효과가 surrogate 품질에 의해 조절됨을 시사함.
4. 경량화되고 재학습 불필요한 개선안 제시: 제안된 방법은 계산 비용이 낮고 추가 재학습이 필요 없으며, outer-loop 최적화 전략과 orthogonal하여 기존 파이프라인에 쉽게 통합 가능함.

How

Figure 3. JAK2 surrogate score trajectory (mean ± SEM) under

• Fitness function으로 D-MPNN ensemble(Chemprop v2.0)을 사용하여 GB-GA 후보 분자를 스코어링하고, 별도의 단일 D-MPNN을 attribution 모델로 사용해 scoring 모델과 attribution 모델을 분리함으로써 단일 모델 아티팩트에 attribution이 좌우되는 것을 방지함.
• 각 attribution 방법이 분자 내 모든 원자에 대해 점수를 계산: (1) single-atom ablation은 원자 특징을 0으로 만들고 예측 변화를 측정(Ai = f(G̃i) − f(G)), n번의 forward pass 필요; (2) Shapley value attribution은 모든 부분집합에 대한 평균 한계 기여도를 계산; (3) subgraph removal은 더 넓은 구조적 맥락을 고려함.
• Attribution score가 가장 긍정적인(즉 fitness 기여도가 가장 낮다고 추정되는) ⌊n/10⌋개의 원자를 candidate operation site로 지정하고, bias probability b=0.5로 이 candidate set에서, 그렇지 않으면 전체 원자 집합에서 site selection을 수행하여 탐색적 능력을 보존함.
• DOCKSTRING 벤치마크의 세 가지 target(JAK2 등 포함)에 대해 favorable initialization과 random initialization 두 조건에서 top-25 surrogate score 및 oracle 평가를 통해 성능을 비교함.

Originality

• Attribution 기법을 해석(interpretability) 도구가 아닌 분자 최적화 과정의 능동적 개입(operational guidance) 요소로 재해석(reframing)한 최초의 시도로, atom attribution을 이용해 atom-level site selection을 직접 유도한 첫 사례임을 주장.
• 기존 model-guided 최적화 연구(Fu et al. 2022의 RL 기반 분자/템플릿 선택, Lee et al. 2023의 단백질 서열 최적화용 Shapley value 기반 GA, Wüthrich et al. 2021의 fragment boundary 정의)와 달리, 순서가 없는 분자 그래프의 원자 단위에서 직접 작동하며 추가 학습 없이 표준 GA 연산자를 그대로 보존하는 modular한 접근을 취함.
• 특정 attribution 방법에 종속되지 않는(agnostic) 설계로, 세 가지 서로 다른 attribution 방법(single-atom ablation, Shapley value, subgraph removal)에 대해 모두 일관된 성능 향상을 실증적으로 보여 방법론의 일반성을 뒷받침함.

Limitation & Further Study

• Atom threshold(⌊n/10⌋)와 bias probability b=0.5는 별도의 최적화 없이 heuristic하게 선택되었으며, 저자들도 adaptive scheduling이나 attribution magnitude에 비례하는 continuous probability weighting과 같은 개선 여지를 언급함.
• Single-atom ablation은 n번의 forward pass가 필요하고 first-order contribution만 포착하며, Shapley value attribution은 조합적 계산 비용이 크다는 근본적 한계가 있어 대규모 분자나 대규모 population에는 확장성 문제가 있을 수 있음.
• 성능 개선의 폭이 surrogate 모델의 신뢰도에 의존한다는 결과는, surrogate가 부정확한 경우(예: 새로운 화학 공간, low-data regime) attribution guidance의 효용이 제한될 수 있음을 시사하며, 이에 대한 체계적 검증이 부족함.
• 평가가 DOCKSTRING의 세 가지 target에 국한되어 있어, 더 다양한 target class, 더 크고 복잡한 분자, 실제 실험적 검증(wet-lab validation)으로의 일반화 가능성은 추가 연구가 필요함.

Evaluation

총평: 단순하지만 참신한 아이디어를 통해 attribution 기법의 활용 범위를 해석에서 능동적 최적화 guidance로 확장하였으며, 재학습이 필요 없고 기존 파이프라인과 orthogonal하다는 실용적 장점을 갖춘 견고한 연구이나, 벤치마크 범위와 하이퍼파라미터 최적화 측면에서 추가 검증이 필요하다.