Alan Lee

计算 LSTM 的参数量

2020/06/04 Share

理论上的参数量

之前翻译了 Christopher Olah 的那篇著名的 Understanding LSTM Networks,这篇文章对于整体理解 LSTM 很有帮助,但是在理解 LSTM 的参数数量这种细节方面,略有不足。本文就来补充一下,讲讲如何计算 LSTM 的参数数量。

建议阅读本文前先阅读 Understanding LSTM Networks 的原文或我的译文

首先来回顾下 LSTM。一层 LSTM 如下:

lstm layer

这里的 $x_t$ 实际上是一个句子的 embedding(不考虑 batch 维度),shape 一般为 [seq_length, embedding_size]。图中的 $A$ 就是 cell,$x_t$ 中的词依次进入这个 cell 中进行处理。可以看到其实只有这么一个 cell,所以每次词进去处理的时候,权重是共享的,将这个过程平铺展开,就是下面这张图了:

unrolled lstm layer

实际上我觉得这里 $x_t$ 并不准确,第一个 $x_t$ 应该指的是整句话,而第二个 $x_t$ 应该指的是这句话中最后一个词,所以为了避免歧义,我认为可以将第一个 $x_t$ 重命名为 $x$,第二个仍然保留,即现在 $x$ 表示一句话,该句话有 $t+1$ 个词,$x_t$ 表示该句话的第 $t+1$ 个词,$t \in [0, t]$。

始终要记住这么多 $A$ 都是一样的,权重是一样的,$x_0$ 到 $x_t$ 是一个个词,每一次的处理都依赖前一个词的处理结果,这也是 RNN 系的网络难以像 CNN 一样并行加速的原因。同时, 这就像一个递归过程,如果把求 $h_t$ 的公式展开写,$A$ 里的权重记为 $W$,那么就会发现需要 $t$ 个 $W$ 相乘,即 $W^t$,这是非常恐怖的:

一个不那么小的数被多次相乘之后会变得很小,一个不那么大的数被多次相乘之后会变得很大。所以,这也是普通 RNN 容易出现梯度消失/爆炸的问题的原因

扯远了点。

那么 LSTM 的参数很明显了,就是这个 $A$ 中的参数。这个 $A$ 内部具体是这样的:

cell.png

从这张图来理解参数的数量你可能有点懵逼,一步一步来看,实际上这里面有 4 个非线性变换(3 个 门 + 1 个 tanh),每一个非线性变换说白了就是一个两层的全连接网络。重点来了,第一层是 $x_i$ 和 $h_i$ 的结合,维度就是 embedding_size + hidden_size,第二层就是输出层,维度为 hidden_size,所以该网络的参数量就是:

1
(embedding_size + hidden_size) * hidden_size + hidden_size

一个 cell 有 4 个这样结构相同的网络,那么一个 cell 的总参数量就是直接 × 4:

1
((embedding_size + hidden_size) * hidden_size + hidden_size) * 4

注意这 4 个权重可不是共享的,都是独立的网络。

所以,一般来说,一层 LSTM 的参数量计算公式是:

其中 4 表示有 4 个非线性映射层,$d_h + d_x$ 即 Understanding LSTM Networks 中的 $[h_{t-1}, x_t]$ 的维度,后面的 $d_h$ 表示 bias 的数量。所以,LSTM 层的参数数量只与输入维度 $d_x$ 和输出维度 $d_h$ 相关,和普通全连接层相同。

那么显而易见,一层双向 LSTM 的参数量就是上述公式 × 2。

TensorFlow 中的实现

在 TensorFlow 中,这些 $d_x$、$d_h$ 如何与代码对应上呢?

我们可以如下实现一个简单的以 LSTM 为核心的网络:

1
2
3
4
5
6
7
8
import tensorflow as tf

model = tf.keras.model.Sequential(
tf.keras.layers.Embedding(1000, 128),
tf.keras.layers.LSTM(units=64),
tf.keras.layers.Dense(10)
)
model.summary()

输入如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
Model: "sequential"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
embedding_1 (Embedding) (None, None, 128) 128000
_________________________________________________________________
lstm_1 (LSTM) (None, 64) 49408
_________________________________________________________________
dense_1 (Dense) (None, 10) 650
=================================================================
Total params: 178,058
Trainable params: 178,058
Non-trainable params: 0
_________________________________________________________________

可以看到 TF 给出的 LSTM 层参数量是 49408。我们来根据上面的公式验证下。

  • $d_x$:输入维度,在这里就对应于 128,就是词向量维度。
  • $d_h$:输出维度,在这里就是 LSTM 的参数 64,在 TF 这里叫 units

所以,参数量就是 $4 \times \left[64 \times \left(64+128\right) + 64 \right] = 49408$,和 TF 给出的一样。

另外,tf.keras.layers.LSTM() 的默认输出大小为 [batch_size, units],就是只使用最后一个 time step 的输出。假如我们想要得到每个 time step 的输出($h_0,\cdots,h_t$)和最终的 cell state($C_t$),那么我们可以指定另外两个参数 return_sequences=Truereturn_state=True

1
2
3
4
5
inputs = tf.random.normal([64, 100, 128])  # [batch_size, seq_length, embedding_size]
whole_seq_output, final_memory_state, final_carry_state = tf.keras.layers.LSTM(64, return_sequences=True, return_state=True)(inputs)
print(f"{whole_seq_output.shape=}")
print(f"{final_memory_state.shape=}")
print(f"{final_carry_state.shape=}")

输出:

1
2
3
whole_seq_output.shape=TensorShape([32, 100, 64])  # 100 表示有 100 个词,即 100 个 time step
final_memory_state.shape=TensorShape([32, 64])
final_carry_state.shape=TensorShape([32, 64])

OK,LSTM 的参数量应该挺清晰了,欢迎在评论区留下你的想法。😋

Reference

END

CATALOG
  1. 1. 理论上的参数量
  2. 2. TensorFlow 中的实现
  3. 3. Reference
  4. 4. END