pLSTM: parallelizable Linear Source Transition Mark networks
Authors: name: Korbinian Pöppel, name: Richard Freinschlag, name: Thomas Schmied, name: Wei Lin, name: Sepp Hochreiter

Modern recurrent architectures, such as xLSTM and Mamba, have recently
challenged the Transformer in language modeling. However, their structure
constrains their applicability to sequences only or requires processing
multi-dimensional data structures, such as images or molecular graphs, in a
pre-defined sequential order. In contrast, Multi-Dimensional RNNs (MDRNNs) are
well suited for data with a higher level structure, like 2D grids, trees, and
directed acyclic graphs (DAGs). In this work, we extend the notion of
multi-dimensionality to linear RNNs. We introduce parallelizable Linear Source
Transition Mark networks (pLSTMs) using Source, Transition, and Mark gates that
act on the line graph of a general DAG. This enables parallelization in analogy
to parallel associative scans and the chunkwise-recurrent form of sequential
linear RNNs, but for DAGs. For regular grids (1D and 2D), like images, this
scheme can be efficiently implemented using einsum operations, concatenations,
and padding in logarithmic time. pLSTMs tackle the vanishing/exploding
activation/gradient problem for long distances in DAGs via two distinct modes:
a directed propagation mode (P-mode) and a diffusive distribution mode
(D-mode). To showcase the long-range capabilities of pLSTM, we introduce
arrow-pointing extrapolation as a synthetic computer vision task that contains
long-distance directional information. We demonstrate that pLSTMs generalize
well to larger image sizes, whereas Transformers struggle to extrapolate. On
established molecular graph and computer vision benchmarks, pLSTMs also show
strong performance. Code and Datasets are available at:
https://github.com/ml-jku/plstm_experiments.