トランスフォーマーに代わるアーキテクチャ「SSM」とは?

GPTモデルなど生成AI市場で幅を利かせる大規模言語モデル(LLM)は、ほとんどが「トランスフォーマー」と呼ばれるアーキテクチャをベースとしている。アーキテクチャとは、AIモデルの基本的な構造や設計のことを指す。しかし、このトランスフォーマーアーキテクチャにはいくつかの課題があり、それを克服するための研究開発が日々進められている。

トランスフォーマーモデルが抱える課題の1つが長文処理だ。トランスフォーマーモデルは、その構造上、テキスト内の各単語(またはトークン)を他のすべての単語と比較してコンテキストを理解しようとするため、コンテキストウィンドウ(モデルが一度に処理できる文章の長さ)が大きくなるにつれて計算能力とメモリ利用が急増する傾向にあるのだ。そのため、リソースが適切にスケーリングされない場合、推論速度が低下し、一定の長さを超えるテキストを処理できなくなる可能性がある。最近は大幅に改善されたが、以前のChatGPTでも、長文プロンプトを入力するとエラーが頻繁に起こっていた。もしかすると、この構造上の問題が影響していた可能性が考えられる。

この課題を克服するために提案されたアプローチの1つが、カーネギーメロン大学とプリンストン大学の研究者らが2023年12月に提案した「Mamba SSM」アーキテクチャだ。SSMとは「State Space Model(状態空間モデル)」の略で、単語を処理しながら「状態」を継続的に更新することで文脈を理解する仕組み。ここでいう「状態」とは、モデルが文章を読み進めながら蓄積する情報のこと。この手法により、トランスフォーマーのような注意機構を使わずに、長い文章を効率的に処理できるという。

Mamba SSMの特徴は「選択メカニズム」にある。これは、人間が文章を読むときに重要な部分に注目するのと似たような働きをする。たとえば、長い小説を読むとき、私たちは物語の展開に重要な部分に注目し、それほど重要でない細かい描写は軽く読み流すことがある。Mamba SSMの選択メカニズムも、これと似たようなことを行っている。

この仕組みにより、Mamba SSMは長い文章(例えば本1冊分)を効率的に処理することが可能となった。従来のAIモデルでは、文章が長くなるほど必要な計算量やメモリが急激に増えてしまう問題があったが、Mamba SSMにはこの問題が発生しない。

さらに、Mamba SSMには「線形でスケーリングする」という特性がある。これは、処理する文章が長くなっても、必要な計算時間がゆるやかにしか増えないということだ。たとえば、100ページの本を読むのに10分かかるとすれば、200ページの本は20分、300ページの本は30分というように、ページ数に比例して読む時間が増えていくイメージである。

一方、従来のトランスフォーマーモデルは「二次関数的にスケーリング」する。これは、文章が長くなるほど計算時間が急激に増加することを意味する。100ページの本を10分で読めても、200ページの本は40分、300ページの本は90分かかるというように、ページ数の二乗に比例して読む時間が増えていくイメージだ。

この特性により、Mambaは長いコンテキストを扱う際に優れたパフォーマンスを発揮できるという。100万トークン(日本語で約100万語)を超えるデータでも性能が向上し続けることが確認されたという。対照的に、多くのモデルは長いコンテキストで性能が低下する傾向にある。

アブダビ発の最新AIモデル、Falcon Mamba 7Bとは?

上記のSSMアーキテクチャをベースに開発されたのが「Falcon Mamba 7B」だ。

これはアラブ首長国連邦アブダビを拠点とする政府系研究機関Technology Innovation Institute(TII)が開発したモデル。SSMの進化版「Mamba State Space Language Model (SSLM)」アーキテクチャを採用している。

Falcon Mambaの特筆すべき点は、Mambaモデルとしては初めての汎用モデルであるということだ。これは幅広い種類の言語タスクに対応できるAIモデルであることを意味する。

これまでもいくつかのState Space Language Model(SSLM)が開発されてきたが、文章の理解や生成の質という点では、従来の最先端トランスフォーマーモデルに及ばなかった。言わば、SSLMは「長い文章を速く読める」が「内容の理解は今一つ」という状態だった。これに対し、「速く読める」という特性を保ちつつ、「内容もしっかり理解できる」能力を大きく高めることに成功したのがFalcon Mambaということになる。

Falcon Mambaの設計は、上記のSSMをベースにしているが、大規模な学習をより安定させるために、データの偏りを抑える仕組み(RMS正規化層)が追加されている。この設計により、Falcon Mambaはまず、メモリを増やさなくても、どんなに長い文章でも処理できるようになった。それほど高性能ではないGPU(A10 24GB)1台でも十分であるという。

また、文章の長さに関係なく、一定のスピードで文章を生成することも可能だ。IITは、1つの指示(プロンプト)から最大13万語を生成する実験を行ったが、Falcon Mambaは始めから終わりまで一定の速度で文章を生成し、メモリ使用量も増えなかったという。一方、従来のモデルでは、生成する語数が増えるにつれてメモリ使用量が増加し、生成速度も遅くなったと報告されている。

これらの結果は、Falcon Mambaが長い文章の処理や生成において、従来のモデルよりも効率的であることを示している。この特性は、大量の文書を扱う業務や、長文の自動生成が必要なアプリケーションにおいて重宝することになるかもしれない。

Falcon Mamba 7Bの実力、メタの最新モデル超えも

Hugging Faceで公開されたベンチマーク結果から、Falcon Mamba 7Bの性能を見ていきたい。

新しいリーダーボードでは、Falcon Mamba 7Bが純粋なSSMモデルの中で最高スコアを記録。IFEval(推論能力を測定)で33.36、BBH(人間レベルの推論を要する課題)で19.88、GPQA(大学院レベルの質問応答)で8.05など、同サイズの競合モデルを大きく上回る記録を叩き出した。たとえば、メタのLlama3.1 8Bは、IFEvalが12.70、GPQAが6.15など。一方BBHでは、25.29とFalcon Mamba 7Bを超える性能を示した。

また、ARC(小学生レベルの科学的常識を問うテスト)で62.03%、TruthfulQA(モデルの誠実さや真実性を評価)で53.42%、GSM8K(数学的推論能力を測定)で52.54%を記録。これらのスコアは、70億パラメータ規模のモデルだけでなく、より大規模なモデルをも上回る成績だ。特にARCテストではLlama3 8B(60.24%)やFalcon2 11B(59.73%)を上回った。

MMLU(多領域にわたる知識の評価)やHellaSwag(常識的推論能力の評価)では他のトランスフォーマーモデルに若干劣る結果となったものの、全体平均では64.09%を記録。これはLlama3 8B(62.62%)やMistral7B v0.1(60.97%)を上回る数値だ。また平均で、メタのLlama3.1 8Bの62.28を上回った点は特筆に値する。

純粋なSSMモデルがトランスフォーマーモデルと互角以上の性能を発揮できることを示したのは、これが初とのこと。今後、さらに研究開発が進めば、SSM/SSLMベースの強力なAIモデルが続々登場するかもしれない。

なお、Falcon Mamba 7BはTII Falcon Mamba 7B License 1.0の下でリリーされており、Hugging Faceのエコシステム内で研究や応用目的で利用することができる。

文:細谷元(Livit