6144個のTPU、5400億個のパラメータ、バグ修正、ジョーク解釈など、GoogleはPathwaysを使用して大規模なモデルをトレーニングしました

Google の次世代アーキテクチャ Pathways は、大規模なモデルのトレーニングに使用されています。

規模が大きくなるにつれて、複数のタスクを処理する際のモデルのパフォーマンスが徐々に向上し、新しい機能が継続的に実現されます。

既存の AI モデルの限界について議論していたとき、Google の AI 責任者であるジェフ・ディーン氏は、今日の AI システムは常に新しい問題をゼロから学習すると述べました。最終的には、何千もの個別のタスクに対して何千ものモデルを開発することになりました。この方法で新しいタスクを学習すると、時間がかかるだけでなく、より多くのデータが必要になるため、非常に非効率的です。

ジェフ・ディーンの見解では、理想的な開発モデルは、モデルをトレーニングして何千ものことを実行することであるべきだ。このビジョンを実現するために、彼のチームは昨年、「Pathways」と呼ばれる汎用 AI アーキテクチャを提案しました。 Jeff Dean 氏は、Pathways は 1 つのアーキテクチャで複数のタスクを同時に処理するように設計されており、新しいタスクを迅速に学習して世界をよりよく理解する能力を備えていると紹介しました。しばらく前に、チームはついに Pathways 論文を発表しました。

この論文では、PATHWAYS は新しい非同期分散データフロー設計を使用していると述べられています。この設計により、PATHWAYS は単一のコントローラー モデルを採用できるようになり、複雑な新しい並列パターンをより簡単に表現できるようになります。実験結果によると、2048 TPU で SPMD (単一プログラム複数データ) 計算を実行すると、PATHWAYS のパフォーマンス (アクセラレータの使用率は 100% に近くなります) は SOTA システムに匹敵します。

Google Pathways システム アーキテクチャの概要。

強力なシステムを備えた次のステップは、モデルをトレーニングすることです。

Google は、先日発表された論文「PaLM: Pathways による言語モデルのスケーリング」の中で、5,400 億のパラメータを持つ大規模な言語モデル PaLM (Pathways Language Model) をトレーニングするために Pathways システムを使用したことを発表しました。

論文リンク: https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf

これはデコーダー専用の Dense Transformer モデルです。このモデルをトレーニングするために、Google は 6144 個の TPU を使用し、Pathways が 2 つの Cloud TPU v4 Pod で PaLM をトレーニングできるようにしました。

強力なシステムとコンピューティング能力への投資により、驚くべき成果がもたらされました。研究者らは、数百の言語理解および生成タスクで PaLM を評価した結果、ほとんどのタスクで SOTA の少数ショット学習パフォーマンスを達成し、ジョークの解釈、バグ修正、絵文字から動画を推測するなどの言語およびコードタスクでも優れたパフォーマンスを発揮できることを発見しました。

あるネットユーザーは「GoogleがTPUを開発している目的がようやく分かった。このレベルの自然言語理解がアプリケーションで活用され、十分に効率的で広く使われるようになれば、Googleのあらゆる活動に革命が起こるだろう。この目標を達成できるチップを持つことは非常に価値があり、Nvidiaのチップの購入やレンタルにこだわる企業は不利になるだろう」と嘆いた。

PaLM アーキテクチャの概要

PaLM はデコーダーのみで構成されており (各タイム ステップは自身と過去のタイム ステップのみに対応できます)、標準の Transformer アーキテクチャに次の変更を加えます (Vaswani ら、2017)。

SwiGLU アクティベーション

MLP 中間アクティベーションには SwigLU アクティベーション (Swish(xW) · xV) を使用します。これは、標準の ReLU、GeLU、または Swish アクティベーションと比較して品質が大幅に向上することが研究で示されているためです。 MLP では 2 回ではなく 3 回の行列乗算が必要になることに注意してください。ただし、Shazeer (2020) は計算上同等の実験で品質の向上を実証しています。

並列レイヤー

研究者らは、標準的な「シリアル」アプローチではなく、各 Transformer モジュールで「パラレル」アプローチを使用しました。具体的には、標準的な方法は次のように記述できます。

並列法は次のように書ける。

MLP と注意入力行列の乗算を融合できるため、ここでの並列アプローチにより、大規模なトレーニングを 15% 高速化できます。アブレーション実験では、8Bスケールでは品質低下は小さいが、62Bスケールでは品質低下が見られないことが示されており、研究者らは、540Bスケールでは平行層の影響が品質中立に達すると推測している。

マルチクエリアテンション

標準的な Transformer アプローチでは、k 個のアテンション ヘッドが使用されます。各タイム ステップの入力ベクトルは、形状 [k, h] の「クエリ」、「キー」、および「値」テンソルに線形投影されます。ここで、h はアテンション ヘッドのサイズです。ここで、キー/値の投影は各ヘッドで共有されます。つまり、「キー」と「値」は [1, h] に投影されますが、「クエリ」は引き続き形状 [k, h] に投影されます。これまでの研究では、これがモデルの品質とトレーニング速度に中立的な影響を与える一方で、自己回帰デコード時間の大幅なコスト削減をもたらす可能性があることが示されています。これは、自己回帰デコード中、キー/値テンソルがインスタンス間で共有されず、一度に 1 つのトークンのみがデコードされるため、標準のマルチヘッド アテンションがアクセラレータ ハードウェア上で非効率的であるためです。

RoPE埋め込み

研究者は、絶対位置や相対位置の埋め込みではなく、RoPE 埋め込みを使用しました。これは、RoPE 埋め込みが長いシーケンス長でより優れたパフォーマンスを発揮することが示されているためです。

共有入出力埋め込み

研究者らは、入力と出力の埋め込み行列を共有しました。これは、過去の研究でよく行われていましたが、普遍的ではありませんでした。

偏見なし

高密度カーネルやレイヤーノルムではバイアスを使用しませんでした。これにより、大規模モデルのトレーニングの安定性が向上することがわかりました。

用語集

研究者らは、トレーニングコーパスで複数の言語をサポートするために（過度なトークン化なしで）選択された、256kトークンのSentencePiece語彙を使用しました。語彙はトレーニングデータから生成され、研究者はこれによってトレーニングの効率が向上することを発見しました。

Pathways で 5400 億パラメータの言語モデルをトレーニングする

PaLM は、Google が Pathways システムを初めて大規模に使用して、トレーニングを 6,144 チップに拡張したもので、これは、これまでのトレーニングに使用された最大の TPU ベースのシステム構成です。研究者らは、各ポッド内では標準のデータとモデルの並列処理を使用しながら、2 つの Cloud TPU v4 ポッド全体でポッド レベルのデータ並列処理を使用してトレーニングを拡張しました。これは、これまでのほとんどの LLM と比較して、規模が大幅に増加したものです。

PaLM は、ハードウェア FLOP 使用率 57.8% のトレーニング効率を達成します。これは、この規模で LLM が達成した最高の効率です。このレベルを達成するために、研究者は並列戦略と Transformer ブロックの再設計の結果を組み合わせました。これにより、アテンション層とフィードフォワード層の並列化が可能になり、TPU コンパイラの最適化によってもたらされる高速化が実現しました。

PaLM は、高品質の Web ドキュメント、書籍、Wikipedia、会話、GitHub コードなどの英語および多言語データセットを使用してトレーニングされます。研究者らはまた、すべての空白（特にコードにとって重要）を保持し、語彙外の Unicode 文字をバイトに分割し、数字を数字ごとに 1 つずつ個別のトークンに分割する「ロスレス」語彙も作成した。

画期的な言語、推論、コード生成機能

PaLM は、言語の理解、生成、推論、コードに関連するタスクを含む、多くの非常に困難なタスクにおいて画期的な機能を発揮します。

言語理解と言語生産

29 の英語ベースの NLP タスクにおいて、PaLM 540B は以前の SOTA の結果を改善しました。

PaLM は、英語の NLP タスクに加えて、トレーニング コーパスの 22% のみが英語以外であるにもかかわらず、翻訳を含む多言語 NLP ベンチマークでも優れたパフォーマンスを示しています。

研究者らはまた、最近リリースされた 150 を超える新しい言語モデリング タスクのスイートである Beyond the Imitation Game Benchmark (BIG-bench) で PaLM の新機能と将来の機能を調査しまし た。このプロセスでは、これらのタスクの 58 個の共通サブセットで PaLM のパフォーマンスが Gopher および Chinchilla のパフォーマンスと比較されます。

興味深いことに、研究者らは、規模の関数としての PaLM のパフォーマンスが以前のモデルと同様の対数線形挙動を示し、規模によるパフォーマンスの向上がまだ頭打ちになっていないことを示唆していると指摘しています。 PaLM 540B 5-shot は、同じタスクにおいて平均的な人間のパフォーマンスを上回ります。

58 個の BIG-bench タスクのサブセットにおける PaLM のスケーリング動作。

PaLM は、いくつかの BIG-bench タスクで優れた自然言語理解および生成機能を実証します。たとえば、このモデルは原因と結果を区別し、文脈における概念の組み合わせを理解し、さらには絵文字から映画を推測することもできます。

PaLM 540B の BIG ベンチタスクにおける 1 ショットのパフォーマンス: 因果関係のラベル付け、概念の理解、絵文字からの映画の推測、同義語と反事実の検出。

推論

モデルのスケーリングと思考連鎖の促進を組み合わせることで、PaLM は、複数ステップの算術演算や常識的な推論を必要とする推論タスクにおいて画期的な機能を発揮します。これまで、Gopher などの大規模言語モデルは、パフォーマンスの向上という点では、モデルの規模によるメリットがほとんどありませんでした。

小学校の算数問題の例における標準的なプロンプトと思考の連鎖によるプロンプト。思考連鎖プロンプトは、人間が処理する方法と同様に、複数ステップの推論問題のプロンプトを中間ステップ (黄色の部分) に分解します。

研究者たちは、3 つの算術データセットと 2 つの常識推論データセットで、思考連鎖プロンプトを使用した PaLM 540B の優れたパフォーマンスを観察しました。たとえば、8 ショットのプロンプトを使用して、PaLM は、何千もの難しい小学校レベルの数学の問題を含むベンチマークである GSM8K の問題の 58% を解きました。これは、外部計算機と検証機能を組み合わせた 7,500 の問題セットでトレーニングされた GPT-3 175B 微調整モデルが達成したこれまでの最高スコア 55% を上回りました。

この新しいスコアは、問題集の対象者である 9 歳から 12 歳の児童の 60% の問題解決レベルに近いため注目に値します。研究者たちは、PaLM 語彙における数字の個別のエンコードがこれらのパフォーマンスの向上に貢献していると仮説を立てています。

驚くべきことに、PaLM は、複数ステップの論理的推論、世界認識、および深い言語理解の複雑な組み合わせを必要とするシナリオでも、明示的な説明を生成できます。たとえば、Web では入手できない新しいジョークに対して、高品質な説明を提供できます。

PaLM は、ツーショットプロンプトを使用してオリジナルのジョークを説明します。

コード生成

大規模言語モデルは、自然言語記述を与えられたコードの記述 (テキストからコードへ)、ある言語から別の言語へのコードの翻訳、コンパイル エラーの修正 (コードからコードへ) などのコーディング タスクにうまく一般化できることが示されています。

PaLM 540B は、事前トレーニング データセットのコードが 5% しかないにもかかわらず、単一のモデルでコーディングと自然言語のタスク全体で優れたパフォーマンスを示します。具体的には、PaLM 540B の少数ショットのパフォーマンスは顕著で、微調整された Codex 12B と同等でありながら、Python トレーニング コードを 50 分の 1 しか使用していません。この結果は、他のプログラミング言語や自然言語データからの学習をより適切に転送できるため、大規模なモデルの方が小規模なモデルよりも効果的であるという以前の調査結果を裏付けています。

テキストからコードへのタスク (GSM8K - Python や HumanEval など) およびコードからコードへのタスク (Transcoder など) における PaLM 540B 微調整モデルの例。

さらに、純粋な Python コード データセットで PaLM を微調整することで、モデルのパフォーマンスがさらに向上し、チームはこれを PaLM-Coder と名付けました。下の図に示すように、PaLM-Coder には DeepFix と呼ばれるサンプル コード修復タスクが与えられました。このタスクの目的は、最初は壊れていた C プログラムを修正して、正常にコンパイルできるようにすることです。PaLM-Coder 540B は、82.1% のコンパイル率を達成し、以前の SOTA 結果の 71.7% を上回る素晴らしいパフォーマンスを発揮しました。これにより、ソフトウェア開発中に発生するより複雑なバグを修正する機会が提供されます。

DeepFix コード修復タスクの例。微調整された PaLM-Coder 540B は、コンパイル エラー (左) をコンパイル可能なバージョンのコード (右) に修正します。

倫理的な観点から、最近の研究では、オンラインテキストでトレーニングされた LLM に関連するさまざまな潜在的なリスクが強調されています。これらの潜在的な悪影響のリスクを、モデル カードやデータ シートなどの透明な成果物を通じて分析し、文書化することが重要です。これらの成果物には、意図された使用方法やテストに関する情報も含まれています。この目的のために、Google の論文では、データ テーブル、モデル カード、Responsible AI ベンチマークの結果が提供され、バイアスとリスクを検出するためのデータセットとモデル出力の包括的な分析が報告されています。分析はモデルの潜在的なリスクの一部を明らかにするのに役立ちますが、起こり得る損害を真に調整し、状況に即して軽減するには、ドメインおよびミッション固有の分析が不可欠です。これらのモデルのリスクと利点をさらに理解すること、および言語モデルの悪意のある使用を防ぐためのスケーラブルなソリューションを開発することは、現在も研究が進められているテーマです。