从零实现双向RNN（BRNN）：反向传播实现困惑求助

嘿，我完全懂你从零实现双向循环神经网络（BRNN）时卡在反向传播上的头疼——前向传播逻辑很直观，但反向传播因为涉及到反转序列的梯度对齐，确实容易绕晕。咱们一步步把它拆解清楚，你就能明白了。

先回顾下你的前向传播：你用了两个并行的RNN流——一个处理原始输入序列，另一个处理反转后的输入，最后把两个的隐藏状态拼接得到输出。反向传播的关键就是分别处理这两个RNN分支的梯度，再把对应部分的梯度合并。

反向传播的分步实现 #

1. 拆分输出梯度 #

前向传播最后是把hidden_state（原始RNN的隐藏状态）和hidden_state_reversed（反转RNN的隐藏状态）拼接成了output。所以当拿到输出的梯度d_output时，第一步要把它拆成两个部分：

• d_hidden_state：对应原始RNN隐藏状态的梯度
• d_hidden_state_reversed：对应反转RNN隐藏状态的梯度

拆分的维度取决于你拼接的方式——如果是在隐藏单元维度拼接（比如每个时间步的隐藏状态是[batch_size, hidden_dim]，拼接后变成[batch_size, 2*hidden_dim]），直接把d_output在这个维度上切成前半段和后半段就行。

2. 处理反转RNN的梯度 #

反转RNN的前向传播用的是反转后的输入，所以它的梯度处理要多一步时间步对齐：

• 先把d_hidden_state_reversed反转——因为反转RNN的时间步是倒过来的，它的最后一个时间步对应原始输入的第一个时间步，所以梯度也要反转回去，才能和反转RNN的前向缓存对应上。
• 调用你已经实现的RNN反向传播函数，传入反转后的梯度和反转RNN的前向缓存，得到反转输入的梯度d_reversed_input和反转RNN的参数梯度d_params_reversed。
• 再把d_reversed_input反转一次——因为d_reversed_input是对应反转输入序列的梯度，反转后才能和原始输入序列的时间步对齐，得到d_input_from_reversed。

3. 处理原始RNN的梯度 #

这部分就简单了：直接调用你的RNN反向传播函数，传入d_hidden_state和原始RNN的前向缓存，得到原始输入的梯度d_input_from_forward和原始RNN的参数梯度d_params_forward。

4. 合并最终梯度 #

• 输入的总梯度：把d_input_from_forward和d_input_from_reversed相加——因为原始输入同时被两个RNN分支用到（一个直接用，一个反转后用），所以两个分支的梯度要叠加。
• 参数的总梯度：这里要分两种情况：
  • 如果你的BRNN是正反方向共享同一套RNN参数，那把d_params_forward和d_params_reversed相加即可；
  • 如果是正反方向用独立的参数（这是更常见的情况，因为要学习不同方向的序列依赖），那就分别保存这两组参数梯度。

伪代码示例 #

def brnn_backward(d_output, forward_cache, reversed_cache, hidden_dim, share_params=False):
    # forward_cache: 原始RNN前向传播时保存的缓存（用于反向传播）
    # reversed_cache: 反转RNN前向传播时保存的缓存
    
    # 1. 拆分输出梯度
    split_dim = 1  # 假设在隐藏单元维度拼接，根据你的实际情况调整
    d_hidden_state, d_hidden_state_reversed = d_output.split(split_size_or_sections=hidden_dim, dim=split_dim)
    
    # 2. 处理反转RNN的反向传播
    # 反转梯度以匹配反转RNN的时间步顺序
    d_hidden_state_reversed_reversed = d_hidden_state_reversed.flip(dims=[0])  # 假设时间步是第0维
    # 调用RNN反向传播
    d_reversed_input, d_params_reversed = RNN.backward(d_hidden_state_reversed_reversed, reversed_cache)
    # 反转回原始输入的时间步顺序
    d_input_from_reversed = d_reversed_input.flip(dims=[0])
    
    # 3. 处理原始RNN的反向传播
    d_input_from_forward, d_params_forward = RNN.backward(d_hidden_state, forward_cache)
    
    # 4. 合并梯度
    d_input = d_input_from_forward + d_input_from_reversed
    if share_params:
        d_params = {k: d_params_forward[k] + d_params_reversed[k] for k in d_params_forward}
    else:
        d_params = (d_params_forward, d_params_reversed)
    
    return d_input, d_params

关键注意事项 #

• 缓存必须保存：你的brnn_forward函数要把原始RNN和反转RNN的前向缓存（比如每个时间步的输入、隐藏状态等）都返回并保存，否则RNN的反向传播没法计算梯度。
• 时间步对齐要仔细：反转操作的维度一定要对应正确，建议实现时多打印各个张量的形状，确保每一步的维度都符合预期——这是最容易踩坑的地方。
• 参数共享与否要明确：提前确定你的BRNN是共享参数还是用独立参数，这直接影响参数梯度的合并方式。

内容的提问来源于stack exchange，提问作者Ivan Telnov