タンパク質言語モデルの注意機構を紐解く
3つの要点
✔️ タンパク質言語モデルにより事前学習したTransformerのアテンションを分析
✔️ 言語モデルによる学習を行うだけでタンパク質フォールディングなどの複雑な生物学的特徴を獲得できることが判明
✔️ 上記の現象を複数のアーキテクチャ(TAPE etc.)およびデータセットで確認
BERTology Meets Biology: Interpreting Attention in Protein Language Models
written by Jesse Vig, Ali Madani, Lav R. Varshney, Caiming Xiong, Richard Socher, Nazneen Fatema Rajani
(Submitted on 26 Jun 2020 (v1), last revised 28 Mar 2021 (this version, v3))
Comments: ICLR 2021.
Subjects: Computation and Language (cs.CL); Machine Learning (cs.LG); Biomolecules (q-bio.BM)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
タンパク質（アミノ酸配列）に対して自然言語処理における言語モデルを適用したものをタンパク質言語モデルといいます。タンパク質が生体内でどのように機能しているのか不明な点が数多くあるため、大量のデータとBERTに代表される事前学習を用いることで新たな知見が得られることが期待されています。BERTで利用されているTransformerは自然言語処理において広く用いられているアーキテクチャです。近年の研究ではこのTransformerにおけるアテンション機構を解析することでモデルの解釈性を向上しようという試みが増えています。特にBERTの解釈を行うものはBERTologyと呼ばれています。
今回紹介する論文ではタンパク質の配列情報を用いた事前学習モデルにおいてアテンション機構の解析を行うことでモデルの解釈性向上に取り組んでいます。
モデルの解釈手法
本論文ではTransformerベースの事前学習モデルに対して主に二つの解析手法を適用しています。
アテンション機構の解析
アミノ酸残基間のアテンション重みが閾値を超えたもののうち、実際に生物学的特徴を有している組み合わせがどの程度含まれているのかを確認します。
具体的には以下のような式で計算されます。適合率(Precision)に近い評価指標となっています。
プローブタスク
プローブタスクとは自然言語処理において事前学習モデルの解釈性向上に用いられる補助タスクです。事前学習モデルにおいて獲得された内部表現がどのような情報を含んでいるのか確認します。
具体的には、モデルから得られる表現を分類器に入力してなんらかのタスクを解かせ、そのタスクに有用な情報がエンコードされているかを確かめます。Transformerベースのモデルに対するプローブタスクには埋め込みプローブとアテンションプローブの二種類が存在し、埋め込みプローブでは各層の出力、アテンションプローブではアテンションの重みが対象となります。
アテンション機構の解析結果
タンパク質の二次構造
タンパク質の二次構造における重要な特徴としてコンタクトマップがあります。コンタクトマップとはタンパク質フォールディングにおいて空間的に近いアミノ酸の組み合わしたものです。
このコンタクトマップをアミノ酸ペアに対する特徴とし、上述のアテンション解析手法を適用したところ、各事前学習モデルでの$p_{\alpha}(f)$は44.5%から63.2%の値を取っていました。また、最もコンタクトマップを反映しているアテンションヘッドはより深い層に存在していました。
コンタクトマップにおけるバックグラウンド確率が1.3%であることを考えるとタンパク質言語モデルにおける事前学習はコンタクトマップを反映した高次の表現を獲得していると言えそうです。
結合部位
結合部位とはタンパク質が他の分子と相互作用する部位のことです。結合部位はタンパク質の機能に非常に重要な特徴となっています。
この結合部位を参照しているアテンションの割合$p_{\alpha}(f)$は45.6%～50.7%であり、結合部位のバックグラウンド確率が4.8%であることを考えると非常に高い値となっています。
また、アテンションヘッドの大半が結合部位を高い割合で参照するという結果になっていました。
結合部位が外部の分子との相互作用を反映する特徴であるにもかかわらずこのようにモデルに重要視されている理由として、タンパク質の機能に直接関わる特徴であるために構造モチーフが高度に保存されている可能性が示唆されています。
翻訳後修飾
翻訳後修飾とはタンパク質がmRNAから翻訳された後に施される変更のことであり、代表的な翻訳後修飾としてリン酸化が挙げられます。
翻訳後修飾はタンパク質の構造と機能に大きな役割を果たしていることが知られています。翻訳後修飾における$p_{\alpha}(f)$は64%であり、翻訳後修飾のバックグラウンド確率が0.8%であることを考えると非常に高い値となっています。
しかし、翻訳後修飾部位を参照するアテンションヘッドの数は少数でした。
プローブタスクの結果
事前学習モデルの各層において埋め込みプローブとアテンションプローブを行った結果が以下の図です。
オレンジ色のプロットが埋め込みプローブ、青色のプロットがアテンションプローブでの結果となっています。用いた評価指標（y軸）は二次構造予測ではF1スコア、結合部位予測とコンタクト予測では適合率(Precision)となっています。 ヘリックス、ターン・ベンド、ストランドなどの二次構造予測は比較的低い層における出力を用いても精度が出ていることがわかります。
埋め込みプローブでは精度が基本的に右肩上がりとなり、層を積み重ねるごとに少しずつ有用な情報が蓄積されていくことがわかります。一方で、アテンションプローブでは最終層で急にプローブタスクの精度が上昇しており、埋め込みとアテンションで表現のされ方が異なることがわかります。
まとめ
いかがだったでしょうか。本論文は自然言語処理における事前学習モデルの解釈手法をタンパク質言語モデルにも適用してみたという内容でした。アミノ酸という２０種類のトークンを用いた事前学習を行うだけで二次構造やコンタクトマップなどの高次構造に関しての内部表現が獲得されるというのは興味深い内容です。タンパク質の機能に重要な構造情報は進化の過程でも保存されているため、大量のデータでMLMタスクを解くことで隠れたシグナルを見つけることができるのだと考えられます。
今回は既知の生物学的特徴に関する評価がメインでしたが、事前学習モデルには我々が知らない生物学的知見が眠っている可能性があります。生命の謎を機械学習を用いて解き明かすというのはとてもワクワクしますが、機械学習における解釈性の向上がボトルネックになってきそうです。
今後の展開に注目です！ 
この記事に関するカテゴリー
