InstaDeep が Cloud TPU v4 を使用してサステナブルな農業を支援した方法
Google Cloud Japan Team
※この投稿は米国時間 2022 年 11 月 30 日に、Google Cloud blog に投稿されたものの抄訳です。
「食べ物が体をつくる」とは誰もが耳にしたことがある言葉です。ただし、実際には、私たちが食べるものは、少なくとも遺伝的には、私たち人間よりも複雑であることが多いのです。米粒を取り上げてみましょう。米をつくる植物である稲には、人間の 2 倍にのぼる 40,000 から 50,000 の遺伝子があります。しかし、私たちは、植物の生命よりも、ヒトゲノムの構成に関する知識の方を多く有しているのが現状です。特に、世界中の食料安全保障が気候変動によって悪化する可能性がある中、80 億人分の食料を供給するという緊急性のある課題に対処するには、この知識のギャップを早急に埋める必要があります。
このようなことから、AI 企業である InstaDeep は、Google Cloud と協力し、穀物と食用野菜の基準ゲノムのデータセットを使用して 200 億以上のパラメータを持つ大規模な AI モデルをトレーニングしました。このトレーニングでは、大規模なトレーニングの効率性を向上するのに特に適している Google の最新世代 Tensor Processing Unit（Cloud TPU v4）が使用されました。私たちの目的は、ゲノム配列から植物の農業形質を分析および予測できるツールを作成して、食料安全保障とサステナブルな農業の改善を図ることです。このツールを使用すると、作物の栄養価を上げ、栽培効率を向上させ、害虫、病気、干ばつに対する回復力と耐性を上げる遺伝子を特定できるようになります。
サステナブルな農業のためのゲノム言語モデル
農耕が始まって以来、私たち人類は直接的にも間接的にも、より良質で回復力が高い作物の収穫量を増やし、運が良ければ味も良い作物を栽培しようと努めてきました。何千年もの間、このために試行錯誤を重ね、毎年作物を育てながら、進化における突然変異から自然に生まれた最も有益な形質を特定し、それのみを保持するよう試みてきました。植物のゲノム配列が明らかになった現在、有益な遺伝子を直接特定して、新たな突然変異の影響を予測することが期待されます。
しかし、植物ゲノムは複雑であるため、多くの場合、有益な変異種を特定するのは容易ではありません。機械学習（ML）の革新的な進歩により、DNA 配列と分子表現型との関係性の把握が容易になりました。つまり、正確で費用対効果の高い予測手法により、遺伝情報と観察可能な形質との間のギャップを埋めることが容易になったのです。このような予測により、機能的な変異を特定して、どの遺伝子がどの形質に関連しているかを把握しやすくなるため、より良い作物を選択しやすくなります。
さらに、作物の遺伝子配列の膨大なライブラリがあるおかげで、何百もの植物ゲノムを用いて大規模なモデルをトレーニングすることができます。こうして、徹底的に研究された種から、あまり理解されていないものの特に開発途上国の食料生産にとって重要な種に、知識を伝達させることができるのです。さらに、これをデジタルで行うことで、一般的な作物品種と希少な作物品種の両方のゲノムを AI ですばやくマッピングし、これらのゲノムのアノテーションを行うことができます。
植物ゲノムにおいて従来の ML 手法が限界となる主な理由の一つとして、ほとんどが教師あり学習という手法に依存していることが挙げられます。教師あり学習にはラベル付きデータが必要です。そのようなデータは非常に少なく、収集するにも費用がかかるため、植物ゲノムにおける従来の ML 手法には大きな限界があります。Transformer アーキテクチャや BERT スタイルのトレーニング（Bidirectional Encoder Representations from Transformers）など、最近の自然言語処理（NLP）の進歩により、意味のある表現を学習するために、サイエンティストが生のテキストデータを使用して大規模な言語モデルをトレーニングできるようになりました。この「教師なし学習」という手法が流れを変えたのです。一度学習した表現は、ラベル付きデータが不足している場合でも、複雑な回帰タスクや分類タスクに対処するために利用できます。
InstaDeep が Google Cloud と提携し、TPU を使用してゲノム学用の新世代の AI モデルをトレーニング
プロテオミクスにおいては、大規模な言語モデルが特に効果的であることが研究者により実証されています。この仕組みを理解するために、アミノ酸を単語として、また、タンパク質を文として読み取ることを頭に浮かべてみてください。生のゲノムデータ（配列形式）が豊富にあるため、InstaDeep と Google Cloud は同様の技術をヌクレオチドに応用することにしました。この場合はヌクレオチドを単語として、また、大量のゲノム情報を文として読み解きます。
さらに、システムが学習した表現は、モデルとデータセットのサイズに比例して向上することが、NLP の調査研究で明らかになっています。この発見から、InstaDeep の研究者は、10 億から 200 億ものパラメータを持つゲノムのデータセットを使用して、より大規模な言語モデルをトレーニングすることにしました。
10 億のパラメータと 50 億のパラメータがあるモデルを、果物、穀物、野菜などの複数の食用植物の基準ゲノムで構成されるデータセットを使用してトレーニングしました。ヌクレオチドの数は合計で 750 億です。
トレーニング データセットは、モデルの容量と同じ割合で増やす必要があることが、最近の研究で示されています。そこで、ヒト、動物、非食用植物、細菌のゲノムなど、National Center for Biotechnology Information（NCBI）のデータベースで公開されているすべての基準ゲノムを集めて、より大規模なデータセットを作成しました。このデータセットを使用して、200 億のパラメータがある Transformer モデルをトレーニングしました。このデータセットは 7,000 億のトークンで構成されており、Common Crawl や Wikipedia のデータセットなど、NLP アプリケーションで通常使用されるほとんどのデータセットを上回る大きさになりました。
両チームは、植物ゲノムの研究をさらに推進するために、10 億のパラメータがあるモデルを科学コミュニティと共有することを発表しました。
これらのモデルが学習した、簡潔で意味のあるヌクレオチド配列の表現は、分子表現型の予測に伴う問題を解決する際に利用できます。その力量を示すために、さまざまな食用植物種の遺伝子機能と遺伝子オントロジー（遺伝子の属性）を予測するモデルをトレーニングしました。
このモデルでこれらの特性を高精度で予測できることが早期に実証されました。そこで、これらのモデルの意義を詳しく調べることにしました。その結果、多くの発展途上国にとって非常に重要な 3 つの植物種（キャッサバ、サツマイモ、ヤマイモ）のゲノムのアノテーションを行うことにしました。私たちは、こうしたアノテーションを科学コミュニティが自由に行うことができるように取り組んでいます。そして、自由なアノテーションにより新たなゲノム研究が導かれ、促進されることを期待しています。
大量のモデルとデータセットに伴うスケーリングの課題を Cloud TPU で克服
200 億のパラメータがあるモデルを数十億のトークンを使用してトレーニングするためのコンピューティング要件は膨大です。最新のアクセラレータでは、チップ 1 個あたりのピーク時のパフォーマンスは優れていますが、このパフォーマンスを発揮させるには、多くの場合、密結合したハードウェアとソフトウェアの最適化が必要になります。さらに、この効率性を維持しながら数百個のチップにスケールするとなると、システム設計上の別の課題が生じます。Cloud TPU の密接に結合したハードウェア スタックとソフトウェア スタックは、このような課題に対処するのに特に適しています。Cloud TPU のソフトウェア スタックは、XLA コンパイラをベースとしており、すぐに使える最適化機能（演算と通信のオーバーラップなど）と、並列処理を表現するための簡単なプログラミング モデルを備えています。
私たちは Google Tensor Processing Units（TPUv4）を利用することで、ゲノムの大規模モデルのトレーニングに成功しました。コードは JAX フレームワークを使用して実装されています。JAX は、関数型プログラミング ベースのアプローチにより、計算を関数として表現します。この関数は、XLA を利用した JAX API により簡単に並列化できます。これにより、システム設計上の課題に取り組むことなく、単一ホスト（チップ 4 個）構成からマルチホスト構成にスケールすることができました。TPU の費用対効果の高い相互通信機能および内部通信機能により、チップ数とトレーニング時間の間でほぼ線形のスケーリングが実現されました。これにより、1024 個の TPUv4 コア（チップ 512 個）のグリッドで迅速かつ効率的にモデルをトレーニングすることができました。
まとめ
最終的に、ディープ ラーニング モデルによって予測される変異したゲノムの機能的特徴は、ゲノムの編集と分析に大きく依存する新時代の農業に不可欠なものになると考えています。インシリコ変異誘発法（ゲノム領域で起こりうるすべての変化をコンピュータ モデルで評価する方法）などの新しいアプローチは、植物の適応性を向上させる変異の優先順位を決め、作物の改良を導くうえで有益なものになると考えています。ウェットラボの実験で同様の作業を行おうとしても、スケーリングが難しく、実質的にほぼ不可能です。私たちは、現在と将来のアノテーション プロセスを研究コミュニティに公開することで、育種技術の民主化を支援したいと思っており、その技術が世界の農業全体に利益をもたらすことを願っています。
関連情報
Cloud TPU v4 のハードウェア スタックおよびソフトウェア スタックの独自の機能の詳細については、Cloud TPU v4 の発表をご覧ください。また、スケーリング特性の詳細については、こちらのベンチマークをご覧ください。最後に、Cloud TPU で JAX と SPMD の並列処理を始めるには、PJIT の概要を読むことをおすすめします。
この研究は、Google の TPU Research Cloud（TRC）プログラムのサポートにより、この作業に不可欠な Cloud TPUv4 チップを使用できたことで実現しました。
