ロジットを使ってエントロピーを計算する

方策勾配で強化学習を行う際に、方策が決定論的になっていないか監視するために、ログにエントロピーを出力するようにしたい。

エントロピーは、
$\displaystyle H({\bf p})=\sum_i p_i \log p_i$
で計算できるが、確率がほぼ0の場合、 $\log(p_i)$ が-infになるため、この式のままでは計算できない。

$p_i$ の最小値をクリップして、

F.log(F.clip(p, 1e-32, 1.0))

のようにすると、-infになることを防ぐことができるが、ニューラルネットワークの出力をロジットにしている場合、よりスマートに実装できる。

ロジットを使ったエントロピーの計算

ロジットを $y_i$ とすると、ソフトマックスの定義から、
$\displaystyle \begin{align} \log p_i &= \log \frac{e^{y_i}}{\sum_j e^{y_j}} \\ &= y_i - \log\sum_j e^{y_j} \end{align}$
となる。
ここで、 $\sum_j e^{y_j}$ がオーバーフローする可能性があるため、指数部から ${\bf y}$ の最大値を引いて
$\displaystyle \begin{align} \log p_i &= y_i - \log(\sum_j e^{y_j - y_{max}} \cdot e^{y_{max}}) \\ &= y_i - (\log\sum_j e^{y_j - y_{max}} + y_{max}) \end{align}$
のように変形する。

これを使うと、Chainerでは、エントロピーの計算を以下のように実装できる。

    p = F.softmax(y)
    #entropy = F.sum(- p * F.log(p), axis=1)
    y_max = F.max(y, axis=1, keepdims=True)
    log_p = y - (F.log(F.sum(F.exp(y - y_max), axis=1, keepdims=True)) + y_max)
    entropy = F.sum(- p * log_p, axis=1)

※コメントしている行のまま計算すると結果が、nanになる場合がある。

シグモイドの場合

結果がシグモイドの場合は、
$\displaystyle p = \frac{1}{1+e^{-y}}$
となるので、 $\log p$ と $\log(1-p)$ は、それぞれ、
$\displaystyle \begin{align} \log p &= \log \frac{1}{1+e^{-y}} \\ &= y - \log(e^y+1) \end{align}$
$\displaystyle \begin{align} \log (1 - p) &= \log(1- \frac{1}{1+e^{-y}}) \\ &= \log\frac{1}{e^y + 1} \\ &= -\log(e^y + 1) \end{align}$
となる。
これを使って、Chainerでは、エントロピーの計算を以下のように実装できる。

    p = F.sigmoid(y)
    #entropy = -(p * F.log(p) + (1 - p) * F.log(1 - p))
    log1p_ey = F.log1p(F.exp(y))
    entropy = -(p * (y - log1p_ey) + (1 - p) * -log1p_ey)

※コメントしている行のまま計算すると結果が、nanになる場合がある。