「Attention 就是一個核平滑器。」
第一次看到這句話的時候我以為是某種比喻，畢竟「核方法」聽起來像是 SVM 那個年代的事，跟 Transformer 沒什麼關係。後來自己把公式展開一次，才發現它不是在打比方。
給定查詢矩陣 Q ∈ R n × d Q \in \mathbb{R}^{n \times d} Q ∈ R n × d K ∈ R m × d K \in \mathbb{R}^{m \times d} K ∈ R m × d V ∈ R m × d v V \in \mathbb{R}^{m \times d_v} V ∈ R m × d v
Attn ( Q , K , V ) = softmax ( Q K ⊤ d ) V . \operatorname{Attn}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right) V . Attn ( Q , K , V ) = softmax ( d Q K ⊤ ) V . 看整體沒感覺，把它逐項展開。第 i i i
o i = ∑ j = 1 m exp ( q i ⊤ k j d ) v j ∑ j = 1 m exp ( q i ⊤ k j d ) . o_i = \frac{\sum_{j=1}^{m} \exp\!\left(\frac{q_i^{\top} k_j}{\sqrt{d}}\right) v_j}{\sum_{j=1}^{m} \exp\!\left(\frac{q_i^{\top} k_j}{\sqrt{d}}\right)} . o i = ∑ j = 1 m exp ( d q i ⊤ k j ) ∑ j = 1 m exp ( d q i ⊤ k j ) v j . 左邊的分子是一堆 v j v_j v j q i q_i q i k j k_j k j
如果把 κ ( q , k ) = exp ( q ⊤ k / d ) \kappa(q, k) = \exp\!\left(q^{\top}k / \sqrt{d}\right) κ ( q , k ) = exp ( q ⊤ k / d )
o i = ∑ j κ ( q i , k j ) v j ∑ j κ ( q i , k j ) , o_i = \frac{\sum_j \kappa(q_i, k_j)\, v_j}{\sum_j \kappa(q_i, k_j)} , o i = ∑ j κ ( q i , k j ) ∑ j κ ( q i , k j ) v j , ——非參數回歸裡的 Nadaraya–Watson 估計。一模一樣。
差別只在於：這裡的「樣本」是模型自己學出來的，不是給定的資料集。
不是為了數學上的優雅。是因為它會告訴你哪裡可以砍 。
如果 κ \kappa κ κ ( q , k ) ≈ ϕ ( q ) ⊤ ϕ ( k ) \kappa(q,k) \approx \phi(q)^{\top}\phi(k) κ ( q , k ) ≈ ϕ ( q ) ⊤ ϕ ( k )
o i ≈ ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) v j ⊤ ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) . o_i \approx \frac{\phi(q_i)^{\top} \sum_j \phi(k_j) v_j^{\top}}{\phi(q_i)^{\top} \sum_j \phi(k_j)} . o i ≈ ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) v j ⊤ . 右邊的兩個累加量可以事先算好，每個查詢只要查表。複雜度從 O ( n m ) O(nm) O ( nm ) O ( n + m ) O(n + m) O ( n + m )
標準 attention : O ( n m d ) 線性 attention : O ( n d 2 + m d 2 ) \begin{aligned}
\text{標準 attention} &: \quad O(n m d) \\
\text{線性 attention} &: \quad O\!\left(n d^2 + m d^2\right)
\end{aligned} 標準 attention 線性 attention : O ( nm d ) : O ( n d 2 + m d 2 ) 這就是 linear attention 那一整家族在做的事。
這裡有個很容易被忽略的地方：誤差最大的地方不是分子，是分母。
分子是「加權和」，近似得差一點，輸出只是稍微偏掉；分母是「歸一化常數」，它一錯，整個輸出會跟著縮放。序列越長，ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) \phi(q_i)^{\top}\sum_j \phi(k_j) ϕ ( q i ) ⊤ ∑ j ϕ ( k j ) ∑ j κ ( q i , k j ) \sum_j \kappa(q_i,k_j) ∑ j κ ( q i , k j )
所以我們在評估這類方法時，看的第一個數字不是整體 perplexity，而是分母的相對誤差。
檢查 為什麼 分母的相對誤差 比分子更致命，因為它會整體縮放輸出 特徵維度 d ϕ d_\phi d ϕ 太少則近似粗糙，太多則省不下計算 序列長度分佈 短序列幾乎沒有收益，長序列才值得換
# 概念示意：先把 K、V 的累加量算好
kv = einsum("m d, m e -> d e", phi(K), V) # O(m d^2)
z = phi(K).sum(dim=0) # O(m d)
out = einsum("n d, d e -> n e", phi(Q), kv) / einsum("n d, d -> n", phi(Q), z)
如果你正在為長序列煩惱，我會建議先把 attention 當成核平滑器看一遍，再決定要不要換模型。
因為換成 linear attention 不是「更快的 attention」，而是用特徵映射換掉精確的歸一化 ——你付出的是準確度，不是免費的午餐。
知道自己在付什麼，比知道自己在用什麼重要。
