Additive Attention

The problem, from first principles

Before the transformer, attention was invented to fix a bottleneck in translation: an encoder squeezed a whole sentence into one vector. Bahdanau's idea was to keep all encoder states and let the decoder look back at them with a learned relevance score. The score between a decoder query q and each encoder state k_t comes from a tiny one-hidden-layer network: v^T tanh(W_q q + W_k k_t). A softmax over the scores gives attention weights and the context vector is the weighted average of the encoder states. The transformer replaced this additive score with a cheaper dot product, but the weighted-average idea is identical.

From theory to code

Implement additive_attention.

Constraints

• query has shape (d_q,), keys has shape (T, d_k). W_q is (d_q, h), W_k is (d_k, h) and v is (h,).
• Scores are tanh(query @ W_q + keys @ W_k) @ v, shape (T,). Weights are a numerically stable softmax over the scores.
• The context vector is weights @ keys, shape (d_k,). Return (context, weights).

Hints