ソフトマックス関数

ソフトマックス関数（Softmax Function）とは、K個の実数値ベクトルを受け取り、各値を指数変換して合計で割ることで、合計が1.0になる確率分布に変換する関数です。e^xi / Σe^xj という式で計算され、最も大きな値が最も高い確率を持ちつつ、他の値もゼロにはならない「ソフトな」最大値選択を行います。

LLMにおける2つの重要な役割として、(1)トークン確率の計算：LLMは最終層で語彙全体（数万トークン）に対するロジット（スコア）を出力し、ソフトマックスでそれを確率分布に変換します。この確率から次トークンをサンプリングします。温度パラメータ（temperature）はソフトマックス前のロジットをスケールし、分布のシャープさを制御します、(2)アテンション重みの計算：スケールドドットプロダクトアテンションでは Q・Kの内積を次元数の平方根でスケール後にソフトマックスを適用し、各トークンへの注意の重みとして使用します。

temperatureとの関係として、temperature=1.0は通常のソフトマックス、temperature < 1.0（例: 0.3）はロジットを拡大して分布をシャープ化（確定的・創造性低）、temperature > 1.0はロジットを縮小して分布を平坦化（多様・創造性高）する効果があります。

数値的安定性として、指数関数は大きな値でオーバーフローしやすいため、実装では「max値を引いてからexp」するLog-Sum-Exp トリックが標準的に使われます。FlashAttentionもこの数値安定性の課題に対応した高効率実装です。