File size: 3,431 Bytes
7b61fa8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
In this section, we study the limitations of state-space duality from its two ends:
(i) on the attention side, we show the impossibility of extending SSD to practical softmax attention \cite{vaswani2017attention};
(ii) on the SSM side, we show that even SSMs with low state dimension or low-rank state matrices may not always admit an SSD extension.

\paragraph{Impossibility of Extending SSD to Softmax Attention.}
We present a simple example showing that softmax attention \cite{vaswani2017attention} does not admit state-space duality.
Consider a matrix $V \in \mathbb{R}^{T \times T}$ with entries $V_{i,j} = i \times j$ for all $i, j \in [T]$.
The matrix $V$ has rank $1$ since every column is a scalar multiple of $(1, 2, \dots, T)^\top \in \mathbb{R}^T$.
However, $\Softmax(V)$ has rank $T$ by the Vandermonde determinant, and in fact, every submatrix of $\Softmax(V)$ is full rank.

This shows that even when the attention matrix $QK^\top$ has very low rank, $\Softmax(QK^\top)$ typically expands to full rank $T$.
Moreover, any attention matrix with a state-space dual must be $N$-semiseparable, where $N$ is the state dimension of the corresponding SSM.
Hence, softmax attention does not have a state-space model dual.

\paragraph{Impossibility of Extending SSD to General SSM with Low State Dimension.}

We provide an example for this impossibility in the following proposition.
It shows that a general SSM, \textit{even} with low state dimension or low-rank state matrices, \textit{may not} admit a $1$-SS masked-attention dual.

\begin{proposition}
    Consider the SSM layer with state dimension $N \ge 2$ defined by \eqref{eqn:recurrence}, there exist $A^{1:T+1},b_{1:T+1},c_{1:T+1}$
    such that the recurrence relation does not have an attention dual.
\end{proposition}

\begin{proof}
    According to \cref{lem:N_semiseparable_equals_NSSS}, 
    there exist $A^{1:T+1},b_{1:T+1},c_{1:T+1}$
    such that the recurrence relation \eqref{eqn:recurrence} has representation
    $Y = M \cdot X$, where $M := I_{T \times T} + E^{T,1}$ is a $2$-SS matrix.
    Here
    \begin{align*}
        E^{T,1}_{j, i} := 
        \begin{cases}
            1, \quad & ~ j = T ~ \text{and} ~ i = 1;\\
            0, \quad & ~ \text{otherwise}.
        \end{cases}
    \end{align*}

    We claim that $M$ does not have the representation of $L \odot (QK^\top)$ where 
    $Q,K \in \R^{T\times N}$ and $L$ is a $1$-SS matrix.
    
    Otherwise if $M = L \odot (QK^\top)$, suppose 
    \begin{align*}
        L_{j, i} = \begin{cases}
        a_{j}\cdots a_{i+1}, &\quad\text{for}\quad j \ge i;\\
        0, &\quad\text{for}\quad j < i.
        \end{cases}
    \end{align*}
    for $i,j \in [T]$.

    Since $M_{T,1} =1$, $L_{T,1} = a_2\cdots a_T$ is none-zero, i.e. each of $a_2,a_3, \cdots, a_T$ is none-zero.
    
    Thus, we have
    \begin{align*}
        (QK^\top)_{i.j} = 0,
        \quad\text{for all}\quad
        1 \le j < i \le T-1.
    \end{align*}
    Since each diagonal element of $M$ is none-zero, each diagonal element of $QK^\top$ is also none-zero.
    Given that $(QK^\top)_{i.j} = 0$ for all $1 \le j < i \le T-1$,
    we deduce that $QK^\top$ has rank at least $T-1$.
    This is a contradiction to $\rank(QK^\top) \le \rank(Q) \le N$.

    This completes the proof.
\end{proof}
Therefore, general SSM does not always have $1$-SS masked attention dual, even when the SSM has very low state dimension.