[논문 리뷰] Regression Transformer enables concurrent sequence regression and generation for molecular language modelling

Jannis Born, Matteo Manica|arXiv (Cornell University)|2022. 02. 01.

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

한 줄 요약

Regression Transformer (RT)은 회귀를 조건부 시퀀스 모델링 작업으로 공식화하여 동시에 시퀀스 회귀와 조건부 시퀀스 생성을 수행하는 혁신적인 다중 작업 프레임워크를 도입한다. 이는 소분자, 단백질, 화학 반응 분야에서 단일 통합 아키텍처를 사용해 소수의 분자, 단백질, 화학 반응에 걸쳐 강력한 zero-shot 일반화 능력을 보이며, 성능 예측 및 특성 기반 분자의 생성에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Despite significant progress of generative models in the natural sciences, their controllability remains challenging. One fundamentally missing aspect of molecular or protein generative models is an inductive bias that can reflect continuous properties of interest. To that end, we propose the Regression Transformer (RT), a novel method that abstracts regression as a conditional sequence modeling problem. This introduces a new paradigm of multitask language models which seamlessly bridge sequence regression and conditional sequence generation. We thoroughly demonstrate that, despite using a nominal-scale training objective, the RT matches or surpasses the performance of conventional regression models in property prediction tasks of small molecules, proteins and chemical reactions. Critically, priming the same model with continuous properties yields a highly competitive conditional generative model that outperforms specialized approaches in a substructure-constrained, property-driven molecule generation benchmark. Our dichotomous approach is facilitated by a novel, alternating training scheme that enables the model to decorate seed sequences by desired properties, e.g., to optimize reaction yield. In sum, the RT is the first report of a multitask model that concurrently excels at predictive and generative tasks in biochemistry. This finds particular application in property-driven, local exploration of the chemical or protein space and could pave the road toward foundation models in material design. The code to reproduce all experiments of the paper is available at: https://github.com/IBM/regression-transformer

연구 동기 및 목표

• 분자 및 단백질 생성 모델에서 연속적 성질에 대한 인덕티브 바이어스의 부족을 해결하기 위해.
• 생물화학 분야에서 예측 모델링과 생성 모델링 간의 의미적 격차를 해소하기 위해, 단일 아키텍처에서 회귀와 조건부 생성을 통합하기 위해.
• 예측 및 생성 양쪽에서 뛰어난 성능을 발휘하는 단일 모델을 통해 성질 기반의 국소적 탐색을 가능하게 하기 위해.
• 동일한 모델이 공유된 시퀀스 모델링 목표를 사용하여 회귀 및 생성 작업 간을 원활하게 전환할 수 있도록 하는 훈련 방식을 개발하기 위해.

제안 방법

• RT는 입력 시퀀스와 목표 수치 값을 모두 조건으로 삼아 회귀를 조건부 시퀀스 모델링 문제로 공식화한다.
• 사전 훈련 중에 마스크된 수치 토큰 예측(회귀)과 마스크된 시퀀스 토큰 생성(생성)을 번갈아가며 수행하는 새로운 교차 훈련 방식을 사용한다.
• 두 작업 모두에 동일한 파라미터를 공유하는 공유된 Transformer 인코더-디코더 아키텍처를 사용하여 파라미터 효율성과 공동 최적화를 달성한다.
• MoleculeNet, Boman, TAPE, 화학 반응의 수율 예측 벤치마크 등 다양한 데이터셋에서 미세조정을 수행한다.
• 학습된 임베딩을 사용한 SMILES 및 단백질 시퀀스 토크나이제이션을 활용하고, 회귀 및 생성 스트림 양쪽에 마스크된 언어 모델링 목표를 적용한다.
• 지속적인 성질 값(예: 용해도, pLogP)을 프리밍하여 원하는 성질을 가진 분자의 조건부 생성을 유도한다.

실험 결과

연구 질문

• RQ1 단일 신경망 아키텍처가 분자 및 단백질 모델링에서 동시에 시퀀스 회귀와 조건부 시퀀스 생성을 효과적으로 수행할 수 있는가?
• RQ2 회귀를 조건부 시퀀스 모델링으로 공식화하는 것이 별도의 모델 대비 일반화 및 성능 향상에 기여하는가?
• RQ3 통합 모델이 서브스트럭처 제약 조건 하에서 전용 모델을 능가하는 성질 기반 분자 생성에서 승리할 수 있는가?
• RQ4 교차 훈련 방식이 동시에 두 작업인 회귀와 생성 학습을 효과적으로 가능하게 하는가?
• RQ5 RT가 소분자, 단백질, 화학 반응을 포함한 다양한 생물화학 도메인으로 일반화되는 정도는 어느 정도인가?

주요 결과

• RT는 소분자, 단백질, 화학 반응 분야의 성질 예측 작업에서 기존의 회귀 모델과 동일하거나 이를 초월하며, MoleculeNet 데이터셋에서 최신 기술 수준의 성능을 달성한다.
• 성질 최적화 벤치마크에서 RT는 시드 구조와 유사성을 유지하면서 최대 pLogP를 갖는 분자를 생성하는 데 있어 전용 조건부 생성 모델을 능가한다.
• 서브스트럭처 제약 조건 하에서 pLogP 값이 3.0 이상인 분자를 생성하는 데 성공률가 92.3%에 달하며, 기준 모델 대비 15% 이상 높은 성능을 기록한다.
• 조건부 생성에서 RT는 화학적으로 타당하고 구조적으로 다양한 성질을 가진 분자(예: QED 용해도, Boman 지수 안정성)를 생성한다.
• 교차 훈련 방식은 모델이 두 작업인 회귀 및 생성을 효과적으로 학습하도록 하며, 어느 스트림에서도 성능 저하 없이 작동한다.
• RT는 자연어 처리 작업, 예를 들어 원하는 유머 점수를 가진 텍스트 생성 등으로 일반화되어 있으며, 화학 분야를 넘어서도 넓은 적용 가능성을 보여준다.