🎓 A1Learner
🔍

PyTorch Attention Variants

Attention variants are changed forms of attention that score fewer pairs or score them more cheaply. Each variant trades something away, like reach or exactness, to buy speed on long sequences.

Count the pairs saved

Six positions score 36 pairs under full attention. A window of one neighbour on each side scores only 16 pairs.

import torch

positions = torch.arange(6)
distance = (positions.unsqueeze(0) - positions.unsqueeze(1)).abs()
local = distance <= 1
print("Full pairs:", 6 * 6)
print("Local pairs:", local.sum().item())

A sparse pattern

This mask lets every position score the first position and itself. Every other pair carries minus infinity and drops out of the mix.

sparse = torch.full((5, 5), float("-inf"))
sparse[:, 0] = 0.0
sparse.fill_diagonal_(0.0)
print(sparse)

What the window gives up

The last position holds its biggest score on the first position, far outside its window. The masked weights give that pair exactly zero.

scores = torch.tensor([5.0, 1.0, 0.0, 1.0, 0.0])
mask = torch.tensor([float("-inf"), float("-inf"), 0.0, 0.0, 0.0])
weights = torch.softmax(scores + mask, dim=0)
print("Weights:", weights)

Example

import torch

scores = torch.tensor([[2.0, 1.0, 0.0, 1.0, 0.0],
                       [1.0, 2.0, 1.0, 0.0, 1.0],
                       [0.0, 1.0, 2.0, 1.0, 0.0],
                       [1.0, 0.0, 1.0, 2.0, 1.0],
                       [0.0, 1.0, 0.0, 1.0, 2.0]])

positions = torch.arange(5)
distance = (positions.unsqueeze(0) - positions.unsqueeze(1)).abs()
local_mask = torch.where(distance <= 1, 0.0, float("-inf"))

weights = torch.softmax(scores + local_mask, dim=1)

print("Local mask:")
print(local_mask)
print("Local-window weights:")
print(weights)
print(f"Pairs scored per row: {(distance <= 1).sum(dim=1).tolist()}")

Key points

  • Full attention scores every pair, so its cost grows with the square of the length
  • Local attention lets each position score only a small window of neighbours
  • A local window cannot use a clue that sits outside the window
  • Sparse attention scores a fixed pattern of pairs, like the first position and each position itself
  • Linear attention rewrites the maths so the full score table is never built
  • Linear attention pays for that speed with an approximation of the softmax mix
  • In short, every variant saves work by giving up some pairs or some exactness

In this example, a local mask with a window of one neighbour is built from positions alone. The distance table holds how far apart each pair of positions sits. Pairs at distance 1 or less keep a mask value of 0, and farther pairs get minus infinity. Adding the mask to the scores and running softmax gives the printed weights. Every row has zeros outside its window, and each row scores at most 3 pairs. The last line prints the pair counts per row, which are 2 or 3 instead of 5.

Note

On short sequences full attention is already cheap, so variants earn their keep only as sequences grow

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does local attention give up in exchange for scoring fewer pairs?