[논문 리뷰] BERTology Meets Biology: Interpreting Attention in Protein Language Models

Jesse Vig, Ali Madani외 4명|arXiv (Cornell University)|2020. 06. 26.

Machine Learning in Bioinformatics인용 수 24

한 줄 요약

이 논문은 단백질 언어 모델의 주의 메커니즘을 분석하기 위한 해석 가능성 방법을 제안하며, 주의가 계층을 거쳐 3차원 단백질 구조, 기능적 결합 부위, 생물물리적 특성과 같은 핵심 생물학적 특징을 포착함을 보여준다. 두 데이터셋에서 BERT, ALBERT, XLNet 모델을 대상으로 한 분석에서 주의는 접촉 맵(p < 0.00001), 결합 부위, 번역 후 수정(PTM)과 유의미하게 일치하며, 자기지도 학습을 통한 사전 훈련 과정에서 주의가 의미 있는 생물학적 추론을 학습하고 있음을 시사한다.

ABSTRACT

Transformer architectures have proven to learn useful representations for protein classification and generation tasks. However, these representations present challenges in interpretability. In this work, we demonstrate a set of methods for analyzing protein Transformer models through the lens of attention. We show that attention: (1) captures the folding structure of proteins, connecting amino acids that are far apart in the underlying sequence, but spatially close in the three-dimensional structure, (2) targets binding sites, a key functional component of proteins, and (3) focuses on progressively more complex biophysical properties with increasing layer depth. We find this behavior to be consistent across three Transformer architectures (BERT, ALBERT, XLNet) and two distinct protein datasets. We also present a three-dimensional visualization of the interaction between attention and protein structure. Code for visualization and analysis is available at https://github.com/salesforce/provis.

연구 동기 및 목표

• 단백질 전용 트랜스포머 모델(예: BERT, ALBERT, XLNet)의 주의 메커니즘을 해석하여 이들이 생물학적 지식을 어떻게 인코딩하는지 이해하는 것.
• 이러한 모델의 주의가 단백질의 구조적 및 기능적 특징, 예를 들어 3차원 공간적 근접성과 결합 부위를 얼마나 잘 포착하는지 조사하는 것.
• 은닉 상태 표현과의 비교를 통해 주의 기반 해석 가능성의 일관성을 여러 아키텍처와 데이터셋 간에 평가하는 것.
• 모델의 주의가 알려진 생물학적 특징(예: 접촉 맵, 치환 행렬, 번역 후 수정)과 연결되도록 하여 과학적 발견을 지원하는 것.
• 생물학적 통찰을 지원하기 위해 3차원 단백질 구조에 주의 무게를 오버레이하는 데 사용할 수 있는 오픈소스 시각화 도구를 제공하는 것.

제안 방법

• 저자는 사전 훈련된 단백질 트랜스포머(TapeBert, ProtBERT, ProtALBERT, ProtXLNet)의 주의 무게를 여러 계층과 헤드에서 분석한다.
• 주의 무게가 실제 생물학적 주석(공간적 근접성, 결합 부위, 번역 후 수정(PTM), 아미노산 빈도 등)과 얼마나 일치하는지 정량화한다.
• 통계적 유의성은 95% 신뢰구간과 배경 빈도에 대한 보너페리-보정된 가설 검정을 통해 평가한다.
• 3차원 시각화 도구는 NGL Viewer를 사용하여 단백질 구조에 주의 무게를 오버레이함으로써 공간적 해석을 가능하게 한다.
• 허수 모델(랜덤화된 주의 헤드)을 사용하여 기준선을 설정하고 관측된 주의 패턴이 우연이 아님을 확인한다.
• 주의 패턴이 BLOSUM 치환 행렬과 얼마나 유사한지 평가하기 위해 주의 패턴을 BLOSUM 치환 행렬과 비교한다.

실험 결과

연구 질문

• RQ1 단백질 언어 모델의 주의 헤드는 접촉 맵로 측정된 3차원 단백질 구조에서의 공간적 근접성과 일치하는가?
• RQ2 특정 주의 헤드는 HIV-1 프로테아제와 같은 효소의 알려진 기능적 부위인 결합 부위를 타겟으로 하는가?
• RQ3 주의가 네트워크 계층을 거쳐 깊어질수록 고차원의 생물물리적 또는 기능적 특징을 점차로 인코딩하는가?
• RQ4 주의 분포는 BLOSUM 치환 행렬과 같은 알려진 생물학적 사전 지식과 어떻게 비교되는가?
• RQ5 주의 패턴은 BERT, ALBERT, XLNet 등의 다양한 트랜스포머 아키텍처와 단백질 데이터셋 간에 일관된가?

주요 결과

• 단백질 트랜스포머의 주의 헤드는 3차원 접촉 맵과 유의미하게 일치한다(p < 0.00001), 특히 서열상으로는 떨어져 있지만 3차원적으로 가까운 아미노산 쌍에 집중하는 상위 헤드가 존재한다.
• 특정 주의 헤드는 알려진 기능적 결합 부위를 타겟으로 한다. 예를 들어, ProtXLNet의 헤드 7-1은 HIV-1 프로테아제의 잠재적 약물 결합 부위인 잔류물 27G에 집중한다.
• 주의 패턴은 BLOSUM 치환 행렬과 강하게 상관되며, 이는 주의가 생물학적으로 의미 있는 아미노산 관계를 학습하고 있음을 시사한다.
• 번역 후 수정(PTM)에 집중하는 주의 비율은 배경 빈도보다 유의미하게 높다(p < 0.00001), 일부 헤드는 PTM 부위에 70% 이상의 비율로 주의를 기울인다.
• 아미노산 별 주의 분포는 매우 비균형적이다. 예를 들어, 일부 헤드에서 글리신과 프로린은 각각 최대 98.3%와 98.1%의 주의를 받으며 배경 빈도보다 훨씬 높다.
• 표준 모델의 상위 주의 헤드는 랜덤화된 대조군보다 일관되게 뛰어나다(예: TapeBert에서 알라닌의 경우 25.5% 대비 12.1%), 이는 비랜덤이자 생물학적으로 의미 있는 주의임을 확인한다.