[論文レビュー] DenseMamba: State Space Models with Dense Hidden Connection for Efficient Large Language Models

Wei He, Kai Han+5 more|arXiv (Cornell University)|Feb 26, 2024

Topic Modeling被引用数 5

ひとこと要約

DenseMambaはDenseSSMを導入します。これはState Space Models (SSMs) の dense hidden connection 機構で、層間の情報フローを改善し、最小限のパラメータと計算オーバーヘッドでRetNet/Mamba様式のアーキテクチャの性能を向上させます。

ABSTRACT

Large language models (LLMs) face a daunting challenge due to the excessive computational and memory requirements of the commonly used Transformer architecture. While state space model (SSM) is a new type of foundational network architecture offering lower computational complexity, their performance has yet to fully rival that of Transformers. This paper introduces DenseSSM, a novel approach to enhance the flow of hidden information between layers in SSMs. By selectively integrating shallowlayer hidden states into deeper layers, DenseSSM retains fine-grained information crucial for the final output. Dense connections enhanced DenseSSM still maintains the training parallelizability and inference efficiency. The proposed method can be widely applicable to various SSM types like RetNet and Mamba. With similar model size, DenseSSM achieves significant improvements, exemplified by DenseRetNet outperforming the original RetNet with up to 5% accuracy improvement on public benchmarks. code is avalaible at https://github.com/WailordHe/DenseSSM

研究の動機と目的

• 効率的なシーケンスモデル（SSM）をTransformersの代替として大規模言語モデルに適用する動機づけ。
• SSMにおける層間の隠れ状態劣化と情報フローへの影響を特定する。
• DenseSSMを提案し、浅い隠れ状態を深い層に濃密に接続しつつ、学習/推論の効率を維持する。
• DenseSSMがRetNetおよびMamba派生モデルの公的ベンチマークで性能を改善することを示す。
• 選択的遷移モジュールと隠れ融合モジュールの実用的でハードウェアへ優しい実装戦略を提供する。

提案手法

• DenseSSMを、前の m 層から現在のレイヤへ密に隠れ状態を接続するために、選択的遷移モジュール φ と隠れ融合モジュール Fuse を介して実装する。
• φに projection+gate 機構を用いて浅い隠れ状態をターゲットサブ空間へ投影し、有用な成分を選択する。
• 現在の層の隠れ状態と選択された浅い状態を加法融合（h_t^L = h_t^L + sum φ(h_t^{l-i})）で結合する。
• DenseSSMを再帰型（DenseRetNet）および畳み込み/GAU風ブロック変種の両方に拡張し、並列可能な学習と効率的な自己回帰推論を維持する。
• DenseSSMをRetNetおよびMambaアーキテクチャに適用し、DenseRetNetとDenseMambaを同様のモデルサイズで性能を向上させる。
• 56GB The Pile-derived data、15B tokens、LLaMA tokenizer を用い、多様な QA/常識タスクでゼロショットおよび 4-shot 能力を評価する。

実験結果

リサーチクエスチョン

• RQ1 SSM層間に濃密な隠れ状態接続を導入することで、従来の層制限の隠れ状態伝播と比較して情報フローを維持または改善できるか。
• RQ2 DenseSSMはRetNet-およびMambaベースのモデルの精度と困惑度を、学習/推論の効率を犠牲にせずに改善できるか。
• RQ3 選択的遷移モジュールと隠れ融合モジュールは性能とパラメータ効率にどのように影響するか。
• RQ4 ゼロショットおよび4-shotの commonsense QAおよび推論ベンチマークで、改善は一貫して見られるか。
• RQ5 DenseSSMは最小限のアーキテクチャ変更で、異なるSSM変種（例：RetNet、Mamba）に適用可能か。

主な発見

• DenseSSMは、浅層の隠れ状態を深層へ濃密に接続することで情報フローを強化し、選択的投影とゲート機構を用いた。
• DenseRetNetおよび DenseMamba は、公開ベンチマークでゼロショットおよび少数-shot 設定の下、ベースラインのRetNet、Mamba、Transformerベースのモデルより性能が向上。
• DenseRetNetは、公開ベンチマークで元のRetNetに対して最大5%の精度向上を達成（要約欄に報告された）。
• DenseSSMは、パラメータ増加はわずかで、学習の並列性と自己回帰推論の効率を維持する。
• アブレーション研究は、選択的遷移のIdentity projectionとMLP、そして加法融合を用いると、性能とパラメータ数のトレードオフが良好になることを示唆する。