Long-context Extension: 長文脈拡張
Olmo 3 では、ベースモデル（8K トークンのコンテキスト長）から 65K トークンへと長文脈拡張を実施した。この拡張により、長文書の理解や複雑なタスクに対応できるモデルとなっている。
拡張の概要
長文脈拡張では、以下の規模でトレーニングを行った。
- 7B モデル: 50B トークンで学習
- 32B モデル: 100B トークンで学習
- コンテキスト長: 8K トークン → 65K トークンへ拡張
この拡張は、特定のデータミックス（Dolma 3 Longmino Mix）と複数の技術的手法を組み合わせて実現されている。
Dolma 3 Longmino Mix の構成
長文脈拡張に使用されたデータセットは、以下の3つの主要コンポーネントから構成されている（Table 11 参照）。
1. olmOCR PDFs
PDF から抽出された長文書データで、様々な長さのバケットに分類されている。
| Length Bucket | Documents | Tokens |
|---|---|---|
| 8K-16K | 1,090,349 | 13.1B |
| 16K-32K | 508,354 | 11.0B |
| 32K-64K | 142,983 | 6.1B |
| 64K-128K | 54,992 | 4.5B |
| 128K-256K | 20,893 | 3.2B |
| 256K-512K | 8,130 | 2.4B |
| 512K-1M | 3,394 | 1.7B |
| 1M+ | 1,172 | 1.8B |
2. 合成データ
長文脈能力を強化するための合成的に生成されたデータである。
- CWE (Common Word Extraction): 7.4B トークン
- REX (Rewriting Expressions): 1.5B トークン
3. Midtraining Data Mix
Midtraining phase で使用されたデータの 66% を含めることで、一般的な能力の維持を図っている。
- Midtraining data mix: 34.9B トークン（66% の割合）
主要技術コンポーネント
長文脈拡張の実現には、Figure 13 に示される 5つの技術コンポーネントが使用されている。
1. RoPE 拡張 (YaRN)
RoPE (Rotary Position Embedding) を拡張するために YaRN (Yet another RoPE extensioN) を採用した。
YaRN は Full attention layers のみに適用されている。Sliding window attention layers では、元の RoPE 設定を維持している。
2. Document Packing
効率的なトレーニングのために、複数の文書を1つのシーケンスにパッキングする。
- 手法: Best-fit packing アルゴリズム
- 目的: GPU メモリの効率的な利用とトレーニングスループットの向上
3. Intra-document Masking
パッキングされた文書間で情報が漏れないように、文書内マスキングを適用する。
Document packing を使用する場合、異なる文書間でアテンションが計算されないようにマスキングすることが重要である。これにより、各文書が独立して処理される。
4. Model Souping
複数のチェックポイントを平均化することで、モデルの安定性と性能を向上させる。
- 手法: 異なるトレーニングステップで保存されたチェックポイントの重みを平均化
- 効果: より汎用的で安定したモデルの獲得
5. Token Budget
長い拡張により多くのトークンを割り当てることで、より良い性能を達成する。
- 7B モデル: 50B トークン
- 32B モデル: 100B トークン
合成データ生成パイプライン
長文脈能力を効果的に向上させるため、2つの合成データ生成手法を使用している。
CWE (Common Word Extraction)
文書内で共通して出現する単語を抽出するタスクを生成する。これにより、モデルは長文書全体を参照する能力を獲得する。
REX (Rewriting Expressions)
REX では、12 種類の vignettes（小場面）を用いて、様々な長文脈タスクをシミュレートする。
REX は多様なタスク形式を含み、要約、情報抽出、質問応答など、実際のユースケースに近い合成データを生成する。これにより、モデルは様々な長文脈タスクに適応できるようになる。
評価結果
長文脈拡張モデルの性能は、RULER（開発スイート）と HELMET（Held-out 評価）で評価されている（Table 12 参照）。
RULER 評価結果
RULER は、様々なコンテキスト長での性能を測定する開発用評価スイートである。
| Model | 4K | 8K | 16K | 32K | 64K | 128K | Average |
|---|---|---|---|---|---|---|---|
| Olmo 3 7B | 92.7 | 91.7 | 88.1 | 82.5 | 70.3 | - | 85.1 |
| Olmo 3 32B | 95.8 | 94.9 | 92.8 | 89.4 | 82.1 | - | 91.0 |
HELMET 評価結果
HELMET は、実際のユースケースに近い held-out 評価セットである。
Olmo 3 の長文脈モデルは、同規模の他のオープンモデルと比較して競争力のある性能を示している。特に、32B モデルは多くのベンチマークで高いスコアを達成している。
主要な発見
長文脈拡張の評価から、以下の知見が得られた。
- Token budget の効果: より多くのトークンで学習することで、長文脈能力が大幅に向上
- 合成データの重要性: CWE と REX の合成データが、実際のタスクでの性能向上に寄与
- Model souping の効果: 複数チェックポイントの平均化により、安定した性能を獲得
これらの技術を組み合わせることで、Olmo 3 は 65K トークンの長文脈を効果的に扱えるモデルとなっている。