本当は奥が深いSoftmaxとそのお仲間達

17日目書かせていただきます．よろしくお願いします．

機械学習やってます！という方は，きっとsoftmaxに日頃お世話になっているだろう．しかし多くの方は，「softmaxを使えば，ニューラルネットとかで分類問題解けたりするんでしょ？」程度の理解ではないだろうか？私もそうだ．使えれば何でも良い．しかしこの記事では，softmaxについて少々深掘りをしてみたいと思う．**本記事の流れとしては，argmaxを導入し，softmaxをラグランジュの未定乗数法で導出し，出力がスパースになるsparsemax,それらの一般化であるentmaxを紹介したいと思う**． おまけとして，温度付きsoftmaxとGumbel-softmaxも入れたので，適当に興味あるところだけでも読んでくれると嬉しい．

この記事では，なんちゃらmaxに入れる前の値をロジットと呼び，$\mathbf{z}$で表す．また，なんちゃらmax後の出力を$\mathbf{p}$(各値は非負で足して1)とする．また，表記等がテキトーなところもあるが，全く気しないスタイルでいく．

- [argmax](#argmax) 
- [softmax](#softmax) 
- [sparsemax](#sparsemax) 
- [entmax](#entmax) 
- [おまけ](#おまけ)
- [まとめ](#まとめ)

# TL;DR
いろんな(.+)-maxがあるけど，繋がってるんだよね
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/d8db429b-396a-ca00-dbc9-c088cc942c1e.png)
[参考資料](https://docs.google.com/presentation/d/13PZfY6uLHdHVr9c9nDjHLDCkhdkX-JR3gKWGSokPwr8/edit#slide=id.p)

# 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条件](http://marupeke296.com/MT_No2_lagurange2.html)から$\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関数と同等である．見慣れたグラフになってますね．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/2ff08668-bec8-2c07-b782-4a44a51f3b35.png)
論文[[1](#論文)]より

見るとわかるように，softmaxではtが大きくなっても小さくなっても，完全に1(または0)にはならないがわかる．softmaxはソフトで優しいやつで，どんなにtが小さくても見捨てて0にはしない奴なのだ．

対してsparsemaxはこの例の場合，tが1以上なら出力が[1,0]になるし，tが-1以下なら出力が[0,1]となるわけだ．sparsemaxは値が小さいとすぐ見捨てる奴なわけだが，スパースな出力をするモデルが欲しい場合，sparsemaxは非常に役に立つ．

sparsemaxをどう計算するか知りたい人は原著論文または[この記事](https://data-analytics.fun/2021/08/17/understanding-sparsemax/)を読んでほしい．ちなみにPyTorchなら[これ](https://github.com/KrisKorrel/sparsemax-pytorch)が良いと思っている（私はこれを使っている）．

# entmax
softmaxとsparsemaxの違いは，実はエントロピー項のみということがわかった．また，softmaxはソフトな優しいやつ，sparsemaxはすぐ見捨てるやつ，と表現した．だとすると，以下のような疑問が浮かぶはずだ．

**softmaxとsparsemaxの補間を考えられないのか？**
**優しさはあるけど，捨てる時は捨てる丁度いいやつはいないのか？**

結論から言うと，答えはYesだ．エントロピー項を一般化すれば良いわけで，実は1988年にツァリスという人がツァリスのエントロピーというのを発表している([ref](https://link.springer.com/article/10.1007/BF01016429))．これを利用して発表されたのが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})$は以下である（うまく数式が反映されないのでスクショ載せてます）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/bae391e7-df85-492b-e495-df3681650ee5.png)

$\mathrm{H}_{\alpha}^{\mathrm{T}}(\mathbf{p})$はツァリスのエントロピーと呼ばれ，式を見てみると

- $\alpha=1$の時にsoftmaxの式（シャノンのエントロピー）
- $\alpha=2$の時にsparsemaxの式（ジニエントロピー）

が出てくるのがわかる．従って，一般化になっており，パラメータ$\alpha$を調整することによって間の行き来が可能だ．先ほど同様，二次元の入力が入ってきた時の各パラメータの出力の違いを示す．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/5a359885-ffaf-6d37-93e7-4d8ea048edbb.png)
(論文[[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に近づくのがわかる．
![tempup.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/6b5f8ff5-fa32-f55f-b2a5-9130265c33b8.gif)

逆に上げていったシミュレーションが以下だ．上がるほど等確率に近づくのがわかる．知識の蒸留においては，温度を上げることによって，教師モデルの出力の最大値以外も考慮することができるわけだ．
![tempdown.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/7487104a-c07d-93d6-0be3-4ccc0b9763bb.gif)

## おまけ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分布](https://ja.wikipedia.org/wiki/%E3%82%AC%E3%83%B3%E3%83%99%E3%83%AB%E5%88%86%E5%B8%83)からのサンプルだ．これによって確率モデルになっているわけだが，なぜ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$と，サンプルされた回数から算出した確率がほぼ一致する．詳しくは[このわかりやすい記事](https://data-analytics.fun/2021/04/06/understanding-gumbel-max-trick/)や，以下のgistを見てほしい．上式を微分可能とするために，softmaxへと変えたのがGumbel-Softmaxと言える．

https://gist.github.com/matsutakk/6fa145c6e3274dfb2ce9dadb08aa6ed8

温度付きSoftmaxとは違い確率モデルなので，必ずしも値の大きい場所が常にone-hotになるわけではない．$\mathbf{z}=[3, 2, 1]$として，$T=0.01$固定でGumbel-Softmaxを100回実行したアニメーションは以下のようになる．

![gumbel.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/223338/72c825d4-e50e-2d1f-faaf-b03c352c7bf3.gif)
温度付き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.