Linear Self-Attention

原始题目：LeetGPU - Linear Self-Attention

题目描述

实现论文 "Transformers are RNNs" 中提出的线性注意力。给定 $Q$、$K$、$V$（均为 $M \times d$），计算：

$$\text{LinearAttention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)\left(\sum_j \phi(K_j)\right)}$$

其中 $\phi(\cdot) = \text{elu}(\cdot) + 1$（elu 加 1 确保非负）。与标准 softmax attention（$O(M^2)$）不同，线性注意力利用 $\phi(Q)(\phi(K)^T V)$ 的计算顺序，将复杂度降为 $O(M \cdot d^2)$。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

约束条件

• $1 \le M \le 10{,}000$，$1 \le d \le 128$。
• 性能测试在 $M = 5{,}000$ 下进行。

解题思路

线性注意力的核心技巧是改变计算顺序：先算 $\phi(K)^T V$（$d \times d$），再与 $\phi(Q)$ 相乘——当 $d \ll M$ 时，这比 $O(M^2)$ 的 standard attention 快得多。$\phi$ 函数（elu+1）保证输出非负。但线性注意力缺少 softmax 的归一化效应，在建模长程依赖方面不如标准注意力。欢迎在 GitHub Discussions 分享你的解法。