Transformers and Attention Questions

The short version

Default to a Transformer. It processes all positions of a sequence at once during training, connects any two positions in a single step, and benefits from large-scale pretraining. An RNN (recurrent neural network: reads one token at a time and carries a fixed-size summary, the hidden state) or LSTM (an RNN with gates that control what to keep) is still reasonable when you need constant memory per step on a stream, run on a tiny device, or train from scratch on very little data with no pretrained model available.

Comparison along the four axes asked

Measured: gradient reaching the first input

The program takes a plain tanh RNN (default PyTorch initialization, width 32) and one self-attention layer (random projections, width 32), and measures the size of the gradient of the last position's output with respect to the first input vector, for sequence lengths 10, 50 and 200, averaged over 10 seeds. It then prints counted quantities (not timings): sequential steps, the n by n attention entries, and the size of the RNN's hidden state, which is 32 numbers whatever the sequence length. The comparison is deliberately narrow: a plain tanh RNN at initialization, not an LSTM and not a trained network, against one random attention layer. It shows how a gradient signal behaves when it must pass through many repeated multiplications versus one step; it does not show what a trained LSTM can remember.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(1)

D = 32

def rnn_grad(T, seed):
    torch.manual_seed(seed)
    rnn = nn.RNN(D, D, nonlinearity="tanh", batch_first=True)
    x = torch.randn(1, T, D, requires_grad=True)
    out, _ = rnn(x)
    out[0, -1].sum().backward()
    return x.grad[0, 0].norm().item()

def attn_grad(T, seed):
    torch.manual_seed(seed)
    Wq, Wk, Wv = (torch.randn(D, D) / D ** 0.5 for _ in range(3))
    x = torch.randn(1, T, D, requires_grad=True)
    out = F.scaled_dot_product_attention(x @ Wq, x @ Wk, x @ Wv)
    out[0, -1].sum().backward()
    return x.grad[0, 0].norm().item()

print("T     RNN d out_last/d x_0   attention d out_last/d x_0   (mean of 10 seeds)")
for T in (10, 50, 200):
    r = sum(rnn_grad(T, s) for s in range(10)) / 10
    a = sum(attn_grad(T, s) for s in range(10)) / 10
    print(f"{T:<5} {r:<22.2e} {a:.2e}")

print("\nT      sequential steps (RNN)  sequential steps (attention layer)  attn entries n^2  RNN hidden-state numbers")
for T in (512, 4096, 32768):
    print(f"{T:<6} {T:<23} {1:<35} {T * T:<17} {D}")

Output:

T     RNN d out_last/d x_0   attention d out_last/d x_0   (mean of 10 seeds)
10    2.54e-03               5.93e-01
50    2.47e-15               1.19e-01
200   0.00e+00               2.23e-02

T      sequential steps (RNN)  sequential steps (attention layer)  attn entries n^2  RNN hidden-state numbers
512    512                     1                                   262144            32
4096   4096                    1                                   16777216          32
32768  32768                   1                                   1073741824        32

Reading it: in the RNN the gradient from the last step back to the first input shrinks from about 2.5e-3 at length 10 to about 2.5e-15 at length 50 and underflows (becomes too small for float32 to represent, so it prints as exactly zero) at length 200. This is the vanishing-gradient effect at this default initialization (other initializations can make it explode instead); LSTM gates were designed to slow it down, not remove it. In the attention layer the path from first to last position is one step, so the gradient shrinks only slowly (0.593, 0.119, 0.0223), roughly as the weight on any one token drops when the weights spread over more tokens. The counted columns show the price: attention entries grow as n^2 (about 1.07 billion at n = 32,768), while the RNN's hidden state stays at 32 numbers and its sequential steps grow linearly. For scale, the attention gradient at length 200 (0.0223) is still usable, while the plain RNN's at length 50 (2.5e-15) is far below anything an optimizer can act on; LSTM gates raise how far a signal survives, and how far in a trained model depends on the task and training, so measure it on your data. This experiment is about credit assignment over long distances; the small-data point in the next section is a separate argument and is not tested here.

Limited labelled data and long documents

• Small labelled sets. A Transformer's strength comes largely from pretraining (training on a large amount of unlabelled text first) and then fine-tuning (continuing training on your labelled task). If a pretrained checkpoint exists for your language and domain, fine-tuning it is usually the strongest option even with few labels. Trained from scratch on a few thousand examples, a Transformer has weak inductive bias (built-in assumptions about the data), whereas an RNN assumes order matters and processes one step at a time, which can make it more sample efficient (reaching a given accuracy with fewer labelled examples) in that setting, though not always; a simple baseline such as TF-IDF with logistic regression (word counts weighted by how distinctive each word is, fed to a linear classifier) should also be tried. Decide on your own validation set.
• Long documents. Full attention costs O(n^2) memory, while a recurrent state is fixed size. Mitigations on the Transformer side: chunk and aggregate, sliding-window or sparse attention, memory-efficient exact kernels, or retrieval of relevant passages.

When an RNN is still reasonable

1. Streaming input (sensor data, on-device speech) where a fixed-size state and constant per-step cost matter.
2. Very tight memory or compute budgets and short, simple sequences.
3. Small datasets with no pretrained model, as a baseline you should beat before adding complexity.
4. Research directions such as state-space and linear-recurrent layers (layers that update a fixed-size state with a linear rule, so training can run in parallel while inference runs step by step) combine a recurrent view for inference with parallel training; they are examples, not a replacement you can assume.

Common mistakes

• Saying Transformers are always better; the answer depends on data size, latency and memory.
• Claiming LSTMs solve vanishing gradients; they reduce it.
• Forgetting that Transformer decoding is not constant per step.
• Ignoring that parallelism is a training property; generating still goes one token at a time.

Running the code

Save the program as demo.py and run python demo.py. The output above was produced with Python 3.12, numpy 2.1.3 and torch 2.5.1 (CPU). The seed is fixed in the program.