[人工智能]新一代 Kaldi 中的 LSTM GradientFilter

本文介绍新一代 Kaldi 中的 LSTM GradientFilter：?

?相关代码：https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py#L115

1. LSTM 梯度问题

https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html

我们先通过上图简单回顾下 LSTM cell 的内部结构和运算机制。作为 RNN 的变种之一， LSTM 旨在解决 RNN 存在的梯度消失或梯度爆炸的问题。

虽然在大多数情况下，LSTM 可以解决梯度问题。然而，如果矩阵??的值非常大，梯度爆炸问题仍然会发生，原因在于：在反向传播中，梯度会与矩阵??相乘；随着序列的长度增大， 相乘次数越多，梯度会无限制地随之增大。

我们在 icefall 的 PR （https://github.com/k2-fsa/icefall/pull/558 ）中讨论了训练 LSTM 模型时遇到的梯度问题。通过使用 icefall 中提供的模型分析工具（https://github.com/k2-fsa/icefall/blob/master/icefall/diagnostics.py），发现在第 10 层发生了梯度爆炸，由于梯度的反向传播，导致前面的层输出值对应的梯度都超过了正常范围。

LSTM diagnostics

我们使用的 LSTM 模型结构相对简单，包括 12 层的 encoder-layer。如下图所示，每个 encoder-layer 包括一个 ScaledLSTM、一个 ScaledFeedForward 和 一个 BasicNorm。

lstm encoder-layer

感兴趣的同学可以查看关于 LSTM 的 recipe 相关代码:https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/lstm_transducer_stateless3/lstm.py

2. Pytorch?中的 clip_grad_norm_

尽管 Pytorch 中的梯度裁剪函数?torch.nn.utils.clip_grad_norm_?可以在一定程度上解决梯度爆炸的问题，然而，该函数作用于对整个模型梯度反向传播结束之后，核心代码如下所示。

#?see?https://github.com/pytorch/pytorch/blob/435e78e5237d9fb3e433fff6ce028569db937264/torch/nn/utils/clip_grad.py#L10
total_norm?=?torch.norm(
??torch.stack(
????[torch.norm(p.grad.detach(),?norm_type).to(device)?for?p?in?parameters]
??),?norm_type
)
clip_coef?=?max_norm?/?(total_norm?+?1e-6)
clip_coef_clamped?=?torch.clamp(clip_coef,?max=1.0)
for?p?in?parameters:
????p.grad.detach().mul_(clip_coef_clamped.to(p.grad.device))

阅读代码，可发现其潜在的问题为：

• 裁剪之后，那些发生爆炸的模块会占主导，其它模块的梯度将变得特别小；

• 爆炸的梯度向前面层反向传播没有太大意义。

3. GradientFilter

我们希望在梯度的反向传播过程中，及时裁剪发生爆炸的梯度，而不是在整个模型的梯度反向传播完成之后。因此，我们在 icefall 实现了用于训练 LSTM 的 GradientFilter 模块。

该模块基于 torch.autograd.Function 实现，作用于每个 LSTM 层的内部，如：

#?see?https://github.com/k2-fsa/icefall/blob/9b671e1c21c190f68183f05d33df1c134079ca18/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py#L594
input,?*flat_weights?=?self.grad_filter(input,?*flat_weights)

如下面代码块所示，GradientFilter 在前向计算过程不进行任何操作，在反向传播过程中以 mask 的形式：

• 过滤每个 batch 中发生梯度爆炸的那些元素（序列）；

• 并根据梯度裁剪的程度，对应地放缩 LSTM 模块参数的梯度。

#?see?https://github.com/k2-fsa/icefall/blob/9b671e1c21c190f68183f05d33df1c134079ca18/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py#L115
eps?=?1.0e-20
dim?=?ctx.batch_dim
norm_dims?=?[d?for?d?in?range(x_grad.ndim)?if?d?!=?dim]
norm_of_batch?=?(x_grad?**?2).mean(dim=norm_dims,?keepdim=True).sqrt()
median_norm?=?norm_of_batch.median()

#?filter?gradients?of?batch?elements
cutoff?=?median_norm?*?ctx.threshold
inv_mask?=?(cutoff?+?norm_of_batch)?/?(cutoff?+?eps)
mask?=?1.0?/?(inv_mask?+?eps)
x_grad?=?x_grad?*?mask

#?filter?gradients?of?module?parameters
avg_mask?=?1.0?/?(inv_mask.mean()?+?eps)
param_grads?=?[avg_mask?*?g?for?g?in?param_grads]

4. 实验结果

下面展示了 12 层的 LSTM 模型 （参数量为86M）在 Full LibriSpeech 数据集上训练了 40 个 epoch，使用不同解码方法时，在 test-clean & test-other 上的 WER：

可以发现， GradientFilter 可以稳定 LSTM 的训练过程，有助于函数收敛，并取得一定的性能提升。

5.?实验结果

本文介绍了新一代 Kaldi 中用于训练 LSTM 的 GradientFilter，如若大家有相关的优化建议，欢迎向 icefall 提交 PR。