【DL輪読会】LongNet: Scaling Transformers to 1,000,000,000 Tokens

DEEP LEARNING JP LongNet: Scaling Transformers to 1,000,000,000 Tokens [DL Papers] ⾼城 頌太（東京⼤学 ⼯学系研究科 松尾研 M2） http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 タイトル： LongNet: Scaling Transformers to 1,000,000,000 Tokens https://arxiv.org/abs/2307.02486 著者： Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, Furu Wei Microsoft Research Jiaotong University 概要： 10億トークンまでのテキストを処理できるLongNetというTransformerベースのモデルを提案． 公式実装： https://github.com/microsoft/unilm/tree/master 2

• https://arxiv.org/abs/2307.02486
• https://github.com/microsoft/unilm/tree/master
• https://arxiv.org/search/cs?searchtype=author&query=Ding, J
• https://arxiv.org/search/cs?searchtype=author&query=Ma, S
• https://arxiv.org/search/cs?searchtype=author&query=Dong, L
• https://arxiv.org/search/cs?searchtype=author&query=Zhang, X
• https://arxiv.org/search/cs?searchtype=author&query=Huang, S
• https://arxiv.org/search/cs?searchtype=author&query=Wang, W
• https://arxiv.org/search/cs?searchtype=author&query=Zheng, N
• https://arxiv.org/search/cs?searchtype=author&query=Wei, F

概要 • 「シーケンス⻑の増⼤」「Transformerの計算量の削減」という2つの課題に対し て，Dilated Attentionというアーキテクチャを提案 • このアーキテクチャを導⼊したLongNetにより，10億トークン(GPT-4の約25万 倍)を線型オーダーで処理できる – Vanila TransformerはN^2オーダー 3

Sparse Attention • Attentionを向ける先を絞ることで計算量を削減する • Sparse attention patern によってどこにattentionを向けるかを決定 する 4

Dilated Attention • Sparse attention paternをsegment length, dilated rateという⼆つの変数を持っ たパターンで分割する – Segment length(w): セグメントの幅 – dilated rate(r): 各セグメントの中でどの程度隙間を開けるか 5

Dilated Attention • それぞれのセグメント は並列に処理され，最終的な出⼒がconcatされる • 分割後のDilated AtentionはDense Atentionに変換することができるため，Flash Attentionなどの通常のTransfomerで⽤いられる最適化が使⽤できる 6

Dilated Atention • セグメント⻑は⼤域性をどれほど考慮するか，拡張率は計算効率をどれほど良くするか – 短期，⻑期の情報を効率的に処理するためには，複数のセグメント⻑，拡張率を組み合わせて実装 する – Local atentionは正確性が必要だが，global atentionはだいたいで良いので，セグメント⻑が⼤き くなるほど，拡張率を⼤きくしていく ※ はソフトマックスの分⺟の値 7

Multi-Head Dilated Attention • マルチヘッドにする際は，offsetを変更する – offsetの値は 8

計算量について • セグメントごとのQuery, Key, Valueの次元は • セグメント⻑と拡張率は なので，最終的な計算量は と近似できる 9

分散アルゴリズム • Dilated Attentionによって理論的な計算オーダーは削減されたが，計算資源とメ モリの制約によって，単⼀のGPUで10億トークンの計算は不可能 • そこで，LongNetではトークン⻑が⼤規模になった際の新しい分散アルゴリズム を提案 10

分散アルゴリズム 1. ⼊⼒トークンを分割(下の例だと2つに) 2. それぞれの分割されたXでQ, K, Vを作成 3. Atentionの計算 – If セグメント⻑ <= ローカルデバイスのシーケンス⻑: • – そのまま各デバイスでdilated attentionを計算し，最後に出⼒をconcat elif セグメント⻑ > ローカルデバイスのシーケンス⻑: • Sparse化したK, Vをconcatしてからそれぞれの出⼒Oを求める 11

10億トークンまでスケール • 通常のAttentionとDilated Attentionの順伝播の実⾏速度を⽐較 – 8K~1Bまでトークン⻑を変更 – 最⼤3つのセグメント⻑を採⽤，10回の実⾏時間の平均を計算 – メモリ節約，計算速度向上のためどちらもFlashAttentionを適応している • Dilated Attentionでは線型に実⾏時間が増えていく 12

性能確認: 実験設定 • バックボーンのアーキテクチャとしてはMAGNETOを使⽤ – Transformerの亜種の⼀つ – Position encodingにはxPOSを使⽤ • 学習データセットはThe Stack datasetを使⽤ – BigCodeと呼ばれるプロジェクトで作成された，300ものプログラミング⾔語で構成されたソ ースコード集 – Githubの公開データを⽤いて作成 – 1.5TBのデータがある 13

実験結果 • LongNetは全ての場合においてもPerplexityが低い(良い) – – – – 各バッチごとのトークン⻑は⼀定に設定 モデルサイズよりも⼊⼒トークン⻑が⻑い場合はBlockwise Causal Attention(BCA)を適応 Segment length = (2048, 4096, 8192, 16384, 32769) Dilated Ration = (1, 2, 4, 6, 8) • 少ない計算量でTransfomerと同等の Perplexityを達成 14

モデルサイズごとの計算量 & コンテキストサイズを変更した際の性能 • 125Mパラメータから2.7Bパラメータで訓練させた結果，スケーリング則が適⽤ できる(といえる？？) – 2.7Bモデルは300B token, その他は40B tokenで学習 • コンテキストサイズを⼤きくするとテストロスが減少 – (それは当たり前では．．？) 15

Thank you. 16