Clarification regarding (truncated) backpropagation through time

1

Hi, I am unsure how the backpropagation works/the gradient is computed in a scenario with truncated backpropagation through time. Currently, I do something akin to

outputs = []
for t in range(max_steps):
     y_t, hidden  = custom_rnn(x[t], hidden.detach())
     outputs.append(y_t)

y = torch.stack(outputs)
my_loss = criterion(y,x)
return y, my_loss

In this code, when the derivative of hidden with respect to the rnn’s parameters are computed ends at every time step, because due to the .detach() it acts as a constant in every time step.

Now, what about this:

outputs = []
for t in range(max_steps):
     y_t, hidden  = custom_rnn(x[t], hidden)
     if t % 10 == 0 and t > 0:
          hidden = hidden.detach()
     outputs.append(y_t)

y = torch.stack(outputs)
my_loss = criterion(y,x)
return y, my_loss

When the derivative of the loss is computed with respect to the rnn’s parameters, are only the last <= 10 time steps relevant for the chain-rule derivative of hidden (with respect to the rnn’s parameters)? Or is the derivative somehow restarted inbetween?
If it is the former (only the last <= 10 time steps are relevant), what about this:

outputs = []
segments = []
segments.append(torch.tensor(0, requires_grad=False))
for t in range(max_steps):
     y_t, hidden  = custom_rnn(x[t], hidden)
     if t % 10 == 0 and t > 0:
          hidden = hidden.detach()
          segments.append(t)
     outputs.append(y_t)

y = torch.stack(outputs)
total_loss = 0.0
for t in segments:
     total_loss += criterion(y[t:t+10], x[t:t+10]) # ignore the last chunk possibly being shorter
my_loss = total_loss/len(segments)
return y, my_loss

Now here I compute the loss independently for each segment between the detach operations. Is the gradient flowing properly for each chunk/segment?

Now in practice I have the following scenario:
My whole model is roughly:

class MyModel(nn.Module):
 # some defs 
def forward(self, x):
    y = self.submodel_A(x)
    z, rnn_loss = my_rnn(y)
    return z, rnn_loss

and then I have an outer criterion:

my_model = MyModel()
for (input, target) in dataloader:
    output, my_loss = my_model(input)
    loss = outer_criterion(output, target) + my_loss
    loss.backward()

How does the loss.backward() interact with the segment losses I compute inside the rnn? Are the derivates correctly computed for each segment? Or just for the last one? Also, for the outer_critertion, does the derivative of hidden with respect to the rnn’s parameters stop at the last hidden.detach()? Or does the gradient somehow propagate through the entire sequence the rnn processes?

It would be really helpful to get some clarification on this as I am still unsure what is going on under the hood.
Thank you :slight_smile:

2

Did you get this sorted? Essentially what you need to be doing is computing ten steps forward (no detach), computing the loss function, calling backward on the graph that contains all ten steps, then calling detach on the hidden to free everything up. Then you go around again.

There is example code at RNN — PyTorch 2.12 documentation, you can strip it of all the batch_first, layer and bias code and it gets readable and understandable.

3

Your intuition is correct on all three variants. The single rule that governs everything: .detach() creates a hard boundary in the autograd graph — gradient cannot cross it. Everything else follows.

Variant 1 : no gradient flow through the recurrent connection at all. Each timestep sees hidden as a constant input. You are still learning something (the RNN parameters get gradients from the direct dependency of y_t on them at that timestep), but the recurrent structure contributes nothing.

Variant 2: you are right. When loss.backward() runs on the concatenated y, gradient flows back through each y_t to its corresponding hidden. Because hidden gets detached every 10 steps, the gradient chain through the recurrent connection reaches at most 10 steps back before hitting the barrier. Different y_ts reach different depths - y_39 reaches back to step 30, y_45 reaches back to step 40, and so on. Each timestep also contributes gradient to the RNN parameters through its own direct forward computation; that part is not truncated.

Variant 3: gradients are correctly computed per segment. The stack preserves each y_t’s individual graph, and slicing then selects which timesteps’ losses to compute. Each segment’s loss produces gradients that flow at most to its own start (where the previous detach happened). Total gradient is the sum of per-segment gradients. This is mathematically equivalent to the “compute loss per segment, backward per segment with retain_graph, then detach” pattern, just fused into one backward pass.

4

Thank you so much for the clarification. So the gradients in Variant 3 are also computed for the RNN’s weights with respect to the outer criterion? The gradient with respect to these weights of the outer criterion really does not stop at the last detach operation?