[논문 리뷰] In-Context Learning for Few-Shot Molecular Property Prediction

Christopher Fifty, Jure Leskovec외 1명|arXiv (Cornell University)|2023. 10. 13.

Machine Learning in Materials Science인용 수 5

한 줄 요약

이 논문은 소수의 예시에서 분자의 성질을 직접 추론할 수 있는 새로운 인코어 컨텍스트 학습 프레임워크인 Context Aware Molecule Prediction (CAMP)를 소개한다. 이는 파라미터 미세조정 없이 (분자, 성질) 쌍의 예시 컨텍스트로부터 분자의 성질을 추론할 수 있다. CAMP는 분자 및 레이블 임베딩을 순서대로 처리하기 위해 트랜스포머 인코더를 사용하며, 소규모 지원 크기에서 최신 기술 수준의 성능을 달성하고, 프로토타입 네트워크보다 약 4.5배 빠른 추론 속도를 제공한다.

ABSTRACT

In-context learning has become an important approach for few-shot learning in Large Language Models because of its ability to rapidly adapt to new tasks without fine-tuning model parameters. However, it is restricted to applications in natural language and inapplicable to other domains. In this paper, we adapt the concepts underpinning in-context learning to develop a new algorithm for few-shot molecular property prediction. Our approach learns to predict molecular properties from a context of (molecule, property measurement) pairs and rapidly adapts to new properties without fine-tuning. On the FS-Mol and BACE molecular property prediction benchmarks, we find this method surpasses the performance of recent meta-learning algorithms at small support sizes and is competitive with the best methods at large support sizes.

연구 동기 및 목표

• 자연어에 국한되어 있던 인코어 컨텍스트 학습을 약물 발굴에 핵심적인 소수의 예시 학습이 필요한 분자 성질 예측 분야로 확장하기 위해.
• 파라미터 미세조정이나 기울기 업데이트 없이도 새로운 분자 성질에 신속하게 적응할 수 있는 방법을 개발하기 위해.
• 기존의 소수의 예시 학습 방법이 분자 과학에서 가지는 한계를 극복하고, 대규모 걸러내기 작업에서 저지연, 고정확도 추론을 가능하게 하기 위해.
• 시퀀스 모델링을 사용하여 인코어 컨텍스트 학습 원리가 언어 외 영역인 분자 과학 분야로 성공적으로 이식될 수 있는지 조사하기 위해.
• 위치 임베딩과 예시 순서가 분자 소수의 예시 학습에서 모델 일반화에 미치는 영향을 평가하기 위해.

제안 방법

• CAMP는 분자 인코더(예: MPNN)를 사용해 쿼리 및 예시 분자를 인코딩하고, 레이블 인코더를 사용해 성질 측정치를 인코딩한다.
• 분자 및 레이블 임베딩을 연결하여 시퀀스로 만들고, 이를 트랜스포머 인코더에 입력하여 쿼리 분자의 성질을 예측한다.
• 모델은 단일 전방향 프로세스를 통해 추론을 수행하며, 파라미터 미세조정이나 중심점 계산 없이 예시 컨텍스트로부터 직접 레이블을 예측한다.
• 아키텍처는 예시 순서에 대해 불변이며, 분석 결과 위치 임베딩은 성능을 떨어뜨리고 과적합을 증가시킴을 보여준다.
• LLM 기반의 인코어 컨텍스트 학습과 달리, 이 방법은 소규모 소수의 예시 데이터셋(~5 GB)에서 무작위 초기화로부터 학습된다.
• 분석 연구에서는 Naive ICL(교차 배치된 임베딩)와 CAMP를 비교하여, 별도의 임베딩 연결 방식이 학습 안정성과 성능을 향상시킴을 보여준다.

실험 결과

연구 질문

• RQ1 인코어 컨텍스트 학습 원리가 언어 외 영역인 분자 성질 예측에 성공적으로 적용될 수 있는가?
• RQ2 분자 환경에서 기울기 기반 및 메트릭 기반 소수의 예시 학습 베이스라인에 비해 분자 및 레이블 임베딩을 사용하는 시퀀스 모델링 접근 방식이 우수한 성능을 내는가?
• RQ3 컨텍스트 내 예시 순서가 모델 성능에 미치는 영향은 무엇이며, 위치 인코딩은 유익한가, 아니면 해로운가?
• RQ4 소규모 데이터셋에서 무작위 초기화로부터 학습된 모델이 파라미터 미세조정 없이도 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 임베딩 연결 및 위치 인코딩과 같은 아키텍처 선택 사항이 학습 안정성과 일반화에 미치는 영향은 무엇인가?

주요 결과

• CAMP는 FS-Mol 및 BACE 벤치마크에서 소규모 지원 크기에서 최근의 메타학습 알고리즘을 능가하며, 뛰어난 소수의 예시 일반화 성능을 입증한다.
• 대규모 지원 크기에서 CAMP는 프로토타입 네트워크를 포함한 최고의 베이스라인과 경쟁하는 성능을 달성하며, 추론 시 약 4.5배 더 빠른 속도를 기록한다.
• 위치 임베딩이 없는 모델은 위치 임베딩이 있는 변형보다 더 낮은 검증 손실(0.5746)을 기록하여, 위치 인코딩이 성능에 악영향을 준다는 것을 시사한다.
• 위치 임베딩을 사용한 학습은 더 빠른 과적합을 유도하며, 최종 학습 손실이 0.3181로 위치 임베딩이 없는 변형(0.3374)보다 낮다.
• 분자 및 레이블 임베딩을 교차 배치하는 Naive ICL은 CAMP의 별도 연결 시퀀스 방식에 비해 수렴이 불안정하고 성능이 劣 劣하다.
• CAMP의 성능은 예시 순서에 대해 강건하여, 분자 예측에서 인코어 컨텍스트 학습이 순열 불변이어야 한다는 가설을 지지한다.