Function std::linear_attention_forward

#include <std/ml/attention.h> source · definition

&FloatTensor linear_attention_forward(
    const &FloatTensor Q,
    const &FloatTensor K,
    const &FloatTensor V,
)

Kernel-trick reformulation (Katharopoulos et al., "Transformers are RNNs"): with feature map phi(x) = elu(x)+1 (elementwise, always positive), softmax attention's O(seqLenkvLenheadDim) cost becomes O((seqLen+kvLen)*headDim^2) by computing phi(K)^T @ V once ([headDim,headDim]) instead of the full [seqLen,kvLen] score matrix:

out_i = (phi(Q_i) . sum_j phi(K_j) (x) V_j) / (phi(Q_i) . sum_j phi(K_j))

Q: [seqLen,headDim], K/V: [kvLen,headDim] -> result: [seqLen,headDim].