Jamba 1.5が採用した「トランスフォーマー×Mamba」のハイブリッドアプローチとは？

Jamba 1.5の基礎の1つとなっているMambaは、カーネギーメロン大学とプリンストン大学の研究者らが2023年12月に提案したトランスフォーマーを代替する新しいアーキテクチャだ。Mambaは構造化状態空間モデル（SSM）の一種で、従来のトランスフォーマーモデルが抱える長文処理の課題を解決するために開発された。 Mambaの特徴は「選択メカニズム」にある。これは、人間が文章を読むときに重要な部分に注目するのと似たような働きをする。長い小説を読むとき、私たちは物語の展開に重要な部分に注目し、それほど重要でない細かい描写は軽く読み流すことがある。Mambaの選択メカニズムも、これと似たようなことを行っている。 この仕組みにより、Mambaは長い文章を効率的に処理することが可能となった。従来のAIモデルでは、文章が長くなるほど必要な計算量やメモリが急激に増えてしまう問題があったが、Mambaにはこの問題が発生しない。さらに、Mambaには「線形でスケーリングする」という特性がある。これは、処理する文章が長くなっても、必要な計算時間がゆるやかにしか増えないことを意味する。一方、従来のトランスフォーマーモデルは「二次関数的にスケーリング」する特性を持っており、文章が長くなるほど計算時間が急激に増加するのだ。 この構造の違いにより、Mambaは長いコンテキストを扱う際に優れたパフォーマンスを発揮できる。たとえば、100万トークン（日本語で約100万語）を超えるデータでも性能がほとんど落ちないことが報告されている。 しかし、Mambaにも課題は存在する。AI21の説明によると、Mambaは長文処理には優れているものの、全体のコンテキストを考慮する注意機構がないため、特に想起関連のタスクで既存の最高性能モデルと同等の出力品質を達成するのが難しいという。ここでいう「想起」とは、過去に学習した情報を適切に思い出し、利用する能力を指す。Mambaは、文章全体のコンテキストを十分に考慮できないため、こうしたタスクで既存の最高性能モデルと同等の性能を発揮するのが難しかったのだ。 この課題を克服するために、AI21はMambaとトランスフォーマーのハイブリッドアーキテクチャを開発した。これがJambaの核心部分となる。Jambaは、トランスフォーマー、Mamba、そしてMixture-of-Experts（MoE）層で構成されており、メモリ、スループット、パフォーマンスを同時に最適化することを目指したモデルだ。 この新しいアーキテクチャにより、Jamba 1.5は長文処理の効率性と高い出力品質を両立させることに成功。AI21によれば、Jamba 1.5は長いコンテキストで同サイズの他モデルの3倍のスループット（スピード）を達成し、単一のGPUで14万トークンのコンテキストを処理できるという。