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は後でやります）
```math
\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の出力結果と手組みの出力結果を比較することで確かめます。

まず入力値をランダムに生成します。
```python
# 時刻ステップ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'として値を入れます。
```python
# 内部メモリは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つの行列が結合されていることに注意）
```python: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 ]]

```python: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]]

```python: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
この部分は結果に合うようにリバースエンジニアリングした結果です。また、元々の行列のまま計算する方法もあると思います。
:::

```python
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を定義しておきます。
```python
def sigmoid(x):
    return 1/(1+np.exp(-x))
```

最後に、各サンプルごとに、GRUの数式に従って計算し、各時刻ステップ毎に結果を出力します。このとき、行列積の順序に注意します。（入力値や重み行列が数式とは転置になっているため）
```python
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の計算式は以下のように表されることがわかりました。
```math
\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の結果は次のようになります。（レイヤーを作り直しているので、重み行列も再初期化されています）
```python
# 内部メモリは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)

上記と同様にして、手組みの場合の結果を一気に計算します。
```python
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のパラメータ数の計算方法](https://qiita.com/yuichi0625/items/29cf86b911013e3f9476)