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

ソフトマックスの問題点

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

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

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

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

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

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

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

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

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

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

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