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の出力で確認してみます。
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の授業では、$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の出力で確認してみます。
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）でもLSTMでも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の出力で確認してみます。
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のChapter 15: Processing Sequences Using RNNs and CNNs
- CourseraのSequence Models by DeepLearning.AIのWeek 1
- How many parameters are in a gated recurrent unit (GRU) recurrent neural network (RNN) layer?