17日目書かせていただきます．よろしくお願いします．
機械学習やってます！という方は，きっとsoftmaxに日頃お世話になっているだろう．しかし多くの方は，「softmaxを使えば，ニューラルネットとかで分類問題解けたりするんでしょ？」程度の理解ではないだろうか？私もそうだ．使えれば何でも良い．しかしこの記事では，softmaxについて少々深掘りをしてみたいと思う．本記事の流れとしては，argmaxを導入し，softmaxをラグランジュの未定乗数法で導出し，出力がスパースになるsparsemax,それらの一般化であるentmaxを紹介したいと思う． おまけとして，温度付きsoftmaxとGumbel-softmaxも入れたので，適当に興味あるところだけでも読んでくれると嬉しい．
この記事では，なんちゃらmaxに入れる前の値をロジットと呼び，$\mathbf{z}$で表す．また，なんちゃらmax後の出力を$\mathbf{p}$(各値は非負で足して1)とする．また，表記等がテキトーなところもあるが，全く気しないスタイルでいく．
TL;DR
いろんな(.+)-maxがあるけど，繋がってるんだよね
参考資料
argmax
まずはargmaxを導入しよう．argmaxは以下の式で表現される．
$$\operatorname{argmax}(\mathbf{z}):=\underset{\mathbf{p} \in \Delta^{d}}{\max} \mathbf{p}^{\top} \mathbf{z}$$
つまりargmax関数は，与えられたベクトル$\mathbf{z}$との内積が最大となるようなベクトル$\mathbf{p}$(ただし全て値は非負で足して1，つまり単体の要素)を返すような関数である．例えば，$\mathbf{z}=[100,20,-3]$とかだったら，$\mathbf{p}=[1,0,0]$となるわけだ．一番大きいとこを取ってくるわけだね．
softmax
それで，softmaxはどうなるのかというと以下の式で表される．
$$\operatorname{softmax}(\mathbf{z}):=\underset{\mathbf{p} \in \Delta^{d}}{\max} \mathbf{p}^{\top} \mathbf{z}+\mathrm{H}^{\mathrm{S}}(\mathbf{p})$$
ここで，$\mathrm{H}^{\mathrm{S}}(\mathbf{p}):=-\sum_{j} p_{j} \log p_{j}$である．情報理論で習うシャノンのエントロピーだね．先ほどのargmaxの式と違い，$\mathrm{H}^{\mathrm{S}}(\mathbf{p})$という項が加わっている．つまり，**「softmaxというのは，argmaxの式にシャノンのエントロピー項を加えたもの」**なのだ.右辺をラグランジュの未定乗数法で解くことで，見慣れたsoftmaxの式$\frac{e^{z_{i}}}{\sum_{i} e^{z_{i}}}$が出てくる．確認してみよう．
まずラグランジュ関数は以下のようになる．
$$
L(\mathbf{p}, \boldsymbol{\mu}, \tau) = \mathbf{p}^{\top}\mathbf{z}+\mathrm{H}_{\alpha}^{\mathrm{S}}(\mathbf{p})-\boldsymbol{\mu}^{\top} \mathbf{p}+\tau\left(\mathbf{1}^{\top} \mathbf{p}-1\right)
$$
第３項は$\mathbf{p}\geq 0$という制約から，第4項は$\mathbf{p}$が足して１になる制約から出てきている.$\boldsymbol{\mu}$と$\tau$はラグランジュ乗数だね．ラグランジュ関数が作れたら，$\mathbf{p}$で微分しよう！わかりやすくするために，$p_{i}$だけを考慮してみる．$\mathrm{H}^{\mathrm{S}}(\mathbf{p}):=-\sum_{j} p_{j} \log p_{j}$だったため，
$$\frac{\partial\mathrm{H}^{\mathrm{S}}(\mathbf{p})}{\partial p_{i}}=-\log \left(p_{i}\right)-1$$
なので，
$$\frac{\partial L(\mathbf{p}, \boldsymbol{\mu}, \tau)}{\partial p_{i}}=-\log \left(p_{i}\right)-1+z_{i}-\mu_{i}+\tau$$
となる．微分=0を考えれば良いので．
$$-\log \left(p_{i}\right)-1+z_{i}-\mu_{i}+\tau=0$$
だね．これを$p_{i}$に関して整理すると，
$$p_{i}=\exp \left(z_{i}+\tau-\mu_{i}-1\right)$$
と変形できる．ここで，ラグランジュの未定乗数法（不等式制約）におけるKKT条件から$\mu_{i} p_{i}=0$が成り立つことと，$\log \left(p_{i}\right)$より$p_{i} \neq 0$なので，$\mu_{i}=0$が言える（わからんかったらとりあえず成り立つと認めると）．したがって，
$$p_{i}=\exp \left(z_{i}+\tau-1\right)$$
ベクトル$\mathbf{p}$は全部足すと1だったのを思い出すと，以下が成り立つよね．
$$\sum_{i} p_{i}=\sum_{i} e^{z_{i}+\tau-1}=1$$
整理すると，
$$e^{\tau}=\frac{e}{\sum_{i} e^{z_{i}}}$$
ここで，$p_{i}=\exp \left(z_{i}+\tau-1\right)$から，$e^{\tau}=p_{i}e^{-z_{i}}e$
が成り立つから，これを上式に代入してあげると，
$$p_{i} = \frac{e^{z_{i}}}{\sum_{i} e^{z_{i}}}$$
voila! 見慣れたsoftmaxの式が本当に出てきたね！！
sparsemax
softmaxはargmaxの式にシャノンのエントロピーが加わっている最大化問題を解けば導けることがこれでわかった．では，sparsemaxについて今度は紹介する．sparsemaxは論文[1]で提案された．sparsemaxは以下の式で表される（ただし，以下の式で表現したのは論文[2]）．
$$\operatorname{sparsemax}(\mathbf{z}):=\underset{\mathbf{p} \in \Delta^{d}}{\max} \mathbf{p}^{\top} \mathbf{z}+\mathrm{H}^{\mathrm{G}}(\mathbf{p})$$
ここで，$\mathrm{H}^{\mathrm{G}}(\mathbf{p}):=\frac{1}{2} \sum_{j} p_{j}\left(1-p_{j}\right)$であり，ジニエントロピーと呼ぶ．決定木を勉強したことがあるひとは，ジニ不純度という言葉で聞いたのではないだろうか．経済学でも出てくるようだ．知らんけど．
sparsemaxの正体は上記の式で，実はsoftmaxとの違いはエントロピー項のみということだ．では，これらの補間が出るのか？という問いに答えたのが次に紹介するentmax関数なのである．
entmaxの紹介前に，２次元の入力に対するsoftmaxとsparsemaxの違いを見てみよう．以下は入力[t,0]をそれぞれに入れた時の図で，横軸がt，縦軸がtに対する出力である．つまりこの場合のsoftmaxはsigmoid関数と同等である．見慣れたグラフになってますね．
論文[1]より
見るとわかるように，softmaxではtが大きくなっても小さくなっても，完全に1(または0)にはならないがわかる．softmaxはソフトで優しいやつで，どんなにtが小さくても見捨てて0にはしない奴なのだ．
対してsparsemaxはこの例の場合，tが1以上なら出力が[1,0]になるし，tが-1以下なら出力が[0,1]となるわけだ．sparsemaxは値が小さいとすぐ見捨てる奴なわけだが，スパースな出力をするモデルが欲しい場合，sparsemaxは非常に役に立つ．
sparsemaxをどう計算するか知りたい人は原著論文またはこの記事を読んでほしい．ちなみにPyTorchならこれが良いと思っている（私はこれを使っている）．
entmax
softmaxとsparsemaxの違いは，実はエントロピー項のみということがわかった．また，softmaxはソフトな優しいやつ，sparsemaxはすぐ見捨てるやつ，と表現した．だとすると，以下のような疑問が浮かぶはずだ．
softmaxとsparsemaxの補間を考えられないのか？
優しさはあるけど，捨てる時は捨てる丁度いいやつはいないのか？
結論から言うと，答えはYesだ．エントロピー項を一般化すれば良いわけで，実は1988年にツァリスという人がツァリスのエントロピーというのを発表している(ref)．これを利用して発表されたのがentmaxである．entmaxは以下の式で表され，パラメータ$\alpha$を持つ．
$$
\text{entmax}(\mathbf{z}):=\underset{\mathbf{p} \in \Delta^{d}}{\max} \mathbf{p}^{\top} \mathbf{z}+\mathrm{H}_{\alpha}^{\mathrm{T}}(\mathbf{p})
$$
ここで，$\mathrm{H}_{\alpha}^{\mathrm{T}}(\mathbf{p})$は以下である（うまく数式が反映されないのでスクショ載せてます）
$\mathrm{H}_{\alpha}^{\mathrm{T}}(\mathbf{p})$はツァリスのエントロピーと呼ばれ，式を見てみると
- $\alpha=1$の時にsoftmaxの式（シャノンのエントロピー）
- $\alpha=2$の時にsparsemaxの式（ジニエントロピー）
が出てくるのがわかる．従って，一般化になっており，パラメータ$\alpha$を調整することによって間の行き来が可能だ．先ほど同様，二次元の入力が入ってきた時の各パラメータの出力の違いを示す．
(論文[2]より)
$\alpha$がデカくなっていけばいくほど，ステップ関数に近づくのがわかるだろう．面白いね．
おまけ
これら以外にも(.+)-maxは存在する．以下では２つ紹介する．温度付きSoftmaxとGumbel-Softmax(Concrete Distribution)だ．
おまけ１ 温度付きSoftmax
Softmax関数に温度パラメータ$T$を導入したもので，以下の式で表される．これは知識の蒸留とかで使われる．詳しくはHinton氏の論文[3]を読んでみると良いだろう．
$$\text{温度付きSoftmax} := \frac{e^{\frac{z_{i}}{T}}}{\sum_{i} e^{\frac{z_{i}}{T}}}$$
$\mathbf{z}=[3, 2, 1]$として，温度を下げていったアニメーションが以下だ．温度を下げていけばいくほど，one-hotに近づくのがわかる．
逆に上げていったシミュレーションが以下だ．上がるほど等確率に近づくのがわかる．知識の蒸留においては，温度を上げることによって，教師モデルの出力の最大値以外も考慮することができるわけだ．
おまけ2 Gumbel-Softmax
これは文献[4,5]で同時に提案された．片方はGoogleから，片方はDeepMindからだ．以下ではGumbel-Softmaxと呼ぶが，[4]ではConcrete Distributionと呼んでいる．
Gumbel-Softmaxは，離散潜在変数に対する微分可能な確率モデリングを可能にする．式は以下で表される．
$$p_{i}=\frac{\exp \left(\left(\log \left(z_{i}\right)+g_{i}\right) / T\right)}{\sum_{j=1}^{k} \exp \left(\left(\log \left(z_{j}\right)+g_{j}\right) / T\right)}$$
ほぼ温度付きSoftmaxと同じように見えるが，$g_{i}$というのが加わっている．これはGumbel分布からのサンプルだ．これによって確率モデルになっているわけだが，なぜGumbel分布からのサンプルが入っているのだろうか？より詳細に理解するには，Gumbel Max Trickというものの理解が必要となる．Gumbel Max Trickを式で表すと以下のようになる．
$$z=\underset{i}{\arg \max }\left[g_{i}+\log \pi_{i}\right]$$
これは簡単に言ってしまえば，パラメータ$\pi$に従うカテゴリカル分布から効率よくサンプルを得たければ．$\log{\pi_{i}}に$ガンベル分布からのサンプル値$g_{i}$を足して一番大きいところを選択カテゴリとすれば良いということだ．実際にシミュレーションしてみると，真のパラメータ$\pi$と，サンプルされた回数から算出した確率がほぼ一致する．詳しくはこのわかりやすい記事や，以下のgistを見てほしい．上式を微分可能とするために，softmaxへと変えたのがGumbel-Softmaxと言える．
温度付きSoftmaxとは違い確率モデルなので，必ずしも値の大きい場所が常にone-hotになるわけではない．$\mathbf{z}=[3, 2, 1]$として，$T=0.01$固定でGumbel-Softmaxを100回実行したアニメーションは以下のようになる．
温度付きSoftmaxの場合は１番左が大きな値を必ず取るが，そうなっていないのがわかる．カテゴリカル分布からのサンプリングになってるように見えるが，重要なポイントは，Gumbel-Softmaxは微分可能という点だ．したがって勾配法でEnd-to-End学習が可能となる．
まとめ
- softmaxというのは
- argmax + シャノンのエントロピー
- sparsemaxというのは
- argmax + ジニエントロピー
- entmaxというのは
- argmax + ツァリスのエントロピー
- (.+)-maxで他に有名なのは
- 温度付きSoftmax
- Gumbel-Softmax(Concrete Distribution)
Transformerで使われるAttentionだって，中身はsoftmax．当たり前のように使っていたがなかなか奥が深かった．メリークリスマス！
参考文献
Softmaxの面白さがわかる資料
- https://docs.google.com/presentation/d/13PZfY6uLHdHVr9c9nDjHLDCkhdkX-JR3gKWGSokPwr8/edit#slide=id.p
- http://www.wdong.org/equivalence-between-entropy-regularization-and-the-softmax-function.html
- https://www.cs.utah.edu/~tli/posts/2019/01/blog-post-1/
論文
[1] Martins, Andre, and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. International conference on machine learning. PMLR, 2016.
[2] Peters, B., Niculae, V., & Martins, A. (2019). Sparse Sequence-to-Sequence Models. In Proc. ACL.
[3] Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. "Distilling the knowledge in a neural network." arXiv preprint arXiv:1503.02531 (2015).
[4] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel- softmax. In International Conference on Learning Representations (ICLR), 2017.
[5] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A contin- uous relaxation of discrete random variables. In 5th International Conference on Learning Representations (ICLR), 2017.