Adaptive Softmaxとは

言語モデルでは必ずといっていいほどソフトマックス（Softmax）が登場します。というのも言語モデルの学習では次に来るトークンを予測する訓練を行うからです。しかし、これが実はクセ者です。

ソフトマックスの問題点

以下の図では、言語モデルが文章の中で「犬と猫が」という言葉の後に何が続くかを計算している様子を表しています。

次に来るトークンの確率を予測する

次のトークンを予測するために、言語モデルは過去の文章を学習して、その文脈や単語の出現頻度を考慮します。例えば、「犬と猫が」の後に続く言葉として、「走っている」という単語が多く出現している場合、言語モデルはそのような予測をします。こうした地道な訓練が文章生成や文章の自動翻訳などに活用され、私たちの生活に欠かせない技術の1つとなりました。

人間だったら次に来る言葉を直感的に選んでいるのかもしれませんが、コンピューターはあらゆる言葉に対して次のトークンとなる確率を計算する必要があります。

仮に100個のトークンを扱えるモデルであれば、ソフトマックスを文字通りに実行すると100個のトークン一つ一つに対して確率を計算し、その中からより確率の高いものを選んだりするわけです。

もちろん100個であれば大したことはありません。

近年、言語モデルが扱えるトークンの総数（語彙サイズ）は大きくなるばかりです。扱う言語やデータセットやトークン化の仕方によりますが、数万から数十万というのはざらにあるわけです。

なお、扱える語彙を集めたものをコーパスとも呼びます。このコーパスは巨大化してきています。普通にソフトマックスを実行していたら膨大な計算量になってしまいます。

言語モデルが扱うコーパスの巨大化は、トランスフォーマーが登場する以前から問題になっていました。このためソフトマックスをより効率的に実行するための研究は早くから行われていました。

特に、今回紹介する適応ソフトマックス（Adaptive Softmax）は2016年に当時のFacebook AI（現Meta AI）によって論文が発表されています。また、2023年2月に発表された Meta の巨大言語モデル LLaMA でも使われています。

適応ソフトマックスは標準のソフトマックスと比べると倍以上に高速でしかもGPUをより効率よく使うことができます。それはちょっとした工夫の賜物でもあります。

この記事では適応ソフトマックスの性能とその仕組みを解説します。

適応ソフトマックスの性能

標準のソフトマックスは巨大なコーパスに対して非常に非効率です。なぜなら、語彙の中の多くのトークンがほぼゼロの確率を持つにも関わらず、すべてのトークンの確率を計算する必要があるからです。コーパスの中の語彙数が増えれば増えるほど無駄な計算がふえスピードが遅くなります。

これと比較して、適応ソフトマックスは2倍から10倍の高速化が可能です。巨大言語モデル（LLM）が膨大なコーパス（語彙）を扱う際に、時間と電力の両方を節約します。しかも、言語モデルの性能・精度を犠牲にすることもほぼありません。

実際にかかった時間で測ると適応ソフトマックスは標準のソフトマックスより早く訓練を終えることができるのですが、その理由として確率の高い語彙のグループに対してだけに集中して計算を行うことがあげられます。

次にその仕組みを見ていきます。

適応ソフトマックスの仕組み

言語モデリングは、単純に言うと、文章の中で次に来る単語を予測するタスクです。目的は、文の続きとして最も適切な単語を選ぶことです。

つまり、与えられたトークンのシーケンス（順番に並んだもの）の次に来るトークンを予測するのですが、ここで確率の計算が必要となります。言い換えると、言語モデリングの目標は与えられた文章に続く次のトークンとして確率が高いものを選ぶことです。

アンバランスな使用頻度

自然言語のトークン（単語や記号）は、使用される頻度に大きなばらつきがあります。つまり、言語の中にはよく使われる単語や記号がある一方で、あまり使われないものもたくさん存在します。

具体的には、日本語には「わたし」や「あなた」、「食べる」、「行く」といった頻出単語があります。これらの単語は、よく文章で使用されるため、データセットにも頻繁に登場します。

一方で、「蕾つぼみ」や「茗荷みょうが」、「榊さかき」などの単語は、比較的まれにしか使われないため、データセットにはあまり出現しません。

つまり、自然言語のトークンが使われる頻度にはかなりのアンバランスが生じています。一部の単語や記号は頻繁に使われる一方、他の単語や記号はほとんど使われません。

例えば、ウォール・ストリートジャーナルなどの記事からの語彙を集めたPenn TreeBank というデータセットでは、おおよそ20％の語彙が87％の使用率をカバーするそうです。

このことから、すべてのトークンに対してソフトマックスを使って確率を計算するのは非効率なのがわかります。

なぜなら、よく使われるトークンの確率はある程度高くなりますが、あまり使われないトークンの確率はほとんどゼロに近く、その計算は無駄な作業になることが多いからです。

従って、計算時間を節約するために、重要な（よく使われる）トークンだけに焦点を当ててソフトマックス計算を行う方が効率的です。

そこで、トークンをグループ化することを考えます。

２つのクラスターの例

わかりやすい例として、２つのトークンのグループ（クラスター）を考えます。

クラスター１には頻繁に使われるトークンが、クラスター２にはあまり使われないトークンが含まれています。仮にクラスター１が20％の語彙を持ち、クラスター２が80％の語彙を持っているとしましょう。

文章が与えられたときに次のトークンを予測する際、最初に行うことは、次のトークンがどちらのクラスター（クラスター1またはクラスター2）に含まれるかを予測することです。この予測により、計算が必要な範囲を絞り込み、効率的に次のトークンを特定することができます。

どちらのクラスターに次のトークンが入っているか？

つまり、次に来るトークンがクラスター１に属する確率とクラスター２に属する確率を計算します。ここでソフトマックスの計算を使いますが、対象はトークンではなくクラスターになります。

次に、予測されたクラスター（例えばクラスター１）が選ばれた後、そのクラスター内で最も確率が高いトークンを選ぶために、再びソフトマックス計算を行います。この２段階の計算によって、効率的に次のトークンを予測することができます。

まとめると、最初のソフトマックス計算でクラスター１に属する確率が高いと判断された場合、クラスター１内のトークンに対してソフトマックスの計算を行い、最も適切なトークンを予測します。もしクラスター２が選ばれた場合、同様にクラスター２内でソフトマックスの計算を行います。

ただし、クラスター２は選ばれる頻度が低いため、全体の平均として計算量が減ります。この２段階の計算方法によって、効率的に次のトークンを予測しながら、計算リソースを節約することができます。

以上により、2つのクラスターの例を用いて、段階的なソフトマックス計算を適用することで、効率が向上し、予測精度を維持しながら計算リソースを節約することが可能になることが理解できます。

この方法がGPUにとって効率的な理由は、GPUにはメモリの制限があるためです。適応ソフトマックスはメモリ使用量が少なく、計算負荷を減らすことができるため、GPUに適した計算方法となります。これにより、高速で効率的な言語モデリングが実現できます。

ただし、2つのクラスターの例では、常に2段階の処理が必要になるため、頻繁に使われるトークンに対しては効率が良くありません。

そこでより一般的な適応ソフトマックスでは更なる工夫を加えます。

より一般的な例

より一般的な適応ソフトマックスでは、頻繁に使われるトークンについての効率を向上させることを考えます。

そこで、頻繁に使われるトークンをそれぞれ独立したクラスターのように扱うことにします。

これにより、よく使われるトークンに対しては１段階で予測が可能となり、計算の効率がさらに向上します。この方法により、適応ソフトマックスは、言語モデリングをさらに効率的に行うことができるようになります。

順を追って解説します。

まず、最初のソフトマックスでは、頻繁に使われるトークンと、あまり使われないトークンが入ったクラスターの中から、確率が高いものを選びます。これが図の緑と白の部分に相当します。

この段階で使用されるメモリは、コーパス全体と比較してかなり少なくなります。このため、より効率的に次のトークンかクラスターを予測することができ、同時に計算リソースの節約も可能になります。

最初のソフトマックスで頻繁に使われるトークンが選ばれた場合、その確率をもって処理が終了します。多くの場合、頻繁に使われるトークンが選ばれるので、１段階目のソフトマックス計算だけで処理が終わり、効率が良くなります。これが緑の部分に相当します。

一方で、稀にあまり使われないトークンが属するクラスターが選ばれた場合、そのクラスター内で２段階目のソフトマックス計算が行われます。これが青い部分になります。各クラスターのサイズはコーパス全体よりも小さいため、メモリ使用量が少なく、計算を速く実行できます。

以上より、この一般的な適応ソフトマックスのアプローチでは、より効率的な言語モデリングが可能になり、計算リソースを節約しながら高速な処理が実現できます。

トークンの分配によるコスト

頻繁に使われるトークンを小さなグループに、あまり使われないトークンを大きなグループに分けることで、効率的に確率計算ができます。その理由は、小さなグループが平均的によくアクセスされるため、計算が早くなるからです。

しかし、どのように効率的な分配を行えば良いのでしょうか。

効率的な分配のアプローチ

効率的な分配を行うために適応ソフトマックス（Adaptive Softmax）を提案した論文「Efficient softmax approximation for GPUs」では、この方法が詳しく説明されています。

この論文では以下のアプローチが提案されています。

1. トークンの頻度を計算します。これは、トークンがコーパス内でどれだけ頻繁に現れるかを数えることです。

2. トークンを頻度順に並べます。これにより、最も頻繁に使われるトークンから最も稀に使われるトークンまでの順序が得られます。

3. トークンをクラスターに割り当てます。割り当ては、より頻繁に使われるトークンをより小さなクラスターに、頻度が低いトークンを大きなクラスターに割り当てることです。この割り当て方法により、平均的に小さなクラスターがより頻繁にアクセスされることになり、効率が向上します。

4. 最適なクラスター数と各クラスターのサイズを決定します。これは、実際の計算効率と予測精度のトレードオフを考慮して決定されます。適切なクラスター数とサイズを選ぶことで、効率と精度のバランスが最適化されます。

このアプローチでは、適切なクラスター数とサイズを選べることを前提としています。そのためには、ソフトマックスによる計算コストを知る必要があります。

次にソフトマックスによる計算コストの定義から始まってどのようにコストを最小化していくのかを論文からの解説に従って説明します。

トークン分配の全体像

すでに解説した通り、適応ソフトマックスでは、最初の段階で頻繁に使われるトークンを単独で扱い、あまり使われないトークンをクラスターにまとめます。

論文では、最初の段階のクラスター分けを $${V_h}$$（V head）と呼んでいます。この$${V_h}$$には、単独のトークンと2段階目で使われるクラスターを指すものが混ざっています。

つまり、$${V_h}$$はこの最初の段階で選択される可能性のある単独のトークンと、2段階目で選択されるクラスターの集合を表します。前述の図で言うと緑と白を全て合わせた部分になります。

論文の図３では、次のように個別トークンを青色で、クラスターを白で表現しています。

論文の図３

図を見てわかるように、$${V_h}$$は最初の段階のクラスター分けであり、単独のトークンと第二段階目で使われるクラスターを指すものが混じっています。

なお、図では第二段階目のクラスターが$${V_1, V_2, V_3}$$の３つしかないですが、一般に$${V_1, V_2, \dots, V_J}$$と$${J}$$個あるとします。

つまり、$${V_h}$$を含めると、全体として$${V_h, V_1, V_2, \dots, V_J}$$のクラスターがあることになります。

ここで各クラスターに対する計算コストを$${C_h, C_1, C_2, \dots, C_J}$$とすると全体のコストは各コストの総和なので、次のように定義されます。

$$
C = C_h + \sum\limits_i^J C_i
$$

次に、第一段階目のコストを計算し、その次に第二段階目のコストを計算します。

第一段階目のコスト

まず、$${C_h}$$の計算方法を考えます。第一段階目のソフトマックスでは$${k_h}$$個の個別トークンがあるとします。そのほかに$${J}$$個のクラスターがあるので合計$${J + k_h}$$個のクラスターを対象としてソフトマックスの計算になります。

第一弾目のソフトマックス

さらに、バッチサイズを$${B}$$とすると$${C_h}$$を次のように定義できます。

$$
C_h = g(J + k_h, B)
$$

$${g}$$はソフトマックスの計算コストで、それがクラスター$${V_h}$$の要素数$${J + k_h}$$とバッチサイズ$${B}$$に依存していることになります。

もちろん、トークンの分散表現である埋め込みベクトルのサイズも計算量に影響を与えますが、モデルに対して固定の値となるので暗に含まれるものとして無視します。

なお、トークンの分散表現などについては「トランスフォーマーを理解する」を参照してください。

第二段階目のコスト

次に２段階目のソフトマックスのコストを考えます。これはクラスターごとに計算します。各クラスターのコストを$${C_1, C_2, \dots, C_J}$$として、次のように定義します。

$$
C_i = g(k_i, p_i B)
$$

ここで$${k_i}$$はクラスター$${i}$$に含まれるトークンの数です。$${p_i}$$はこのクラスターが選ばれる確率で、$${p_i B}$$は選ばれる確率を考慮したバッチサイズになります。これにより各クラスターがどれだけ計算に貢献するかを示します。

例えば、バッチサイズが32で、クラスターが選ばれる確率が5%の場合、実質的なバッチサイズは 32 x 0.05 = 1.6 になります。この場合、一回の計算でバッチサイズがあたかも1.6であるかのような影響しか与えないことになります。

つまり、頻度が低いクラスターの選ばれる確率$${p_i}$$​は小さくなるため、実質的なバッチサイズ$${p_i B}$$も小さくなります。これにより、頻度が少ないクラスターの計算量が少なくなることが数式で示されています。

逆に言うと、頻度が少ないトークンをなるべく大きなクラスターにバランス良く入れることで計算コストを低く抑える分配が可能になるのが理解できます。

しかし、「バランス良く」というところがミソです。

合計のコストの計算

これまでの議論をまとめると第一段階と第二段階のソフトマックスの計算量の合計は以下になります。

$$
C = g(J + k_h, B) + \sum\limits_i^J g(k_i, p_i B)
$$

「バランス良く」するには$${J}$$や$${k_h}$$や$${k_i}$$などを調節して$${C}$$を最小にする必要があります。

なので、まずはコスト関数$${g(k, B)}$$をどのように計算するかを決める必要があります。

そこで論文では、$${g(k, B)}$$が関数としてどのような振る舞いを持つのかをGPUを使って実験的に確かめています。

論文の図1では、ソフトマックスの計算コストが対象となるトークン（埋め込みベクトル）の数$${k}$$に応じてどのように変化するかを示しています。K40とM40は、それぞれ異なるタイプのGPUを表しています。

論文の図１

また、図1からわかることは、異なるGPUにおいても、トークン（埋め込みベクトル）の数に対するソフトマックスの計算コストの傾向が似ていることです。

どのGPUでもトークンの数が50ぐらい（$${2^6 = 64}$$より少し下）まではコストが一定であり、それ以上ではトークンの数に比例してコストが増えています。

このような傾向を利用すれば、GPUの種類に関わらず適応ソフトマックスが効率的に機能することを示しています。

よって、コスト関数を次のように定義します。

$$
g(k, B) = \max(c + \lambda k_0 B, c + \lambda k B)
$$

ここで$${\lambda}$$は直線の傾きの値で、$${c}$$は切片の値です。ただし、$${k_0 \approx 50}$$までは、コストが最低値$${c + \lambda k_0 B}$$になっているというわけです。

トークンの数が$${k_0}$$より少ないとコストが一定なのでGPUを効率よく利用しているとは言えません。逆に、すべてのクラスターに含まれるトークンの数が$${k_0}$$以上であるようにすれば、より効率的にGPUを利用できることになります。

よって、$${k \ge k_0}$$であるという条件をつけ、GPUを効率よく利用できるようにします。こうすることで、コスト関数を次のように簡略化できます。

$$
g(k, B) = c + \lambda k B
$$

ここで第一段階と第二段階のクラスターに、それぞれ$${J + k_h \ge k_0}$$と$${k_i \ge k_0}$$という条件をつけることになります。

まとめると二段階で行う適応ソフトマックス全体のコストは次のようになります。

$$
\begin{aligned}
C &= g(J + k_h, B) + \sum\limits_i^J g(k_i, p_i B) \\
  &= c + \lambda (J + k_h)B + \sum\limits_i^J (c + \lambda k_i p_i B) \\
  &= (J + 1)c + \lambda B \left[ J + k_h + \sum\limits_i^J p_i k_i \right]
\end{aligned}
$$

この式で我々が決められるのはクラスターの数$${J}$$と各クラスターへのトークンの分配数$${k_h, k_1, \dots, k_J}$$になります。

そして、各トークンをどのクラスターに割り当てるかを調整して、全体のコストを最小化するようにします。この調整によって、より効率的な適応ソフトマックスが実現されることが期待できます。

合計のコストの最小化

繰り返しますが、合計のコスト$${C}$$を最小にするようにトークンを「バランス良く」分配する必要があります。

仮に、$${J}$$の値が固定されており、$${k_h, k_1, \dots, k_J}$$も決まっているとします。ここでクラスターのサイズは順番通りに増加するようになっているとします。この場合は、クラスターの数と各クラスターに含まれる個数が決まっているので、確率が高い順番にトークンを割り当てるだけで良くなります。

つまり、よく使われるトークンは$${V_h}$$​に、あまり使われないトークンは$${V_1, \dots, V_J}$$​の中で適切に分割されます。

動的計画法

次に条件を緩めて、$${J}$$の値だけが固定されているとすると、各クラスターへ割り当てるトークン数$${k_h, k_1, \dots, k_J}$$を決める必要があります。論文では、動的計画法（Dynamic Programminng）をアルゴリズムとして使うことで計算できるとしています。

論文には、その実装までは書いてはないのですが、動的計画法を使うということなので次のようなアルゴリズムではないかと推定されます。

簡単な例として、5個のトークンを二つのクラスターに分けることを考えます。データセットから各トークンの出現頻度が分かっているので、それぞれの確率を計算できます。それを$${p_1, p_2, \dots, p_5}$$とします。

下図のように各トークンを確率の高い順位並べます。

各トークンを確率が高い順に並べる。

そして、累積の確率を計算します。

$$
\begin{aligned}
f_1 &= p_1 \\
f_2 &= p_1 + p_2 \\
f_3 &= p_1 + p_2 + p_3 \\
f_4 &= p_1 + p_2 + p_3 + p_4 \\
f_5 &= p_1 + p_2 + p_3 + p_4 + p_5
\end{aligned}
$$

これを使うとクラスターごとの確率が計算できます。最後のトークンまで到達すると累積確率は100％になるので、$${f_5 = 1}$$になります。

例えば、トークン１と２をクラスター１に入れ、残りをクラスター２に入れるとします。すると、クラスター１の確率は$${f_2}$$となり、クラスター２の確率は$${f_5 - f_2}$$となります。

つまり、トークンの順列のどこでクラスターを分けるか、カットする場所を指定すれば、各クラスターの確率が簡単に計算できます。

これを使ってすべての可能な区分けに対してコストの総和を計算すれば、最小のコストになる分け方がわかります。

クラスターの数が増えても、考え方は同様です。二つのクラスターが三つになれば区切るところが二つに増えるだけです。

もちろん、たくさんのトークンがあると計算量は増えますが、一度きりの計算なので数時間かかったとしても問題ありません。

クラスターの数を決める

クラスターの数$${J}$$が決まっていれば動的計画法で最小コストになる分配が決められます。よって、残りの問題はどのようにクラスターの数$${J}$$を決めるかです。

以下は論文の図４の引用です。Bulgarian Europarlというデータセットで全体のクラスター数と実際の計算時間の関係を示しています。

論文の図４

これによると、10～15のクラスター数で計算時間が最小になりますが、5以上のクラスター数では時間の減少が大きくなく、数ミリ秒程度の差しかありません。

なお、クラスター数やサイズの調整で計算効率は向上しますが、適切なバランスが大切です。クラスター数を増やしすぎると、クラスターの予測が難しくなり、トークンの予測精度が低下します。

適応ソフトマックスの目標は、計算コストを抑えつつ、予測精度を維持することです。実験では、2～5のクラスター数で十分な効率が得られることがわかっており、その範囲が無難です。

もちろん、データセットの偏りによって、クラスター数やサイズを調整する必要があるのでケースバイケースではありますが、目安として2～5のクラスター数でも十分な成果が得られることを覚えておくと良いでしょう。

まとめ

本記事では、適応ソフトマックス（Adaptive Softmax）について解説しました。

適応ソフトマックスは、トークンの使用頻度によって、個別に対応するトークンとグループに分けて扱うことで、より効率的にGPUを利用することができます。

論文では、適応ソフトマックスが2倍から10倍のスピードを達成しつつ、性能・精度を下げることがないことが実証されています。

また、適応ソフトマックスは、巨大言語モデルであるMetaのLLaMAでも使われています。

関連記事