Affineレイヤの逆伝播: 2通りの勾配計算手法

# はじめに
書籍「[ゼロから作る Deep Learning](https://www.oreilly.co.jp/books/9784873117584/)」で紹介されているAffineレイヤ$\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}+\boldsymbol{b}$の逆伝播について、省略されていた以下の計算過程を2通りの方法で導出します。1つ目の方法は連鎖律を用いるもので、2つ目はフロベニウス内積と微分を利用します。
```math
\displaylines{
\begin{align}
  &\frac{\partial L}{\partial \boldsymbol{x}} = \frac{\partial L}{\partial \boldsymbol{y}}\boldsymbol{W}^T \\
  &\frac{\partial L}{\partial \boldsymbol{W}} = \boldsymbol{x}^T \frac{\partial L}{\partial \boldsymbol{y}}
\end{align}
}
```


# 表記規則
ベクトルや行列で微分を行う場合、表記法について注意が必要です。例えば、$n$次元行ベクトル$\boldsymbol{y}$を$m$次元行ベクトル$\boldsymbol{x}$で微分する場合、$\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}}$の行列は$n\times m$または$m\times n$の2通りの表記記法が存在します。参考資料「[Matrix calculus](https://en.wikipedia.org/wiki/Matrix_calculus)」によれば、微分の表記方法には分子レイアウト記法(numerator-layout notation)、分母レイアウト記法(denominator-layout notation)、およびこれらの混在した表記法の3つがあります。どの記法を用いるかは筆者によって異なり、統一されていません。そこで、本稿では、Affineレイヤの入力($\boldsymbol{x}$)と重み($\boldsymbol{W}$)の形状をそれぞれの勾配の形状と一致させる、以下の混合型レイアウトを採用します。この表記法を用いることで、SGDなどのパラメータ更新プログラムにおいて、パラメータと勾配の変数の形状を気にする必要がなくなります。
```math
\displaylines{
\begin{align}
  &\boldsymbol{x} = \bigl[x_1\cdots x_m\bigr],\hspace{0.3cm}
  \boldsymbol{y} = \bigl[y_1\cdots y_n\bigr],\hspace{0.3cm}
  \boldsymbol{W} =
    \begin{pmatrix}
      w_{11} & \cdots & w_{1n} \\
      \vdots & \ddots & \vdots \\
      w_{m1} & \cdots & w_{mn} \\
    \end{pmatrix},\\
  &\frac{\partial L}{\partial \boldsymbol{x}} = \Biggl[\frac{\partial L}{\partial x_1}\cdots \frac{\partial L}{\partial x_m}\Biggr],\hspace{0.3cm}
  \frac{\partial L}{\partial \boldsymbol{W}} =
    \begin{pmatrix}
      \frac{\partial L}{\partial w_{11}} & \cdots & \frac{\partial L}{\partial w_{1n}} \\
      \vdots & \ddots & \vdots \\
      \frac{\partial L}{\partial w_{m1}} & \cdots & \frac{\partial L}{\partial w_{mn}} \\
    \end{pmatrix},\\
  &\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}} =
    \begin{pmatrix}
      \frac{\partial y_1}{\partial x_{1}} & \cdots & \frac{\partial y_1}{\partial x_{m}} \\
      \vdots & \ddots & \vdots \\
      \frac{\partial y_n}{\partial x_{1}} & \cdots & \frac{\partial y_n}{\partial x_{m}} \\
    \end{pmatrix},\hspace{0.3cm}
  d\boldsymbol{W} =
    \begin{pmatrix}
      dw_{11} & \cdots & dw_{1n} \\
      \vdots & \ddots & \vdots \\
      dw_{m1} & \cdots & dw_{mn} \\
    \end{pmatrix}
\end{align}
}
```


# 連鎖律
$\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}+\boldsymbol{b}$、$\boldsymbol{x}\in \mathbb{R}^m$、$\boldsymbol{y},\boldsymbol{b}\in \mathbb{R}^n$、$\boldsymbol{W}\in \mathbb{R}^{m\times n}$に対して、連鎖律を用いて勾配を計算します。
```math
\displaylines{
\begin{align}
  \Biggl(\frac{\partial L}{\partial \boldsymbol{x}}\Biggr)_j
    &=\Biggl(\frac{\partial L}{\partial \boldsymbol{y}} \frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}}\Biggr)_j
    =\sum_{k=1}^n \frac{\partial L}{\partial y_k} \frac{\partial y_k}{\partial x_j}
    =\sum_{k=1}^n \frac{\partial L}{\partial y_k} \frac{\partial \Bigl(\sum_{l=1}^m x_l w_{lk}+b_k\Bigr) }{\partial x_j} \\
    &=\sum_{k=1}^n \frac{\partial L}{\partial y_k} w_{jk}
    =\sum_{k=1}^n \frac{\partial L}{\partial y_k} (w^T)_{kj}
    =\Biggl(\frac{\partial L}{\partial \boldsymbol{y}} \boldsymbol{W}^T\Biggr)_j \\
  \Biggl(\frac{\partial L}{\partial \boldsymbol{W}}\Biggr)_{ij}
    &=\Biggl(\frac{\partial L}{\partial \boldsymbol{y}} \frac{\partial \boldsymbol{y}}{\partial \boldsymbol{W}}\Biggr)_{ij}
    =\sum_{k=1}^n \frac{\partial L}{\partial y_k} \frac{\partial y_k}{\partial w_{ij}}
    =\sum_{k=1}^n \frac{\partial L}{\partial y_k} \frac{\partial \Bigl(\sum_{l=1}^m x_l w_{lk}+b_k\Bigr) }{\partial w_{ij}} \\
    &=\frac{\partial L}{\partial y_j} x_i
    =(x^T)_i\frac{\partial L}{\partial y_j}
    =\Biggl(\boldsymbol{X}^T \frac{\partial L}{\partial \boldsymbol{y}}\Biggr)_{ij} \\
\end{align}
}
```


# フロベニウス内積と微分
$\boldsymbol{Y}=\boldsymbol{X}\boldsymbol{W}+\boldsymbol{B}$、$\boldsymbol{X}\in \mathbb{R}^{1\times m}$、$\boldsymbol{Y},\boldsymbol{B}\in \mathbb{R}^{1\times n}$、$\boldsymbol{W}\in \mathbb{R}^{m\times n}$に対して、フロベニウス内積と微分を用いて勾配を計算します。ここで、$\boldsymbol{X}$、$\boldsymbol{Y}$、$\boldsymbol{B}$は、それぞれ$\boldsymbol{x}$、$\boldsymbol{y}$、$\boldsymbol{b}$のベクトルを行列に拡張しています。フロベニウス内積は次のように定義されます。
```math
\boldsymbol{A} : \boldsymbol{B} = \sum_{i=1}^m \sum_{j=1}^n a_{ij} b_{ij} = \text{Tr}(\boldsymbol{A}^T \boldsymbol{B})
```
$\text{Tr}(\boldsymbol{P+Q})=\text{Tr}(\boldsymbol{P})+\text{Tr}(\boldsymbol{Q})$と、$\text{Tr}(\boldsymbol{PQ})=\text{Tr}(\boldsymbol{QP})(\boldsymbol{P}\in\mathbb{R}^{p \times q}$、$\boldsymbol{Q}\in\mathbb{R}^{q \times p})$より
```math
\displaylines{
\begin{align}
  &\boldsymbol{A} : (\boldsymbol{B}+\boldsymbol{C}) = (\boldsymbol{A} : \boldsymbol{B}) + (\boldsymbol{A} : \boldsymbol{C}) \\
  &\boldsymbol{A} : \boldsymbol{BC} = \boldsymbol{B}^T\boldsymbol{A} : \boldsymbol{C} = \boldsymbol{A}\boldsymbol{C}^T : \boldsymbol{B}
\end{align}
}
```
を満たします。したがって、
```math
\displaylines{
\begin{align}
  dL &= \sum_{i=1}^1 \sum_{j=1}^n \Biggl(\frac{\partial L}{\partial \boldsymbol{Y}}\Biggr)_{ij} \Bigl(d\boldsymbol{Y}\Bigr)_{ij}\hspace{0.7cm} \text{(全微分)} \\
     &= \frac{\partial L}{\partial \boldsymbol{Y}} : d\boldsymbol{Y} \\
     &= \frac{\partial L}{\partial \boldsymbol{Y}} : \Bigl(d\boldsymbol{X}\boldsymbol{W} + \boldsymbol{X}d\boldsymbol{W} + d\boldsymbol{B}\Bigr)\hspace{0.7cm} \text{(行列微分)}  \\
     &= \biggl(\frac{\partial L}{\partial \boldsymbol{Y}}\boldsymbol{W}^T : d\boldsymbol{X}\biggr) + \biggl(\boldsymbol{X}^T \frac{\partial L}{\partial \boldsymbol{Y}} : d\boldsymbol{W}\biggr) + \biggl(\frac{\partial L}{\partial \boldsymbol{Y}} : d\boldsymbol{B}\biggr)
\end{align}
}
```
を得ます。ここで、$\boldsymbol{W}$と$\boldsymbol{B}$が定数の場合、$d\boldsymbol{W}=d\boldsymbol{B}=0$なので、
```math
\displaylines{
\begin{align}
  dL &= \frac{\partial L}{\partial \boldsymbol{Y}}\boldsymbol{W}^T : d\boldsymbol{X}
      = \sum_{i=1}^1 \sum_{j=1}^m \Biggl(\frac{\partial L}{\partial \boldsymbol{Y}}\boldsymbol{W}^T\Biggr)_{ij} \Bigl(d\boldsymbol{X}\Bigr)_{ij} \\
     &= \sum_{i=1}^1 \sum_{j=1}^m \Biggl(\frac{\partial L}{\partial \boldsymbol{X}}\Biggr)_{ij} \Bigl(d\boldsymbol{X}\Bigr)_{ij}\hspace{0.7cm} \text{(全微分)}\\
  &\frac{\partial L}{\partial \boldsymbol{X}} = \frac{\partial L}{\partial \boldsymbol{Y}}\boldsymbol{W}^T \\
\end{align}
}
```
となります。同様に、$\boldsymbol{X}$と$\boldsymbol{B}$が定数の場合、$d\boldsymbol{X}=d\boldsymbol{B}=0$なので、
```math
\displaylines{
\begin{align}
  &dL = \boldsymbol{X}^T \frac{\partial L}{\partial \boldsymbol{Y}} : d\boldsymbol{W} \\
  &\frac{\partial L}{\partial \boldsymbol{W}} = \boldsymbol{X}^T \frac{\partial L}{\partial \boldsymbol{Y}}
\end{align}
}
```
となります。以上が勾配の計算過程です。
説明では、$\boldsymbol{X}$、$\boldsymbol{Y}$、$\boldsymbol{B}$を1行として扱いましたが、$n$行に拡張しても、形状が異なるだけで勾配の表記は変わりません。さらに、$dL = (* : d\boldsymbol{X}) + (* : d\boldsymbol{W}) + (* : d\boldsymbol{B})$が得られた時点で、それぞれの勾配が同時に計算されています。通常、計算過程では、全微分と記載した式は省略されます。その部分を除くと、フロベニウス内積、行列微分、並び替えの3ステップで勾配を得ることができます。また、トレースの性質により、フロベニウス内積は、$\boldsymbol{A}:\boldsymbol{B}=\boldsymbol{B}:\boldsymbol{A}=\boldsymbol{A}^T:\boldsymbol{B}^T$の関係も満たします。これらの特性などを理解することで、微分を求める他の問題にも適用できる可能性があります。


# おわりに
Affineレイヤの逆伝播における勾配を2つの方法で導出しました。連鎖律を用いる場合、ベクトルを行列で微分($\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{W}}$)するため、3次元テンソルを含む複雑な計算式が出現します。一方、フロベニウス内積と微分を用いると、成分添字が不要で計算式が行列内で完結するため、より簡潔で扱いやすい形になります。