🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch Why Attention

Attention is a way for each step of a sequence to look back at every earlier step directly. An RNN instead carries the whole past inside one hidden vector, and that vector is fixed in size.

One vector carries 32 steps

An RNN reads 32 steps here, one value per step. The final hidden state is a single vector of 8 numbers.

import torch

torch.manual_seed(0)
rnn = torch.nn.RNN(1, 8, batch_first=True)
x = torch.rand(1, 32, 1)
out, hidden = rnn(x)
print("Steps read:", x.shape[1])
print("Final hidden shape:", hidden.shape)

The averaging baseline

The baseline guesses the first value with the mean of the sequence. Its squared error is large on this sequence.

seq = torch.tensor([0.80, 0.10, 0.40, 0.20])
guess = seq.mean()
print(f"First value: {seq[0].item():.2f}")
print(f"Mean guess: {guess.item():.4f}")
print(f"Squared error: {((guess - seq[0]) ** 2).item():.4f}")

Looking back directly

A weight is a share that says how much of a value to take. A weight of 1 on the first step and 0 elsewhere picks the first value out. The squared error drops to zero.

weights = torch.tensor([1.0, 0.0, 0.0, 0.0])
picked = (weights * seq).sum()
print(f"Picked value: {picked.item():.2f}")
print(f"Squared error: {((picked - seq[0]) ** 2).item():.4f}")

Example

import torch

def make(length, count, seed):
    gen = torch.Generator().manual_seed(seed)
    x = torch.rand(count, length, 1, generator=gen)
    y = x[:, 0, :].clone()
    return x, y

class RecallRNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = torch.nn.RNN(1, 8, batch_first=True)
        self.head = torch.nn.Linear(8, 1)

    def forward(self, x):
        out, _ = self.rnn(x)
        return self.head(out[:, -1, :])

def train_and_score(length, steps):
    torch.manual_seed(1)
    train_x, train_y = make(length, 32, 10)
    test_x, test_y = make(length, 64, 20)
    model = RecallRNN()
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    loss_fn = torch.nn.MSELoss()
    for step in range(steps):
        loss = loss_fn(model(train_x), train_y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        rnn_loss = loss_fn(model(test_x), test_y).item()
        avg_loss = loss_fn(test_x.mean(dim=1), test_y).item()
    return rnn_loss, avg_loss

short_rnn, _ = train_and_score(5, 60)
long_rnn, long_avg = train_and_score(32, 60)
print(f"Window 5 - RNN loss: {short_rnn:.4f}")
print(f"Window 32 - RNN loss: {long_rnn:.4f}")
print(f"Window 32 - Averaging loss: {long_avg:.4f}")

Key points

  • Hidden state: the one vector an RNN updates at each step and carries forward
  • The hidden vector keeps a fixed size, however long the sequence runs
  • Early values must survive every later update inside that one vector
  • Attention skips the carrying and reads the earlier steps directly
  • Averaging baseline: a guess that predicts the mean of the sequence, with no learning at all
  • In short, attention trades one cramped memory for a direct look back at every step

In this example, the task is to report the first value of a sequence at the very end. make builds random sequences, and the target for each one is its first value. RecallRNN reads a sequence with an RNN and predicts from the final hidden state. train_and_score trains for 60 steps and then scores on fresh sequences. At window 5 the RNN loss is 0.0208, so the first value survives four updates well. At window 32 the RNN loss climbs to 0.1692. The averaging baseline scores 0.0850 on the same long windows. The trained RNN falls behind a guess that learns nothing, because 31 updates overwrite the first value.

Note

Attention still processes the steps in order, but no step must be remembered through all the steps between

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the RNN loss grow on the long window?