1. softmax 関数 (ソフトマックス関数)とは [概要]
softmax 関数とは，機械学習の「多クラス識別モデル」において，確率的なベクトル出力をする際によく用いられる，Maxをソフトに行う出力が行える活性化関数である．softmax関数は，合計が1になるように，ベクトル中の各次元の変数を正規化する関数であり，(2クラス0~1での分類向けである)シグモイド関数の，多変数向けの一般化である．つまりは「シグモイド関数の多変数版」が，softmax 関数である．
この記事では，まずCNNの識別モデルにおけるsoftmax関数について，1節で説明をおこなう．次に2節では，より一般化された，機械学習では昔から用いられてきた「ベイズ識別モデルにおける，softmax関数の使用」について述べる．最後の3節では，深層学習でその他の「softmax関数を用いる場面の例」を列挙して紹介して整理しておく，アテンション機構，Transformerのマルチヘッドアテンション, Deep Metric Learning, 自己教師あり表現学習など )．
softmax関数は『one-hotベクトル符号化に用いる[argmax関数]の滑らかな近似』に相当しており，その意味で「softargmax関数」としても知られている．つまり「ソフトなMax関数」である．argmaxしてしまうとNMSのようにその1点のみが値が残って，他がゼロ値になってしまうが，soft(arg)maxだと，最大値も元の値よりも大きい値へと活性化しながら周辺もある程度0でない値に生き残らせることができる．
1.1 softmax の主用途：出力のマルチクラス確率分布化
入力のベクトルの要素和を1で正規化し，マルチヌーイ確率分布として確率的な最終出力に変換する際に使われる．ロジスティック回帰や，ニューラルネットワークの「確率分布出力の作成」目的で，以前から頻繁に用いられてきた．
ディープラーニングが登場すると，クラス識別の物体認識CNNや，自然言語処理のRNNにおいて出力層1-of-Kの単語表現(最近はトークン表現)を出力する際に用いられることになった．よって，ディープラーニング全般でも，softmax層は，非常に使用頻度の高い，重要な活性化層となっている．
softmaxでは，以下のように，$K$クラス識別モデルの出力の，クラス$i$の確率をそれぞれ独立に計算する：
\[
\text{softmax}(\bm{z})_i = \frac{\exp(z_i)}{\sum_{k=1}^{K} \exp (z_k)}\tag{1.1}
\]
ここで，入力の$\bm{z}$は，softmax関数に入力する直前の，隠れベクトル表現である．ニューラルネットワークだと，softmax層の，直前の隠れ層が出力するベクトル$\bm{z} \in K$に相当する．
式(1.1) の分母は，$\bm{z}$の各要素である$z_i$の指数和により，「$\bm{z}$全体を正規化する項」である．この正規化機能により，出力 $\bm{y}$ 全体の合計値は1となるので，$\bm{y}$をマルチヌーイ分布に変換することができている（つまりsoftmaxは，正規化指数関数であるとも言える）
1.2 softmax の機能: ソフトな argmax
名前の通り，softmax （またはsoftargmax）は，シグモイド関数のような「$\exp$を用いた関数の活性化」を通して，argmax関数の近似として働いて，入力ベクトルを変形する．argmaxの機能を持ちつつも，微分可能関数であるので，（ディープ）ニューラルネットワークの多クラス識別出力層に，良く用いられる（図1）（※ 2クラス識別ネットワークや2項ロジスティック回帰では，シグモイド関数を出力に用いる）．
softmaxで活性化すると，入力ベクトル$\bm{z}$のうち，特定の要素$z_k$だけが，最大値として大きくなり，他の要素の出力（確率）は小さく出力される．ただし「soft」なので，極端に要素間の差はつかない（シグモイドの多クラス版であることから，softmaxの関数形はイメージしてもらいたい）．
1.3 活性化関数としての分類．
Softmaxは，活性化関数の1種でもあるが「確率的識別モデルの『出力層用』の活性化関数である」と把握しておくとわかりよい (親記事の活性化関数でも，そう分類している)．
古くは，ロジスティック回帰や三層MLPのニューラルネットワークにおいて，，「クラス識別モデルの出力」として，シグモイド関数と共にsoftmax関数は用いられてきた．
近年のディープラーニングでは，画像識別CNNの出力層や，その他モデルの多クラス識別層に用いられる．RNN を用いた自然言語処理では，RNN言語モデル や，系列対系列変換の出力層で，単語・トークンの次フレームの「クラス予測」に，softmax層が用いられる．
2. ベイズ識別モデルにおける softmax 出力
2節では，1節でみた式(1.1)を，実際に識別モデル（特にニューラルネット）の出力の活性化関数に用いる様子を，ベイズ推定モデル的な識別モデルの目線で説明する．これにより，識別モデルにおけるsoftmaxの役割を，もう少し直感的に，わかりやすく分解して把握できるようにしたい．
2.1 ベイズ識別モデルの出力としての解釈
教師あり学習で学習をおこなう，確率的出力の「条件付き識別モデル$P(\bm{y}|\bm{x})$」の出力分布を，そのターゲット変数がone-hot表現ベクトルであることから，マルチヌーイ分布によりモデル化することを考える(図1)：
\[ P(\bm{y}|\bm{x}) ＝ \text{Multinoulli}(\bm{y}) \tag{2.1}\]
ここで，識別モデルの入出力は以下の通りとする（$k=1,2,\ldots$K$を各クラスの添え字）：
- $\bm{x} = (x_1, x_2, \ldots, y_D)$（入力）：$D$個の特徴量 $x_i$ から構成される，$D$次元ベクトル．
- $\bm{y} = (y_1, y_2, \ldots, y_K)$（出力）：one-hotエンコーディングされた，1-of-Kベクトル．$ \bm{y} \in {0,1}^K $．
識別モデルでは，予測結果$\hat{\bm{y}}$も，$\bm{y}$から最大の確率のみを1にして，他は0になった「one-hotエンコーディング化されたベクトル」を，argmax処理により出力したい:
\[
\hat{\bm{y}} = \arg\max{P(\bm{y}|\bm{x})}
\]
このargmax関数の近似として，softmaxを用いる．図1 に示したように，このsoftmax=argmax近似を最後に行うことで，出力が確率分布化して，one-hotエンコーディングしたマルチヌーイ分布表現と同等の表現に変換できる．$\bm{y}$と$\hat{\bm{y}}$のあいだで，交差エントロピー損失を計算できるようになり，識別モデル・ネットワークを学習できるようになる．
ここで，softmax関数の，直前のベクトル$\bm{z} \in \mathbb{R}^K$は，識別モデル（ロジスティック回帰やニューラルネット）により，$\phi(\bm{x})$から（softmax後の）K次元ベクトル$\bm{z}$へ変換される（図1 上部）：
\[
\bm{z} = \phi(\bm{x}) \tag{2.2}
\]
そして，最後に(1.1)のsoftmax 関数を用いて，$\bm{z}$を入力として，各クラス$k$の出力確率値を計算し，識別モデルの出力$\bm{y}$が得られる（図1 下部）：
\[ y_k = \text{softmax}(\bm{z})_k = \frac{\exp(z_i)}{\sum_{k=1}^K \exp(z_{k}) } \tag{2.3}\]
以上のように，(2.1),(2.2),(2.3)の処理に分解して考えると（図1），識別モデルにおけるsoftmaxの役割が全体像として理解しやすいと思う．
ただし，softmaxを用いても，ベイズ推定モデルのような確率出力ベクトルが得られるだけで，決定的関数であることに注意である．つまり，実際にベイズ学習するわけではないし，出力もランダムサンプリングするわけではないことに注意．
2.2 シグモイド関数と softmax 関数の類似部分
式(1.1)および式 (2.3)のsoftmax関数は，「シグモイド関数を$K$個用意し，それらの出力全体を正規化することで，和が必ず1になって，確率出力ができるようにした関数」と捉えることができる．
多次元空間でK次元ベクトル上でのsoftmax出力の挙動を，うまく可視化はできないのでイメージはしづらいかもしれない．とはいえ，（概要でも述べたように）softmaxは「シグモイド関数の多変数への一般化」であり「シグモイド関数のK変数版である」と覚えておき，数式的にも，そのようにイメージすると理解しやすい．
また対数尤度の最大化を行う際に，softmaxでも「指数関数の使用」がシグモイド関数と同様に効いてくる
詳しくはDeep learning book の6.2.2.3の 式(6.30) 近辺を参照．
3. 深層学習における，他の主な softmax 活用例
1,2節でみてきた「教師あり学習の識別モデルにおける，出力の活性化関数」としての使用以外でも，ディープラーニングではsoftmaxが用いられることが多い．
ここでは「アテンション（3.1）」，「Deep Metric Learning（3.2）」とその応用の１つの「自己教師あり学習（3.3）」について使用例を見ておきたい．
3.1 アテンション機構での使用
アテンション機構 では，softmax関数が，アテンション係数の算出に用いられる．また，並列化したアテンションであるマルチヘッドアテンションをメイン部品で多数用いるTransformerでも用いらている．
入力集合の各ベクトルに対する，現フレームの入力ベクトルの関連度を，学習済みの「スコア関数」もしくは単なる「内積」からて「アテンション係数」として計算する．そして，そのアテンション係数で，入力系列の各ベクトルを重みづけして和を取った「コンテキストベクトル」も予測に使うことで，各入力ベクトル（とのその中の各特徴）が重要度で重みづけされた状態で予測できるようになる仕組みである．
系列対系列変換向けのアテンション機構では，入力のベクトル系列 $\bm{S} = \{ \bm{s}_1,\bm{s}_2,\ldots, \bm{s}_N \}$ に対して，出力系列のフレーム$i$における入力ベクトル$\bm{h}_i$から，フレーム $i$ における $\bm{S}$の 各ベクトル $\bm{s}_j$ に対応する重み計数 $a_j(i)$ を，アテンション係数として推定する．
各ベクトルのアテンション係数は，なんからのスコア関数 $\text{score}(\bm{s}_j,\bm{h}_i)$ を学習して，2ベクトル間の関連度・類似度を，学習済みのネットワーク $\text{score}(\cdot,\cdot)$ により推定する．もしくは，スコア関数は学習せずに，単に $\bm{s}_j,\bm{h}_i$ 間の内積などを通じてスコア算出をおこなう．
しかし，アテンション係数はスコア関数や内積で計算すると，値の範囲がバラバラになって扱いづらい．そこで，以下のようにスコア関数で計算した値を，softmax関数で活性化し（= 確率ベクトル化してmax値も和らげソフト化し），確率ベクトル化による正規化を通じてスケールが揃った値を，最終的なアテンション係数に用いる：
\[
a_j(i) = \frac{ \exp(\text{score}(\bm{s}_j, \bm{h}_i))}
{\sum_{j=1}^N \exp(\text{score}(\bm{s}_j,\bm{h}_i))}
\]
分母の正規化項の計算は，$j$方向に行っている点に注意されたい．
このアテンション機構の例のように，クラス識別出力が欲しいわけではなく，ベクトルの各要素を合計1にして正規化したいときにも，機械学習・ラーニングではSoftmaxを活用することが多い．アテンションの詳しい定式化については，アテンション機構とseq2seq with attentionの記事を参考のこと．
3.2 Deep Metric Learningでの使用
画像認識で Deep Metric Learning（DML）を行う際，2対のCNNから構成される「Siamese Network」構造により，対照損失(contrastive loss)や，その3つ組版であるトリプレット損失(triplet loss)を用いた学習をおこなう．画像サンプル同士のpos / negラベルを頼りにするサンプルペアを用いた，これらの損失によって，CNNを「埋め込み関数」として学習する．
Deep Metric Learningを用いると，CNNバックボーンで抽出できる特徴ベクトルの「埋め込み空間」を，クラスラベルやIDラベルにランクラベルなどから学習し，特徴ベクトル間の「距離関係度合い」を，埋め込み空間への射影ネットワークによって学習することが可能である．つまりは，類似画像検索や画像認証・照合モデルを，ディープラーニングを通じて学習できる．
また，「画像+テキスト」や「動画 + 音楽」ペアなどの「異なる2モーダル間のペア」から，クロスモーダルな埋め込み（表現）学習をおこなうことも多い
※ クロスモーダルの話も説明すると長くなるので，ここでは割愛する．このクロスモーダル研究は，著者の専門の1つであるが，このサイトでもしばらく詳しい話までは行わない予定である．
さて，そのDMLでは，対照損失・トリプレット損失に追加する合成損失（あるいは正則化項）として，「Softmax + 交差エントロピー損失」も使用することで，安定した収束をうながすことが多い（※ 詳しい話まではしないが，当初はPerson Re-ID界隈などで発展してきた発想）．「トリプレット損失」は，3サンプルペア間の関係なのでまだよいが，「対照損失」は2サンプル間の距離のみで学習するゆえ，初期値次第ではうまく収束しづらい，学習が難しいデータセットも出てくる．よって，対照損失に，交差エントロピー損失も(正則化項として)追加すると，収束を安定化させて，なおかつ過学習を防いだ埋め込みを作成しやすくなる．
ただし，クラスタリング的で中央値(centroid)を形成しやすい DMLとは異なり，「softmax+交差エントロピー損失」は，識別境界の学習しかできない点に注意である．あくまで，学習の序盤にパラメータ検索空間が不安定になりがちになった際に「学習安定化」を助ける目的で，「softmax + 交差エントロピー損失」がDMLで追加使用できると捉えておきたい．
3.3 自己教師あり表現学習での使用
最近のコンピュータビジョンや周辺分野では，DML・対照学習を用いての，自己教師あり視覚表現学習(self-supervised visual representation learning)の研究も盛んである．
その自己教師あり学習による画像表現学習の際にも，教師なしデータセットからランダムに作成した「Postive/Negagtiveサンプルのペア」を用いた教師なし対照学習が行われることが多い．よって，このケースを対照表現学習(contrastive representation learning)と呼ぶことも多い．その際に，損失関数を確率分布化（あるいはボルツマン分布化）する必要があるゆえ，（温度パラメータ付きの）softmax を使用することが多くなる．
例としては，InfoNCE（Noise Contrasitve Estimation の相互情報量最大化版）を用いる各手法があげられる．代表例として，まずInfoNCEが提案された CPC (Contrastive Predictive Coding) [Oord et al., 2018]や，画像認識向け事前学習のMomentum Contrast [He et al., 2019], [He et al., 2020] (メモリーバンク系手法)があげられる．一方，他の対照表現学習手法の SimCLR [Chen et al., 2020] でも，(対照学習なので)softmax を用いる．
4. まとめ
softmax関数が，多クラス識別モデルの出力としても用いられる「ソフトなMax関数」であることを紹介した(1節)．ディープニューラルネットワークの出力層がクラス識別である際，softmax層がファーストチョイスとして頻繁に用いられている．
ベイズ推定モデルの目線から考えると，モデルの出力をsoftmaxで処理した場合は，マルチヌーイ分布(カテゴリカル分布)として，出力分布をモデル化していることを紹介した(2節)．ただし，softmax関数は「決定的出力」であり，ランダムサンプリングを行う確率モデル化まではしていないことに注意である．
更に，ディープラーニングでは，古典的な教師ありクラス識別問題以外でも，softmaxの出番が多くなっていることを紹介した(3節)，それは，アテンション機構や，対照学習・自己教師あり表現学習などである．
関連書籍
- Pythonで学ぶ画像認識 (機械学習実践シリーズ), 田村 雅人, 中村 克行, インプレス, 2023.
- 深層学習 改訂第2版 (機械学習プロフェッショナルシリーズ) 岡谷貴之，講談社，2022.
- 2.4.3 多クラス分類 (p21)
References
- [Chen et al., 2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020.
- [Horiguchi et al.,2019] Shota Horiguchi, Daiki Ikami, and Kiyoharu Aizawa. Significance of Softmax-based Features in Comparison to Distance Metric Learning-based Features. IEEE transactions on pat- tern analysis and machine intelligence, 2019.
- [He et al., 2019] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual repre- sentation learning. arXiv preprint arXiv:1911.05722, 2019. 2, 3, 5, 6
- [He et al., 2020] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual rep- resentation learning. In, CVPR 2020.