Backpropagation Through Time › Gradient Flow Through Hidden States

def gradient_decay_by_position(
    seq_len, vocab_size=27, hidden_size=64, batch_size=4
):
    """
    Compute gradient norms as a function of distance from the end of the sequence.
    Uses manual RNN forward pass to retain gradients at each hidden state.
    """
    torch.manual_seed(0)
    m = SimpleRNN(vocab_size, hidden_size, vocab_size)
    x, targets = make_batch(batch_size, seq_len, vocab_size)

    # Initial hidden state: requires_grad=True so retain_grad() works
    h_prev = torch.zeros(batch_size, hidden_size, requires_grad=True)
    all_h = []

    # Manual forward to retain all hidden state grads
    for t in range(seq_len):
        x_t = x[:, t, :]  # (batch, input_size)
        h_new = torch.tanh(
            x_t @ m.rnn.weight_ih_l0.T
            + m.rnn.bias_ih_l0
            + h_prev @ m.rnn.weight_hh_l0.T
            + m.rnn.bias_hh_l0
        )  # (batch, hidden_size) - requires_grad=True because it depends on parameters
        h_new.retain_grad()
        all_h.append(h_new)
        h_prev = h_new

    # Apply the loss only at the final timestep. This isolates the gradient
    # propagated backward through the recurrent chain instead of adding a
    # direct loss contribution at every hidden state.
    logits = m.fc(all_h[-1])
    loss = criterion(logits, targets[:, -1])
    loss.backward()

    grad_norms = []
    for h_t in all_h:
        if h_t.grad is not None:
            grad_norms.append(h_t.grad.norm().item())
        else:
            grad_norms.append(0.0)

    return grad_norms


seq_len_decay = 30
grad_decay = gradient_decay_by_position(seq_len_decay)