tf.kerasにおけるRNNのパラメータ数の計算方法

# 1. はじめに
kerasではモデルを構築したあと`model.summary()`と入力するだけで、モデルの概要が確認できます。その際、右列の`Param #`に各層のパラメータ数が表示されますが、毎回「あれ、何でこんな値になるんだっけ？」と頭の中が初期化されていました。
そこで、備忘録もかねて、RNNにおける各層のパラメータ数の計算方法をまとめます。

# 2. 環境
- Python 3.8.7
- TensorFlow 2.5.0

# 3. パラメータ数の計算方法
## 3-1. SimpleRNN
SimpleRNN層の出力は以下の式で表されます。

$y_{t} = \tanh (W_{x}\ x_{t}\ +\ W_{y}\ y_{t-1}\ +\ b)$

- $y_{t}$はタイムステップ$t$における出力$y$
- $x_{t}$はタイムステップ$t$における入力$x$

下で説明するLSTMやGRUと同様に$h_{t}$を用いて説明されることもありますが、その場合は結局$y_{t} = h_{t}$なので、ここでは書いていません。

| 重み | 説明 | 形状 |
| :-: | :-: | :-: |
| $W_{x}$ | 入力$x_{t}$に対応する重み | $n_{inputs}$ $\times$ $n_{outputs}$ |
| $W_{y}$ | 出力$y_{t-1}$に対応する重み | $n_{outputs}$ $\times$ $n_{outputs}$ |
| $b$ | バイアス | $n_{outputs}$ |

よってSimpleRNN層のパラメータ数は以下のように算出されます。
$
n_{inputs}\ \times\ n_{outputs}\ +\ n_{outputs}\ \times\ n_{outputs}\ + \ n_{outputs}
$

この結果を、実際にtf.keras.layers.SimpleRNNの出力で確認してみます。

```python
from tensorflow import keras

model = keras.Sequential([
    keras.layers.SimpleRNN(units=20, input_shape=[None, 10], return_sequences=True),
    keras.layers.SimpleRNN(units=5, return_sequences=True),
    keras.layers.SimpleRNN(units=2)
])

model.summary()
```

```
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
simple_rnn (SimpleRNN)       (None, None, 20)          620       
_________________________________________________________________
simple_rnn_1 (SimpleRNN)     (None, None, 5)           130       
_________________________________________________________________
simple_rnn_2 (SimpleRNN)     (None, 2)                 16        
=================================================================
Total params: 766
Trainable params: 766
Non-trainable params: 0
```

| 層の名前 | $n_{inputs}$ | $n_{outputs}$ | 計算結果 |
| :-: | :-: | :-: | :-: |
| simple_rnn | input_shape=10 | units=20 | 10$\times$20+20$\times$20+20=620 |
| simple_rnn_1 | 一つ前の層の出力=20 | units=5 | 20$\times$5+5$\times$5+5=130 |
| simple_rnn_2 | 一つ前の層の出力=5 | units=2 | 5$\times$2+2$\times$2+2=16 |

## 3-2. LSTM
LSTM層の出力は以下の式で表されます。

$
g_{t} = \tanh (W_{xg}\ x_{t}\ +\ W_{hg}\ h_{t-1}\ +\ b_{g}) \\\
f_{t} = \sigma (W_{xf}\ x_{t}\ +\ W_{hf}\ h_{t-1}\ +\ b_{f}) \\\
i_{t} = \sigma (W_{xi}\ x_{t}\ +\ W_{hi}\ h_{t-1}\ +\ b_{i}) \\\
o_{t} = \sigma (W_{xo}\ x_{t}\ +\ W_{ho}\ h_{t-1}\ +\ b_{o}) \\\
c_{t} = f_{t}\ \otimes\ c_{t-1}\ +\ i_{t}\ \otimes\ g_{t} \\\
y_{t} = h_{t} = o_{t}\ \otimes\ \tanh (c_{t})
$

- LSTMは、$h_{t}$で短期間の情報（短期記憶）を、$c_{t}$で長期間の情報（長期記憶）を伝えるイメージです。
- $g_{t}$はSimpleRNNの式と同じです。
- LSTMでは$g_{t}$に加えて、$f_{t}$（forget gate）、$i_{t}$（input gate）、$o_{t}$（output gate）を計算します。$g_{t}$の活性化関数をシグモイド関数に置き換えたものです。
- $f_{t} \otimes c_{t-1}$で長期記憶をどの程度忘れるか、$i_{t} \otimes g_{t}$で現在の情報をどの程度残すか、それらを足すことで長期記憶$c_{t}$を更新します。（$\otimes$は要素ごとの掛け算を表します。）
- $o_{t} \otimes \tanh (c_{t})$で長期記憶をどのくらい短期記憶$h_{t}$として出力するかを決めます。
- 短期記憶$h_{t}$と出力$y_{t}$は同じです。

ちなみに[Andrew Ng先生のCourseraの授業](https://www.coursera.org/learn/nlp-sequence-models)では、$i_{t}$（input gate）は$\Gamma_{u}$（update gate）という名前で説明されていました。（$\Gamma$はアルファベットではGにあたる文字で、gateのgだと思います。）

| 重み | 説明 | 形状 |
| :-: | :-: | :-: |
| $W_{xg}$ $W_{xf}$ $W_{xi}$ $W_{xo}$ | 入力$x_{t}$に対応する重み | $n_{inputs}$ $\times$ $n_{outputs}$ |
| $W_{hg}$ $W_{hf}$ $W_{hi}$ $W_{ho}$ | 一つ前の短期記憶$h_{t-1}$に対応する重み | $n_{outputs}$ $\times$ $n_{outputs}$ |
| $b_{g}$ $b_{f}$ $b_{i}$ $b_{o}$ | バイアス | $n_{outputs}$ |

よってLSTM層のパラメータ数は以下のように算出されます。
$
(n_{inputs}\ \times\ n_{outputs}\ +\ n_{outputs}\ \times\ n_{outputs}\ + \ n_{outputs})\ \times\ 4
$

この結果を、実際にtf.keras.layers.LSTMの出力で確認してみます。

```python
from tensorflow import keras

model = keras.Sequential([
    keras.layers.LSTM(units=20, input_shape=[None, 10], return_sequences=True),
    keras.layers.LSTM(units=5, return_sequences=True),
    keras.layers.LSTM(units=2)
])

model.summary()
```

```
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
lstm (LSTM)                  (None, None, 20)          2480      
_________________________________________________________________
lstm_1 (LSTM)                (None, None, 5)           520       
_________________________________________________________________
lstm_2 (LSTM)                (None, 2)                 64        
=================================================================
Total params: 3,064
Trainable params: 3,064
Non-trainable params: 0
```

| 層の名前 | $n_{inputs}$ | $n_{outputs}$ | 計算結果 |
| :-: | :-: | :-: | :-: |
| lstm | input_shape=10 | units=20 | (10$\times$20+20$\times$20+20)$\times$4=2480 |
| lstm_1 | 一つ前の層の出力=20 | units=5 | (20$\times$5+5$\times$5+5)$\times$4=520 |
| lstm_2 | 一つ前の層の出力=5 | units=2 | (5$\times$2+2$\times$2+2)$\times$4=64 |

## 3-3. GRU
GRU層の出力は以下の式で表されます。

$
g_{t} = \tanh (W_{xg}\ x_{t}\ +\ b_{xg}\ +\ W_{hg}\ (r_{t}\ \otimes\ h_{t-1})\ +\ b_{hg}) \\\
r_{t} = \sigma (W_{xr}\ x_{t}\ +\ b_{xr}\ +\ W_{hr}\ h_{t-1}\ +\ b_{hr}) \\\
z_{t} = \sigma (W_{xz}\ x_{t}\ +\ b_{xz}\ +\ W_{hz}\ h_{t-1}\ +\ b_{hz}) \\\
y_{t} = h_{t} = z_{t}\ \otimes\ h_{t-1}\ +\ (1\ -\ z_{t})\ \otimes\ g_{t}
$

- バイアスが二度足されていることは、一旦無視してください。
- LSTMを簡略化したものです。
- GRUでは、LSTMにあった短期記憶$h_{t}$と長期記憶$c_{t}$は、$h_{t}$に一本化されています。
- $g_{t}$はSimpleRNNやLSTMと似ていますが、$h_{t-1}$が$r_{t} \otimes h_{t-1}$に置き換えられていて、$h_{t-1}$の影響をどのくらい取り入れるかが$r_{t}$によって制御されています。
- $f_{t}$（forget gate）と$i_{t}$（input gate）は、$z_{t}$と$1-z_{t}$に置き換えられています。

SimpleRNNやLSTMと違って、バイアスが二度足されています。これは「PyTorchはそうしている」のと「そうしないとパラメータ数の計算が合わない」という理由から追加したものです。下で確認する通り、こうすると条件を変えてもパラメータ数が正しく計算されるので、おそらく正しいと思っています。
PyTorchでは、[SimpleRNN（torch.nn.RNN）](https://pytorch.org/docs/stable/generated/torch.nn.RNN.html#torch.nn.RNN)でも[LSTM](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html#torch.nn.LSTM)でも[GRU](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html#torch.nn.GRU)でもバイアスが二度足されているようで、なぜtf.kerasがこのような使い分けをしているのかは分かりませんでした...。

| 重み | 説明 | 形状 |
| :-: | :-: | :-: |
| $W_{xg}$ $W_{xr}$ $W_{xz}$ | 入力$x_{t}$に対応する重み | $n_{inputs}$ $\times$ $n_{outputs}$ |
| $W_{hg}$ $W_{hr}$ $W_{hz}$ | $h_{t-1}$に対応する重み | $n_{outputs}$ $\times$ $n_{outputs}$ |
| $b_{xg}$ $b_{hg}$ $b_{xr}$ $b_{hr}$ $b_{xz}$ $b_{hz}$ | バイアス | $n_{outputs}$ |

よってGRU層のパラメータ数は以下のように算出されます。
$
(n_{inputs}\ \times\ n_{outputs}\ +\ n_{outputs}\ \times\ n_{outputs}\ + \ 2\ \times\ n_{outputs})\ \times\ 3
$

この結果を、実際にtf.keras.layers.GRUの出力で確認してみます。

```python
from tensorflow import keras

model = keras.Sequential([
    keras.layers.GRU(units=20, input_shape=[None, 10], return_sequences=True),
    keras.layers.GRU(units=5, return_sequences=True),
    keras.layers.GRU(units=2)
])

model.summary()
```

```
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
gru (GRU)                    (None, None, 20)          1920      
_________________________________________________________________
gru_1 (GRU)                  (None, None, 5)           405       
_________________________________________________________________
gru_2 (GRU)                  (None, 2)                 54        
=================================================================
Total params: 2,379
Trainable params: 2,379
Non-trainable params: 0
```

| 層の名前 | $n_{inputs}$ | $n_{outputs}$ | 計算結果 |
| :-: | :-: | :-: | :-: |
| gru | input_shape=10 | units=20 | (10$\times$20+20$\times$20+2$\times$20)$\times$3=1920 |
| gru_1 | 一つ前の層の出力=20 | units=5 | (20$\times$5+5$\times$5+2$\times$5)$\times$3=405 |
| gru_2 | 一つ前の層の出力=5 | units=2 | (5$\times$2+2$\times$2+2$\times$2)$\times$3=54 |

# 4. 参考
- [Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow](https://www.amazon.co.jp/Hands-Machine-Learning-Scikit-Learn-TensorFlow-ebook/dp/B07XGF2G87/)のChapter 15: Processing Sequences Using RNNs and CNNs
- Courseraの[Sequence Models by DeepLearning.AI](https://www.coursera.org/learn/nlp-sequence-models)のWeek 1
- [How many parameters are in a gated recurrent unit (GRU) recurrent neural network (RNN) layer?](https://stats.stackexchange.com/questions/328926/how-many-parameters-are-in-a-gated-recurrent-unit-gru-recurrent-neural-network)