# xLSTM parallel computation - mismatch in dimensions

In this recent paper, a new architecture is proposed, called xLSTM. I've implemented the sequential version in PyTorch, but it's slower than I would like, so I'm now implementing the parallel version that's explained in the appendix (page 25-26). I feel like this page might contain a mistake, or maybe I'm missing something, so I wanted to check here.

The issue is as follows. We consider an input sequence $$X\in\mathbb{R}^{T\times d}$$ and we obtain two matrices $$F, I\in\mathbb{R}^{T\times T}$$ (see the paper for details), which combine into $$D\in\mathbb{R}^{T\times T}$$. All well and good. Now, it is stated that $$Q, K, V\in\mathbb{R}^{T\times d}$$. I believe this should be $$\mathbb{R}^{T\times e}$$ where $$e$$ denotes the embedding dimension, but that is not the main source of confusion. The point is that we have the equation $$C=QK^T\odot D$$ in the paper, and subsequently, $$H$$ is obtained as $$CV$$. But $$QK^T\in \mathbb{R}^{T\times e\times e}$$, which makes sense since this is a sequence of linear transformations which will transform the sequence $$V\in\mathbb{R}^{T\times e}$$. But then the dimensions of $$QK^T$$ do not match those of $$D$$, so the equation $$QK^T\odot D$$ does not make sense to me.

Have I missed something here, or is there an issue with the equation in the paper?

$$Q \in R^{T \times d}, K \in R^{T \times d} \\ Q \cdot K^T \in R^{T \times d} \times R^{d \times T}\\ Q \cdot K^T \in R^{T \times T}$$
Which is exactly the dimension of $$D$$
• Except that does not make sense. $C$ has to be a sequence of matrices transforming $V$, as mentioned in the original post, so it needs to be of dimension $T\times d\times d$. Strictly speaking what you wrote is the correct way to do the matrix product, but I think it should be read as performing an outer product between $Q_t$ and $K_t^T$ for each $t\in [1, T]$. This is also what is done in the non-parallel case. Commented May 9 at 16:45