Cannot recreate backpropagation of `torch.nn.RNN`

Dylan_Hanna: <p>Hello.</p>
<p>I’m currently trying to understand Backpropagation Through Time (BPTT), and am trying test my understanding by re-creating the backpropagation calculation that <code>torch.nn.RNN</code> performs.</p>
<p>After having read many articles and performed the mathematical proofs, I am 90% certain that not only my theoretical understanding, but my code implementation is also correct – however, the gradients that I am manually calculating do not match the gradients produced by <code>torch</code>, and I have no idea why.</p>
<p>I’ll show you what I mean. Below is a simple <code>torch_rnn</code> object, which is a simple <code>RNN</code> layer:</p>
<pre><code class="lang-auto">import torch.nn as nn

torch_rnn = nn.RNN(input_size=1, hidden_size=2, batch_first=True)
</code></pre>
<p>I’ll also be using the following sequence as the input to our <code>RNN</code>:</p>
<pre><code class="lang-auto">import torch

x = torch.tensor([[1.0],
                  [2.0],
                  [3.0]])
</code></pre>
<p>I’ll now perform a forward pass using <code>torch_rnn</code>, and will store the final hidden state (or the output of the layer) as <code>torch_last_hidden</code>:</p>
<pre><code class="lang-auto">_, torch_last_hidden = torch_rnn(x)
</code></pre>
<p><em>Now</em>, I’ll use my own hard-coded implementation to do the same thing, but will copy the weights and biases <em>directly</em> from the <code>torch_rnn</code> object. I’ll then walk-forward through the RNN calculations step by step, until I reach <code>h_3</code>, the final hidden state. I’ll also be initialising <code>h_0</code> to be all-zero; as is the default behaviour for <code>nn.RNN</code>:</p>
<pre><code class="lang-auto">W_ih = torch_rnn.weight_ih_l0.detach()
b_ih = torch_rnn.bias_ih_l0.detach()
W_hh = torch_rnn.weight_hh_l0.detach()
b_hh = torch_rnn.bias_hh_l0.detach()

h_0 = torch.zeros(1, 64)

tanh = torch.nn.Tanh()

z_1 = (x[0] @ W_ih.T) + b_ih + (h_0 @ W_hh.T) + b_hh
h_1 = tanh(z_1)

z_2 = (x[1] @ W_ih.T) + b_ih + (h_1 @ W_hh.T) + b_hh
h_2 = tanh(z_2)

z_3 = (x[2] @ W_ih.T) + b_ih + (h_2 @ W_hh.T) + b_hh
h_3 = tanh(z_3)
</code></pre>
<p>Let’s make sure that our manual calculation has produced a very close approximation of the <code>torch</code> implementation:</p>
<pre><code class="lang-auto">assert torch.allclose(h_3, torch_last_hidden, atol=1e-6)
</code></pre>
<p>^^ This assert statement doesn’t throw an error, so the forward pass has been successfully recreated.</p>
<p>Next, is the backward pass. For this example, I will try and calculate the gradients of <code>W_ih</code> with respect to <code>h_3</code>. The theory behind how to do this (as I understand it), is to calculate the following:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2a313406c911b8b307e0760fcba0e1cefcfba36c.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/2a313406c911b8b307e0760fcba0e1cefcfba36c" title="backpropagation_formula_full"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a313406c911b8b307e0760fcba0e1cefcfba36c_2_574x500.jpeg" alt="backpropagation_formula_full" data-base62-sha1="61ftdfJ7GoABPZlbUMRfOMTSl1q" width="574" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a313406c911b8b307e0760fcba0e1cefcfba36c_2_574x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a313406c911b8b307e0760fcba0e1cefcfba36c_2_861x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a313406c911b8b307e0760fcba0e1cefcfba36c_2_1148x1000.jpeg 2x" data-dominant-color="F8F8F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">backpropagation_formula_full</span><span class="informations">1390×1210 117 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’ll now try applying this to our above problem and manually perform the backpropagation for <code>W_ih</code>. I’ll also allow <code>torch</code> to perform its own calculation of the <code>W_ih</code> gradients for comparison, and save the result to <code>torch_W_ih_gradients</code>:</p>
<pre><code class="lang-auto">import torch.optim as optim

optimiser = optim.SGD(torch_rnn.parameters(), lr=0.001)

optimiser.zero_grad()

torch_last_hidden.sum().backward()
torch_W_ih_gradients = copy.deepcopy(torch_rnn.weight_ih_l0.grad)

dh_3_dW_ih_3 = (1 - tanh(z_3)**2)*x[2]

dh_3_dW_ih_2 = (((1 - tanh(z_3)**2) @ W_hh) * 
                ((1 - tanh(z_2)**2) * x[1]))

dh_3_dW_ih_1 = (((1 - tanh(z_3)**2) @ W_hh) * 
                ((1 - tanh(z_2)**2) @ W_hh) * 
                ((1 - tanh(z_1)**2) * x[0]))

dh_3_dW_ih = dh_3_dW_ih_1 + dh_3_dW_ih_2 + dh_3_dW_ih_3

assert torch.allclose(dh_3_dW_ih.T, torch_W_ih_gradients, atol=1e-6)
</code></pre>
<p>The above code will throw an <code>AssertionError</code>, and a manual inspection of the gradients reveals that they are in fact different:</p>
<pre><code class="lang-auto">dh_3_dW_ih.T -------------&gt; tensor([[0.5162],
                                    [0.5511]])
                                    
torch_W_ih_gradients -----&gt; tensor([[0.5053],
                                    [0.5356]]))
</code></pre>
<p>The difference is small, but significant enough to cause concern.<br>
Have I made a mistake with my theoretical understanding of the backpropagation in the RNN?<br>
Is there a mistake with my code implementation?<br>
Why are my calculations not yielding the same gradients as the <code>torch</code> implementation?</p>