The FFN as a key–value memory, next to attention

The FFN as a key–value memory, next to attention A workflow diagram generated by Archify. 01 / Attention: keys and values come from other tokens 02 / FFN: keys and values are fixed parameters Query q · projected from this token · Attention: keys and values come from other tokens Query q projected from this token q · Kᵀ · keys of other tokens · Attention: keys and values come from other tokens q · Kᵀ keys of other tokens softmax · weights summing to 1 · Attention: keys and values come from other tokens softmax weights summing to 1 · V · mixes context values · Attention: keys and values come from other tokens · V mixes context values Hidden state x · one token, d_model dims · FFN: keys and values are fixed parameters Hidden state x one token, d_model dims x · K · compared with d_ff keys · FFN: keys and values are fixed parameters · K = W_in x · K compared with d_ff keys K = W_in f(·) · ReLU, GELU: keep matches · FFN: keys and values are fixed parameters f(·) ReLU, GELU: keep matches · V · sums the active values · FFN: keys and values are fixed parameters · V = W_out · V sums the active values V = W_out Write to stream · boosts a group of tokens · FFN: keys and values are fixed parameters Write to stream boosts a group of tokens match also writes to the stream match match score coefficients Legend query from the token computation steps residual stream keys and values are parameters