タンパク質における「言語」を紐解く～ProteinBERT～

Bioinformatics 2022年03月22日

3つの要点
✔️ タンパク質言語モデルの事前学習タスクに従来のMLMタスクに加えてGOアノテーションを予測するタスクを追加した
✔️ 局所的な特徴と全体的な特徴を分離して扱うことにより、従来のニューラルネットワークよりも小さく高速なアーキテクチャを提案した
✔️ 構造予測や翻訳後修飾予測を含むベンチマークで従来手法を上回る性能を記録した

ProteinBERT: A universal deep-learning model of protein sequence and function
written by Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, View ORCID ProfileMichal Linial
(Submitted on 25 May 2021)
Subjects: Bioinformatics
Copyright:The copyright holder for this preprint is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made available under a CC-BY 4.0 International license.
code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

タンパク質は20種類のアミノ酸が連結されたものであることから、一種の文字列として考えることができます。このことから自然言語処理の技術をタンパク質へと応用する研究が増えてきています。

このような研究の一例としてタンパク質言語モデルが挙げられます。自然言語処理におけるBERTのような文脈を考慮した単語埋め込みの計算手法をタンパク質に適用することで、二次構造予測や翻訳後修飾などの下流タスクに有用な事前学習モデルを獲得する研究が進められています。タンパク質言語モデルを解釈することによって今まで見つかっていないような生物学的知見を発見することができる可能性があり、注目を集めています。

しかし、タンパク質と自然言語では以下に挙げるような大きな違いがあります。

タンパク質には単語や文などの明確な構成要素が存在しない
タンパク質は自然言語よりも配列長におけるバリエーションが大きい
タンパク質は立体構造などの影響で、配列上は離れていても相互作用する可能性がある

従って、このようなタンパク質の特徴を上手くモデリングする方法を見つけることが求められています。

今回ご紹介するProteinBERTではタンパク質におけるこれらの特徴を考慮するため、局所的な表現とグローバルな表現を分離して考えるアーキテクチャを考案しています。これにより、従来の自然言語処理で用いられてきたアーキテクチャよりも小さなネットワークサイズでより高速な訓練を可能にしています。

事前学習タスク

ProteinBERTでは自然言語処理におけるBERTに倣い、Masked Language Modeling(MLM)と遺伝子オントロジー(GO)アノテーション予測の二つのタスクにより事前学習を行います。

ここで、GOとはある遺伝子が細胞内で果たす機能や細胞内局在の情報に基づいた階層的分類のことであり、アノテーションとは遺伝子とGOの間の対応関係を構築することです。本論文では入力のアミノ酸配列とGOアノテーションラベルに対してノイズを加え、深層学習モデルに元の入力を復元させることで学習を行っています。これら二つのタスクは同時に学習され、モデル全体の損失を各部位のトークンを予測する際の交差エントロピーとGOアノテーションラベルを予測する際の二値交差エントロピーの和としています。

各部位のアミノ酸を予測するMLMタスクはローカルな特徴の学習、GOアノテーションの予測はグローバルな特徴の学習と捉えることができます。

深層学習モデルのアーキテクチャ

ProteinBERTにおいて用いられた深層学習モデルのアーキテクチャは以下の図のようになっています。事前学習タスクで述べたようにProteinBERTでは入力を復元する自己教師あり学習が行われることから、入力と出力のテンソルのサイズがそれぞれ同じであることがわかります。

ネットワークアーキテクチャはBERTに倣ってTransformerブロックが６層重なった構造をしており、ローカルな特徴量の計算として一次元畳み込み層が採用され、グローバルな特徴量の計算として全結合層が採用されています。図中のグローバルアテンションを通してローカルな特徴量がグローバルな特徴量に反映される仕組みとなっており、ブロードキャスト全結合層を通してグローバルな特徴量がローカルな特徴量に影響を与える仕組みとなっています。

このように、ProteinBERTではローカルな表現の学習とグローバルな表現の学習がネットワークアーキテクチャ上で明示的に分けれられていることが特徴的です。

potein bert architecture

下流タスクにおける性能

事前学習の効果を確認する方法はファインチューニングによる下流タスクでの性能変化の確認です。ファインチューニングでは事前学習により得られたネットワークパラメータを初期値として事前学習とは異なるタスクで学習を行います。

先行研究において提案されていたタンパク質にまつわるベンチマークを用いてProteinBERTの評価を行いました。ベンチマークとしてTAPEで提案された二次構造予測、ホモロジー予測、蛍光予測、安定性予測の４つのタスクを利用しています。

比較手法としてTransformerEncoderを用いたTAPE TransformerやLSTMを用いています。これらのモデルは最大3800万個のパラメータで構成され巨大であるのに対し、ProteinBERTは最大1600万個と比較的小さなモデルとなっています。

protein bert downstream

この結果から、従来のタンパク質言語モデルの研究で得られた結果と同様に事前学習が下流タスクの性能向上に有益であることがわかります。また、ProteinBERTが従来手法と同等もしくはそれ以上の性能を記録していることがわかります。

さらに、ProteinBERTでは長さへのオーバーフィッティングを避けるためモデルの訓練時に定期的に配列長を変更する工夫を施しており、多くの場合でタンパク質の配列長が長いほど下流タスクの性能が低下する現象が確認されたものの、大幅な低下ではなく、配列の長さに関しても汎化性能を確認できたと報告されています。

また、一部のベンチマークタスクでは配列長が長いほど性能が高くなる場合もあったことから、このような性能の変化は配列長ではなく異なる原因によるものと推測しています。

アテンション機構の理解

深層学習モデルのアーキテクチャで述べたグローバルアテンションはローカルな特徴量をグローバルな特徴量に反映する役割を担っています。従って、このアテンション重みを解析することで下流タスクが配列のどの部分に注目しているのかというパターンを見つけることができます。

本論文によればアテンション重みのパターンはタンパク質によって大きくばらつくものの、共通のパターンも存在することがわかっており、一部のアテンションヘッドでは配列の上流部分など一定の箇所に注目する傾向があるなどの報告があります。また、ファインチューニング前後でアテンション重みを比較すると、モデルの最終層において大きく変化することがわかりました。

おわりに

いかがだったでしょうか。本論文ではタンパク質に対する新しい事前学習手法ProteinBERTを提案していました。本論文の貢献はローカルな特徴の学習とグローバルな特徴の学習を明示的に分けて扱うネットワークアーキテクチャを導入することで、より効率的な事前学習を可能にした点だと考えられます。これによりたった一つのGPUを用いるだけで従来手法と同程度かそれ以上の性能を出すことができるようです。驚異的ですね。

タンパク質言語モデルのこれからの進展が楽しみです。