ICML 2024: Paper Review #1

Vassilis Papadopoulos, Jérémie Wenger, Clément Hongler

Large Language Models (LLMs) typically model the probability of observing a token given the previous tokens, generating text autoregressively. The product of these conditional probabilities forms the joint probability over the sequence, representing the likelihood of observing the entire sequence of tokens.

The authors of the paper observe that this probability distribution can also be learned in a right-to-left fashion, by predicting the previous token instead of the next one. This raises the question: does modelling left-to-right produce a better joint distribution estimate? If so, why?

The authors denote this phenomenon as a “forward arrow of time” when the left-to-right model consistently outperforms the right-to-left model. Their main result is the identification of a consistent forward arrow of time in language models across various architectures, model sizes and languages, with the performance gap increasing as model size and capability increase.

They argue that natural language inherently possesses a certain degree of sparsity. In the forward direction, this sparsity is less pronounced due to the natural progression of language being more predictable and structured. This reduced sparsity makes it easier for language models to predict the next token, leading to better performance (lower perplexity). Conversely, predicting previous tokens in the backward direction encounters higher sparsity and less predictable structures, resulting in higher perplexity.

To illustrate this, the authors construct artificial datasets where one direction is easier to model. For instance, they create a dataset of random prime numbers followed by their product. Modelling this sequence in reverse, starting with the product, would require the model to perform prime factorization, a computationally complex task. This example demonstrates a strong forward arrow of time, emphasising the computational complexity aspect.

The authors hope these insights can inform improved training procedures that explicitly account for the arrow of time present in the dataset. This could lead to more efficient and effective language models.

Tri Dao and Albert Gu

The main contributions of this paper are two-fold: a new framework for reasoning about sequence models, and a new architecture, Mamba-2, which improves on Mamba.

The authors’ new framework, termed State Space Duality (SSD), demonstrates that Transformers and SSMs can be understood as specific instances of a broader class of models. This framework represents sequence models as a single matrix multiplication on a sequence, where the “sequence-mixing” matrix may depend on the input.

The authors show that SSMs, viewed in this light, produce a sequence-mixing matrix that is structured and semiseparable. It is lower triangle and every submatrix in this lower triangle is of low-rank. This property allows the matrix to be factorised and efficiently multiplied with sub-quadratic complexity. This is the basis of an efficient algorithm for training such models and leads to greatly improved computational performance.

Part of the improved performance, relies on the fact that previous SSMs utilised scan algorithms, which are harder to implement efficiently on GPUs. The SSD framework, which generalises previous SSMs, allows for the model to be computed via batched matrix multiplications instead. Furthermore, NVIDIA GPUs have specialised so called “tensor cores” which are used for fast tiled matrix multiplications, and as a result, this reparameterisation leads to much better throughput.

The authors show this empirically as the SSD implementation is up to 8x faster and is particularly well-suited to large state sizes which are important for some tasks. Moreover, the SSD framework allows for more natural distributed formulations that use fewer collective operations per layer, further improving the performance of large-scale training.

The authors note that for sufficiently short sequence lengths, around 2048, fused Attention implementations (such as FlashAttention) are still slightly faster as they take advantage of specific features introduced by NVIDIA in the Hopper generation of GPUs. They expect this gap to decrease as the SSD implementation improves, and, crucially, for large sequences, the sub-quadratic complexity makes SSDs much faster already. For example, at a sequence length of 32k, the performance gap is more than 10x between SSDs and FlashAttention.

In summary, this paper extends the previous SSM literature with some favourable computational properties and introduces a general framework for reasoning about sequence models.