NeurIPS 2025

pLSTM: parallelizable Linear Source Transition Mark
networks
Korbinian Pöppel1,2
Richard Freinschlag1
Thomas Schmied1
Wei Lin1
Sepp Hochreiter1,3
1 ELLIS Unit Linz, LIT AI Lab, Institute for Machine Learning
Johannes Kepler University Linz, Austria
2 Zuse School ELIZA
3 NXAI GmbH
poeppel@ml.jku.at
Abstract
Modern recurrent architectures, such as xLSTM and Mamba, have recently chal-
lenged 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 linegraph 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 dif-
fusive distribution mode (D-mode). To showcase the long-range capabilities of
pLSTM, we introduce arrow-pointing extrapolation as a synthetic computer vi-
sion task that contains long-distance directional information. We demonstrate
that pLSTMs generalize well to larger image sizes, whereas Transformers strug-
gle to extrapolate. On established molecular graph and computer vision bench-
marks, pLSTMs also show strong performance. The complete code is available at
https://github.com/ml-jku/plstm_ex...