OLMo: 言語モデル科学の強化

言語モデルの開発と進歩は、NLP研究だけでなく商業的提供と実世界のアプリケーションにおいても、ほぼすべての場所でその存在を示している。ただし、言語モデルの商業的需要の増加は、ある程度、コミュニティの成長を妨げている。これは、多くの最先端で有能なモデルが、プロプライエタリインターフェイスの後ろに隠されており、開発コミュニティがトレーニングアーキテクチャ、データ、開発プロセスの重要な詳細にアクセスすることができないためである。トレーニングと構造の詳細は、潜在的なリスクと偏見を含む研究研究に不可欠であるため、研究コミュニティが真正にオープンで強力な言語モデルにアクセスする必要性が生じている。

この要件を満たすために、開発者はOLMoを開発しました。OLMoは、最先端で真正にオープンな言語モデルフレームワークであり、研究者が言語モデルを構築して研究することができる。最先端の言語モデルのほとんどは、インターフェイスコードとモデル重みのみをリリースしているが、OLMoフレームワークは真正にオープンソースであり、評価コード、トレーニング方法、トレーニングデータが公開されている。OLMoの主な目的は、オープン研究コミュニティと言語モデルの継続的な開発を強化することである。

この記事では、OLMoフレームワークについて詳細に説明し、そのアーキテクチャ、方法論、現在の最先端フレームワークとの比較を検討する。

OLMo: 言語モデル科学の強化

言語モデルは、AIとMLコミュニティだけでなく、テクノロジー業界全体で、過去数年間で最もホットなトレンドであり、その人間のようなパフォーマンスで実世界のタスクを実行する能力が注目されている。ChatGPTは、言語モデルが持つ可能性の良い例であり、テクノロジー業界の大手企業が言語モデルを自社製品に統合することを検討している。

NLP（自然言語処理）は、言語モデルを広く採用している業界の一つである。ただし、業界が人間の注釈を使用して大規模な事前トレーニングを行い始めて以来、言語モデルは商業的な実用性が大幅に向上し、最先端の言語モデルとNLPフレームワークの多くが制限されたプロプライエタリインターフェイスを持っているため、開発コミュニティは重要な詳細にアクセスできない。

言語モデルの進歩を確実にするために、OLMoは、開発者が言語モデルを構築、研究、開発するためのフレームワークを提供する。さらに、研究者がトレーニングコード、トレーニング方法、トレーニングデータ、トレーニングログ、モデルのチェックポイントにアクセスできる。既存の最先端モデルは、さまざまな程度のオープン性を持っているが、OLMoモデルは、トレーニングからデータ、評価ツールまでの全フレームワークをリリースしており、LLaMA2モデルなどの最先端モデルと比較したときのパフォーマンスギャップを狭めている。

モデリングとトレーニングのために、OLMoフレームワークには、トレーニングコード、フルモデル重み、アブレーション、トレーニングログ、トレーニングメトリクスがインターフェイスコードとして含まれている。また、Weights & Biasesログも含まれている。分析とデータセットの構築のために、OLMoフレームワークには、AI2のDolmaとWIMBDモデルで使用されるフルトレーニングデータと、トレーニングデータを生成するコードが含まれている。評価のために、OLMoフレームワークには、ダウンストリーム評価のためにCatwalkモデルと、パープレクシティベースの評価のためにPalomaモデルが含まれている。

OLMo: モデルとアーキテクチャ

OLMoモデルは、Neural Information Processing Systemsに基づくデコーダのみのトランスフォーマーアーキテクチャを採用しており、1億パラメータと7億パラメータの2つのモデルを提供している。現在、65億パラメータのモデルが開発中である。

OLMoフレームワークのアーキテクチャは、バニラトランスフォーマーコンポーネントを含む他のフレームワークよりも、いくつかの改善点がある。最近の最先端モデル、たとえばOpenLM、Falcon、LLaMA、PaLMと比較したときのOLMoモデルのパフォーマンスは、以下の図に示されている。

OLMoフレームワークは、ハードウェアでのトレーニングスループットを最適化しながら、スローダイバージェンスとロススパイクのリスクを最小限に抑えるために、ハイパーパラメータを選択する。OLMoフレームワークがバニラトランスフォーマーアーキテクチャと異なる主な変更点は以下のとおりである。

バイアスなし

Falcon、PaLM、LLaMAなどの言語モデルとは異なり、OLMoフレームワークにはトレーニングの安定性を高めるためにバイアスが含まれない。

ノンパラメトリックレイヤーノルム

OLMoフレームワークは、ノンパラメトリックなレイヤーノルムの公式を採用している。ノンパラメトリックレイヤーノルムには、ノルム内にアダプティブゲインやバイアスは存在せず、さらにセキュリティが高く、パラメトリックレイヤーノルムよりも高速である。

SwiGLU活性化関数

PaLMやLLaMAなどの多くの言語モデルと同様に、OLMoフレームワークは、ReLU活性化関数ではなくSwiGLU活性化関数を採用しており、隠れ活性化サイズを128の最も近い倍数に増やしてスループットを向上させている。

RoPEまたはロータリーポジショナルエンベッディング

OLMoモデルは、LLaMAとPaLMモデルと同様に、絶対的なポジショナルエンベッディングをRoPEまたはロータリーポジショナルエンベッディングに置き換えている。

ドルマを使用した事前トレーニング

開発コミュニティは、モデルパラメータへのアクセスが強化されたが、事前トレーニングデータセットへのアクセスはまだ制限されている。事前トレーニングデータは、クローズドモデルやオープンモデルとともにリリースされていない。さらに、データに関する技術文書は、モデルを完全に理解し、再現するために必要な重要な詳細が欠けていることが多い。OLMoフレームワークは、事前トレーニングデータセットであるDolmaを構築してリリースし、言語モデル事前トレーニングに関するオープンな研究を促進する。Dolmaデータセットは、7つの異なるソースから収集された5億のドキュメントにわたる3兆トークンの多源および多様なコレクションであり、大規模言語モデルが事前トレーニングに一般的に使用するソースからアクセス可能である。Dolmaデータセットの構成は、以下の表にまとめられている。

Dolmaデータセットは、言語フィルタリング、品質フィルタリング、コンテンツフィルタリング、マルチソースミキシング、重複排除、トークン化の5つのコンポーネントからなるパイプラインを使用して構築される。OLMoは、設計原則と構築の詳細に関するDolmaレポートも提供しており、コンテンツの要約も含まれている。また、OLMoは、高性能データキュレーションツールをオープンソース化しており、事前トレーニングデータコーパスのキュレーションを容易かつ迅速に行える。評価は、トレーニング中の決定に使用するオンライン評価と、モデルチェックポイントからの集約評価の2段階の戦略で行われる。オフライン評価のために、OLMoは、幅広いデータセットとタスク形式にアクセス可能な公開評価ツールであるCatwalkフレームワークを使用する。フレームワークは、Catwalkをダウンストリーム評価と新しいパープレクシティベンチマークであるPalomaでの言語モデリングの評価の両方に使用する。OLMoは、固定評価パイプラインを使用して、ダウンストリーム評価とパープレクシティ評価の両方で、他の公開モデルと比較する。

OLMoは、モデルアーキテクチャ、初期化、オプティマイザ、学習率スケジュール、データの混合について、トレーニング中にいくつかの評価メトリクスを実行する。これは、OLMoの「オンライン評価」と呼ばれ、トレーニングの各1000ステップ（または約4Bトレーニングトークン）でモデル品質の早期および継続的な信号を提供する。評価の設定は、オフライン評価に使用されるコアタスクと実験設定に大きく依存する。OLMoは、OLMo-7Bと他のモデルとの比較だけでなく、科学的評価をより包括的に行うことを目的としている。OLMo-7Bは、パープレクシティ評価のために明示的な除汚を実行する最大の言語モデルである。

OLMoトレーニング

OLMoフレームワークのモデルは、ZeROオプティマイザ戦略を使用してトレーニングされており、これはFSDPフレームワークを通じてPyTorchによって提供され、GPUメモリの消費を大幅に削減する。7Bスケールで、トレーニングは、各GPUあたり4096トークンのマイクロバッチサイズで実行できる。OLMo-1Bと-7Bモデルのトレーニングフレームワークは、約4Mトークン（2048インスタンス、各インスタンスのシーケンス長は2048）のグローバル定数バッチサイズを使用する。現在トレーニング中のOLMo-65Bモデルでは、バッチサイズのウォームアップを使用し、約2Mトークン（1024インスタンス）から開始し、100Bトークンごとに2倍に増やし、約16Mトークン（8192インスタンス）まで増やす。

スループットを向上させるために、FSDPの組み込み設定とPyTorchのampモジュールを通じて、混合精度トレーニングを使用する。後者は、ソフトマックスなどの特定の操作が常にフル精度で実行されることを保証する。また、特定の設定では、各GPUにローカルなシャーディングモデル重みとオプティマイザ状態はフル精度で保持される。各トランスフォーマーブロック内の重みは、トレーニングのフォワードパスとバックワードパス中に各GPUで完全なサイズのパラメータが実現化されたときにのみ、bfloat16形式にキャストされる。勾配はフル精度でGPU間で減算される。

オプティマイザ

OLMoフレームワークは、以下のハイパーパラメータでAdamWオプティマイザを使用する。

すべてのモデルサイズについて、学習率は最初の5000ステップ（約21Bトークン）で線形に増加し、最大値に達した後、ステップ数の逆二乗根で線形に減少する。ウォームアップ期間後、モデルは、パラメータの勾配の合計l-ノルムが1.0を超えないように、勾配をクリッピングする。以下の表は、7Bスケールでのオプティマイザ設定を、他の最近のLLMと比較する。

トレーニングデータ

トレーニングには、文とBPEトークナイザを使用して文をトークン化し、各ドキュメントの末尾に特殊なEOSトークンを追加し、2048トークンの連続するチャンクをグループ化してトレーニングインスタンスを形成する。トレーニングインスタンスは、各トレーニング実行で同じ方法でシャッフルされる。データの順序と各トレーニングバッチの正確な構成は、リリースされたアーティファクトから再構築できる。リリースされたOLMoモデルはすべて、少なくとも2Tトークン（トレーニングデータの1エポック）までトレーニングされており、一部のモデルは、データの異なるシャッフル順序で2番目のエポックを開始することによって、さらにトレーニングされている。繰り返されるデータの量は小さいため、影響は無視できる。

結果

OLMo-7Bの評価に使用されるチェックポイントは、Dolmaデータセットで2.46Tトークンまでトレーニングされており、事前に述べた線形学習率減少スケジュールを使用する。さらに、Dolmaデータセットで1000ステップ、線形に減少する学習率でこのチェックポイントをファインチューニングすると、パープレクシティとエンドタスク評価スイートのモデルパフォーマンスが向上する。最終評価のために、開発者はOLMoを、LLaMA-7B、LLaMA2-7B、Pythia-6.9B、Falcon-7B、RPJ-INCITE-7Bなどの公開モデルと比較する。

ダウンストリーム評価

コアダウンストリーム評価スイートは、以下の表にまとめられている。

すべてのケースで、ランク分類アプローチによるゼロショット評価を実行する。候補テキストの完了（たとえば、異なる多肢選択肢）が、通常は正規化因子で正規化された可能性によってランク付けされ、予測精度が報告される。

Catwalkは、各データセットに選択された正規化戦略を使用する。正規化戦略には、トークンごとの正規化、キャラクタごとの正規化、回答の無条件的な可能性などが含まれる。具体的には、arcとopenbookqaタスクでは正規化は行われず、hellaswag、piqa、winograndeタスクではトークンごとの正規化が行われ、boolq、copa、sciqタスク（シングルトークン予測タスクに近い形式のタスク）では正規化は行われなかった。

以下の図は、9つのコアエンドタスクの精度スコアの進捗を示している。OLMo-7Bをさらにトレーニングするにつれて、ほとんどのタスクで精度が一般的に上昇していることがわかる。最後から2番目のステップと最後のステップの間で、多くのタスクの精度が急上昇することは、最後の1000トレーニングステップで学習率を0に線形に減少させることの利点を示している。たとえば、内在的な評価の場合、Palomaは、各ドメインのパフォーマンスの検査から、ドメインの組み合わせの上位レベルの結果まで、分析のシリーズを通じて議論する。

最終的な考え

この記事では、OLMoについて話し合った。OLMoは、開発者が言語モデルを構築、研究、開発するためのフレームワークを提供する、最先端で真正にオープンな言語モデルである。既存の最先端モデルは、さまざまな程度のオープン性を持っているが、OLMoモデルは、トレーニングからデータ、評価ツールまでの全フレームワークをリリースしており、LLaMA2モデルなどの最先端モデルと比較したときのパフォーマンスギャップを狭めている。