TECHNOLOGY NLP AI

【日本語訳】LONGNET: Scaling Transformers to 1,000,000,000 Tokens【10億トークンをサポートしたTransformer】【Microsoft】

悩んでいる人

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの日本語訳を教えて!

こういった悩みにお答えします.

$$\newcommand{\1}{\mbox{1}\hspace{-0.25em}\mbox{l}}$$

本記事の信頼性

  • リアルタイムシステムの研究歴12年.
  • 東大教員の時に,英語でOS(Linuxカーネル)の授業.
  • 2012年9月~2013年8月にアメリカのノースカロライナ大学チャペルヒル校(UNC)コンピュータサイエンス学部で客員研究員として勤務.C言語でリアルタイムLinuxの研究開発.
  • プログラミング歴15年以上,習得している言語: C/C++PythonSolidity/Vyper,Java,Ruby,Go,Rust,D,HTML/CSS/JS/PHP,MATLAB,Assembler (x64,ARM).
  • 東大教員の時に,C++言語で開発した「LLVMコンパイラの拡張」,C言語で開発した独自のリアルタイムOS「Mcube Kernel」GitHubにオープンソースとして公開
  • 2020年1月~現在はアメリカのノースカロライナ州チャペルヒルにあるGuarantee Happiness LLCのCTOとしてECサイト開発やWeb/SNSマーケティングの業務.2022年6月~現在はアメリカのノースカロライナ州チャペルヒルにあるJapanese Tar Heel, Inc.のCEO兼CTO.
  • 最近は自然言語処理AIイーサリアムに関する有益な情報発信に従事.
    • (AI全般を含む)自然言語処理AIの論文の日本語訳や,AIチャットボット(ChatGPT,Auto-GPT,Gemini(旧Bard)など)の記事を50本以上執筆.アメリカのサンフランシスコ(広義のシリコンバレー)の会社でプロンプトエンジニア・マネージャー・Quality Assurance(QA)の業務委託の経験あり.
    • (スマートコントラクトのプログラミングを含む)イーサリアムや仮想通貨全般の記事を200本以上執筆.イギリスのロンドンの会社で仮想通貨の英語の記事を日本語に翻訳する業務委託の経験あり.

こういった私から学べます.

AIのプログラミング言語「C++/Python言語」を学べるおすすめのWebサイトを知りたいあなたはこちらからどうぞ.

独学が難しいあなたは,AIを学べるオンラインプログラミングスクール3社で自分に合うスクールを見つけましょう.後悔はさせません!

国内・海外のAIエンジニアのおすすめ求人サイトを知りたいあなたはこちらからどうぞ.

国内・海外のプロンプトエンジニアのおすすめ求人サイトを知りたいあなたはこちらからどうぞ.

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの日本語訳を紹介します.

Microsoftによる10億トークンをサポートしたTransformer「LONGNET」がわかります.

※図表を含む論文の著作権はLONGNET: Scaling Transformers to 1,000,000,000 Tokensの著者に帰属します.

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの目次は以下になります.

  • Abstract
  • 1章:Introduction
  • 2章:LONGNET
  • 3章:LONGNET as a Distributed Trainer: Scaling up to 1B Tokens
  • 4章:Experiments on Language Modeling
  • 5章:Conclusion and Future Work
  • Acknowledgement
  • References
  • 付録A:Hyperparameters

LONGNET: Scaling Transformers to 1,000,000,000 Tokensを解説しつつ,私の考えも語ります.

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの概要と私の日本語訳は以下になります.

Scaling sequence length has become a critical demand in the era of large language models.
大規模言語モデルの時代において,シーケンス長のスケーリングは重要な要求となっている.

However, existing methods struggle with either computational complexity or model expressivity, rendering the maximum sequence length restricted.
しかし,既存の手法は,計算の複雑性とモデルの表現力のどちらにも苦戦しており,最大シーケンス長が制限されている.

In this work, we introduce LONGNET, a Transformer variant that can scale sequence length to more than 1 billion tokens, without sacrificing the performance on shorter sequences.
本研究では,より短いシーケンスに対する性能を犠牲にすることなく,シーケンス長を10億トークン以上に拡張できるTransformerのVariantであるLONGNETを紹介する.

Specifically, we propose dilated attention, which expands the attentive field exponentially as the distance grows.
具体的には,距離が長くなるにつれて指数関数的にAttentionフィールドを拡張するDilated Attentionを提案する.

LONGNET has significant advantages: 1) it has a linear computation complexity and a logarithm dependency between tokens; 2) it can be served as a distributed trainer for extremely long sequences; 3) its dilated attention is a drop-in replacement for standard attention, which can be seamlessly integrated with the existing Transformer-based optimization.
LONGNETには大きな利点がある:1)計算複雑度が線形で,トークン間の依存性が対数であること,2)非常に長いシーケンスの分散訓練士として機能すること,3)Dilated Attentionは標準的なAttentionの完全互換であり,既存のTransformerベースの最適化とシームレスに統合できること.

Experiments results demonstrate that LONGNET yields strong performance on both long-sequence modeling and general language tasks.
実験の結果,LONGNETは長いシーケンスのモデリングと一般的な言語タスクの両方で強力な性能を発揮することが実証された.

Our work opens up new possibilities for modeling very long sequences, e.g., treating a whole corpus or even the entire Internet as a sequence.
我々の研究は,非常に長いシーケンスをモデル化するための新しい可能性を開くものである.例えば,コーパス全体やインターネット全体をシーケンスとして扱うことができる.

https://arxiv.org/abs/2307.02486

私の日本語訳の注意点は以下になります.

  • 概要は英語と日本語を両方掲載しましたが,本文は私の日本語訳のみを掲載していること(英語で読みたいあなたは原文を読みましょう!)
  • 基本的には原文の直訳ですが,わかりにくい箇所は意訳や説明を追加している箇所があること
  • 原文の「Acknowledgements」(謝辞)は省略していること
  • 本文中に登場する表記「[BMR+20, KMH+20, ZKHB22, CND+22, DDM+23]」などは参考文献ですので,興味がある方は本記事の参考文献を参照されたいこと

それでは,LONGNET: Scaling Transformers to 1,000,000,000 Tokensの本文を読みすすめましょう!

1章:Introduction(はじめに)

LONGNET Figure1
図1:Transformerのシーケンス長のトレンド.

近年,ニューラルネットワークをスケールアップする傾向が見られる[BMR+20, KMH+20, ZKHB22, CND+22, DDM+23].

深さは主に指数的な表現力のためにスケールアップされ,多くの強力なディープネットワーク[HZRS16, HCB+19, WMD+22]を生み出している.

次に,スパースMoEモデル[LLX+21, FZS21, ZBK+22]とモデル並列化アプローチ[SPP+19, KCL+22]が効率的に隠れ次元を拡大する.

シーケンス長は,ニューラルネットワークの最後の原子次元として,無制限であることが望ましい.

シーケンス長の制限を破ることは大きな利点をもたらす.

第一に,モデルに大きなメモリと受容野を提供し,人間や世界と相互作用するのに実用的である.

第二に,より長いコンテキストは,モデルが訓練データで利用できる,より複雑な因果関係と推論経路を含んでいる.

これに対して,短い依存関係は,より多くのスプリアス相関を持ち,これは汎化に有害である.

第三に,コンテキスト内学習の限界を探ることができる.

これは,非常に長いコンテキストがモデルの破滅的忘却を軽減するのに役立つ可能性があるため,Many-Shot学習のパラダイムシフトとなる可能性を秘めている.

シーケンス長をスケールアップする際の大きな課題は,計算の複雑性とモデルの表現力の間で適切なバランスを取ることである.

RNNスタイルのモデルは,主に長シーケンス化のために実装される.

しかし,そのシーケンシャルな性質により,長シーケンスモデリングに不可欠な訓練中の並列化が制限される.

最近では,状態空間モデル[GGR22, SWL23, FDS+23, PMN+23]がシーケンスモデリングに魅力的である.

訓練時にはCNNとして動作し,テスト時には効率的なRNNに変換することができる.

長距離のベンチマーク[TDA+21]では良い性能を発揮するが,通常の長さでの性能はTransformersほど良くなく,主にモデルの表現力によって制限される[FPB+23].

シーケンス長をスケーリングするもう1つの筋は,Transformerの複雑性,つまりSelf-Attentionの2次的な複雑性を減少させることである.

スライディングウィンドウや畳み込みモジュールをAttentionの上に実装することは,複雑性をほぼ線形にする簡単な方法である.

とはいえ,この方法では初期のトークンを思い出す能力が犠牲になり,シーケンスの一番最初のプロンプトを忘れてしまう.

Sparse AttentionはAttention行列をスパース化することで計算量を減らし,遠くの情報を呼び出す可能性を維持する.

例えば,[CGRS19]は固定のスパースパターンで\(\mathcal{O}(N \sqrt{N}d)\)の時間複雑度を得ている.

発見的なパターン[ZGD+20, BPC20]の他に,学習可能なパターン[KKL20, ALdJ+23]がSparse Attentionに有用であることが証明されている.

また,Low-Rank Attention[WLK+20, WCL+20],Kernelベースの手法[KVPF20, CLD+21, QHS+22],ダウンサンプリングアプローチ[LLK+19, JGB+21, MKW+21],リカレントモデル[DYY+19, BKB23],検索ベースの手法[WRHS22, WDC+23]など,Transformerベースの効率的なVariantsもある.

しかし,10億トークンに拡張されたものはない(図1参照).

※訳注:原文では表1を参照していませんが,こちらに記載しておきます.

LONGNET Table1
表1:各手法の計算量の比較.Nはシーケンス長,dは隠れ次元.

この研究では,シーケンスの長さを10億トークンまで拡張することに成功した.

我々のソリューションはLONGNETであり,Vanilla TransformersのAttentionをDilated Attentionという新しい要素に置き換えたものである.

一般的な設計原理は,トークン間の距離が大きくなるにつれて,Attentionの割り当てが指数関数的に減少するというものである.

我々は,それが線形計算複雑性とトークン間の対数依存性を得ることを証明する.

これは,限られたAttentionリソースと全てのトークンへのアクセス可能性との間の矛盾に対処するものである.

実装において,LONGNETはDense Transformerに変換することができ,Transformerのための既存の最適化(例:カーネル融合,量子化,分散訓練)をシームレスにサポートする.

LONGNETは線形複雑性を利用して,ノード間で訓練を並列化することができ,分散アルゴリズムによって計算とメモリの両方の制約を打ち破ることができる.

これにより,シーケンス長を1Bトークンまで,ほぼ一定の実行時間で効率的にスケールアップすることができる(図5参照).

一方,Vanilla Transformerは2次関数的な複雑性に悩まされる.

2章:LONGNET

2.1節:Preliminary(準備)

Transformers[VSP+17]のコアはSelf-Attentionであり,クエリと,キーとバリューの集合を出力にマッピングする.

入力\(Q, K, V \in \mathbb{R}^{N*d}\)が与えられると,以下のように出力Oを計算する.

$$ O = softmax(QK^T)V $$

Self-Attentionはシーケンス長に2次依存するため,長いシーケンスでは苦労する.

一回のクエリで全てのキーとバリューにAttentionすることになり,計算効率が悪くなる.

Sparse Attentionは,クエリのアクセスをキーとバリューのサブセットに制限することで,この問題を軽減する.

Sparse AttentionのキーはSparse Attentionパターン\(S \in \{0,1\}^{N*N}\)であり,クエリQがAttentionできる特定のキーとバリューを決定する.

$$ O = softmax(QK^T \odot \1_S)V $$

例えば,Sparse Transformer[CGRS19]の固定パターンは,ローカルパターンとストライドパターンで構成される.

シーケンスは長さlのブロックに分割される.

ローカルパターンは1つのクエリが同じブロック内のトークンにAttentionすることを可能にし,ストライドパターンは1つのクエリが各ブロックの最後のc個のトークンにAttentionすることを可能にする.

形式的には,ローカルパターン\(S_i^{(1)} = \{ j \ | \ \lfloor j / l \rfloor = \lfloor i / l \rfloor \}\)であり,ストライドパターン\(S_i^{(2)} = \{ j | j {\rm\ mod\ } l \in \{ t, t + 1, …, l \} \}\)である.

2.2節:Dilated Attention

LONGNET Figure2
図2:LONGNETで使われているDilated Attentionのビルディングブロック.これは,短距離依存と長距離依存の両方をモデル化するための一連のAttentionパターンから構成される.Attentionパターンの数はシーケンスの長さに応じて拡張することができる.

図2はDilated Attentionの概要を示している.

Dilated Attentionは,入力\((Q, K, V)\)をセグメント\(\{(\widetilde{Q}_i, \widetilde{K}_i, \widetilde{V}_i)\}^{\frac{N}{w}}\)個に,セグメント長wで均等に分割する.

各セグメントは,間隔rで行を選択することで,シーケンス次元に沿ってスパース化される.

この計算は次のように書ける.

\begin{eqnarray}
\widetilde{Q}_i &=& [Q_{iw}, Q_{iw+r}, Q_{iw+2r}, ..., Q_{(i+1)w-1}] \tag{1} \\
\widetilde{K}_i &=& [K_{iw}, K_{iw+r}, K_{iw+2r}, ..., K_{(i+1)w-1}] \\
\widetilde{V}_i &=& [V_{iw}, V_{iw+r}, V_{iw+2r}, ..., V_{(i+1)w-1}] \tag{2}
\end{eqnarray}

スパース化されたセグメント\(\{\widetilde{Q}_i, \widetilde{K}_i, \widetilde{V}_i) \}^{\frac{N}{w}}\)は,並列にAttentionに供給され,その後,散在して連結され,出力Oとなる.

\begin{eqnarray}
\widetilde{O}_i &=& softmax(\widetilde{Q}_i, \widetilde{K}_i^T) \widetilde{V}_i \\
\hat{O}_i &=& \{ \widetilde{O}_{i,j} | j {\rm\ mod\ } r = 0; 0 | j {\rm\ mod\ } r \neq 0 \} \\
O &=& [\hat{O}_0, \hat{O}_1, ..., \hat{O}_{\frac{N}{w}-1}] \tag{3}
\end{eqnarray}

実装において,Dilated Attentionは,入力\((Q, K, V)\)に対する収集操作と出力\(\widetilde{O}_i\)に対する散乱操作の間でDense Attentionに変換することができるので,Vanilla Attention(例:Flash Attention[DFE+22])に対する任意の最適化を直接再利用することができる.

Dilated Attentionは,Vanilla Attentionに比べ\(\frac{N}{w} r^2\)のファクターで計算コストを大幅に削減することができる.

実際には,セグメントサイズwはAttentionのグローバル性を効率と引き換えにし,サイズrによるDilationはAttention行列を近似することで計算コストを削減する.

長距離と短距離の両方の情報を効率的に捕捉するために,異なるセグメントサイズとDilation Rates \(\{r_i,w_i\}^k\)を持つDilated Attentionの混合を実装する.

\begin{eqnarray}
O &=& \sum_{i=1}^k \alpha_i O |_{r_i,w_i} \\
\alpha_i &=& \frac{s_i}{\Sigma_j s_j}
\end{eqnarray}

ここで\(s_i\)は,\(O |_{r_i,w_i}\)のAttentionソフトマックスの分母を示す.

※訳注:\(\alpha_i\)の分子に\(s_i\)があるため,約分されるからだと思います.

\(\{O |_{r_i,w_i} \}^k\)の計算は並列である.

なぜなら,これらの間には計算依存関係がないからである.

実験によると,Attentionソフトマックスの分母によって計算される動的な重みは,学習可能な固定の重みよりも優れている.

クエリが異なるDilated Attention内のキーにAttentionする場合,Dilated Attentionを混合する我々の方法は,異なる部分のキーを集めて一緒にソフトマックスを計算することと等価である.

直感的には,ローカルなAttentionは正確に計算されるべきだが,グローバルなAttentionは近似的に計算できる.

したがって,より大きな\(r_i\)でより大きな\(w_i\)を設定する.

さらに,最大長NまたはAttentionパターン数kに達するまで,各Attentionに対して\(w_i\)を徐々に増加させる.

\begin{eqnarray}
w &=& \{ w_0, w_1, w_2, ..., N \}^k & (w_i < w_{i+1} < N ) \tag{4} \\
r &=& \{ 1, r_1, r_2, ..., r_k \}^k & (1 < r_i < r_{i+1} ) \tag{5}
\end{eqnarray}

実際には,wとrを指数関数的なAttentive Fieldの幾何級数としている.

2.3節:Multi-Head Dilated Attention

LONGNET Figure3
図3:複数のHeadsによるDilated Attention.順次位置をずらすことで,AttentionパターンがHeadsによって異なる.

図3に示すように,クエリ・キー・バリューのペアの異なる部分をスパース化することで,異なるHeads間で計算が異なる.

具体的には,j番目のHeadでは,\((Q, K, V)\)を選択する際にオフセット\(s_j = j {\rm \ mod \ } r\)を設定する.

\begin{eqnarray}
\widetilde{Q}_i &=& [Q_{iw+s_j}, Q_{iw+s_j+r}, Q_{iw+s_j+2r}, ..., Q_{(i+1)w+s_j-1}] \\
\widetilde{K}_i &=& [K_{iw+s_j}, K_{iw+s_j+r}, K_{iw+s_j+2r}, ..., K_{(i+1)w+s_j-1}] \\
\widetilde{V}_i &=& [V_{iw+s_j}, V_{iw+s_j+r}, V_{iw+s_j+2r}, ..., V_{(i+1)w+s_j-1}]
\end{eqnarray}

Vanilla Multi-Head Attentionに続いて,異なるHeadsの出力が最終出力に連結される.

残りの計算は2.2節のSingle-Headと同じである.

2.4節:Computational Complexity and Token Dependency(計算複雑性とトークン依存性)

セグメントサイズとDilation Rateが\((r, w)\)のDilated Attentionが与えられた場合,各クエリ・キー・バリューのペアは\((Q, K, V) \in \mathbb{R}^{N*d}\)から\((Q, K, V) \in \mathbb{R}^{\frac{w}{r} * d}\)にスパース化されるため,Attention計算のFLOPs数は次のように推定される.

$$ FLOPs = \frac{2N}{w} (\frac{w}{r})^2 d = \frac{2Nwd}{r^2} $$

さらに,これを複数のセグメントサイズとDilation Ratesを持つDilated Attentionに拡張する.

FLOPsは次のように書ける.

$$ FLOPs = 2Nd \sum_{i=1}^k \frac{w_i}{r_i^2} $$

式4と式5のセグメントサイズとDilation Ratesを用いて,FLOPsは次式で与えられる.

$$ FLOPs = 2w_0 Nd \sum_{i=0}^{k-1} \frac{1}{\alpha_i} \leq \frac{2 \alpha}{\alpha - 1} w_0 Nd \ \ \ (\alpha > 1) $$

ここで,\(w_0\)はあらかじめ定義された定数であり,\(\alpha\)は幾何学的系列wとrの共通比率である.

したがって,Dilated Attentionの計算複雑度は近似的に\(\mathcal{O}(Nd)\)となる.

さらに,各トークンの情報は最大距離Dまで伝播することができる.

$$ D = \sum_{i=0}^{l-1} w_i = w_0 \sum_{i=0}^{l-1} \alpha^i \approx \frac{w_0}{\alpha - 1} \alpha^l $$

ここで,lは伝播パスの長さである.

したがって,N個のトークンを持つシーケンスの最大パス長は以下のように推定できる.

$$ L \approx \log_\alpha \frac{N(\alpha - 1)}{w_0} \ \ \ (\alpha > 1) $$

これは,トークン依存性が\(\mathcal{O}(\log N)\)に近似していることを証明している.

3章:LONGNET as a Distributed Trainer: Scaling up to 1B Tokens(分散訓練士としてのLONGNET:1Bトークンへのスケールアップ)

Dilated Attentionの計算複雑度は\(\mathcal{O}(Nd)\)と大幅に削減されたが,計算とメモリの制約により,単一のGPUデバイスでシーケンス長を100万レベルにスケールすることは不可能である.

モデル並列[SPP+19],シーケンス並列[LXLY21, KCL+22],パイプライン並列[HCB+19]などの大規模モデル訓練用の分散訓練アルゴリズムがある.

しかし,LONGNETでは,特にシーケンスの次元が非常に大きい場合には不十分である.

3.1節:Distributed Algorithm(分散アルゴリズム)

LONGNET Figure4
図4:2つのGPUデバイス上でのLONGNETの分散訓練.シーケンス次元を分割することで訓練を並列化する.計算と通信のコストはデバイスの数が増えてもほぼ一定である.

我々は,シーケンス次元の分散学習において,LONGNETの線形計算複雑度を利用している.

一般性を損なうことなく,図4は2つのGPU上での分散アルゴリズムを示すが,これはさらに任意の数のデバイスに拡張することができる.

まず入力シーケンスをシーケンス次元に沿って分割する.

各シーケンスは別々に1つのデバイスに置かれる.

$$ X = [X_1, X_2] $$

そして,それらは2つのデバイス上のクエリ,キー,バリューに投影される.

$$ [Q_1, K_1, V_1] = [W_Q, W_K, W_V] X_1, \ \ \ [Q_2, K_2, V_2] = [W_Q, W_K, W_V] X_2 $$

セグメント長\(w_i \leq l\)(lはローカルデバイスのシーケンス長)の場合,式1から式3を用いてローカルにAttentionを計算する.

セグメント長\(w_i > l\)の場合,キーとバリューは異なるデバイスに分散している.

したがって,Attentionを計算する前に,キーとバリューのペアを収集する.

式1から式2を使用して,\(\{ Q, K, V \} \)を\(\{ \widetilde{Q}, \widetilde{K}, \widetilde{V} \}\)にスパース化する.

キーとバリューのペアを収集するために,全集合演算を実行する.

$$ \widetilde{K} = [\widetilde{K}_1, \widetilde{K}_2], \ \ \ \widetilde{V} = [\widetilde{V}_1, \widetilde{V}_2] $$

後方での全集合操作は,reduce-scatter操作になることに注意されたい.

Vanilla Attentionとは異なり,\(\widetilde{K}_i\)と\(\widetilde{V}_i\)のサイズはいずれもシーケンス長Nに依存しないため,通信コストは一定となる.

最後に,ローカルクエリ\(\widetilde{Q}_i\)とグローバルなキーと値のペア\(\{ \widetilde{K}, \widetilde{V} \} \)とのCross Attentionを計算する.

定式化は次のようになる.

$$ \widetilde{O}_1 = softmax(\widetilde{Q}_1, \widetilde{K}^T) \widetilde{V}, \ \ \ \widetilde{O}_2 = softmax(\widetilde{Q}_2, \widetilde{K}^T) \widetilde{V} $$

異なるデバイス間の出力を連結したものが,最終的なAttention出力となる.

$$ \widetilde{O} = [\widetilde{O}_1, \widetilde{O}_2] $$

上述した分散アルゴリズムは,バッチ次元を分割するデータ並列,隠れ次元を分割するモデル並列,レイヤを分割するパイプライン並列など,他の並列性と直交している.

3.2節:Scaling up to 1B Tokens(1Bトークンへのスケールアップ)

最新の分散システムで10億トークンへのスケーリングの実現可能性を検証する.

8Kから始めて,GPUメモリの限界までシーケンス長を徐々にスケーリングする.

それに応じてバッチサイズを縮小し,バッチあたりのトークン数を10億に維持する.

異なるシーケンス長の各モデルは,最大3つのセグメント長(2,048,デバイスあたりのトークン数,シーケンス長)を持つ.

10回の異なる実行について.順伝播の平均速度を計算する.

LONGNET Figure5
図5:Dilated AttentionとVanilla Attentionの実行時間.どちらもFlashAttention[DFE+22]を搭載している.

図5は,Vanilla AttentionとDilated Attentionの実行時間である.

どちらもFlashAttention Kernelで実装されており,メモリの節約とスピードの向上が図られている.

これは,Dilated Attentionがほぼ一定の待ち時間でシーケンス長をスケールアップできることを示している.

シーケンス次元を分割することにより,分散システムを活用してシーケンス長を10億トークンまでスケールすることができる.

これに対して,Vanilla Attentionはシーケンス長に対する2次依存性に悩まされる.

シーケンスの長さが長くなると,その待ち時間は劇的に長くなる.

さらに,Vanilla Attentionにはシーケンス長の制限を破る分散アルゴリズムは存在しない.

このことは,LONGNETの線形複雑性と分散アルゴリズムの優位性を証明している.

4章:Experiments on Language Modeling(言語モデリングの実験)

4.1節:Setup(セットアップ)

我々はLONGNETを言語モデリング上に実装した.

バックボーンとなるアーキテクチャはMAGNETO[WMH+22]であり,XPOS[SDP+22]の相対位置エンコーディングを用いる.

ただし,標準的なAttentionをDilated Attentionに置き換えることを除く.

我々はMAGNETOのベースサイズ構成を使用し,これは768の隠れ次元,12のAttention Heads,12のデコーダ層を持つ.

300以上のプログラミング言語のソースコード集であるThe Stackデータセット[KLA+22]でモデルを事前訓練する.

データはtiktoken tokenizerで前処理され,cl100k_baseエンコーディングされる.

モデルは0.5Mトークンのバッチサイズで300Kステップ訓練される.

ハイパーパラメータの詳細は付録を参照されたい.

すべての実験はtorchscale[MWH+22]コードベースに基づいて行われている.

4.2節:Results(結果)

LONGNETをVanilla TransformerとSparse Transformerの両方と比較する.

各アーキテクチャの違いはAttention層であり,その他は同じである.

これらのモデルのシーケンス長を2Kから32Kにスケールし,バッチサイズを小さくしてバッチあたりのトークン数を一定に保つ.

LONGNETでは,セグメント長をw = {2048, 4096, 8192, 16384, 32768}とし,Dilated Ratiosをr = {1, 2, 4, 6, 12}とする.

我々は[CGRS19]と同様に,複数のHeads異なるサブブロックにAttentionするSparse Attentionの固定パターンを実装する.

ブロックサイズは2048に設定される.

LONGNETとの比較が公平になるように,計算FLOPs数に合わせてスパース比率を調整する.

Vanilla TransformersのAttention層は高密度で完全に接続されているため,計算コストはかなり高くなる.

計算の制約から,我々は32Kシーケンス長までしかスケールアップしない.

Attention Variantsの全ての実装は,訓練の効率化のためにFlashAttentionに基づいている.

我々はFlash Attention KernelをSparse AttentionとDilated Attentionの両方にカスタマイズしている.

LONGNET Table2
表2:LONGNETとベースラインの言語モデルのパープレキシティ.

表2はStackデータセットにおけるこれらのモデルの結果をまとめたものである.

評価指標としてパープレキシティを用いる.

モデルは2Kから32Kまでの異なるシーケンス長でテストされた.

モデルがサポートする最大長よりも入力が長い場合,言語モデル推論のための最先端の外挿手法であるBlockwise Causal Attention(BCA)[SDP+22]を実装する.

また,絶対位置エンコーディングを除去する.

主な結果として,訓練中のシーケンス長を長くすることで,より良い言語モデルが得られることが実証された.

次に,推論におけるシーケンス長の外挿法は,シーケンス長がモデルのサポートする長さよりはるかに大きい場合には適用されない.

最後に,LONGNETは一貫してベースラインモデルを上回っており,言語モデリングにおける有効性を証明している.

4.3節:Scaling Curves of Sequence Length(シーケンス長のスケーリング曲線)

先行研究[KMH+20]では,言語モデルはパラメータや訓練トークンを増やすことで,いくつかのスケーリング則に従うことが示されている.

我々は,訓練中にコンテキストの長さをスケールアップした場合の言語モデルの性能に興味がある.

我々は1Kから32Kまでの様々な長さの入力を混合して損失をテストする.

推論中にBCAを用いることで,シーケンス長の汎化を改善する.

LONGNET Figure6
図6:LONGNETとDense Transformersの訓練時のシーケンス長を変えた場合のテストパープレキシティ.LONGNETはDense Transformersよりも低いパープレキシティと非常に少ない計算量の観点で優れている.

図6は,Vanilla TransformersとLONGNETの両方について,シーケンス長のスケーリング曲線をプロットしたものである.

行列乗算の総FLOPs数を計算することで,計算量を推定している.

この結果から,Vanilla TransformersとLONGNETの両方が,訓練中にコンテキスト長を大きくすることで恩恵を受けることがわかる.

しかし,LONGNETはより効率的にコンテキスト長をスケールアップすることができ,より少ない計算量でより低いテスト損失を達成することができる.

これは,外挿よりも長い訓練入力の優位性を示している.

結論として,我々の実験は,LONGNETが言語モデルのコンテキスト長をスケールアップするより効率的な方法であることを示している.

これはLONGNETがより長い範囲の依存関係をより効果的に学習できるためである.

4.4節:Scaling up Model Size(モデルサイズのスケールアップ)

LONGNET Figure7
図7:(左)モデルサイズの増加に伴うLONGNETのテスト損失.スケーリング曲線はVanilla Transformersと同様の法則に従っている.(右)異なるコンテキストウィンドウを用いたLONGNETのテスト損失.コンテキストウィンドウを長くすることで,より良い言語モデリングが得られる.

大規模言語モデルの重要な性質として,損失は計算量に比例してべき乗則でスケールするというものがある.

LONGNETが同様のスケーリング則に従っているかどうかを検証するために,1億2500万から27億のパラメータまで,異なるモデルサイズで一連のモデルを訓練する.

27億個のモデルは3000億個のトークンで訓練し,残りのモデルは約400億個のトークンで訓練する.

図7(a)はLONGNETの計算量に関するスケーリング曲線をプロットしたものである.

同じテストセットで計算量を計算する.

計算量は,訓練中の行列乗算の総FLOPs数を計算することで推定している.

その結果,LONGNETはべき乗則に従うことがわかった.

このことは,Dense Transformerは言語モデルをスケーリングするための前提条件ではないことを意味している.

さらに,スケーラビリティと効率はLONGNETによって得られる.

4.5節:Long Context Prompting(長いコンテキストのプロンプト)

プロンプトは言語モデルを誘導し,追加情報を提供するために不可欠な手法である.

我々は,LONGNETがプロンプトのための長いコンテキストウィンドウから利益を得られるかどうかを検証するための実験を行う.

具体的には,接頭辞の一部をプロンプトとして予約し,その接尾辞のパープレキシティをテストする.

プロンプトの長さは2Kから32Kまで徐々にスケールさせた.

公平に比較するため,接尾辞の長さは同じにし,接頭辞の長さをモデルの最大長まで増加させる.

図7(b)はテストセットでの結果である.

これを見ると,コンテキストウィンドウが大きくなるにつれて,LONGNETのテスト損失が徐々に減少していることがわかる.

これは,言語モデルを改善するために長いコンテキストを十分に活用するLONGNETの優位性を示している.

5章:Conclusion and Future Work(結論と今後の課題)

我々はTransformerのVariantであるLONGNETを発表する.

LONGNETはシーケンスの長さを10億トークン以上にスケールすることができ,より短いシーケンスでも損失はない.

LONGNETのコアはDilated Attentionであり,計算の複雑性を2次関数から線形関数に減少させる.

LONGNETは複数のGPUデバイスにまたがってシーケンスの訓練を並列化する分散訓練士として利用することができる.

実験によれば,LONGNETは長いシーケンスと短いシーケンスの両方をモデル化する際に,強力なベースラインよりも優れた性能を発揮する.

将来的には,LONGNETをより多くのタスク,例えばマルチモーダル大規模言語モデリング[HDW+23, PWD+23],BEiT事前訓練[BDPW22, PDB+22, WBD+23],ゲノムデータモデリングをサポートするように拡張する予定である.

参考文献

  1. [ALdJ+23] Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David C. Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, Yun-Hsuan Sung, and Sumit Sanghai. CoLT5: Faster long-range transformers with conditional computation. CoRR, abs/2303.09752, 2023.
  2. [BDPW22] Hangbo Bao, Li Dong, Songhao Piao, and FuruWei. BEiT: BERT pre-training of image transformers. In International Conference on Learning Representations, 2022.
  3. [BKB23] Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer to 1m tokens and beyond with RMT. CoRR, abs/2304.11062, 2023.
  4. [BMR+20] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurIPS 2020, 2020.
  5. [BPC20] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020.
  6. [CGRS19] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv, abs/1904.10509, 2019.
  7. [CLD+21] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
  8. [CND+22] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Many Others, Jeff Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling language modeling with Pathways. ArXiv, abs/2204.02311, 2022.
  9. [DDM+23] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Many Others, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. CoRR, abs/2302.05442, 2023.
  10. [DFE+22] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In NeurIPS, 2022.
  11. [DYY+19] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc Viet Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Anna Korhonen, David R. Traum, and Lluís Màrquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pages 2978–2988. Association for Computational Linguistics, 2019.
  12. [FDS+23] Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
  13. [FPB+23] Mahan Fathi, Jonathan Pilault, Pierre-Luc Bacon, Christopher Pal, Orhan Firat, and Ross Goroshin. Block-state transformer. CoRR, abs/2306.09539, 2023.
  14. [FZS21] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR, abs/2101.03961, 2021.
  15. [GGR22] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
  16. [HCB+19] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In NeurIPS 2019, pages 103–112, 2019.
  17. [HDW+23] Shaohan Huang, Li Dong, Wenhui Wang, Yaru Hao, Saksham Singhal, Shuming Ma, Tengchao Lv, Lei Cui, Owais Khan Mohammed, Qiang Liu, Kriti Aggarwal, Zewen Chi, Johan Bjorck, Vishrav Chaudhary, Subhojit Som, Xia Song, and Furu Wei. Language is not all you need: Aligning perception with language models. ArXiv, abs/2302.14045, 2023.
  18. [HZRS16] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778, 2016.
  19. [JGB+21] Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and João Carreira. Perceiver: General perception with iterative attention. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 4651–4664. PMLR, 2021.
  20. [KCL+22] Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. CoRR, abs/2205.05198, 2022.
  21. [KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
  22. [KLA+22] Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: 3 TB of permissively licensed source code. CoRR, abs/2211.15533, 2022.
  23. [KMH+20] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08361, 2020.
  24. [KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5156–5165. PMLR, 2020.
  25. [LLK+19] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 3744–3753. PMLR, 2019.
  26. [LLX+21] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR 2021, 2021.
  27. [LXLY21] Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Making 4d parallelism possible. CoRR, abs/2105.13120, 2021.
  28. [MKW+21] Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. In Marc’Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 2441–2453, 2021.
  29. [MWH+22] Shuming Ma, Hongyu Wang, Shaohan Huang, Wenhui Wang, Zewen Chi, Li Dong, Alon Benhaim, Barun Patra, Vishrav Chaudhary, Xia Song, and Furu Wei. TorchScale: Transformers at scale. CoRR, abs/2211.13184, 2022.
  30. [PDB+22] Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. BEiT v2: Masked image modeling with vector-quantized visual tokenizers. 2022.
  31. [PMN+23] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. CoRR, abs/2302.10866, 2023.
  32. [PWD+23] Zhiliang Peng, Wenhui Wang, Li Dong, Yaru Hao, Shaohan Huang, Shuming Ma, and Furu Wei. Kosmos-2: Grounding multimodal large language models to the world. ArXiv, abs/2306, 2023.
  33. [QHS+22] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pages 7025–7041. Association for Computational Linguistics, 2022.
  34. [SDP+22] Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A length-extrapolatable transformer. CoRR, abs/2212.10554, 2022.
  35. [SPP+19] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. CoRR, abs/1909.08053, 2019.
  36. [SWL23] Jimmy T. H. Smith, AndrewWarrington, and ScottW. Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
  37. [TDA+21] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
  38. [VSP+17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS 2017, pages 5998–6008, 2017.
  39. [WBD+23] Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorck, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, and Furu Wei. Image as a foreign language: BEiT pretraining for vision and vision-language tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
  40. [WCL+20] Genta Indra Winata, Samuel Cahyawijaya, Zhaojiang Lin, Zihan Liu, and Pascale Fung. Lightweight and efficient end-to-end speech recognition using low-rank transformer. In 2020 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2020, Barcelona, Spain, May 4-8, 2020, pages 6144–6148. IEEE, 2020.
  41. [WDC+23] WeizhiWang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Augmenting language models with long-term memory. CoRR, abs/2306.07174, 2023.
  42. [WLK+20] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768, 2020.
  43. [WMD+22] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. DeepNet: Scaling transformers to 1,000 layers. CoRR, abs/2203.00555, 2022.
  44. [WMH+22] Hongyu Wang, Shuming Ma, Shaohan Huang, Li Dong, Wenhui Wang, Zhiliang Peng, Yu Wu, Payal Bajaj, Saksham Singhal, Alon Benhaim, Barun Patra, Zhun Liu, Vishrav Chaudhary, Xia Song, and Furu Wei. Foundation transformers. CoRR, abs/2210.06423, 2022.
  45. [WRHS22] Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
  46. [ZBK+22] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing effective sparse expert models. CoRR, abs/2202.08906, 2022.
  47. [ZGD+20] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
  48. [ZKHB22] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 1204–1213. IEEE, 2022.

付録A:Hyperparameters(ハイパーパラメータ)

LONGNET Table3
表3:表2のモデルに使用したハイパーパラメータ.

LONGNET Table4
表4:図7(a)の実験に使用したハイパーパラメータ.

参考:LONGNET: Scaling Transformers to 1,000,000,000 Tokensの解説動画

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの解説動画です.

まとめ

LONGNET: Scaling Transformers to 1,000,000,000 Tokensの日本語訳を紹介しました.

Microsoftによる10億トークンをサポートしたTransformer「LONGNET」がわかりました.

AIのプログラミング言語「C++/Python言語」を学べるおすすめのWebサイトを知りたいあなたはこちらからどうぞ.

独学が難しいあなたは,AIを学べるオンラインプログラミングスクール3社で自分に合うスクールを見つけましょう.後悔はさせません!

国内・海外のAIエンジニアのおすすめ求人サイトを知りたいあなたはこちらからどうぞ.

国内・海外のプロンプトエンジニアのおすすめ求人サイトを知りたいあなたはこちらからどうぞ.

-TECHNOLOGY, NLP AI
-, , , , ,