LLM生成の精度を維持しつつ、20％の応答高速化を達成　「Medusa Head+Speculative Sampling」を活用した工夫

「Romi」における過去のAIアップデート

信田春満氏：ちょっと話がずれますが、2023年の「TECH CONFERENCE」からのアップデートとして、「Romi」のAIがけっこう進化しています。

「ChatRomi」と呼んでいるんですが、パブリックなLLMをベースとして、「Romi」の独自会話データでファインチューンすることによって、今の「Romi」はしゃべっています。

現在は「StableLM」を使っているんですが、常に新しいモデルも考えていて、今は「Swallow」に目をつけています。

これによって、ユーザーさんからもすごくいい声をいただいています。出したのは2023年9月とけっこう前ですが、なにも言わずにサイレントリリースしたんですね。

こっそり出しておいて、こっそりアンケートを取って、「なんか最近変わったと思います？」と聞いてみたら、4割の人が変化に気づいていて。変化に気づいた人の中の9割に「いい変化だ」と言ってもらえた感じで。よかったということが、データで証明されています。

GPUインスタンス、ライブラリ変更による高速化

今回からLLMベースに変わりました。LLMってデカいので遅いです。

というわけで、「どうやって高速化しましょうか？」という話がここから始まります。

1つは、より強いGPUインスタンスを使いましょうという話です。当時の我々は「G4」というものを使っていました。「G5」が当時、まぁ今も最新なんですが、これにすると2倍弱ぐらい速くなります。2倍は速くならなかったです。一応AWS公式は、「2倍速くなる」と言っているんですけど。

これはお値段も2倍するので、コスパとしてはちょっと落ちるけれども、とはいえユーザー体験のためには仕方ないかなというところが、私の所感です。

あとは、実際に本番でAIモデルを使っている方だとやっているんじゃないかと思うんですが、「推論モデルをそのまま動かすんじゃなくて、fp16にしてモデルをコンパクトにしましょう」とか、あと、弊社では「TorchScript」を使っているんですが、そういうライブラリを使って推論専用のものに固めてしまえば高速化できるというテクニックを使っています。

ただ、「もうちょっと速くなりません？」みたいなことをチーム内で話していて、モデルの構造でも工夫できる点があるので、残りの時間はそこをメインで話そうかなと思います。

同じコストで2単語分を生成する「Medusa Head」

まず1つ目はMedusa Headっていう方法です。（スライドを示して）この左側の図は、通常のLLMの生成プロセスです。最初に1トークン、1つの単語だったり文字だったり、それでもない何かだったりするんですが、LLMはトークンを出してきます。今回は単語だとしましょう。なので、「バナナ」と出してきます。

1回ぐるっと（生成AIプロセスを）回すのに、「cost: 1」がかかるとしましょう。

次の単語を出すのに、また「cost: 1」がかかります。また「cost: 1」がかかりますという感じで、LLMは文字列を生成していきます。

メドゥーサって頭からたくさん髪が生えていますよね。Medusa Headという方法は、（通常の LLM の出力の他に）もう1個Medusa Headというものがくっついていて、こいつも「バナナ」のさらにその次の単語を予測して出してくるような仕組みなんですね。

ということは、「cost: 1」を1周回すだけで、2単語分を出すことができます。なので、従来の2倍速いんですね。

ですが、「1個先を先読みしましょう」ということは、やはり精度的にだいぶ（クオリティが）落ちてしまいます。

というわけで、これだと精度が犠牲になってしまう。じゃあどうするか。

2つのモデルを回すことで高速化できる「Speculative Sampling」

ここで出てくるのがSpeculative Sampling、投機的なサンプリングという方法です。この仕組みはちょっと複雑なので、かいつまんで、正確じゃない部分もありつつお話をします。正確な説明を見たい方は、論文を読んでください。

Speculative Samplingでは、普通の、遅いけれども高精度なモデル（Target Model）と、先ほどのMedusa Headを搭載したような、速いけれども低精度なモデル（Draft Model）の2つを回すことで高速化ができる方法です。

まずは、速いDraft Modelで単語を生成します。「バナナ」と「を食べ」というものを作りました。「バナナ」は普通のモデルと同じ仕組みで出しているので、十分信頼できるものです。

この次にMedusa Headが出してきた「を食べ」は、合っているかどうか、ちょっと自信がないという感じですね。

ここでまず、「cost: 1」を払いました。

次に、Draft Modelの生成結果を入力として、Target Modelでさらにその先を推論します。「バナナ」を使って「を食べ」を出してきました。さらにMedusa Headが出した「を食べ」を使って、「たい」を出してきました。

この時、Medusa Headが出してきた、自信がなかった「を食べ」と、Target Modelが出してきた「を食べ」は一致しているので、この「たい」も結果的に自信がある答えだということになります。ということは、「cost: 2」を払うことによって3トークンを生成できたので、今までの1.5倍高速であるということになります。

あと補足ですが、「ここで2つ作っているから『cost: 2』じゃないの？」って思うかもしれませんが、「バナナ」から「を食べ」を作るのと、「を食べ」から「たい」を作るのは並列で処理を回すことができるので、ここは「cost: 1」しかかからないんです。そこがミソです。

というわけで、話は戻りますが、「cost: 2」で3単語を作ることができました。

精度は維持しつつ、20パーセントの高速化を達成

「でも失敗することあるよね」という話があると思います。Medusa Headが、「『バナナ』の次は『いちご』でしょ？」と間違った答えを出してきました。

この場合、先ほどと同じ仕組みでいくと、「バナナ」から「を食べ」を出してきて、「いちご」から今度は「みかん」を出してきました。そうすると、Draft Modelの2単語目とTarget Modelの単語の結果が違うので、Target Model側が採用されます。

そして、この「みかん」という単語はDraft Modelの誤った推論を前提としたものなので、「この『ミカン』も使わない」としてあげます。

ただ、この場合でも「cost: 2」を使って「バナナ」「を食べ」の2つの単語を作ったので、損はしていないんですね。厳密には「いちご」を採るのか「を食べ」を採るのかは必ずTarget Modelが選ばれるわけではないんですが、そのあたりの詳細を知りたい方は、論文を読んでください。

というわけで、失敗した場合でも「cost: 2」を払って2単語を作れました。1倍です。ここのDraft Modelの応答が合っていた場合には1.5倍なので、たまに速くできる。「全体としてはちょっと高速化するよ」みたいな方法です。

我々の実験では、Medusa HeadとSpeculative Samplingを使ったことによって、精度は維持しつつも、20パーセントの高速化を達成しています。ちなみに、数学的に「Speculative Samplingを使うと精度は絶対に犠牲にならん」みたいな証明があるらしいです。僕はちゃんと理解できていないですけど。

インフラとAppサーバー、AIの2つで高速化を実現

というわけで、まとめです。インフラとAppサーバーの高速化では、ボトルネック調査とかキャッシュとか、その他の細かいことをいろいろやりました。

AIサーバーではより強いGPUインスタンスを使う、推論用のモデルに変換をする。あと、Medusa Head、Speculative Samplingなどを使って高速化していったという話でした。ご清聴ありがとうございました。