ほとんどの言語モデルは短いテキストシーケンスでトレーニングされ、その後、より長いドキュメントでの追加トレーニング（コンテキスト拡張）によって長い入力を処理できるようになります。Llama 3の事前学習データはプロプライエタリであるため、その拡張の容易さがアーキテクチャによるものか、トレーニングデータによるものかを判断するのは困難でした。研究者は同じ拡張手法が他のアーキテクチャに転送可能であると想定してきましたが、OlmPoolの研究はそれがしばしば当てはまらないことを示しています。
OlmPoolは、Allen Institute for AI（Ai2）によって開発された制御されたモデルスイートで、26の7Bモデルから構成されます。すべてのモデルは同じデータで1400億トークンの事前学習を行い、同じ長コンテキストデータ混合と手順を使用して64Kコンテキストに拡張されました。唯一変化するのはアーキテクチャです。研究は、アテンションメカニズムに影響を与える4つのアーキテクチャ選択に焦点を当てています：QK正規化、グループ化クエリアテンション（GQA）、スライディングウィンドウアテンション、および事前学習コンテキスト長。
研究によると、これらの選択は単独では穏やかな影響しかありません：QK正規化が最も大きな影響を持ち、それを削除し正規化順序を変更するとHELMETベンチマークで6ポイントの改善が見られます。GQAと短い事前学習コンテキスト長はそれぞれ小さな低下を引き起こし、スライディングウィンドウアテンションは単独で約1ポイントの低下をもたらします。しかし、これらが組み合わさると、効果は各部分の和をはるかに超えます。例えば、すでにGQAを持つモデルにスライディングウィンドウアテンションを追加すると、平均で9ポイントの低下を引き起こします。最悪の構成は、アテンションの柔軟性を制限する2つ以上の選択を組み合わせたものです。
さらに重要なことに、標準的なトレーニング信号は長コンテキスト性能をほとんど予測できません。トレーニング損失、検証パープレキシティ、および16の短コンテキストベンチマークでは、どのモデルが32Kまたは64Kコンテキストで良いパフォーマンスを発揮するかを識別できません。同じベンチマークの8K分割でのスコアでさえ、拡張後の二桁の変動を予測できません。拡張後にほぼ同一に見えるモデルでも、HELMETで26ポイント以上乖離することがあります。
また、Llama 3のアーキテクチャ構成は強力ですが、すべての場合で最適とは限らないことがわかりました。他のいくつかのモデルがそれを明確に上回っています。これは、Llama 3の長コンテキスト成功が主にアーキテクチャによるものであり、Llamaで検証された拡張手法は他のモデルファミリーに適応させる必要があることを示唆しています。アーキテクチャ起因のギャップはデータを増やしても解消されません。最悪のアーキテクチャは、500億トークンのコンテキスト拡張（トレーニング全体の26％）後でも、Llamaアーキテクチャが10億トークンで達成するパフォーマンスに達しません。
アテンションパターンの分析により、QK正規化がないモデルはより強いアテンションシンクを発達させることが明らかになりました。これは、入力の初期位置（通常最初の数トークン）が、現在の予測に無関係でも一貫して大きなアテンションを受け取る現象です。より強いシンクは長コンテキスト性能の向上と相関しており、QK正規化がないモデルが長距離検索をサポートするためのデフォルト戦略であることを示しています。針の中の針テストでは、QK正規化があるモデルはターゲット情報に割り当てるアテンションが少なく、全体的な長コンテキスト性能の低さと一致しています。
OlmPoolスイートは完全に公開されており、事前学習とコンテキスト拡張の各段階における26モデルすべての38チェックポイントが含まれています。この研究は、アーキテクチャ選択の組み合わせが、実務者が期待するよりもはるかに低い長コンテキスト性能を生み出す可能性があり、この結果は標準的なトレーニング信号からは見えないことを強調しています。研究者は、このリソースがより良いコンテキスト拡張手法の開発や、初期事前学習における他の現象の研究に役立つことを期待しています。