Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[2021] Highly accurate protein structure prediction with AlphaFold [paper-reading] #60

Open
yoheikikuta opened this issue Jul 25, 2021 · 18 comments

Comments

@yoheikikuta
Copy link
Owner

yoheikikuta commented Jul 25, 2021

論文リンク

https://www.nature.com/articles/s41586-021-03819-2
(Supplementary Information も同ページにリンクがある)

公開日(yyyy/mm/dd)

2021/07/15

概要

タンパク質の構造予測コンペである CASP14 において他チームを圧倒した精度を達成した AlphaFold2 の詳細を記した論文と supplementary information。
情報量がすごいので一言でまとめを書くというのは難しいので内容は以降のコメントで書いていくとして、とにかくあらゆるドメイン知識や機械学習知識を総動員して、考えられる最高のモデルを構築しましたという感じの論文。
多くのアイデアが詰め込まれてるので追えてない要素が大量にあるが、自分が興味を持ったところを中心に論文を読んでメモを書いていった。いや〜これは凄いなという気持ちにさせてくれる実に珍しい論文であった。

GitHub repository は https://github.com/deepmind/alphafold

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

DeepMind のブログ https://deepmind.com/blog/article/alphafold-a-solution-to-a-50-year-old-grand-challenge-in-biology が出てからというものそのモデルの詳細が気になっていた AlphaFold2 (出た当初は AlphaFold2 と読んでいたが、ナンバリングせずに新しいものを AlphaFold と呼ぶようにしているみたい。それに合わせて以降は AlphaFold と書く) だが、遂に論文とコードが公開されたので読んでみることにした。

@yoheikikuta
Copy link
Owner Author

Supplementary Information は 62 ページあり、その密度も高い。読み解くには機械学習だけの知識では不十分で biology の知識も必要になる。

これはちょっと全部読んで論文メモを残すのは大変すぎるので、biology に関する議論はほぼスキップしつつ、モデリングで自分が興味を惹かれた部分のみに着目して書いていくことにする。

@yoheikikuta
Copy link
Owner Author

生物の仕組みの理解や創薬などにおいてタンパク質の構造解析は重要。
現在は 10 万くらいの構造が同定されているらしいが、配列が既知のタンパク質は数十億とかあるらしく、一個同定するのに数日とか数ヶ月掛けている現状では解き明かすのにあまりにも時間が掛かり過ぎる。タンパク質同定がスケールできる計算機で実施できれば嬉しい。

物理的な相互作用(熱力学や統計力学)から構造を解き明かすのは理論的には興味深いが、複雑な系なので正確な構造解析が困難というのは想像に難くない。
もう一つの方法として evolutionary history に注目するというものがある。この言い回しは門外漢には馴染みがないが、これは既に構造が分かっているタンパク質を介して未知の構造を同定していく方法論のようだ。アミノ酸配列が類似していればその性質も似ている(類縁性)ことを利用していくもので、Protein Data Bank (PDB) としてその知見が蓄積されていっている。このあたりの知識は Multiple Sequence Alignment (MSA) という形で AlphaFold でも不可欠な要素として現れる。これは 3 つ以上の配列を進化的な対応関係に従って並べ替えることで、以下のようなイメージ。配列が似ていれば性質が似ているということから、似ているところを見つけてタンパク質の機能的に重要な部分を抽出したりすることができる。偉そうに書いたが実際にやったことはないのでよく分かってない。


https://ja.wikipedia.org/wiki/%E5%A4%9A%E9%87%8D%E6%95%B4%E5%88%97 から引用

このような蓄積はあるが、実験で似たような配列のものが同定されていない状況では、どちらの方法を用いても殆どのケースで実験で同定するものより遥かに劣る精度になるという状況だった。

実験に近い精度で構造を同定できる計算機手法を発見しました、というのが AlphaFold になる。
CASP14 https://predictioncenter.org/casp14/ という未発表のタンパク質構造解析コンペでは文字通り圧倒的な精度を見せつけたこと、モデルが公開されてからいろんな人に試されてその性能に偽りなしと認識されたこと、で大きな盛り上がりを見せている。

@yoheikikuta
Copy link
Owner Author

モデルの中身に入る前に性能を先に見ておく。
CASP14 において、median backbone accuracy が 0.96 Å RMSD_{95} だったとのこと。
他のモデルで最も良いものが 2.8 Å RMSD_{95} でそれ移行は団子状態であり、炭素原子のサイズが約 1.4 Å であることを考えると、これはズレが原子レベル以内に収まっているような真に驚くべき精度と言える。

そもそもこの評価指標は何かというと、タンパク質は残基の空間配位が重要で、その残基の空間配位を root-mean-square を計算しているものと思われる。RMSD_{95} の 95 は 95% の残基のカバレッジということらしく、恐らく主要なところを捉えるのが重要で、側鎖でそこだけびよーんと飛び出してるやつは外れ値的な意味もあるので含めずに計算してるのではないかという気がする。
(この辺は正確に定義を追っているわけではないのでちゃんと理解できているわけではない)

全部の原子を考慮したものも計算していて、その場合 AlphaFold が 1.5 Å RMSD で 2 位が 3.5 Å RMSD であり、やはり AlphaFold の性能が際立ったいる。

一番左が CASP の結果。エラーバーは 10,000 の bootstrap sample を使って検証した時の confidence interval ということだが、この bootstrap sample の作り方はよく知らない。全体の構造はあまり変えずに一部のアミノ酸配列だけ置き換えた、系統が同じタンパク質を引っ張ってきているのだろうか? PDB になければそもそも教師データがないので都合よく見つかるのか?という気もするが、誰か知ってる人がいたら教えてもらおう。

左から二番目はあるタンパク質のドメインに対する結果で、モデル予測が実験結果をよく説明していることが理解できる。
タンパク質ドメインとは、独立して進化してある機能を有する単位らしく、雑に言えば主たる構造をきちんと捉えられているかを判断しているということなのだろう。

左から三番目は側鎖についてもよくマッチしていることを示している。
側鎖とは主鎖から枝分かれしている構造で、例えばタンパク質を構成するアミノ酸は、アミノ基とカルボキシル基は共通だけど側鎖の違いによって疎水性か親水性かが分かれたりする。
図のものはタンパク質において重要な亜鉛の結合の部分で、zinc finger とか呼ばれたりするらしい。
そもそもメインの構造部分が精度良く当てられてないとそこから派生する部分もうまく当てられないという話はもちろんあるようだが、AlphaFold の性能はこういうところでも現れてますよというアピール。

一番右は 2180 個もの残基の巨大なタンパク質だが、これくらいの数になっても精度よく当てられているということですごいなぁとなる。

@yoheikikuta
Copy link
Owner Author

モデルの全体像は以下。
supplementary information を読めば分かるが、これはかなり簡約化したものであり、実際は特徴量生成部分などはもっとずっと細かいし、この図には現れない実に様々な工夫がなされている。圧倒的な性能を持つモデルを作るにはこれだけの知見が詰め込まれているというのを肌で感じるのに良い機会なので、興味があれば supplementary information を眺めてみるのをオススメしたい。これが世界を変えるモデリングや!っていうパワーを感じる。

問題設定としては 3D 構造が未知のタンパク質のアミノ酸配列(これは一次元情報)を input にして、3D の構造を output にするモデルを構築したいというものとなる。ポイントだけかいつまんで書くと以下のようになる。

  • input はアミノ酸配列
  • input の情報からモデルに流し込む前の hand-crafted な特徴量を作成
    • 遺伝的なデータベース検索から似た部分配列を持つ配列を検索し、MSA を構築する。ここはドメイン知識がふんだんに使われていて、先述のように似た配列は似た構造をする類縁性というものがあるので、似た配列のものをガッと集めてくればそこにはいま問題としている未知の構造を解き明かすための情報が含まれているだろうということ。これはもちろん PDB というこれまで人類が蓄積してきたデータを有効活用するということになっている。
    • 単なるアミノ酸配列のペアを構築。これは単にアミノ酸配列のペアを作っただけで、入力するアミノ酸配列の情報から 3D 構造を解き明かすための特徴量をモデルに学習させようというもの。ここは機械学習感が全面に出ている。
    • 部分的に構造が似ているものを検索し、それを template として入力に使うという部分もある。詳細は全然把握してないが、軽く読んだ感じ、部分的にここの残基とここの残基がペアになるよという補助情報を入れて局所的な構造を形成するのに役立つ特徴量を提供していそう。
  • これらの特徴量をモデルに流し込む。モデルは大まかにいって 3D 構造を形成するための特徴量を抽出する Evoformer と、その特徴量を使って 3D 構造を形成する Structure module の二つがある。
    • Evoformer では MSA (これは evolutionary history の情報) と配列の pair-representation を混ぜた上でそれぞれの特徴量と同じ次元の特徴量を再構成する。モデルの中で混ぜているというのが重要で、これによって pair-representation という(一次元の)空間的配位に evolutionary history の情報を埋め込むことができる。この Evoformer は残基間の関係性をうまく取り扱うための工夫がなされている。例えばある残基のペアが重要かどうかを判断する際に、別のペアとの結びつき(特徴量)を考慮して三角関係をベースに特徴量をアップデートするようなモデリングにしているとか。これに関しては後でもう少し詳しくみる。
    • いい感じの特徴量が抽出できたらそれを 3D 情報にする。これは具体的には残基間の相互作用を考えずに角度と位置を与えつつ、それらの残基に紐づく側鎖の角度も与えるというモデルを構築する。これだけだと生物学的におかしな構造になったりする(ペプチド結合のところがどうこうとか書いてあったが全然理解してない)ので、fine-tuning のときに分子動力学に基づく loss term を加えてそれを是正するということをやっている。生物学の知識に基づいて geometry がおかしくない構造を探索的に探すというのは大変な感じがするが、ここを loss に置き換えてデータによる学習でカバーする、とかは機械学習の良さを活かしている。ここの structure module にも既存の生物学的な知識がふんだんに使われている(全然追えていないが、原子の種類と torsion angle からありうる local 構造を grouping してその情報をモデルに取り入れたりしている)。
  • モデルに一回通すだけだと複雑な構造を予測する際には精度が不十分らしく、3 回ほど出力をもとの入力に足してもう一度モデルを通すという recycling をしている。
  • 最終的な出力を様々な loss を合わせた loss で学習。学習には FAPE という距離の三角不等式を満たすように構成した loss を使ったり、labelled data だけでなく noisy student を使った unlabelled data での学習などもしている。

@yoheikikuta
Copy link
Owner Author

情報量が多すぎて全部を書いていくのは無理があるので、Evoformer や Structure module で何をしているのか、というところを中心に見ていくことにする。タンパク質の知識とかがないので AlphaFold を理解するのに必要な知識がないので、入力特徴量の細かい話とかモデリングの際に使っているドメイン知識などはいったん諦める。

ドメイン知識をふんだんに使うというのも間違いなく AlphaFold の根幹を成す部分の一つなので、こういうのは詳しい人に聞いてみたいなという気持ちはある。

@yoheikikuta
Copy link
Owner Author

入力のところに言及はしないといったが、supplementary information の図を貼ると以下のようになっている。
先ほどは登場しなかった extra_msa_feat などが登場し(ちなみにこれは cluster center として選ばれなかったものとか書いてあるが、全然なんなのかは分かってない)、様々な特徴量や変換を駆使していることが見て取れる。説明や shape については supplementary information に記載されている。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

Evoformer や Structure module の説明に入る前に、全体の inference の流れについても supplementary information のものを貼っておく。

ざっくりいうと以下。これは特徴量やモデルの詳細を置いておいて、recycling がどのように取り込まれているかが理解しやすい。
supplementary information にはこういう情報がふんだんに記述されていて、なかなかに読み応えがある。すごい。

  • ベースとなる特徴量を準備
  • recycling の情報を加える
  • template の情報を加える
  • Evoformer に入れて 3D 構造構築のための特徴量を抽出
  • これらは ensemble して Structure module への入力として 3D 構造情報を得る
  • recycling を何度か繰り返し、最終的な結果を得る

これだけ見ても理解できないところは多いが、入力特徴量の諸々を無視すれば、Evoformer と Structure Module がモデリングの肝であることが分かるので、それを以降見ていこう。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

Evoformer から。
とは言っても Evoformer 単体でもかなり色々な要素があるので、かいつまんで。

まずは大雑把な全体像は以下。
ここで重要なのは 3 つ。

  • MSA の特徴量の更新部分(上のライン)。MSA はその構築方法から明確に行ごと(これはある一つのタンパク質アミノ酸配列の情報)と列ごと(MSA は異なるタンパク質アミノ酸配列で似た部分配列が揃うように構築しているので進化的な意味合いで対応する部分の情報)の情報が重要である。そのため、row-wise gated self-attention と column-wise gated self-attention を使っている。self-attention の部分はいわゆる axial attention になっていて、gated の部分は行ごと、列ごとに重要とモデルが判断するところのみを使うようにしている、という感じである。
  • 残基の pair representation の更新部分(下のライン)。triangle という文字が見えるが、これが Evoformer のモデリングの肝。残基はペアの関係性が重要で、それはすなわち科学的な結合あるなしに対応する。ただし、単純な残基ペアだけでその化学結合が決まるわけではなく、当然他の残基との関係性が重要になる。Evoformer ではそれを 3 つの残基の関係性に着目してモデルに取り込んでいる。有効グラフとして捉え、pair i,j の edge を更新する際に、別の k との edge の情報を考慮するというところから、triangle という名がついている。これについては後でもう少し詳しく見る。
  • 上下のラインが mix されている縦のラインも重要。これによって進化的な情報と配列の 1 次元的空間情報が混ざり合って情報としてより rich になる。このように性質の異なるものをどう混ぜるのかというのは、機械学習では embedding の値を足したり掛けたりするだけで、あとは学習によってよしなにやってもらうというのは王道の一つだろう。

MSA の方はモデリング的にまあ新しい観点があるわけでない(どうやって pair representation の情報を bias として足し込むかとかのテクニカルな詳細はあるけど)ので、pair representation の triangle 云々の方を見ていく。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

まず、残基の pair representation を配列で表現している部分は、残基を node でその間の関係性(これが残基間の特徴量となり、配列の値となる)を edge とする有効グラフとして解釈することができる。

前述のように、残基のペアの特徴量を考えるときに、それらとは別の残基を考えて、三角関係を考慮してモデリングをする。これはペアだけで扱う場合では得られない相互作用を考慮できる。(もちろんもっと複雑な相互作用を考えることはできるが、それはモデルがより複雑になることを意味するので、始めるならまずは三体の関係性ということだろう)

注目するのは $ i \rightarrow j $ の特徴量をどう更新するかという点で、大きく分けて triangle update と triangle self-attention というものを構築している。triangle update に関しては有効グラフなので outgoing/ingoing の二つの場合があり、triangle self-attention の場合は始点 $ i $ に注目するケースと終点 $ j $ に注目するケースがある。どちらのケースも $ k $ という第三の node との関係性の情報を使って更新していることが見て取れる。

まずは triangle update の方から。これは正式には triangular multiplicative update と呼んでいる。
概念図は以下。

$ i \rightarrow j $ の情報を更新したく、この entry の値を gate として扱いつつ、$ i $ と $ k $ の関係性と $ j $ と $ k $ の関係性(横に長い赤枠の部分)を足し込んでいる。left edges と書いてるのは $ i \rightarrow k $ の特徴量で、right edges と書いてあるのは $ j \rightarrow k $ の特徴量。丸印は要素積で、ポイントとしては left edgesright edges の要素積を取った後の sum r で、ここが第三の node との特徴量の情報を足し上げている部分となる。$ i \rightarrow i $ みたいな自己ループも取り入れてるのが分かる。
グラフ構造として捉えて triangl 云々とか言われると、グラフ理論とか詳しくない人からするとなんかちょっと近寄りがたい感じがするが、やっているのは結局この行列演算だけであるので、これで自信を持って AlphaFold の triangular multiplicative update を理解したと言える。
これは outgoing edge で行方向に和をとっているが、incoming edge を使う場合は列方向に和を取ればよい。

次に triangular self-attention について。
これの概念図は以下。

始点 i に関する edge を更新することを考える($ i \rightarrow j $ を含む $ i \rightarrow $ node のことなので行方向の赤枠が update したい部分)。
dot-product affinities の部分が普通の self-attention になっている部分で、central edges が query part で left edges が key part になっている。これらは最初の三角形のグラフでいうところの $ i \rightarrow j , i \rightarrow k $ に該当する部分となる。これに $ j \rightarrow k $ 部分として right edges の情報を bias として足した後、value part である left edges を掛けて $ k $ 方向に和を取って(さらに gate を通して multihead の部分を concat して)$ i \rightarrow j $ 部分の update (このとき $ j $ 方向は全部の node を走る)が得られる。
言葉で書くとどうしても複雑な感じになるが、shape を意識して図を追っていけば、使っているのは self-attention なので理解は難しくない。これに関しては疑似コードも合わせてみておいた方が理解が容易になるので合わせて貼っておく。
終点 $ j $ に関する edge の更新は適宜行列の読替えをするだけなので、こちらも単純。

supplementary information が実に丁寧に情報を残してくれているので、ちゃんと読めば理解はたやすくできる。
これは作るの大変だっただろうなぁ、圧倒的に成果を出してさらに後続の人たちが理解しやすいように情報もちゃんと整備して...頭が上がらないですね。

一つ一つを追っていけば全く難しいところはなく、Evoformer の pair representation 部分のモデリングに関しては理解したと言えるレベルまで来た。MSA の方は技術的には axial attention が肝で、MSA と pair representation の特徴量をどう混ぜるかの部分に関しても shape をどう合わせるかで注意は必要だが単に bias として足したり外積をとったりするだけで難しいことはないので、割愛。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

続いて Structure module の方を見る。
こっちも一つ一つの要素は理解できないほど難しくはないのだが(ドメイン知識を使っているところは別)、いかんせん要素が多くて複雑なので全部追うのはかなりしんどい。というか全部は追えてない。
とりあえず気になるところだけ追って、それ以外はまた機会があればということにしておこう。

全体の疑似コードは以下。長い。
やりたいことは、Evoformer で抽出したアミノ酸配列部分の特徴量 $ s $ と pair representation 特徴量 $ z $ を入力として、全ての原子の 3 次元配位 $ x $ を得ること(さらに内部的なモデルの正確性ということで信頼度スコアも出しているが、今回は割愛する)。
もう一つ重要な登場人物としてユークリッド空間での変換(回転と平行移動)を表現する $ T $ が存在する。これは local frame から global frame への変換を担うものである。これの意味するところは、3D 構造を予測するときに、回転と平行移動をしても全体の構造は不変だが、正解データの配位と比べるときにうまく比べることができないという問題が発生するので、基準を揃えるための変換ということになる。雑に言えば、正解データの原子の位置を基準として frame を合わせるようにすれば、原子間距離とか比較できるようになるよねという気持ちだろう。細かいことを言うとスケール不変ではないので距離の単位は意識しないといけないとか、変換を求めるのは言うほど簡単ではないのでこれも学習で求めていくとか、そういう話もある。

疑似コードでは学習のことも考慮してるので正解データも含めて loss 計算なども書いているが、一旦それ(と確信度の計算部分)を忘れると大雑把には以下のような手順になっている。

  • アミノ酸配列部分の特徴量 s と pair representation 特徴量 z と変換 T を InvariantPointAttention という構造に入れて、特徴量 s を更新する。
  • その s の情報を使って Backbone、これはすなわち変換 T を更新する
  • s をいくつかの layer を介して特徴量 α として、この α と T で 3D 配位 x を computeAllAtomCoordinates から得る

上述の手順についてもう少し詳しく見ていく。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

まずは Invariant Point Attention (IPA) に関して。
ここでやりたいのは、特徴量 $ s $ の update だが、 3D 配位の不定性を考えると、ある local frame (これは $ T $ でパラメトライズされる) の元で計算可能であるけれど、global なユークリッド変換 $ T_{global} $ の下では不変な計算にしたいという動機が生まれる。これは当然 global なユークリッド変換を施してもタンパク質の持つべき性質は不変であるべきなので、それを取り入れたいというものである。テクニカルな解決方法としては単純で、local frame $ T_i $ が入り込む計算に関しては、例えばノルムを計算することでその以前性を消したり、逆変換 $ T^{-1}_{i} $ を作用させて不変性を保つということをすればよい。

疑似コードで書くと以下のようになる。

よく分からない単語や量がいくつか出てきているが、$ T_i $ が絡む部分は先ほど述べたような構造になっていて、この一連の計算は global なユークリッド変換の元で不変であることが分かる。

概念図は以下のようになる。
結構色々な要素が張り込んできて複雑だが、各計算部分は特に難しいことをしておらず、最終的に対象としている配列特徴量を更新したものが得られるような構造になっていることが分かる。

local frame の部分が backbone frames として記述されているが、これは 3x3 の回転行列と各次元への平行移動 3 自由度が、各残基について与えられているものである。
最終的には全ての原子の 3 次元配列を得たいという話をしたが、基本的にはタンパク質の基本的な構成要素である残基を基本的な unit として取り扱う。

ちょっとめんどくさくなってきたので各計算についてどういう意味かを詳細に記述するのは省くが、pair representation と single representaion と 3D の距離除法などを合わせた特徴量を作った上で、attention を使って特徴量を更新していくというものになっている。

@yoheikikuta
Copy link
Owner Author

yoheikikuta commented Jul 25, 2021

次は Backbone update だが、これは特徴を入力にして、3 次元の回転情報と並進を出力とする Linear layer を使うというだけである。
3 次元の回転は 4 元数でパラメトライズしている。
オイラー角とかでもいいんじゃないかと思うが、回転軸が 2 つ重なって自由度が一個落ちてしまったりする特異的な状況(知らなかったけどジンバルロックと言うらしい)を避けるためということなのかな?ここは詳しく調べてないから知らない。
ここはその Linear layer で $ T_i $ を更新するということなので特に難しいことはない。

@yoheikikuta
Copy link
Owner Author

次は computeAlAtomCoordinates のところ。
ここは正直ちゃんと追えてない... ドメイン知識が入り込みつつ、配位のねじれの情報(7 自由度)でパラメトライズしつつ、座標を求めていくということになるが、さらに literature atom coordinate などが出てきて理解ができてない。

お気持ちとしては、IPA から抽出した特徴量に変換を加えて対象とする残基のねじれ角の情報を取得 → 残基の種類とねじれの自由度でパラメトライズされる ideal literature atom coordinate というもので各原子の相対的配位を獲得(ここがどうやってるか全然追えてない。rigid group という決まったパターンで表現できるものを使っているが、まあ全然分からん) → ねじれ角の情報と local frame の情報を合わせて literature frame という literature atom coordinate から基準となる frame への変換を構築する → ideal literature atom coordinate から基準の frame の配位を得て、欲しかった 3D 配位をゲット、というものになる。

う〜んごちゃごちゃ書いたけど全然分かっていませんと白状しているのと同義な文章だな。
ここはまあ今度やる気が出たときにもっと真面目に追うことにするか。

他にも色々な要素がある(180度の回転対称性があるのでその不定性を取り除くとか)のだが、この辺も全然追えてないので諦めて、とにかくお望みの 3D 配位が得られたものとしよう。そうしたらあとは loss を設計して学習しようという話になる。


ここの 3D 配位を求めるところでは、各残基の位置については明示的に制限を加えることなく、自由に位置を探し求めている。本来は構造があるのでその位置関係は化学的な相互作用を反映していて複雑な拘束条件があったりするはずだが、この段階ではそれを無視している。論文ではその相互作用を無視しているということから residue gas と言ったりしている。これだと化学的にありえない構造になったりしうるので、そこは loss term を入れて fine tuning したりする(後述)。

@yoheikikuta
Copy link
Owner Author

色々なものを無視して強引に進めてきたが、最後に loss function だけ見ておく。
学習は二段階で、 まずは Initial training として上の loss function で学習した後に、fine-tuning として下の loss function で学習する。下の loss function で入れてるのは、詳細は理解してないけど、分子動力学を用いたりして前述の化学的におかしい配位が出たりする場合を是正する loss termi というもの。Amber relaxation とかいうらしく、AMBER99SB force field の最小化をするというものらしい。これはこれでまた別の知識がいるので詳細を追うのはとりあえず置いておいて、構造的な是正も strict にやろうとるするのではなく loss term として表現して数値的に最適化している、ということだけ理解しておくことにする。

ということで上の loss だけをそれぞれ軽く見る。

  • FAPE は Frame Aligned Point Error で、予測結果の配位と正解データの配位をそれぞれの local frame の逆変換で引き戻して、RMSD を計算する(そして無次元量になるように 10 Å で割る)というもの。3 次元配位なので chiral property に注意が必要で 入力とそのパリティ変換をしたものの FAPE が大きな正の値になるとか、三角不等式を満たすとか、色々解説はされているけど、やりたいことは予測結果の配位と正解データの配位の原子の位置の RMSD を計算するということ。
  • auxiliary loss は(各 layer それぞれでの)FAPE とねじれ角の loss の平均としている。ここも詳細は割愛するけど原点に shrink しないようなちょっとしたテクニックが入ったりしていて芸が細かい。
  • distogram loss は予測結果として使うものではないが、残基間の pair representation 特徴量を symmetrize した z_{ij} + z_{ji} を距離に map する layer を作って学習し、そこの差分が小さくなるようにという loss である。これは bin に分けた距離を使って cross entropy として処理している。
  • msa loss は masked MSA prediction の loss で、これは masked language modeling を模して、最終的な MSA 特徴量の情報を使い、アミノ酸のタイプがどれであるかを mask して当てるというものになっている。こういうところでも masked language model の知見が活かされているんだなぁ。そして loss の係数が大きいのでこれは相対的に重要とのことらしい。
  • conf は model confidence prediction の loss とのこと。model confidence のところは今回触れないが、confidence の loss ってなに?という点に関しては pre-residue lDDT-Cα score なるものを計算することが確信度(intrinsic model accuracy)になるとのことでこれを使っているらしい。

loss もかなり色々と工夫して作っていて、ドメイン知識や機械学習の知識の総合格闘技という感じで、いや〜すごいなという感じ。
正直全然追えてない部分が大量にあるんだけど、さすがにしんどいので一旦はこの辺にしておく。

@yoheikikuta
Copy link
Owner Author

最後に ablation study と recycling の重要性に関する結果を貼っておく。

ablation study では modeling での triangles, biasing, gating が重要であることや、recycling が重要であること、そして論文に含まれてる要素はどれも重要であることが示されている。(CASP14 と PDB test set の結果から見るに。特に PDB test set の結果は顕著。)

recycling は青やオレンジの比較的単純な構造に関しては Evoformer block の割と初期に高い精度までいくが、緑のものは複雑らしく、rこれはかなり block を経ないと(モデルの block は 48 なので、それ以上は recycling してることに対応)精度が上がらないことが示されている。これを見ると recycling めっちゃ重要じゃんとなる。

@yoheikikuta
Copy link
Owner Author

ということで抜け落ちてる要素は大量にあるが、自分が興味のある部分を中心に一通り追ってみた。
あれだけの衝撃を与えたモデルということもあり、実にさまざまな要素が含まれていることがわかった。ドメイン知識も存分に使われているし、機械学習知識も存分に使われている。凄いモデルはこうやって作るんや!というパワーを感じることができてとてもよかった。

まだまだ把握してない要素が大量にあるし、コードも Google Colab のものをちょろっと動かしたに過ぎないので、今後も機会があったら少しずつ理解を深めていきたいところだ。

@yoheikikuta
Copy link
Owner Author

Podcast で紹介した: https://anchor.fm/yoheikikuta/episodes/16AlphaFold-e15ubme

@yoheikikuta yoheikikuta changed the title [2021] Highly accurate protein structure prediction with AlphaFold [2021] Highly accurate protein structure prediction with AlphaFold [paper-reading] Aug 20, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant