DeepSeek Sparse Attention
Table of Contents
DeepSeek Sparse Attention (DSA) is introduced along with the DeepSeek V3.2 technical report. It’s designed to reduce computational complexity while preserving model performance in long context scenarios.
1. Prototype of DSA (DeepSeek V3.2)
In vanilla attention, each input query is processed to compute attention score with ALL history KV tokens. Although dense matrix multiplication can be boosted with FlashAttention techniques, the complexity is still \( O(n^2) \) where \(n\) is the length of context.
Reviewing the attention scoring, we found that it’s, in essense, selection of most relevant history tokens where the relevance is measured by the weight computed by softmax.
The intuition here is that there won’t be too much relevant tokens as the context grows. It means that only a few tokens have a meanful weight when performing \(O=PV\), while the weights of a majority of tokens are near \(0\). Thus we want some selector to find out the “a few most relevant tokens”. That is lightning indexer.
1.1. Lightning Indexer
Given a query token \(\mathbf{h}_t\in\mathbb{R}^d\). The lightning indexer computes the index score (which is another measure of relevancy) between the query token and ALL history keys. For some history token \(\mathbf{h}_s\in\mathbb{R}^d\), their indexing score \( I_{t,s} \) is given by
\[ I_{t,s}=\sum_{j=1}^{H^I} w^I_{t,j} \cdot \texttt{ReLU}\Big( \mathbf{q}_{t,j}^I, \mathbf{k}_s^I \Big) \]
- The \( H^I \) denotes the number of lightning indexer heads
- \( \mathbf{k}_s^I \in \mathbb{R}^{d^I} \) is derived from the preceeding token \( \mathbf{h}_s \)
- \( \mathbf{q}_{t,j}^I \in\mathbb{R}^{d^I}, w_{t,j}^I \in\mathbb{R} \) is derived from the query token \( \mathbf{h}_t \)
Other designs include:
- \(\texttt{ReLU}\) is used for throughput consideration
- DSA is implemented on MQA mode of MLA, where each latent vector will be shared across all query heads of the query token.
1.2. Fine-grained Token Selection Mechanism
Now that we have computed the index score for the query token against all history tokens, we next select key-value pairs of top-k history tokens. And apply the vanilla attention mechanism.
\begin{align*} \mathcal{D}' &\gets \set{ (\mathbf{k}_s, \mathbf{v}_s):I_{t,s} \in \texttt{TopK}(I_{t,:})} \\ \mathbf{o}_t &= \texttt{Attn} \Big( \mathbf{h}_t, \mathcal{D}' \Big) \end{align*}