Fast long-context for transformer attention mechanism

A method for use in determining context keys for an attention mechanism of a transformer model, the method comprising:

• generating a plurality of segments, each segment combining a plurality of respective token key mappings generated by mapping the respective token keys using a random feature matrix Ω∈n×d;
• generating a token query mapping by mapping the token query using the random feature matrix Ω;
• calculating attention scores between the token query mapping and each of the plurality of segments;
• and
• returning the token keys of m segments with the highest attention scores.