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

目的

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

対象

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

前提

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

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

結果

reset_after=Falseの場合

天下り的ですが、GRUの計算式は以下のように表されることがわかりました。（注：kerasのデフォルトであるreset_after=Trueは後でやります）

\begin{align}
z_t &= f(W^{(z)} x_t + U^{(z)} h_{t-1} + b^{(z)}) \\
r_t &= f(W^{(r)} x_t  + U^{(r)} h_{t-1} + b^{(r)}) \\
\tilde{h}_t &= g(W^{(h)}x_t + U^{(h)}(r_t\odot h_{t-1}) + b^{(h)}) \\
h_t &= z_t \odot h_{t-1} + (1-z_t)\odot \tilde{h}_t \\
y_t &= h_t
\end{align}

ここで

• $x_t$：入力値（$K$次元ベクトル）
• $z_t$：アップデートゲート
• $r_t$：リセットゲート
• $h_t$：内部メモリ（過去情報の蓄積）（$N$次元ベクトル）（GRUではそのまま出力になる）
• $y_t$：出力値（$N$次元ベクトル）
• $f$：Recurrent Activation（デフォルトはsigmoid）
• $g$：Activation（デフォルトはtanh）
• $W^{(\cdot)}$：入力に対する重み（$N\times K$行列）
• $U^{(\cdot)}$：内部メモリに対する重み（$N\times N$行列）
• $b^{(\cdot)}$：バイアス項（$N$次元ベクトル）
• $\odot$：要素積

となります。

実証

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

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

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

[[[ 0.25023641 -0.109921 ]
[-2.89429182 -0.65836289]
[ 0.0665177 0.59444831]
[-1.24301575 1.5483337 ]]]

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

# 内部メモリはN次元とする
N = 3
gru = tf.keras.layers.GRU(N, bias_initializer='random_normal', return_sequences=True, reset_after=False)
output = gru(input)
print(output)

tf.Tensor(
[[[ 0.03402694 -0.07257576 0.10821893]
[-0.01905339 0.21125174 -0.5427221 ]
[-0.0286486 -0.1144148 -0.39606214]
[-0.1057323 0.0409264 -0.6830395 ]]], shape=(1, 4, 3), dtype=float32)

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

次に、GRUを手組みするために、$W,U,b$の重みをkeras.layerから取得します。（3つの行列が結合されていることに注意）

kernel

W = gru.weights[0].numpy()
print(W.shape)
print(W)

(2, 9)
[[-0.63005173 -0.23849827 0.39704734 -0.4284358 -0.7049546 -0.632142 0.12867337 -0.73239166 0.47501415]
[ 0.5868781 -0.6272168 -0.08583027 0.4702161 -0.28281802 0.2657147 -0.31401938 -0.41231146 -0.4938141 ]]

recurrent_kernel

U = gru.weights[1].numpy()
print(U.shape)
print(U)

(3, 9)
[[ 4.6422434e-01 -7.2382373e-01 -7.3463716e-02 3.8367382e-03 -3.3444062e-01 1.8387362e-01 2.9062083e-01 -9.3553878e-02 -1.2763403e-01]
[-5.4166220e-02 1.7063874e-01 -5.1084316e-01 8.2610264e-02 4.1561580e-01 2.0287602e-01 6.7503279e-01 -1.7492548e-01 2.8214426e-04]
[-2.6494455e-01 -5.3800374e-01 -4.0752131e-01 -2.4554720e-02 3.9878601e-01 -1.1212709e-01 -2.8596750e-01 4.4986060e-01 1.3388421e-01]]

bias

b = gru.weights[2].numpy()
print(b.shape)
print(b)

(9,)
[-0.04370549 -0.00549434 -0.02047178 -0.01586535 -0.01216805 0.00143768 -0.00653381 -0.0085146 0.05727735]

得られた重み行列を数式のそれぞれに割り当てます。
:::note warn
この部分は結果に合うようにリバースエンジニアリングした結果です。また、元々の行列のまま計算する方法もあると思います。
:::

W_z = W[:,N*0:N*1]
W_r = W[:,N*1:N*2]
W_h = W[:,N*2:N*3]
U_z = U[:,N*0:N*1]
U_r = U[:,N*1:N*2]
U_h = U[:,N*2:N*3]
b_z = b[N*0:N*1]
b_r = b[N*1:N*2]
b_h = b[N*2:N*3]

Activationに使うsigmoidを定義しておきます。

def sigmoid(x):
    return 1/(1+np.exp(-x))

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

for x_s in input:
    h_t = np.zeros(N)
    h_hat_t = np.zeros(N)
    for x_t in x_s:
        z_t = sigmoid(x_t @ W_z + h_t @ U_z + b_z) # Recurrent Activationはsigmoidがデフォルト
        r_t = sigmoid(x_t @ W_r + h_t @ U_r + b_r) # Recurrent Activationはsigmoidがデフォルト
        h_hat_t = np.tanh(x_t @ W_h + (r_t * h_t) @ U_h + b_h) # Activationはtanhがデフォルト
        h_t = z_t * h_t + (1 - z_t) * h_hat_t
        print(h_t)

[ 0.03402695 -0.07257576 0.10821895]
[-0.0190534 0.2112517 -0.54272205]
[-0.0286486 -0.11441482 -0.39606211]
[-0.10573233 0.04092636 -0.68303951]

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

reset_after=Trueの場合（kerasのデフォルト）

reset_after=Trueの場合のGRUの計算式は以下のように表されることがわかりました。

\begin{align}
z_t &= f(W^{(z)} x_t + b^{(z,1)} + U^{(z)} h_{t-1} + b^{(z,2)}) \\
r_t &= f(W^{(r)} x_t + b^{(r,1)} + U^{(r)} h_{t-1} + b^{(r,2)}) \\
\tilde{h}_t &= g(W^{(h)}x_t + b^{(h,1)}+r_t\odot (U^{(h)}h_{t-1} + b^{(h,2)})) \\
h_t &= z_t \odot h_{t-1} + (1-z_t)\odot \tilde{h}_t \\
y_t &= h_t
\end{align}

reset_after=Falseのときとの違いは、$\tilde{h}_t$における$r_t$の要素積を取る順番が$U^{(h)}$を掛ける前か後かということと、バイアス項が増えていることです。実際、以下のgru.weights[2]の行列サイズが倍になっています。（バイアス項がなぜ増えるのかについてはよくわかりませんが、$r_t$の要素積を取る前にもバイアスを追加する関係で、行列の次元を揃えるためではと勝手に推測しています）

Kerasの結果は次のようになります。（レイヤーを作り直しているので、重み行列も再初期化されています）

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

tf.Tensor(
[[[-0.07515574 0.03578736 0.01209332]
[-0.02097559 -0.08850388 0.4859622 ]
[-0.02986786 -0.16829771 0.18403143]
[-0.01494701 -0.36375755 0.1712079 ]]], shape=(1, 4, 3), dtype=float32)

上記と同様にして、手組みの場合の結果を一気に計算します。

W = gru.weights[0].numpy() # kernel
U = gru.weights[1].numpy() # recurrent_kernel
b = gru.weights[2].numpy() # bias

W_z = W[:,N*0:N*1]
W_r = W[:,N*1:N*2]
W_h = W[:,N*2:N*3]
U_z = U[:,N*0:N*1]
U_r = U[:,N*1:N*2]
U_h = U[:,N*2:N*3]
b1_z = b[0:1,N*0:N*1]
b1_r = b[0:1,N*1:N*2]
b1_h = b[0:1,N*2:N*3]
b2_z = b[1:2,N*0:N*1]
b2_r = b[1:2,N*1:N*2]
b2_h = b[1:2,N*2:N*3]

for x_s in input:
    h_t = np.zeros(N)
    h_hat_t = np.zeros(N)
    for x_t in x_s:
        z_t = sigmoid(x_t @ W_z + b1_z + h_t @ U_z + b2_z)
        r_t = sigmoid(x_t @ W_r + b1_r + h_t @ U_r + b2_r)
        h_hat_t = np.tanh(x_t @ W_h + b1_h + r_t * (h_t @ U_h + b2_h))
        h_t = z_t * h_t + (1 - z_t) * h_hat_t
        print(h_t)

[[-0.07515575 0.03578737 0.01209332]]
[[-0.0209756 -0.08850388 0.48596215]]
[[-0.02986787 -0.16829772 0.1840314 ]]
[[-0.01494702 -0.36375755 0.1712079 ]]

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

参考

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