パラメータのわずか4%、GPT-3のパフォーマンスに匹敵：開発者がDeepMindのRETROを説明

BERT から GPT-2、そして GPT-3 へと、大規模モデルの規模は拡大し、そのパフォーマンスはますます驚異的なものになってきました。モデルのサイズを大きくすることは改善のための実行可能な道であることが証明されており、DeepMind による最近の研究では、この道はまだ終わっておらず、モデルを拡大し続けることで依然として大きなメリットが得られることが示されています。

しかし同時に、最近のいくつかの研究が示しているように、モデルのサイズを大きくすることがパフォーマンスを向上させる唯一の方法ではないこともわかっています。その中でも代表的な研究としては、DeepMind の RETRO Transformer と OpenAI の WebGPT があります。これら 2 つの研究は、情報を検索/照会する方法でモデルを強化すると、以前は大規模なモデルでしか実現できなかったパフォーマンスを、より小規模な生成言語モデルで実現できることを示しています。

大規模モデルが主流となっている今日の世界では、この種の研究は極めて価値があります。

この記事では、機械学習の視覚化を専門とする著名なブロガーである Jay Alammar が、DeepMind の RETRO (Retrieval-Enhanced TRansfOrmer) モデルを詳細に分析します。このモデルは GPT-3 と同等のパフォーマンスを備えていますが、GPT-3 のパラメーターの 4% しかありません。

RETRO はデータベースから取得した情報を統合し、事実や世界知識の高価な保存からパラメータを解放します。

RETRO 以前にも研究コミュニティでは同様のアプローチを採用した研究がいくつかあったため、この論文ではその新規性ではなくモデル自体について説明します。

言語情報と世界知識情報を分離する

一般的に言えば、言語モデルの役割は空白を埋めることであり、時には次のような事実に関連した情報が必要になる。

しかし、言語に精通している場合は、空白に何を埋めればよいかを推測できる場合もあります。たとえば、次のようになります。

大規模な言語モデルは、知っていることすべてをモデル パラメータにエンコードするため、この区別は重要です。これは言語情報には意味がありますが、事実情報や世界知識情報には効果がありません。検索方法を追加すると、言語モデルを大幅に削減できます。テキスト生成プロセス中に、ニューラル データベースはモデルが必要な事実情報を取得するのに役立ちます。

記憶するトレーニング データの量が減るにつれて、より小さな言語モデルを使用してトレーニングを高速化できます。誰でもこれらのモデルをより小型で安価な GPU に導入し、必要に応じて調整することができます。

アーキテクチャ的には、RETRO はオリジナルの Transformer と同様にエンコーダー/デコーダー モデルです。ただし、データベース検索の助けを借りて入力シーケンスが増加します。モデルはデータベース内で最も可能性の高いシーケンスを見つけ、それを入力に追加します。 RETRO は魔法を使って出力予測を生成します。

モデル アーキテクチャを調べる前に、検索データベースについてもう少し詳しく見てみましょう。

RETROの検索データベース

ここでのデータベースはキーバリューストアデータベースです。キーは標準の BERT 文埋め込みであり、値は 2 つの部分で構成されるテキストです。

1. キーを計算するために使用されるネイバー。
2. 元の文書のテキストの完了、継続。

RETRO のデータベースには、MassiveText データセットに基づく 2 兆個の多言語トークンが含まれています。隣接チャンクと完了チャンクの長さは最大 64 トークンです。

「RETRO データベースの内部」では、RETRO データベース内のキーと値のペアの例を示します。

RETRO 入力プロンプトを複数のチャンクに分割します。簡単にするために、ここでは取得したテキストでチャンクを拡張する方法に焦点を当てます。ただし、モデルは、最初のチャンクを除く入力プロンプトの各チャンクに対してこのプロセスを実行します。

データベース検索

RETRO をクリックする前に、BERT を入力するためのプロンプトを入力します。出力コンテキスト ベクトルを平均化して、文埋め込みベクトルを構築します。このベクトルはデータベースを照会するために使用されます。

BERT を使用して入力プロンプトを処理すると、コンテキスト化されたトークンの埋め込みが生成されます。それらを平均すると、文の埋め込みが生成されます。

この文の埋め込みは、近似最近傍検索に使用されます。最も近い 2 つの近傍が取得され、そのテキストが RETRO への入力の一部になります。

BERT 文埋め込みは、RETRO のニューラル データベースから最も近い近傍を取得するために使用されます。これらは言語モデルの入力に追加されます。

ここで、RETRO への入力は、入力プロンプトと、データベースからのその最も近い 2 つの近傍 (およびそれらの継続) です。

ここから、Transformer ブロックと RETRO ブロックが情報を処理に組み込みます。

取得された近傍は言語モデルの入力に追加されます。ただし、モデル内では若干異なる方法で処理されます。

高レベルRETROアーキテクチャ

RETRO のアーキテクチャは、エンコーダー スタックとデコーダー スタックで構成されています。

RETRO Transformer は、エンコーダー スタック (隣接ノードの処理) とデコーダー スタック (入力の処理) で構成されます。

エンコーダーは、標準の Transformer エンコーダー ブロック (自己注意 + FFNN) で構成されています。 Retro は、2 つの Transformer エンコーダー ブロックで構成されるエンコーダーを使用します。

デコーダー スタックには 2 つのデコーダー ブロックが含まれています。

• 標準トランスフォーマーデコーダーブロック（ATTN + FFNN）
• RETRO デコーダー ブロック (ATTN + チャンク クロス アテンション (CCA) + FFNN)

RETROを構成する3つのTransformerモジュール

エンコーダー スタックは、取得した近傍を処理し、その後注目に使用される KEYS および VALUES マトリックスを生成します。

デコーダー ブロックは、入力テキストを GPT のように処理します。キュー トークンに自己注意を適用し (したがって、前のトークンにのみ焦点を当てる)、FFNN レイヤーを通過します。

RETRO デコーダーに到達した場合にのみ、取得した情報のマージが開始されます。 9 から始まる 3 つおきのブロックは RETRO ブロックです (入力が隣接するブロックに従うことを可能にします)。つまり、9 層目、12 層目、15 層目…32 層目は RETRO ブロックです。

次の図は、取得した情報を参照した後、プロンプトを完了するために必要なノード手順を示しています。