pLSTM: parallelizable Linear Source Transition Mark networks
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.
684fd0cb60b4a34dbe007b70
Korbinian Pöppel
claimed_verified
2025-06-16T08:35:40.461Z
6333650673c07e8aebb2e941
/avatars/bfcc236641671e88c2fe5426740071d3.svg
Korbinian Poeppel
user
korbip
684fd0cb60b4a34dbe007b71
Richard Freinschlag
684fd0cb60b4a34dbe007b72
Thomas Schmied
684fd0cb60b4a34dbe007b73
Wei Lin
684fd0cb60b4a34dbe007b74
Sepp Hochreiter
pLSTM: parallelizable Linear Source Transition Mark networks
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.
pLSTMs are parallelizable linear RNNs designed for DAGs, demonstrating superior performance on long-range tasks and benchmarks compared to Transformers.
xLSTM Mamba Transformer Multi-Dimensional RNNs MDRNNs parallelizable Linear Source Transition Mark networks pLSTMs Source gates Transition gates Mark gates line graph DAGs parallel associative scans chunkwise-recurrent einsum operations concatenations padding vanishing/exploding activation/gradient problem directed propagation mode diffusive distribution mode arrow-pointing extrapolation computer vision task molecular graph performance benchmarks