tf.keras.layers.SimpleRNNの仕組みを手組みで確認する

目的

Tensorflow.Kerasに用意されているRNNレイヤーの構造を、手組みの場合と比較しながら理解します。

対象

本稿では、RNNレイヤーの1つであるSimpleRNNを対象とします。
https://www.tensorflow.org/api_docs/python/tf/keras/layers/SimpleRNN

前提

公式ドキュメント、専門書およびweb上の情報等を参考にリバースエンジニアリングを行いながら、結果を比較することで計算式、および重み行列とKerasレイヤーのパラメータとの対応関係を抽出しました。ソースコードを見て厳密に確認したわけではないため、100%正しい記載になっていない可能性がありますが、理解のためには十分かなと思います。

実行結果は以下の環境で計算した結果となります。
Python: 3.10.6
Tensorflow: 2.12.0 (tensorflow-aarch64)
Keras: 2.12.0
Numpy: 1.23.5

結果

天下り的ですが、SimpleRNNの計算式は以下のように表されることがわかりました。

\begin{align}
s_t &= f(W x_t + U s_{t-1} + b) \\
y_t &= s_t
\end{align}

ここで

• $x_t$：入力値（$K$次元ベクトル）
• $s_t$：内部メモリ（$N$次元ベクトル）（SimpleRNNではそのまま出力になる）
• $y_t$：出力値（$N$次元ベクトル）
• $f$：Activation（デフォルトはtanh）
• $W$：入力に対する重み（$N\times K$行列）
• $U$：内部メモリに対する重み（$N\times N$行列）
• $b$：バイアス項（$N$次元ベクトル）

となります。

実証

SimpleRNNの計算が上式となることを、Kerasの出力結果と手組みの出力結果を比較することで確かめます。

まず入力値をランダムに生成します。

# 時刻ステップTのK次元データをS個生成
S = 3
T = 2
K = 4
input = np.random.randn(S, T, K)
print(input)

[[[-1.00559132 2.19790855 1.22426653 0.43110155]
[ 2.47757924 1.42999603 0.22461884 0.52406924]]

[[ 1.58063093 -0.65496626 0.66998741 -2.05034626]
[ 1.31754524 0.40744175 0.07545818 1.48894891]]

[[-0.46522181 2.16090139 0.71952939 -0.45572353]
[-1.99252205 0.38433699 0.49557174 -2.04901821]]]

Kerasの結果は次のようになります。return_sequences=Trueとして全ての時刻ステップの結果を取得します。また、デフォルトでゼロに初期化されてしまうbiasをbias_initializer='random_normal'として値を入れます。

# 内部メモリはN次元とする
N = 2
rnn = tf.keras.layers.SimpleRNN(N, bias_initializer='random_normal', return_sequences=True)
output = rnn(input)
print(output)

tf.Tensor(
[[[-0.9842801 0.7969481 ]
[ 0.93499136 -0.7715645 ]]

[[ 0.44224197 -0.9944355 ]
[ 0.7529196 -0.6964871 ]]

[[-0.9916976 -0.10067604]
[-0.999102 0.8626111 ]]], shape=(3, 2, 2), dtype=float32)

出力は$(S, T, N)$次元のテンソルとなります。

次に、SimpleRNNを手組みするために、$W,U,b$の重みをkeras.layerから取得します。

W = rnn.weights[0].numpy()
print(W)

[[ 0.83396554 -0.95233345]
[-0.9864013 -0.09958982]
[ 0.21518016 0.05979133]
[ 0.78888416 0.73306966]]

U = rnn.weights[1].numpy()
print(U)

[[ 0.19022751 -0.9817401 ]
[ 0.9817401 0.19022739]]

b = rnn.weights[2].numpy()
print(b)

[-0.01591891 -0.03780531]

最後に、各サンプルごとに、SimpleRNNの数式に従って計算し、各時刻ステップ毎に結果を出力します。このとき、行列積の順序に注意します。（入力値や重み行列が数式とは転置になっているため）

for x_s in input:
    s = np.zeros(N)
    for x_t in x_s:
        s = np.tanh(x_t @ W + s @ U + b)
        print(s)

[-0.98428005 0.79694807]
[ 0.93499136 -0.77156436]

[ 0.44224219 -0.99443547]
[ 0.75291959 -0.69648714]

[-0.99169759 -0.10067606]
[-0.99910198 0.86261118]

kerasの結果と一致していることが確認できました。

参考

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