CTC 教育サービス
[IT研修]注目キーワード Python UiPath(RPA) 最新技術動向 Microsoft Azure Docker Kubernetes
今回は、2022年に公開された論文「PaLM: Scaling Language Modeling with Pathways」を元にして、大規模な自然言語モデル「PaLM」を複数のTPU Podから構成されるPathwaysの分散学習インフラで学習した事例を紹介します。
第133回〜第134回の記事では、Googleのエンジニアが開発した、機械学習モデルの新しい分散学習インフラ「Pathways」について解説しました。その中で、機械学習モデルとしてのPathwaysと分散学習インフラとしてのPathwaysの違いについて次のように説明しました。
・機械学習モデルとしてのPathways:さまざまな役割を持つ「ブロック」を相互接続することで、複数のタスクに対応した機械学習モデルを構成するという考え方
・分散学習インフラとしてのPathways:複数のTPU Podを疎結合した分散学習インフラ
分散学習インフラとしてのPathwaysは、将来的に「機械学習モデルとしてのPathways」、すなわち、複数のタスクに対応した多数のブロックから構成される機械学習モデルの学習に役立つと期待されていますが、現時点では、実際にそのようなモデルが完成しているわけではありません。しかしながら、この新しい分散学習インフラには、従来の機械学習モデルを学習する上でも、複数のTPU Podを利用できるという点で有用性があるはずです。そこで、冒頭の論文では、この観点での実証実験の1つとして、従来型の自然言語モデルをPathwaysの分散学習インフラを用いて学習しています。具体的には、6,144個のTPUチップを用いて、5,400億個のパラメーターを持つモデルの学習に成功しています。この後で紹介するように、論文の中では、「Few-shot learning」と呼ばれる自然言語モデルに特有のタスクにおける興味深い結果が紹介されています。
この論文の中では、「トランスフォーマー」と呼ばれる仕組みを利用した自然言語モデルの学習を行っています。具体的な仕組みは割愛しますが、(途中で途切れた)文章を入力すると、その直後に現れる単語を予測するという機械学習モデルになります。分かりやすい例を示すと、「This is a」という文章を入力すると、「penの確率は〇〇%」「bookの確率は〇〇%」といった具合に、その次に来る単語を確率付きで予測します。さらに、確率が最も高い単語を選択して、それを末尾に加えた文章を再度入力すると、さらにその後に来る単語が予測できます。これを繰り返すことで、ある文章に対して、それに継続する一連の文章を生成するといった事もできます。
トランスフォーマーの仕組み自体はこの論文が公開される以前から知られていましたが、この論文で採用したモデル「PaLM」は、5,400億個という膨大なパラメーター数にその特徴があります。この実験では、このモデルを図1のように、2つのTPU Podに分散配置して学習処理を行っています。
図1 自然言語モデル「PaLM」を2つのTPU Podに分散配置する方法(論文より抜粋)
具体的には、同一のモデルのコピーをそれぞれのTPU Podに配置することで、TPU Pod内部でのモデル分散と、TPU Pod間でのデータ分散を併用しています。さらに、学習処理においては、(A)「学習用のバッチデータからパラメーターの勾配を計算する」という処理と、(B)「得られた勾配を用いてパラメーターを更新する」という2種類の処理を交互に行う必要があります。図1の左にある(A)と(B)はこれらに対応しています。
そして、(A)と(B)では、TPU Pod上に異なる計算処理のコンポーネントを配置する必要があります。ここでは、前回説明したリソーススケジューラーの機能を用いて、動的にコンポーネントを入れ替えます。これにより、(A)と(B)のそれぞれの処理ですべてのTPUコアを無駄なく使用することができます。また、(A)の処理において左右それぞれのTPU Podで計算した勾配の値は、データセンターネットワークを介して、もう一方のTPU Podにも送られます。このように2つのTPU Podを連携する事で、1つのTPU Podの場合に比べて、理想的には2倍の学習速度が得られると期待できます。実際の実験では、1.95倍の学習速度を達成しており、この分散学習インフラの仕組みが理想に近い性能を実現していることがわかります。
先ほど触れたように、この論文では学習済みのモデルによるFew-shot learningの結果が紹介されています。ここでは、まず、Few-shot learningの仕組みについて説明しておきます。前述のように、ここで学習した自然言語モデルは、入力文に対して、それに続く単語、もしくは、一連の文章を予測・生成することができます。そこで、たとえば、「Was George Washington the first American president?(ジョージワシントンはアメリカの初代大統領ですか?)」といった、Yes/Noで答えられる単純な質問文をこのモデルに入力するとどうなるでしょうか? このモデルは、世の中に存在する膨大なテキスト文書を学習データにしていますが、世の中に存在する文書の多くでは、この入力文に対しては、「Yes」で始まる回答文が続いていると想像できます。したがって、このモデルが正しく学習できていれば、この後に続く単語の確率は、「Yes」の方が「No」よりも高くなるはずです。つまり、この後に続く単語が「Yes」である確率と「No」である確率を比較することで、質問の回答が得られます。
一方、Yes/Noだけでは答えられない質問の場合は、それほど簡単ではありません。質問文を入力して、それに継続する文章を生成することはできますが、それが必ずしも期待するような形式の回答になるとは限りません。そこで、はじめに、質問文と回答文のペアを代表例として用意しておき、これに続けて新しい質問文を加えて入力するということを行います。つまり「質問1、回答1、質問2」という形式の文章を入力するのです。すると、機械学習モデルは、「質問1と回答1」のペアを参照することにより、質問2に対しても質問1と同様の形式で回答を出力しようとします。このようにして、本来の質問の前に、参考となる質問と回答のペアを数個付け加える手法を自然言語モデルにおける「Few-shot learning」と呼びます。
そして、図2は、論文の中で紹介されているFew-shot learningの実行例になります。ここでは、小学生レベルの算数の文章題に対する回答を生成する実験を行っており、左右に2種類の結果が示されています。
図2 Few-shot learningの実行例(論文より抜粋)
まず、左側は「失敗例」になります。参考例となる最初の質問に対して、「The answer is 11.」というシンプルな回答例を与えています。機械学習モデルは、それに続く2つ目の質問に対して「The answer is 50.」という回答を生成しており、一見すると最初の回答例と同じ形式で正しく回答しているようにも見えます。しかしながら、問題の内容にあった正しい答えは「9」なので、算数の問題の答えとしては間違っているのです。一方、右側は「成功例」になります。ここでは、最初の質問の回答例として、途中の計算過程を含む文章を与えています。その結果、機械学習モデルは、2つ目の質問に対しても同様の計算過程を含む文章を生成して、最終的に正しい計算結果を得ています。つまり、この機械学習モデルは、入力文として与えられた内容から、期待される回答を得る方法(推論過程)をその場で学習していると考えられます。
冒頭の論文では、この他にも「ジョークの内容を解説する」、あるいは、「プログラミング言語を用いて指示された処理を行うコードを生成する」といった例が紹介されています。いずれも、質問と回答のペアの例を数組与えるというFew-shot learningの手法が利用されており、最初に与える例を工夫することでさまざまな問題に対応できることがわかります。ディープラーニングを利用した自然言語モデルは、このような意味での汎用性を持つことが以前から知られていましたが、今回の実験結果は、その汎用性の高さをさらに高いレベルで実証したものと言えるでしょう。
今回は、2022年に公開された論文「PaLM: Scaling Language Modeling with Pathways」を元にして、大規模な自然言語モデル「PaLM」を複数のTPU Podから構成されるPathwaysの分散学習インフラで学習した事例を紹介しました。論文の中では、本文で触れたFew-shot learningのさまざまな結果が詳しく説明されていますので、興味のある方は、ぜひ論文の方も参照してください。
次回は、Googleのデータセンターネットワークの新たな進化を紹介したいと思います。
Disclaimer:この記事は個人的なものです。ここで述べられていることは私の個人的な意見に基づくものであり、私の雇用者には関係はありません。
[IT研修]注目キーワード Python UiPath(RPA) 最新技術動向 Microsoft Azure Docker Kubernetes