CANN/ops-nn THNN融合LSTM单元梯度

ThnnFusedLstmCellGrad

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn

产品支持情况

功能说明

• 算子功能：LSTMCell中四个门中matmul后剩余计算的反向传播，计算正向输出四个门激活前的值gates、输入cx、偏置b的梯度。

• 计算公式：

变量定义

• 输入梯度：$\delta h_t$ (gradHy)， $\delta c_t$ (gradC)
• 前向缓存：$i，f，g，o$ (各门激活值storage)，$c_{t-1}$ (cx)，$c_t$ (cy)
• 输出梯度：$\delta a_i，\delta a_f，\delta a_g，\delta a_o$ (存入gradInGatesOut)，$\delta c_{t-1}$ (gradCPrevOut)

第一阶段：中间梯度与状态回传

首先计算隐藏状态对细胞状态的贡献，并汇总得到当前时刻细胞的总梯度 $\text{grad_}c_{total}$：

$$ \begin{aligned} gcx &= \tanh(c_t) \ \text{grad_}c_{total} &= \delta h_t \cdot o \cdot (1 - gcx^2) + \delta c_t \ \delta c_{t-1} &= \text{grad_}c_{total} \cdot f \end{aligned} $$

第二阶段：门控分量梯度 (Pre-activation)

根据代码逻辑，各门控在进入激活函数前的梯度 $\delta a$ 计算如下：

$$ \begin{aligned} \delta a_o &= (\delta h_t \cdot gcx) \cdot o \cdot (1 - o) \ \delta a_i &= (\text{grad_}c_{total} \cdot g) \cdot i \cdot (1 - i) \ \delta a_f &= (\text{grad_}c_{total} \cdot c_{t-1}) \cdot f \cdot (1 - f) \ \delta a_g &= (\text{grad_}c_{total} \cdot i) \cdot (1 - g^2) \end{aligned} $$

第三阶段：参数梯度 (db)

**1. 偏置梯度 (db)：**对 Batch 维度（$N$）进行求和：

$$ \delta b = \sum_{n=1}^{N} \begin{bmatrix} \delta a_i \ \delta a_f \ \delta a_g \ \delta a_o \end{bmatrix}_n $$

• 参数说明：

  参数名	输入/输出/属性	描述	数据类型	数据格式
  dhy	输入	表示LSTMCell正向输出隐藏状态的梯度。 shape为[batch，hidden_size]。	FLOAT、FLOAT16	ND
  dc	输入	表示LSTMCell正向输出细胞状态的梯度。 shape为[batch，hidden_size]。	FLOAT、FLOAT16	ND
  cx	输入	表示LSTMCell正向输入细胞状态。 shape为[batch，hidden_size]。	FLOAT、FLOAT16	ND
  cy	输入	表示LSTMCell正向输出细胞状态。 shape为[batch，hidden_size]。	FLOAT、FLOAT16	ND
  storage	输入	表示LSTMCell正向输出四个门的激活值。 shape为[batch, 4 * hidden_size]。	FLOAT、FLOAT16	ND
  dgates	输出	表示LSTMCell正向中四个门预激活值的梯度。 shape为[batch, 4 * hidden_size]。	FLOAT、FLOAT16	ND
  dc_prev	输入	表示LSTMCell正向中输入细胞状态的梯度。 shape为[batch，hidden_size]。	FLOAT、FLOAT16	ND
  db	输入	表示LSTMCell正向中输入偏置的梯度。 shape为[4 * hidden_size]。	FLOAT、FLOAT16	ND
  has_bias	属性	是否需要计算bias梯度。	STRING	-

约束说明

无

调用说明

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn