【論文要約】Jamba: A Hybrid Transformer-Mamba Language Model【メモ】
イントロダクション
今回は以下のJambaアーキテクチャの論文を要約する。Mamba論文も同様に要約した記事を今後書こうと思う。(順序逆なのはすいません。)
研究の目的と背景
<purpose>
本研究の目的は、「Jamba」と名付けられた新しい大規模言語モデルを開発し、提案することである。Jambaは、Transformerレイヤーと「Mamba」と呼ばれる最新のstate-spaceモデルのレイヤーを組み合わせたハイブリッドアーキテクチャと、mixture-of-experts (MoE)を特徴としている。
この研究で達成しようとしている主な目標は以下の通りである:
高いスループットと小さなメモリフットプリントを実現しつつ、標準的な言語モデルのベンチマークおよび長文脈の評価において最先端の性能を達成すること。著者らは、256Kトークンまでの文脈長で強力な結果を示すことを目指している。
TransformerとMambaのレイヤーをどのように組み合わせるか、expertsをどのようにmixするかなど、様々なアーキテクチャ上の意思決定を研究し、大規模モデリングにおいて重要なものを明らかにすること。
訓練とJambaの評価によって明らかになった、これらのアーキテクチャの興味深い特性をいくつか説明し、このアーキテクチャのさらなる探求を奨励するために、様々なアブレーション実験からのチェックポイントを公開すること。
Jambaのような効率的で強力な大規模言語モデルの開発は、自然言語処理の分野にとって重要な意義を持つ。これにより、より高度な言語理解タスクが可能になり、実世界のアプリケーションへの応用が期待できる。
さらに、著者らはJambaの実装の重みをパーミッシブなライセンスで公開することで、研究コミュニティに貢献することを目指している。これにより、他の研究者がJambaを利用して新たな知見を得たり、さらなる改良を加えたりすることが可能になる。
<background>
大規模言語モデルの分野では、Transformerアーキテクチャが広く採用され大きな成功を収めてきた。しかし、Transformerには主に2つの欠点がある。1つは、高いメモリと計算要件のため長文脈の処理が困難なこと、もう1つは、生成されたトークンごとに全文脈に対して計算を行うため推論が遅くスループットが低いことである。一方、従来のRNNモデルは、任意の長さの文脈を単一の隠れ状態にまとめるため、これらの制限はない。しかし、RNNは訓練が高コストで長距離の関係性の捉えが限定的という短所がある。
最近のstate-spaceモデル(SSM)、特にMambaは、RNNよりも訓練が効率的で長距離の関係性の処理が得意だが、同規模のTransformer言語モデルの性能には及ばない。
いくつかの最近の研究でAttentionとSSMモジュールを組み合わせる試みがなされているが、SSMをAttentionとどのように組み合わせるか、実装の規模の点で本研究とは異なっている。特にH3は最大2.7B のパラメータと400Bの学習トークンで実装されたが、純粋なMambaには及ばないことが示されている。
以上のように、Transformerを用いた言語モデルが大きな成功を収めている一方で、メモリ使用量やスループットの点で課題があり、SSMとの組み合わせによる改善が期待されるものの、大規模な実装はこれまで存在しなかった。本研究のJambaは、Transformerの長所とMambaの効率性を兼ね備えた初の大規模ハイブリッドモデルとして、この分野の重要な進展を示すものと位置付けられる。
使用した手法の概要
<methods>
Jambaの中心となる手法は、Transformerレイヤーと「Mamba」と呼ばれるstate-spaceモデル(SSM)のレイヤーを組み合わせたハイブリッドアーキテクチャである。加えて、mixture-of-experts (MoE)モジュールを組み込んでいる。Transformerは、自然言語処理タスクで大きな成功を収めてきたアーキテクチャである。自己注意機構により、系列内の任意の位置間の依存関係を直接的に捉えることができる。しかし、文脈長に対して二次の計算量を要するため、長文脈での処理に課題がある。
Mambaは、最近提案されたSSMの一種である。RNNと同様に、任意の長さの入力系列を単一のベクトル表現にまとめることができ、長文脈の処理に適している。また、Transformerと比べて計算効率が良い。
MoEは、多数のexpert FFNのうち入力に応じて適切なものを選択する機構である。これにより、パラメータ数を増やしつつ計算量を抑えることができ、モデルの容量を効率的に拡張できる。
Jambaでは、これらの手法を以下のように組み合わせている。1つのJambaブロックは、TransformerレイヤーとMambaレイヤーをa:mの比率で交互に重ねたものである。一部のMLP層をMoEに置き換えている。MoEを適用する層の間隔はeで制御される。1つのMoE層にはn個のexpertがあり、各トークンでK個のexpertが選択される。
本研究では、以下のパラメータ設定のJambaを80GBのGPU1基に載せることを目標とした:
l = 8: 層の数
a : m = 1 : 7: TransformerレイヤーとMambaレイヤーの比
e = 2: MoEを適用する間隔
n = 16: expert総数
K = 2: 各トークンで使用するexpert数
Jambaは最大1Mトークンの文脈長で学習を行い、公開モデルでは256Kトークンまでの長さに対応している。
<comparison>
Jambaは、Transformerの強力な性能とMambaの効率性を組み合わせることで、同規模のTransformerモデルと同等の性能を維持しつつ、よりコンパクトなモデルを実現している。純粋なTransformerと比べ、長文脈でのメモリ使用量を大幅に削減できる。
Attention-SSMのハイブリッドモデルの先行研究と比べ、Jambaはその実装規模が異なる。特にH3は2.7Bまでのパラメータで実装されたが、純粋なMambaには及ばないことが示されている。これに対しJambaは、初の大規模なAttention-SSMハイブリッドモデルであり、最先端のTransformerベースのモデルに匹敵する性能を示している。
Closest are perhaps H3 [14], a specially designed SSM that enables induction capabilities, and a generalization called Hyena [35]. The former proposed a hybrid architecture that replaces the second and middle layers with self-attention, and was implemented with up to 2.7B parameters and 400B training tokens. However, as shown in [16], its perfomance lags that of pure Mamba. Based on Hyena, StripedHyena [36] interleaves attention and SSM layers in a 7B parameter model. However, it lags behind the Attention-only Mistral-7B [22]. All of this renders Jamba the first production-grade Attention-SSM hybrid model.
さらに、MoEを組み込むことで、active(使用)パラメータ数を抑えつつ、モデル容量(使用可能なパラメータ総数)を拡張している。これにより、Transformerの性能とMambaの効率性に加え、MoEによるモデル容量の拡張を同時に実現している点がJambaの大きな特徴である。
得られた主な結果
<main_results>
Jambaモデルの主要な結果は以下の通りである。
標準的な言語モデルのベンチマークおよび長文脈の評価において、同程度のサイズの最新の公開モデルと同等かそれ以上の性能を達成した。具体的には、Llama-2 70BやMixtral-8x7Bと同等の性能を示した。
256Kトークンまでの文脈長で強力な結果を示した。これは、公開されている最先端の言語モデルの中で最長の文脈長である。
長文脈でのスループットが、Mixtral-8x7Bの3倍に達した。また、140Kトークンの入力を処理する際にも、1つの80GB GPUに収まった。
以上の結果は、JambaがTransformerとMambaのハイブリッドアーキテクチャにより、高い性能と効率性を両立できることを示している。
<details>
標準的なベンチマークでの性能比較の詳細は、表2に示されている。例えば、HellaSwagでは87.1%、WinoGrandeでは82.5%の精度を達成し、Llama-2 70BやMixtralと同等かそれ以上の性能を示した。ただし、一部のタスク（ARC-EやHumanEvalなど）では、他モデルに及ばない結果も見られた。
長文脈での性能は、図4のneedle-in-a-haystackタスクと表3のQAタスクで評価された。特にQAタスクでは、平均6K～62Kトークンの長い入力に対して、Mixtralを上回る結果を示した。ただし、これらのタスクは実際のアプリケーションとは異なる可能性があり、結果の一般化には注意が必要である。
スループットの比較は、図3に示されている。単一のGPUでバッチサイズを大きくした場合（図3a）と、複数のGPUで文脈長を大きくした場合（図3b）の両方で、Jambaが他モデルを大幅に上回る結果となった。ただし、これらの結果は相対的に解釈すべきであり、絶対的な値ではない。
<comparison>
表2では、Llama-2 13B、Llama-2 70B、GemmaなどのTransformerベースのモデルと、MixtralやJambaのようなMoEを含むモデルが比較されている。全体的に、MoEを含むモデルの方が高い性能を示す傾向があるが、タスクによっては例外も見られる。
表3では、長文脈のQAタスクにおいて、JambaとMixtralが比較されている。5つのデータセットのうち4つで、Jambaが高いF1スコアを達成しており、全体の平均でもわずかに上回っている。ただし、これらの結果の統計的有意性は明記されていない。
図3では、スループットの比較が行われている。Jambaは、バッチサイズや文脈長を大きくした場合に、他のモデルを大幅に上回る結果を示している。特に、128Kトークンの文脈長では、Mixtralの3倍のスループットを達成している。
結果の解釈や考察
<structure>
本論文では、主に5章「Evaluation」と6章「Ablations and Insights」において、結果の解釈や考察が行われている。
5章では、Jambaモデルの評価結果が報告され、他の公開モデルとの比較が行われている。著者は、性能面でJambaが同程度のサイズのモデルと同等かそれ以上の結果を示すことを強調している。また、長文脈でのスループットの高さにも注目している。
6章では、アーキテクチャの設計選択に関する ablation 実験の結果が示され、それぞれの設計がモデルの性能にどのように影響するかが考察されている。著者は、TransformerレイヤーとMambaレイヤーの組み合わせ方や、MoEの適用方法など、Jambaの主要な特徴について詳細に議論している。
<arguments>
著者は、結果の解釈を通じて以下のような主要な主張を提示している。
TransformerとMambaを組み合わせたハイブリッドアーキテクチャは、大規模言語モデルの性能と効率性を向上させる有効な手段である。
MoEは、Attention-Mambaのハイブリッドモデルの性能をさらに向上させる。MoEを適用することで、モデルの容量を効率的に拡張できる。
pure Mamba モデルは、in-context learning の面で Transformer モデルに劣る可能性があるが、Attention-MambaのハイブリッドモデルはTransformerと同等のin-context learningを示す。
明示的なpositional embeddingは、Jambaアーキテクチャでは必ずしも必要ではない。Mambaレイヤーが暗黙的なpositional informationを提供していると考えられる。
これらの主張は、大規模言語モデルの設計に関する新しい知見を提供するものである。特に、TransformerとMambaの組み合わせや、MoEの適用は、今後の研究における有望な方向性を示唆している。
pure Mamba モデルとJambaモデルの差
論文では、pure Mambaモデルとハイブリッドモデルの性能差について、以下のように言及されています。まず、IMDBデータセットでの結果について、以下のように述べられています。
Looking into these results further, we found out that the pure Mamba model often does not follow the correct format. For instance, in the IMDB dataset, answer choices are "Positive" or "Negative". While the Attention model adheres to this format, the pure Mamba model often produces other answers, such as "Very Good", "Very Positive", "Funny", "Bad", "Poor", and "3/10". While these may be considered correct answers, the difficulty of Mamba to adhere to the format suggests a potential problem.
つまり、pure Mambaモデルは正解フォーマットに沿った出力ができていないことが指摘されています。一方、ハイブリッドモデルはAttentionモデルと同様に正しいフォーマットで出力できています。
さらに著者は、このような性能差について以下のように考察しています。
We hypothesize that this phenomenon points to a limitation of SSMs – a potential difficulty in in-context learning (ICL). Indeed, the ability to perform ICL has been linked to the emergence of so-called induction heads in Transformer language models during training, which perform approximate copying operations that are supportive of ICL [31]. We conjecture that the lack of an attention mechanism in the pure Mamba model makes it difficult for it to learn in-context. While Mamba may learn to copy and perform simple ICL when explicitly trained to do so ([16, 32], it is not clear if ICL is an emergent capability in SSM as is typical of Transformer models. In contrast, the hybrid Attention–Mamba model does perform successful ICL, even when only 1 out of 8 layers is an Attention one.
ここでは、pure Mambaモデルにはin-context learning (ICL)の能力が欠けている可能性が指摘されています。一方、ハイブリッドモデルは、わずか1層のAttentionを含むだけでも、ICLを実現できていると述べられています。
また、ハイブリッドモデルのAttentionレイヤーから、Transformerモデルと同様のinduction headが発現していることが示唆されています。
As anecdotal evidence of an emergent induction mechanism, we visualize in Figure 7 the attention of an example head from a 1.3B Attention-Mamba hybrid model (no MoE), on an IMDB example where the pure Mamba failed and the hybrid succeeded. Clearly, the attention from the last token (":") is focused on the labels from the few-shot examples. We have found 12 such heads in our hybrid model, in all three attention layers (which correspond to layers 4, 12, 20 in the model).
以上のように、著者はpure Mambaモデルの限界と、ハイブリッドモデルの優位性を実験的に示し、考察しています。
関連研究
<citation_context>
以下は、主要な関連研究が引用されている文脈の一部です。
1.Vaswani et al. (2017) [46]: Transformer の紹介
Vaswani et al. (2017) [46] proposed the Transformer architecture, which has become the dominant architecture for language models.
2.Gu and Dao (2023) [16]: Mamba の紹介と ablation 実験の文脈
Mamba is a recent state-space model (SSM) [16].
[16] reports that interleaving Mamba and attention layers is only slightly better than pure Mamba in terms of perplexity, with models up to 1.3B parameters.
3.Fedus et al. (2021) [13], Shazeer et al. (2017) [41]: MoE の紹介
Jamba also includes MoE layers [13, 41], which allow increasing the model capacity (total number of available parameters) without increasing compute requirements (number of active parameters).
4.Fu et al. (2022) [14], Poli et al. (2023) [35,36]: 関連する Attention-SSM ハイブリッドモデルの紹介
Closest are perhaps H3 [14], a specially designed SSM that enables induction capabilities, and a generalization called Hyena [35]. The former proposed a hybrid architecture that replaces the second and middle layers with self-attention, and was implemented with up to 2.7B parameters and 400B training tokens. However, as shown in [16], its perfomance lags that of pure Mamba. Based on Hyena, StripedHyena [36] interleaves attention and SSM layers in a 7B parameter model. However, it lags behind the Attention-only Mistral-7B [22].
5.Zhang and Sennrich (2019) [48]: Mamba レイヤーの正規化手法の紹介
Mamba layers incorporate several normalizations that help stabilize training in large model scales. In particular, we apply RMSNorm [48] in the Mamba layers.
<categorize_studies>
論文中で引用されている関連研究は、以下のように分類できます。
理論的基盤となる研究:
Transformer (Vaswani et al., 2017)
Mamba (Gu and Dao, 2023)
S4, SSM (Gu et al., 2021)
MoE (Fedus et al., 2021; Shazeer et al., 2017)
手法の比較対象となる研究:
H3 (Fu et al., 2022)
Hyena (Poli et al., 2023)
StripedHyena (Poli et al., 2023)
手法の一部として使用される研究:
RMSNorm (Zhang and Sennrich, 2019)
特に重要な関連研究は、Transformer, Mamba, MoE に関する研究であり、これらが Jamba の理論的基盤となっています。
Jambaの全体的なアーキテクチャ
<architecture>
Jambaの全体的なアーキテクチャは、以下のように説明されています。
Jambaは、Transformerレイヤー [46] とMambaレイヤー [16] を組み合わせたハイブリッドアーキテクチャに、mixture-of-experts (MoE) [13, 41] を組み込んだモデルです。これらの要素の組み合わせを「Jambaブロック」と呼んでいます。
Jamba is a hybrid decoder architecture that mixes Transformer layers [46] with Mamba layers [16], a recent state-space model (SSM) [17, 18], in addition to a mixture-of-experts (MoE) module [13, 41]. We call the combination of these three elements a Jamba block.
1つのJambaブロックは、TransformerレイヤーとMambaレイヤーを交互に重ねた構造になっています。具体的には、l個のレイヤーがあり、そのうちa個がTransformerレイヤー、m個がMambaレイヤーです（a:mの比率）。
また、一部のMLPレイヤーがMoEに置き換えられています。MoEは、e層ごとに適用されます。
MoEレイヤーには、n個のexpertがあり、各トークンでK個のexpertが選択されます。
本論文のJamba実装では、以下のパラメータ設定が使用されています。
$${l = 8}$$: 層の数
$${a:m =1:7}$$: TransformerレイヤーとMambaレイヤーの比率
$${e = 2}$$: MoEを適用する間隔
$${n = 16}$$: expert総数
$${K = 2}$$: 各トークンで使用するexpert数
Mambaレイヤーには、RMSNorm [48] が適用されています。
この記事が気に入ったらサポートをしてみませんか？