研究者は、Mambaとアテンションアーキテクチャを組み合わせたブロック拡散言語モデル（BDLM）のハイブリッド版を提案し、以前のハイブリッド拡散モデルに内在するキャッシングの課題に対処しました。逆Mambaスキャンをアクティブなノイズ除去ブロックに限定することで、この手法はブロック間で正確なKVキャッシュの再利用を可能にし、フルシーケンススキャンのプレフィックスのみという制限を克服します。

  • BDLM Mamba-Hアーキテクチャは、逆スキャンをアクティブなノイズ除去ブロックに限定し、ブロックごとの拡散中に正確なキャッシングを可能にします。
  • DCLMでの87Mパラメータのスイープにおいて、BDLM Mamba-HはBDLMアテンションやフルシーケンスベースラインと比較して、C4-en検証のパープレキシティで最良の結果を達成しました。
  • 350Mパラメータでは、モデルはパフォーマンス指標においてBDLMアテンションと競合しています。
  • 長文脈推論において、BDLM Mamba-Hは65KトークンでフルシーケンスのDiffuMamba-Hの19.7倍のスループット、262KでBDLMアテンションの3.7倍のスループットに達しました。

このアプローチは、Mambaハイブリッドが長文脈拡散にとって実用的なアーキテクチャであることを示し、競合するパープレキシティを維持しながら大幅なスループットの改善を提供します。