maskとsoftmaxの整理

概要

softmax計算時に特定のロジットを無視した上でsoftmaxをとりたい場合がある。（例：decoder型transformerのcausal attention計算時の未来のトークンをマスクしたいとき、強化学習のpolicyの無効行動をマスクしたいときなど）

その際にいくつかの手順を見かけると思っていたが, 以下はすべて等価であることを確認した.

1. 元のベクトルについてsoftmaxをとってから、無視したい要素の値を$0$にしてから全体の和で各要素を割って正規化
2. 元のベクトルについて無視したい要素の値を$-\infty$にしてから全体についてsoftmax
3. 考慮したい要素の値のみについてsoftmax

以下簡単に証明する

まず手順1について考える.
ベクトルを$\boldsymbol{a}=\{a_1,\cdots,a_n\}$とおく.
$\boldsymbol{a}$の全体についてsoftmaxをとると$i$番目の要素の値$\tilde{a}_{i}$は

\tilde{a}_i = \frac{\exp(a_i)}{\sum_{k=1}^{n} {\exp(a_k)}}

となる.
ここで$\boldsymbol{a}$の$m~(<n)$番目までの要素のみ考慮し, $m+1$番目以降の要素を無視したいとする.
すなわち$\tilde{a}_i = 0~~(i=m+1, \cdots, n)$とおく.
このとき$\boldsymbol{\tilde{a}}= \{ \tilde{a}_j, \cdots, \tilde{a}_n \}$の和は$m$番目の要素までの和であるから, $i$番目の要素その和で割ると

\begin{align}
\frac{\tilde{a}_i}{\sum_{k=1}^m \tilde{a_k}} &= \frac{\exp(a_i)}{\sum_{k=1}^{n} {\exp(a_k)}} \times \frac{1}{\Sigma_{k=1}^m \frac{\exp(a_k)}{\Sigma_{j=1}^n \exp(a_j)}}\\
&= \frac{\exp(a_i)}{\sum_{k=1}^{n} {\exp(a_k)}} \times \frac{\sum_{j=1}^{n} {\exp(a_j)}}{\sum_{k=1}^{m} {\exp(a_k)}} \\
&= \frac{\exp(a_i)}{\sum_{k=1}^{m} {\exp(a_i)}}
\end{align}

この結果は元のベクトル$\boldsymbol{a}$の$m$番目までの要素のみについてsoftmaxをとったものと同じであり, 手順3と等価であることが示された。
また手順2は$\exp(-\infty)=0$であるため手順3と等価になるのは自明である。

感想

手順2と3が等価なのは自明だが, 個人的には手順1と手順2,3が等価なのは若干の直観とのギャップがある.
手順1では正規化に利用する和の計算過程で, 無視する要素の値が間接的に影響していそうな感覚を抱いてしまった.
そのため無視する要素の値を一切計算過程で利用しない, 手順3との等価性に疑問を持ったが実際にはそうでないことが示された.