Why Attention — Query-Dependent Source Access

From one source state to a query-dependent read

A fixed-state encoder-decoder must route all source information through its terminal state. Attention retains the sequence of encoder representations and lets each decoder step retrieve a different weighted context. This removes the requirement that one shared terminal vector carry everything, not every bottleneck or every long-sequence failure. Each context is still a finite weighted sum; computation, representation quality and training data remain constraints.

Bahdanau, Cho and Bengio's Neural Machine Translation by Jointly Learning to Align and Translate was first submitted to arXiv on September 1, 2014 (16:33:02 UTC), and appeared at ICLR 2015. It predates the September 10 Sutskever submission. The previous whiteboard anecdote, December-upload timeline and universal “60-word sentences as accurate as 10” claim were not supported. The original system uses a bidirectional recurrent encoder and a recurrent decoder; the unidirectional LSTM below is a teaching variant, not its exact reproduction.

Goals: derive additive attention, reproduce the three-state arithmetic, implement padding-safe source access, compare dot/bilinear/additive scoring, inspect alignment without calling it a causal explanation, and situate faster attention and recurrence accurately. Prerequisites: the complete seq2seq lab, vector weighted sums, softmax and a two-layer MLP.

Score, normalize, retrieve and update

Let encoder outputs be h_i of width H_e and the decoder query s of width H_s. An alignment hidden width A need not equal either. Parameters W_h, W_s and v have shapes (A,H_e), (A,H_s), and (A). Optional biases must be specified; the code uses a bias in the additive intermediate projection.

Only valid source indices belong to V. Mask invalid scores to negative infinity before softmax; a score of zero still contributes exp(0)=1. Reject an all-invalid row rather than letting all negative infinities produce NaNs. Feed [embedding(previous token);c] into an LSTMCell with both old hidden and cell states, then project [new hidden;c] to target logits. That is one well-defined wiring; Bahdanau-style scoring does not force every later implementation to use that wiring or “six matrices.”

Softmax is differentiable and gives nonnegative unit-sum weights. For fixed finite value vectors, the context lies in their convex hull and its norm is at most the largest value norm. It need not retain exactly one state's norm or remain unchanged when new values are added. Softmax is not the unique possible attention normalization: sparse alternatives and unnormalized/kernel attention are different designs. Hard selection can be trained by other estimators; its lack of an ordinary argmax derivative does not prove it impossible.

The three-state arithmetic and a mask counterexample

With h1=[1,0], h2=[0,1], h3=[1,1], s=[0.5,0.5], identity projections and v=[1,1], scores are approximately [1.367265411,1.367265411,1.810296507]. Softmax is approximately [0.281103,0.281103,0.437795], so context is approximately [0.718897,0.718897]. Values rounded at three decimals in the old example explain small differences; use full precision in tests.

import math, os
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from seq2seq_lab import (PAD,BOS,EOS,Seq2Seq,make_vocab,make_batch,reverse_pairs,
                         train_model,evaluate_lengths,number_words,load_local_tsv)
 
h=torch.tensor([[1.,0.],[0.,1.],[1.,1.]],dtype=torch.float64)
s=torch.tensor([.5,.5],dtype=torch.float64)
e=(h+s).tanh().sum(-1);alpha=e.softmax(-1);context=alpha@h
assert torch.allclose(context,torch.tensor([.7189,.7189],dtype=torch.float64),atol=1e-4)
assert alpha.argmax()==2 and torch.allclose(alpha.sum(),torch.tensor(1.,dtype=torch.float64))
wrong=torch.tensor([1.,0.]).softmax(-1)
right=torch.tensor([1.,-torch.inf]).softmax(-1)
assert wrong[1]>0 and right[1]==0
print('three-state scores/weights/context',e.tolist(),alpha.tolist(),context.tolist())

The gradient does not guarantee that raising one state's relevance under one loss will monotonically increase its weight next epoch: shared parameters affect all scores and the optimizer changes the rest of the model too. A useful direct derivative is dc/de_i=alpha_i(h_i−c), showing both the weight and how different that value is from the current context.

Complete attention-augmented recurrent model

Save these fences together as attention_lab.py beside seq2seq_lab.py. Source lengths drive packing and the attention mask independently. Packed final states prevent decoder initialization from depending on trailing padding. Unpacking with total_length preserves the padded tensor's width so mask/output alignment stays explicit. Project encoder keys once per encoded source, not every decoder timestep.

class EncoderAll(nn.Module):
    def __init__(self,vocab,emb=16,hidden=24):
        super().__init__();self.embed=nn.Embedding(vocab,emb,padding_idx=PAD)
        self.rnn=nn.LSTM(emb,hidden,batch_first=True)
    def forward(self,src,lengths):
        if (lengths<=0).any() or (lengths>src.size(1)).any():raise ValueError('invalid lengths')
        packed=pack_padded_sequence(self.embed(src),lengths.cpu(),batch_first=True,enforce_sorted=False)
        out,(h,c)=self.rnn(packed)
        out,_=pad_packed_sequence(out,batch_first=True,total_length=src.size(1))
        return out,(h[0],c[0])
 
class SourceAttention(nn.Module):
    def __init__(self,encoder_width,query_width,align_width=24,kind='additive'):
        super().__init__();self.kind=kind;self.query_width=query_width
        if kind=='additive':
            self.keys=nn.Linear(encoder_width,align_width,bias=False)
            self.query=nn.Linear(query_width,align_width,bias=True)
            self.score=nn.Linear(align_width,1,bias=False)
        elif kind=='bilinear':self.keys=nn.Linear(encoder_width,query_width,bias=False)
        elif kind in ('dot','scaled-dot'):
            if encoder_width!=query_width:raise ValueError('dot dimensions must agree')
            self.keys=nn.Identity()
        else:raise ValueError('unknown attention kind')
    def forward(self,values,keys,query,valid):
        if valid.dtype!=torch.bool or valid.shape!=values.shape[:2] or not valid.any(-1).all():
            raise ValueError('each source needs at least one valid key')
        if self.kind=='additive':scores=self.score((keys+self.query(query)[:,None]).tanh()).squeeze(-1)
        else:
            scores=torch.bmm(keys,query[:,:,None]).squeeze(-1)
            if self.kind=='scaled-dot':scores=scores/math.sqrt(self.query_width)
        weights=scores.masked_fill(~valid,-torch.inf).softmax(-1)
        return torch.bmm(weights[:,None],values)[:,0],weights
 
class AttentionSeq2Seq(Seq2Seq):
    def __init__(self,vocab,emb=16,hidden=24,kind='additive'):
        nn.Module.__init__(self)
        self.encoder=EncoderAll(vocab,emb,hidden)
        self.attn=SourceAttention(hidden,hidden,hidden,kind)
        self.embed=nn.Embedding(vocab,emb,padding_idx=PAD)
        self.cell=nn.LSTMCell(emb+hidden,hidden);self.head=nn.Linear(2*hidden,vocab)
        self.recorded=[]
    def start(self,src,lengths):
        outputs,state=self.encoder(src,lengths)
        valid=torch.arange(src.size(1),device=src.device)[None]<lengths.to(src.device)[:,None]
        self.recorded=[]
        return state,(outputs,self.attn.keys(outputs),valid)
    def step(self,previous,state,memory):
        outputs,keys,valid=memory
        context,alpha=self.attn(outputs,keys,state[0],valid)
        state=self.cell(torch.cat([self.embed(previous),context],-1),state)
        self.recorded.append(alpha.detach())
        return self.head(torch.cat([state[0],context],-1)),state
 
torch.manual_seed(35)
attn=SourceAttention(2,2,2).double()
with torch.no_grad():
    attn.keys.weight.copy_(torch.eye(2));attn.query.weight.copy_(torch.eye(2));attn.query.bias.zero_();attn.score.weight.fill_(1)
values=h[None].clone().requires_grad_();query=s[None].clone().requires_grad_()
ctx,weights=attn(values,attn.keys(values),query,torch.ones(1,3,dtype=torch.bool))
assert torch.allclose(ctx[0],context,atol=1e-12)
assert torch.autograd.gradcheck(lambda v,q:attn(v,attn.keys(v),q,torch.ones(1,3,dtype=torch.bool))[0],(values,query))
try:attn(values,attn.keys(values),query,torch.zeros(1,3,dtype=torch.bool))
except ValueError:pass
else:raise AssertionError('all-masked rows must be rejected')

Dot attention is q dot k and requires equal projected widths; bilinear attention learns a matrix between spaces; additive attention learns a nonlinear compatibility function. Luong compares dot/general/concat and global/local attention. Scaled dot attention is the later Transformer convention, not Luong's original definition. Under independent zero-mean unit-variance query/key coordinates, dot-product variance is d; division by square root of d makes it 1. Correlated/nonunit data need not obey that model. Saturation can reduce softmax derivatives everywhere, including the winning coordinate; it is not a guaranteed small-d success/large-d failure threshold.

Choose the score by the compatibility you need and the computation you can batch. Additive scoring permits different encoder/query widths and learns nonlinear compatibility; dot scoring after suitable projections can batch all query–key pairs into a matrix multiplication. Bilinear scoring also reconciles different widths, but with a learned linear interaction rather than tanh. Different widths therefore do not force additive attention. Measure task quality and actual runtime rather than assuming all three are equally accurate or that one is always faster; the tiny comparison below establishes working contracts, not a preferred architecture.

Train, compare and export an honest alignment heatmap

The same imported trainer supplies shifted targets, finite loss checks, clipping, disjoint validation and best-state selection. The default run is intentionally tiny and synthetic, not the old fabricated number-word curve. For number-word training, construct the 0–999 pairs and disjoint splits as in S034. For local translation, load separate training/validation/test TSVs, build vocab only from training and evaluate test once after settings are chosen. If reporting BLEU, state the corpus, reference translations, tokenization and scoring implementation/configuration; a bare score is not a universal threshold for good translation.

def export_heatmap(matrix,source_labels,target_labels,path):
    from html import escape
    a=matrix.detach().cpu()
    if a.shape!=(len(target_labels),len(source_labels)):raise ValueError('label/matrix mismatch')
    w=70;left=100;top=60
    parts=[f'<svg xmlns="http://www.w3.org/2000/svg" width="{left+w*a.size(1)}" height="{top+35*a.size(0)}">',
           '<rect width="100%" height="100%" fill="white"/>']
    for j,label in enumerate(source_labels):parts.append(f'<text x="{left+j*w+3}" y="30" font-size="12">{escape(label)}</text>')
    for i,label in enumerate(target_labels):
        parts.append(f'<text x="4" y="{top+i*35+20}" font-size="12">{escape(label)}</text>')
        for j in range(a.size(1)):
            v=float(a[i,j]);shade=int(255*(1-v))
            parts.append(f'<rect x="{left+j*w}" y="{top+i*35}" width="68" height="33" fill="rgb({shade},{shade},255)"/>')
            parts.append(f'<text x="{left+j*w+3}" y="{top+i*35+20}" font-size="12">{v:.2f}</text>')
    Path(path).write_text(''.join(parts)+'</svg>')
 
if __name__=='__main__':
    pairs=reverse_pairs();train,val,test=pairs[:56],pairs[56:68],pairs[68:]
    itos,stoi=make_vocab(train);reports={}
    for kind in ('fixed','additive','dot','scaled-dot','bilinear'):
        torch.manual_seed(35)
        model=Seq2Seq(len(itos)) if kind=='fixed' else AttentionSeq2Seq(len(itos),kind=kind)
        history=train_model(model,train,val,stoi,epochs=3)
        reports[kind]={'last':history[-1],'lengths':evaluate_lengths(model,test,stoi)}
        if kind=='additive':
            x,lens,y=make_batch([test[0]],stoi);model.eval()
            with torch.no_grad():model(x,lens,y)
            alpha=torch.stack(model.recorded,1)[0]
            assert torch.allclose(alpha.sum(-1),torch.ones(alpha.size(0)),atol=1e-6)
            export=os.environ.get('ATTENTION_FIGURE_PATH')
            if export:export_heatmap(alpha,test[0][0].split(),[itos[t] for t in y[0,1:].tolist()],export)
            xp=F.pad(x,(0,3),value=PAD)
            with torch.no_grad():out1=model(x,lens,y);out2=model(xp,lens,y)
            assert torch.allclose(out1,out2,atol=1e-6)
            assert torch.stack(model.recorded,1)[...,-3:].eq(0).all()
    print('synthetic comparative contracts, no required ranking',reports)

The exported heatmap is teacher-forced for one held-out synthetic reversal example; its target labels include EOS and rows are not claimed to be free-running translation alignment. For generation, record the actually emitted labels and stop collecting on EOS. A one-bright-column heatmap may reveal a mask with only one valid key; a bright padded column is a mask error. Diagonal structure is a hypothesis for monotonic alignments, not a universal check: reversal should seek an anti-diagonal relation, while source encodings are already contextual.

Attention is a mixing trace, not automatically an explanation

The query is the reader's state, keys address candidate sources and values carry retrieved content. Cross-attention draws queries and key/values from different streams; self-attention draws them from one stream through possibly different projections. The mechanism generalizes associative lookup, but a Transformer also needs position information, feed-forward transformations, normalization and residual paths. Causal self-attention cannot read future keys.

A large weight does not imply a source token caused the decision: values may have different magnitudes, already encode other tokens, and pass through downstream nonlinearities. If two values are equal, many weight distributions produce exactly the same context. Jain/Wallace demonstrate problems with explanatory claims; Wiegreffe/Pinter argue that conclusions depend on the definition and rigorous controls. Neither licenses declaring every heatmap useless or every diagonal causal.

# Equal values: radically different attention distributions, identical context.
v=torch.tensor([[2.,-1.],[2.,-1.]])
a=torch.tensor([.99,.01]);b=torch.tensor([.01,.99])
assert torch.equal(a@v,b@v)
# Low weight, large value: ablating the lower-weight value changes output more.
v=torch.tensor([[.01,0.],[0.,100.]])
a=torch.tensor([.9,.1]);base=a@v
high_removed=a@torch.stack([torch.zeros(2),v[1]])
low_removed=a@torch.stack([v[0],torch.zeros(2)])
assert (base-low_removed).norm()>(base-high_removed).norm()
print('attention-weight explanation negative controls: pass')

For a real model, intervene on values, compare uniform attention, random seeds and alternative distributions, and measure logit/KL or task changes. Holding weights fixed during value ablation tests a different causal intervention from deleting a token and recomputing the entire network. Sparse-autoencoder features and attribution methods in Scaling Monosemanticity and Ferrando et al.'s primer offer other lenses; they do not make attention alone a proof of mechanism.

Faster attention, context extensions and recurrence alternatives

FlashAttention tiles exact softmax attention to reduce HBM traffic and avoid storing the full quadratic matrix; pairwise dense arithmetic remains quadratic. FlashAttention-2 improves work partitioning and parallelism; FlashAttention-3 targets Hopper asynchrony and low precision. The paper's hardware utilization figures are measured settings, not a guaranteed speedup for this CPU code. PyTorch 2.9 SDPA selects supported backends; it does not promise every call uses FlashAttention. Functional dropout must explicitly be zero in evaluation.

The key to exact tiling is online normalization, not taking a separate softmax in each tile. For one query, retain the maximum score m, exponential sum l, and unnormalized weighted-value sum u over the keys seen so far. When a tile raises the maximum to m_new, multiply the old l and u by exp(m−m_new), then add the new tile's exp(score−m_new) contributions; the final context is u/l. Masked keys contribute nothing; skip an all-masked tile and reject a row with no valid key overall. For scores [0, log(3)] and scalar values [2,6], the correct result is (2+3×6)/(1+3)=5. Averaging separately normalized one-key tiles gives (2+6)/2=4: local normalization loses the tiles' relative mass. FlashAttention §3.1 uses this rescaling and recomputes intermediates during the backward pass rather than storing the quadratic attention matrix. “Exact” means the same mathematical attention operator, not bitwise equality across floating-point kernels. Backend eligibility depends on device, dtype, shape and other input constraints; PyTorch 2.9's SDPA contract still requires explicit dropout_p=0.0 for evaluation.

Ring Attention distributes source blocks across devices and overlaps communication with blockwise attention. Blockwise Parallel Transformer fuses blockwise attention/FFN computation to reduce memory. Neither primary paper establishes that Claude or Gemini privately uses that implementation. Sharding compute does not prove model quality at any supplied context length.

RoPE rotates queries/keys to encode relative offsets for fixed content vectors. Position Interpolation rescales positions; YaRN treats frequency bands differently and adjusts attention scaling; LongRoPE searches nonuniform interpolation and uses progressive extension with short-context readjustment. These are not universally training-free. MLA caches a compressed content latent plus a separate positional key component; its memory advantage depends on actual dimensions and cache representation, not a universal 10×.

Mamba uses input-dependent selective SSM updates and a hardware-aware recurrent scan. Mamba-2 develops structured state-space duality and different efficient algorithms, not just “an RNN parallelized for free.” Jamba-1.5 is a Transformer–Mamba–MoE hybrid, so its attention layers still have cache costs. RWKV-7 introduces generalized delta-rule state evolution; its paper's models span 0.19–2.9B, not an established universal Llama-3-8B match. xLSTM adds exponential gates and scalar/matrix memory; xLSTM-7B's report is from 2025, not a September-2024 release deduced from the earlier abstract. Griffin combines gated linear recurrence and local attention. Fixed-state recurrence trades explicit retrieval bandwidth for bounded state; exact finite workloads, precision and memory capacity determine whether that trade helps.

These alternatives answer the same information-routing question with different operators. They are not all softmax weighted sums over stored history. The historical chain from seq2seq/attention to Transformer, GPT-style models and chat post-training is useful context, not disclosure of GPT-4/Claude internals. GPT-4's report deliberately omits architecture specifics. Building this small model does not reproduce it.

Practice and worked answers

1. Four operations. Score with v transpose tanh(W_h h+W_s s+b), mask then softmax over valid source positions, compute the weighted sum, and update the decoder from previous token/context/state before its logit projection. LSTM state is a pair.
2. Three-state example. State 3 receives the largest score because both tanh inputs are 1.5. Its weight is about 0.438, producing context about [0.719,0.719], not hard selection of h3.
3. Compare scores. Additive learns nonlinear compatibility, dot is an inner product, general/bilinear inserts a learned matrix. Transformer scales dot products by square root of key width under a variance model; this scaling is not in the original Luong dot definition.
4. Why negative infinity? exp(0) is positive, so zero is not exclusion. Mask before normalization and reject all-invalid sources. Correct key masking alone does not repair a padding-contaminated initial decoder state; use valid-length encoding too.
5. Q/K/V. Query asks, keys score and values contribute content. Self-attention uses one source stream for all three projections; cross-attention uses the decoder query and encoder memory. Residuals, FFNs and positional treatment are additional components.
6. Mamba tradeoff. Input-selective recurrent state can retain useful summaries with state size independent of history length, but cannot access arbitrary stored keys in the same manner as dense attention. Test length extrapolation and content retrieval separately; linear sequence scaling alone proves neither equal quality nor inevitable failure.

The stretch task is two self-attention encoder blocks with valid padding/position treatment feeding a recurrent decoder. Taking only the final pooled vector reintroduces a fixed interface; keep all outputs for genuine cross-attention. Compare parameter count, length bins and wall-clock measurements before describing an improvement. The training and heatmap code above replaces unattached plotting and ablation pseudocode with executable contracts.

Primary sources and optional continuation

Additive attention, Luong, Transformer v7, FlashAttention, Ring, BPT, MLA, Mamba, Jain/Wallace, Wiegreffe/Pinter. Claim-specific passages for the remaining named methods are in the accompanying evidence ledger, not inferred from a link list.

Optional Serrano and CodeEmporium videos retain their resource role. Previous · Next: self-attention.

Execution is bounded synthetic CPU training, gradients, masking, mixtures and an exported numerical heatmap. No translation corpus, pretrained model, GPU FlashAttention, distributed Ring or long-context benchmark was executed. Source-supported conceptual review and end-to-end release acceptance are separate.