기술 요약: 그래프 기반 도구를 활용한 소형 언어 모델의 분자 특성 예측 개선

문제 정의

분자 특성 예측(Molecular Property Prediction, MPP)은 약물 발견의 핵심적인 구성 요소이며, 전통적으로 분자의 위상적 구조를 모델링하는 능력이 탁월한 그래프 신경망(Graph Neural Networks, GNN)이 주도해 왔습니다. 그러나 GNN은 인간이 해석 가능한 근거를 제공하지 못하는 "블랙박스"로 기능하는 경우가 많습니다. 반면, 소형 언어 모델(Small Language Models, SLMs)과 대형 언어 모델(Large Language Models, LLMs)은 제로샷 추론에서 유망한 가능성을 보여주었으나, SMILES(Simplified Molecular Input Line Entry System) 문자열을 통해 분자를 처리할 때 "구조적 맹목성(structural blindness)" 문제를 겪습니다. SMILES는 시퀀스 기반의 추상화 방식으로, 복잡한 2D 브랜칭(branching) 및 공간적 관계를 인코딩하는 데 자주 실패하며, 이로 인해 특정 국소 모티프(local motifs)나 전역 연결성에 의존하는 예측 성능이 일관되지 않게 나타납니다.

본 연구가 다루는 핵심 과제는, 언어 모델을 재학습시키거나 계산 비용이 많이 드는 아키텍처 통합을 거치지 않고도, GNN의 구조적 추론 능력과 SLM의 언어적 추론 능력 사이의 간극을 어떻게 메울 것인가 하는 점입니다.

방법론: 컨텍스트 증강 프롬프팅 프레임워크

저자들은 추론 시점에 에이전트 기반의 도구 사용을 가능하게 하는 모듈형 컨텍스트 증강 프롬프팅(Context-Augmented Prompting) 프레임워크를 제안합니다. 이 설정에서 SLM은 사전 학습된 그래프 기반 도구를 쿼리하여 구조화된 컨텍스트를 생성하고, 이를 다시 프롬프트에 피드백하여 제로샷 예측을 강화하는 추론 컨트롤러 역할을 수행합니다.

에이전트 루프(The Agentic Loop)

입력 SMILES 문자열 $s$가 주어지면, SLM은 세 가지 주요 도구를 호출할 수 있습니다:

1. ExpertPredict($s$): SMILES를 분자 그래프 $G$로 변환하고 사전 학습된 GNN(구체적으로 GIN 모델)에 쿼리하여 예측 라벨 $\hat{y}$와 신뢰도 점수 $\hat{p}$를 반환합니다.
2. ExplainSubgraph($s$): GNNExplainer를 사용하여 예측에 기여하는 가장 영향력 있는 서브그래프(모티프)를 식별합니다. 이는 에지 중요도 마스크(edge-importance mask)로 추출되며, 상위 10%의 에지를 유지하도록 임계값이 설정된 후 다시 서브그래프 SMILES 문자열($s_c$)로 변환됩니다.
3. ChemReason($s, s_c, \hat{y}, \hat{p}$): SLM 자체가 식별된 모티프 $s_c$가 왜 예측된 특성(예: 독성)을 뒷받침하는지에 대한 자연어 근거를 생성하는 지시 블록입니다. 이는 더 큰 교사 모델(teacher models)로부터 발생하는 편향을 피하면서, GNN 유래의 증거를 바탕으로 SLM이 직접 추론을 생성하도록 보장합니다.

프롬프트 구성

본 연구는 베이스라인부터 전체 컨텍스트까지 다섯 가지 프롬프트 구성을 평가합니다:

1. SMILES (Baseline): 원시(raw) SMILES 문자열만 입력값으로 사용합니다.
2. SMILES + SUBGRAPH: 원시 SMILES와 추출된 서브그래프 SMILES($s_c$)를 포함합니다.
3. SMILES + HINT: 원시 SMILES와 GNN 전문가의 예측 및 신뢰도를 포함합니다.
4. SMILES + REASONING: 원시 SMILES와 모티프의 관련성에 대해 SLM이 생성한 근거를 포함합니다.
5. ALL CONTEXT: 원시 SMILES, 전문가 예측/신뢰도, 서브그래프 SMILES, 그리고 SLM이 생성한 근거를 모두 결합합니다.

실험 설정

본 프레임워크는 두 가지 벤치마크 데이터셋에 대해 평가되었습니다:

• MUTAG: 돌연변이 유발성을 나타내는 188개의 니트로 방향족 화합물.
• Tox21: 이진 독성 분류 작업으로 변환된 1,000개의 분자 서브셋.

세 가지 SLM이 제로샷 방식으로 테스트되었습니다(데이터셋에 대한 파인튜닝 없음):

• DeepSeek (Lite-Code): 2.4B 활성 파라미터(16B MoE 아키텍처), 코드 및 구조적 추론에 최적화됨.
• Qwen 2.5 (Code): 3B 파라미터, 단계별 생성을 위해 튜닝됨.
• Llama 3.2: 3B 파라미터, 범용 대화형 베이스라인.

GIN 모델은 "전문가" 도구 역할을 수행하기 위해 훈련 분할(70/30 분할) 데이터로 학습되었으며, SLM은 오직 테스트 분할 데이터에 대해서만 평가되었습니다.

주요 결과

연구 결과, 프롬프트를 그래프 유래 컨텍스트로 풍부하게 만드는 것이 상당한 정확도 향상을 가져오지만, 그 성능은 특정 SLM 아키텍처와 데이터셋의 난이도에 따라 크게 달라짐을 입증했습니다.

• 복잡한 작업에서의 유의미한 이득: 더 까다로운 Tox21 데이터셋에서, ALL CONTEXT 구성은 SMILES 전용 베이스라인 대비 DeepSeek에서 74.03%, Qwen 2.5에서 **44.21%**의 상대적 개선을 달이트했습니다.
• 모델 의존성: DeepSeek와 Qwen 2.5는 모든 구성에서 일관되고 강력한 개선을 보였습니다. 반면, Llama 3.2는 혼재된 양상을 보였습니다. 일부 단일 도구 변형은 도움이 되었으나, ALL CONTEXT 구성은 MUTAG에서 베이스라인보다 개선되지 않았으며 Tox21에서도 낮은 향상 폭을 보였습니다. 저자들은 작은 범용 모델들이 이질적인 증거를 통합할 때 "정보 교란(information distraction)"을 겪을 수 있다고 제안합니다.
• 결합된 컨텍스트의 강건성: 단일 도구 신호(예: 힌트 또는 서브그래프 단독)는 때때로 불안정하거나 오히려 성능을 저하시켰습니다(예: MUTAG에서의 힌트가 포함된 DeepSeek). 그러나 ALL CONTEXT 구성은 일반적으로 더 강건했는데, 이는 전문가의 사전 지식, 구조적 증거, 그리고 의미론적 추론이 서로를 보완하는 효과가 있음을 시사합니다.
• 전문화된 모델과의 지속적인 격차: 이러한 개선에도 불구하고, ALL CONTEXT를 사용한 SLM은 전문화된 GNN 전문가(GIN)를 일관되게 능가하지 못했습니다. 예를 들어, MUTAG에서 GIN은 84.21%의 정확도를 달성한 반면, 가장 우수한 SLM(DeepSeek)은 75.44%에 그쳤습니다. 이는 텍스트 조건부 추론이 그래프 표현 학습의 귀납적 편향(inductive bias)을 완전히 대체할 수는 없음을 강조합니다.

인과적 관련성 검증

필요성 기반의 에지 드롭(necessity-based edge-drop) 개입 실험을 통해 GNNExplainer에 의해 추출된 서브그래프가 단순히 상관관계만을 나타내는 것이 아님을 확인했습니다. 중요도 순위가 높은 에지를 제거했을 때 무작위 에지 제거보다 전문가 정확도가 더 급격히 하락하였는데, 이는 추출된 모티프가 결정적인 의사결정 정보를 포함하고 있음을 나타냅니다.

의의 및 주장

본 논문은 언어 모델을 재학습시키지 않고도 분자 특성 예측을 개선할 수 있는 단순하고, 모듈화되었으며, 접근 가능한 접근 방식을 제공한다고 주장합니다. 주요 기여는 다음과 같습니다:

1. 양식 간 가교(Bridging Modalities): GNN 유래 힌트를 통해 그래프 기반의 구조적 학습과 시퀀스 기반의 SLM 추론을 성공적으로 통합하여, SMILES 전용 입력의 "구조적 맹목성"을 완화했습니다.
2. 효능의 정량화: 그래프 유래 컨텍스트가 예측 정확도를 실질적으로 높이지만, 그 이득이 데이터셋 및 모델에 따라 달라짐을 경험적 증거를 통해 제시했습니다.
3. 에이전트 기반 거버넌스: SLM이 의미론적 컨트롤러 역할을 하는 "2단계 거버넌스" 패턴을 보여줍니다. SLM은 그래프 기반 설명을 해석하고, 불일치(예: GNN 전문가가 틀린 경우)를 식별하며, 도메인 제약 조건에 따라 전문가의 출력을 수용하거나 거부할지 결정할 수 있습니다.
4. 실제적 배포: 계산 비용이 많이 드는 하이브리드 아키텍처(예: MolCA, MolLM)와 달리, 이 방식은 SLM의 고정된 가중치를 보존하므로 과학 분야에서 인과적으로 근거 있는 뉴로-심볼릭 에이전트로 나아가는 실용적인 단계가 됩니다.

저자들은 그래프 기반의 구조적 증거와 텍ext 조건부 추론 사이에 여전히 간극이 존재함을 언급하며, 이는 그래프 귀납적 편향이 분자 과학에서 계속해서 필요함을 역설하며, 텍스트 조건부 추론과 전문화된 그래프 모델 사이의 차이를 강조하며 결론을 맺습니다.