大規模言語モデルを支える 分散学習インフラ Pathways 中井 悦司 / Etsuji Nakai Google Cloud / ソリューションズ アーキテクト このスライドはコミュニティイベント「GCPUG Shonan vol.90」での発表資料です

自己紹介 中井 悦司 / Etsuji Nakai Google Cloud / ソリューションズ アーキテクト

3 Introducing Pathways: A next-generation AI architecture | The Keyword | Google https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/

4 Pathways Language Model (PaLM): Scaling to 540 Billion Parameters for Breakthrough Performance | Blog | Google Research https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html

Pathways のアイデアが 生まれた背景

6 ディープラーニング実用化のきっかけ（2012 年の発表） https://diamond.jp/articles/-/165402?page=2 https://blog.google/technology/ai/usin g-large-scale-brain-simulations-for/

7 ディープラーニングが実用化された要因 ● 学習アルゴリズムの研究 ○ 膨大な数のパラメータを効率的に最適化する手法 ● 学習データの蓄積 ○ 膨大な数のパラメータを最適化するのに十分なデータ量 ● 計算リソース ○ 並列計算技術、GPUを用いた計算技術の発展

Google でのディープラーニングの利用実績 https://www.slideshare.net/embeddedvision/largescale-deep-learning-for-building-intelligent-computer-systems-a-keynote-presentation-from-google Used across products: Unique project directories Time 8

9 Neural scaling laws ● 自然言語モデル／翻訳モデルでは、モデルの基本構造が同じであれば、「計算量、 モデルサイズ、データ量」が予測性能により大きな影響を与えるという経験則（2020 / 2021 年） Scaling Laws for Neural Language Models | arxiv | Cornell University https://arxiv.org/abs/2001.08361

10 Neural scaling laws ● GPT-3 モデル（2020 年）：1,750 億パラメーター ● GLaM モデル（2021 年）：1 兆 2,000 億パラメーター ○ https://ai.googleblog.com/2021/12/more-efficient-in-context-learning-with.html GLaM: Efficient Scaling of Language Models with Mixture-of-Experts | arxiv | Cornell University https://arxiv.org/abs/2112.06905 GPT-3 に対する GLaM の性能向上

11 モデルサイズの拡大による応用範囲の広がり ● LaMDA モデル（2021 年） ○ 会話文を学習データとすることで、「テーマに沿った会 話の繋がり」を学習 ○ 3 つの観点でチューニング 意味がある内容 安全な内容 事実に基づいた内容 ○ GAN の学習手法を応用（内容の安全性などを識別す る識別モデルを利用してチューニング） 「ブログのアイデア」につい てアドバイスする LaMDA モ デル https://blog.google/technology/ai/lamda/ https://ai.googleblog.com/2022/01/lamda-towards-safe-grounded-and-high.html

12 「マルチモーダル」なデータを扱うモデル ● 入力と出力でデータの種類が異 なるパターン ○ 例：「テキスト」→「画像」 ● XMC-GAN（2021 年） ○ 学習方法を工夫する事で、従来 よりも高い性能を達成 テキストから画像を 生成するモデル Cross-Modal Contrastive Learning for Text-to-Image Generation | Blog | Google Research https://ai.googleblog.com/2021/05/cross-modal-contrastive-learning-for.html

13 「マルチモーダル」なデータを扱うモデル ● Imagen (2022) Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding | arxiv | Cornell University https://arxiv.org/abs/2205.11487

14 複数の「命令方法」によるロボット操作 BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning | bc-z https://sites.google.com/corp/view/bc-z/home Place grapes in ceramic bowl 自然言語による指示 実技映像による指示 ● BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning（2021 年） ○ 自然言語、もしくは、実技映像による指示でロボットを操作するモデル 指示に対応した ロボットの動作

15 異なる表現形式を統一的に取り扱う仕組み BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning | OpenReview.net https://openreview.net/forum?id=8kbp23tSGYv 自然言語と実技映像を 同一の潜在空間に マッピング ● 指示の「意味」を表した共通の潜在空間（特徴量空間）を介することで、複数の表現 形式を統合

16 異なる表現形式を統一的に取り扱う仕組み 自然言語と実技映像を 同一の潜在空間に マッピング ● 指示の「意味」を表した共通の潜在空間（特徴量空間）を介することで、複数の表現 形式を統合 パーツを結合するための新たなパーツ BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning | OpenReview.net https://openreview.net/forum?id=8kbp23tSGYv

17 Pathways のモデル構造（予想） マルチモーダルな データを並列に入力 データの種類に応じた （内部）モデル 複数データの情報を 潜在空間でマッピング タスクに応じて処理を分岐 複数タスクの結果を 並列に出力 ● 役割ごとの「パーツ」となるブロック （サブグラフ）を相互接続することで、 マルチモーダルなデータに対して複 数のタスクを並列に実行 ● パーツを共有することで、全体の学 習がより効率的に進むと考えられる

Pathways の実現に向けた 技術開発

19 TPU Pod ● TPU Core を搭載したボードを専用のインターコネクトで相互接続 Google showcases Cloud TPU v4 Pods for large model training | Blog | Google Cloud https://cloud.google.com/blog/topics/tpus/google-showcases-cloud-tpu-v4-pods-for-large-model-training

20 TPU Pod による計算処理 ● モデルの各コンポーネントを TPU Core にマップし て、同期的に計算を実行 ● インターコネクトを介して、TPU Core 同士で計算結 果を直接に交換 ● インターコネクトによる高速なメッシュ接続がスケーラ ビリティを実現するポイント Pathways: Asynchronous Distributed Dataflow for ML | arxiv | Cornell University https://arxiv.org/abs/2203.12533 データセンターネットワーク ホスト TPU コア ML モデル

21 Pathways に対応した計算モデルの提案（2022 年） Pathways: Asynchronous Distributed Dataflow for ML | arxiv | Cornell University https://arxiv.org/abs/2203.12533

22 ● Pathways を構成するパーツごとに TPU Pod 内のスライス（小島）にマッピング ● パーツ内の計算は、インターコネクトを介して高速に実行 ● パーツ間のデータ転送は、高速なデータセンターネットワークを介して行う Pathways に対応した計算モデルの提案（2022 年）

23 Pathways に対応した計算モデルの提案（2022 年） ● コンポーネントがシステム全体にマッピングしきれない場合は、スケジューラーが動的に確保したス ライスにコンポーネントを割り当てていく ● パーツ間の依存性（計算順序）を考慮して、計算の並列度を高めるように配置する

24 TPU Pod の課題と解決策 ● モデルのサイズが大きすぎて TPU Pod 全体にマッピングできない場合、効率的な計算が困難にな る ● インターコネクトによる完全なメッシュ接続は、どこかのタイミングでシステムサイズの限界に達する ● これらの問題を解決するために、データセンターネットワークで疎結合した複数の TPU Pod を連携 させる新しい学習インフラを構築 ● この新しい学習インフラと Pathways のアーキテクチャーは親和性が高いと考えられる

25 Pathways のインフラを用いた学習処理の実施例（2022年） PaLM: Scaling Language Modeling with Pathways | arxiv | Cornell University https://arxiv.org/abs/2204.02311

26 Pathways のインフラを用いた学習処理の実施例（2022年） ● 学習対象のモデルは、Transformer を用いた自然言語モデル ○ 入力文書に対して、次の単語（の確率）を予測するシンプルなモデル ○ 予測した単語を再帰的に入力する事で、後続の文章を生成可能 ● 同一モデルのコピーを 2 つの TPU Pod にマッピングして、データ並列で学習 ○ 6144 個の TPU v4 チップを使用 ● スケジューラーを用いて、勾配計算と勾配適用（パラメーター更新）を交互に実行 複数の TPU Pod を用いた学習システ ムの実効性を確認する実験であり、 Pathways の複数タスクモデルを 実現したわけではありません

27 PaLM の性能 ● 従来よりも大きなサイズのモデルの学習に成 功 ● さまざまな自然言語タスクに対して、従来の最 高性能を更新 これまでのモデルとの パラメーター数の比較 これまでのモデルとの 性能の比較

28 （参考）0-shot, 1-shot, few-shot learning について ● 1-shot：「質問と解答の例」＋「質問文」を入力して、それに続くべき文章を予測させる ○ 回答の形式に対するヒントを与えることでモデルが回答を予測しやすくする ● few-shot：「質問と解答の例」を2〜3個程度に増やす ● 0-shot：質問文を入力して、それに続くべき回答を予測させる ○ Yes / No で回答できる質問文を入力して、それに続く単語が「Yes」である確率と「No」で ある確率をそれぞれ予測させる ○ 質問に対する回答をダイレクトに出力させる

29 PaLM による 1-shot learning の実行例 シンプルな解答例だと正 答に失敗する 推論過程を含めたサンプ ルを与えると正答する

30 PaLM によるプログラムコードの出力例 関数の先頭の コメント部分を入力 コメントに続く関数の本 体を出力

31 まとめ ● モデルサイズの拡大に伴う性能向上は今後も継続すると期待される一方、巨大なモ デルを学習するためのインフラには限界がある ● 複数の TPU Pod を疎結合させた新しい学習システムを提案・実装 ○ 従来のサイズを超える高性能な自然言語モデル（PaLM）の学習に成功 ● 将来は、Pathways が想定する複数タスク・マルチモーダルに対応したモデルの学習 に利用できると期待される

Thank You.