pLSTM: parallelizable Linear Source Transition Mark networks
Authors: Korbinian Pöppel
Richard Freinschlag
Thomas Schmied
Wei Lin
Sepp Hochreiter
Venue: NeurIPS 2025 poster ()
Keywords: LSTM
MDRNN
linear RNN
DAG
multi-dimensional data
graphs
Primary area: deep_learning
TL;DR: Enhancing linear RNNs to multi-dimensional structures, stable and parallelizable.

Abstract: 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 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 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. The complete code is available at https://github.com/ml-jku/plstm_experiments.

PLSTM: parallelizable Linear Source Transition Mark networks
Korbinian Pöppel!2
Richard Freinschlag'
Thomas Schmied
Wei Lin'
Sepp Hochreiter
ELLIS Unit Linz, LIT AT Lab, Institute for Machine Learning
Johannes Kepler University Linz, Austria
Zuse School ELIZA
3 NXAI GmbH
poeppel@m1.jkuu
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 images molecular graphs, in a pre-defined sequential order. In contrast, Multi-Dimensional RNNs (MDRNNs) are well suited for data with higher level structure, like 2D grids, trees, and directed acyclic graphs (DAGs). In this work, we extend the notion of multi-diimenioonality to linear RNNs. We introduce parallelizable Linear Source Transition Mark networks (pLSTMs) using Source, Transition, and Mark gates that act on the linegraph of general DAG. This enables parallelization in analogy to parallel associative scans and the Hounkwise-recurres form of sequential linear RNNs, but for DAGs. For regular grids (ID and 2D), like images, this scheme can be efficiently implemented using einsum operations, concatenations, and padding in logarithmic time. pl SS tackle the vanishing/exploding eetivation/graddintt problem for long distances DAGs via two distinet modes: directed propagation mode (P-mode) and a dif- fusive distribution mode (D-mode). To showcase the long-range capabilities synthetic computer vi-
PLSTM, we introduce arrow-pointing extrapolation 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 hattps://githie com/ml jiu/plstt, wetpmmeeepeeaeaae.eee
1 Introduction
Linear RNNs such as DeltaNet [Schlag et al. 2021 Gated Linear Attention Yang et al. 2023
Mamba [Gu and Dao 2023 Dao and Gu| [2024] and xLSTM (mLSTM) [Beck al. (025b) have recently evolved as a powerful alternative to the Transformer (Vaswaan et al. 2017]. In addition to ggguence-parallllizatee
training, such modern recurrent architectures are more efficient inference time than Transformers. In contrast to classical RNNs, they come with associative memory expansion, enabling more performant sequence modeling. However, these modern recurrent architectures are inherently limited to sequences. While linear RNNs have shown good performance for multi- dimensional data such as images, they enforce a pre-defined sequential traversal structure [Zhu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

et al.. 2024, Alkin et al., 2025]. To illustrate, for images, this corresponds to processing patches in scanline form. However, this is problematic as, within one layer, even (vertical) neighbors in a patch grid are only distantly related, in the (horizontal line-wise) processing. Even when switching between line- and column-wise order across layers, diagonal relationships are not covered sufficiently. This mismatch of short-range spatial and long-distnn lingddistance relations in certain path with precise positioning requirements leads to credit assignment problems [Minsky 1961 Bengio and
Frasconi] [993 Schmidhuber, 2015] and suffers from vanishing activations/gradien Hochreiter
1991 Bengio et al. 1994
Hochreiter al. 2001 Hochreiter and Schmidhuber 1997 Pascanu et al.
2013
Non-Recurrent
Recurrent 1D
Recurrent in 2D
ViM ViL
2DMamba
VIT
PlainMamba
pLSTM (ours)
VMamba
V2M
GroupMamba
Mamba-ND
SSM2D
D-mode
CNN
VSSD
+
Mamba2D
fleeo 45"
m-mode
eernable
Skyral Strength
Inpull dep.
Illustration the receptive fields induced by PLSTM and related architectures (for a single layer). CNNs are locally restricted while ViT's have a global receptive field. Modern recurrent architectures, such as ViM, traverse the 2D grid sequentially. PLSTM effectively extends the receptive field via its combination of D-mode and P-mode.
Figure 1:
Multi-dimennsinal RNNs (MDRNNs) [Graves al. 2007 have demonstrated how non-linear RNNs like LSTMs [Hochreiter] 1991 | dimensional data, such images GGraes al., 2007 van den Oord
Hochreiter and Schmidhuber 1997 can be extended to multi-
2016]. trees [Tai al.) 2015), and DAGs [Peng et al. 2017]. In this work, we translate the findings of MDRNNs to linear RNNs and introduce parallelizable Linear Source Transition Mark Networks (pLSTM). To achieve this, we first translate their input, forget, and output gates to Source, Transition, and Mark gates, respectively. This is make the resulting structure applicable to general DAGs. Second, we derive parallelization scheme enabled by the previous adaptation, leading to oigher-order Source, Transition, and Mark terms forming a parallel associative scan in multiple dimensions and DAGs. This also allows for chunkwise-parrdl 2025a], using the parallelization up to certain level and the recurrent option between chunks/patches. Third, we derive two stabilization options, the P-mode for directional propagation and the D-mode for undirected global distribution. Figure [] we illustrate how the two modes affect the receptive field of pLSTM compared other linear RNNs. We use pLSTM with MLP layers in residual pre-LayerNorm backbone, resembling a Transformer with replaced multi-head attention.
formulation known from linear RNNs (Yang et al. 2023 Beck et al.,
To showcase the long-range capabilities of PLSTM in multiple dimensions, we first introduce the synthetic arrow-pointing extrapolation task (see Figure 2.. As the directional information is not tied to only testing horizontal and vertical cases separately, certain line-wise or column- ordering (even alternation) cannot capture the directional information in the sequential RNN transitions, at least for a limited number of layers. Importantly, PLSTM generalizes well to increasing image resolutions compared the Transformer (see Section[5.1] Moreover, we demonstrate the efficacy of PLSTM on established computer-vision and graph benchmarks. Experiments on ImageNet- lk (see Section 5.2) and on molecular graphs (see Section (54) show promising results compared to baselines and good scaling behavior to larger model sizes.

---

In summary, we make the following contributions:
We translate the findings classic MDRNNs to linear RNNs and introduce PLSTM, which comes with adapted gates and a scalable chunkwise-paraley formulation.
We formally derive the general stabilization of pLSTMs for long-range propagation on general DAGS (including images), establishing the P- and D-mode cases.
We introduce the synthetic Arrow Pointing Task to highlight the theoretical advantage pLSTMs, in which PLSTM shows strong extrapolation abilities, and provide highly scalable implementation of PLSTM.
2
Related work
Modern Recurrent Architectures
This work presents a new form of linear RNNs, where DeltaNet [Schlag et al. 2021 Yang et al. 2224]. , LRU [Orvicto et al. 2023 GLA Yang et al. 2023
Mamba [Gu and Dao 2023 Dao and Gu 2024], and xLSTM [Beck et al. 2025b] (in the ml MLSTM
form) have shown their effectiveness
sequence modeling in the language domain. Recently, this
line of work has been complemented by TTT (Sun et al. (024). Titans [Behrouz et al. 2024). and DeltaProduct (Siems et al., 2025], which motivate this structure context, in line with early work on Fast- Weight Programmers [Schlag et al. 2021]. Due to their ability for parallelization during training, modern recurrent architectures scale large-scale datasets similar to Transformers. Moreover, they come with efficient inference, which is attractive for real-world applications [Schmidinger et al. 2024 Schiff et al. 2024
gradient-based optimization in
Schmied et al. 2024].
Non-Linear
RNNs The first founda- tional extension of non-linear RNNs LSTMs to multiple
dimensions was carried out by Graves et al. [2007]. Subse-
quently, Stacked LSTM was proposed [Graves et al. 2013).
which stacks LSTM layers on top of each other. MDRNNs are hard parallelize, while re-arranging the traditional cuboid order of computations in MDLSTMs in pyramidal fashion led to PyraMiD LSTM (Stollenga et al. 2015], which can be par- allelized. Grid LSTM [Kalchbrenner et al. 20155 operates in multiple processing dimensions simultancously, thus general- izing one-dimensional LSTM to more dimensions. Tree-LSTM [Tai et al. [2015] and DAG-LSTM [Zhu et al.] 2016 Peng et al. Chen et al. 2017 extend MDRNNs to tree structures and DAGs. PixeIRNN Van Den Oord et al. 2016) operates directly on pixel-level and improves MDRNNs for images. Due to their strictly recurrent nature, these architectures are not paralleliz- able and therefore unsuitable for modern hardware. Moreover, they lack the associative memory component (state expansion) of modern linear RNNs, making them less powerful.
Figure Illustration of the Ar- row pointing task. The model has to classify whether an arrow is pointing towards a circle (top left). Models with global receptive fields, such Vision Transformers (ViTs), can solve this task by leveraging di- rectional information encoded via positional embeddings (top right), but they often struggle generalize to higher resolutions. In contrast, PLSTMs
Linear RNNs on
data VisionMamba (ViM) Zhuu al. 2024) and Vision-LSTM (ViL) AAlin et al.) 2025) applied linear RNNs to the vision domain to challenge the common Vision Transformer (ViT) with its quadratic scal- ing (Dosovitskiy et al. 2021]. These works rely on traversing the 2D plane in predefined order accommodate their sequen- tial nature and often employ flipping the traversal order across layers. This was also investigated further in Mamba-ND [Li et al. [025]. J. Similarly, several works use Mamba on graph data, covering the graph with multiple paths (Wang et al.] [024a]
effectively solve this task in both diffusive distribution (D-mode, bottom left) and directed propagation (P-mode, bottom right) modes, enabling long-range reason- ing and better scalability.
or forming sequence in other ways [Behrouz and Hashemi 2024), with an overview presented in Atitallah et al. [2024]. To
the best of our knowledge, all of these rely on sequential representation apply Mamba at some point, and therefore not use the graph's structural information for the recurrence. Recent works
extend the Mamba architecture to two dimensions. Most related to work are V2M [Wang et al.)
3

---

2024b] and 2DMamba Zhang et al. 2024a]. which cover only the D-mode of PLSTM, at the loss of directional propagation. Mamba2D Baty al., 2024] and 2-D SSM [Baron et 2024 cover the P-mode of pL.ST'M, but in contrast to PLSTM, they are limited to the diagonal line instead of input-directed propagation along any line because of their fixed normalization factor of 1/2 (see Appendix p.2p.
Relation to GNNs and MPNNS Since pLSTM propagates recurrent states along (directed acyclic) graph, it can be seen as form of Graph Neural Network (GNN) [Scarselli et al. 2009 that covers whole connected component of DAG per layer instead of just a one-hop neighborhood.
3 Background
3.1 Multidimensional RNNs
The most general form of MDRNNs(Gravers et al. 2007 was introduced as DAG LSTM by Zhu ett al. (0016), as it can be applied to all directed acyelic graphs P(N, F). We denote n (e), n. (), incoming and outgoing node of an edge E E, and (n),,, (n) the sets incoming and outgoing edges of a node n E N. The core element of (Multi-Dimensiona) LSTMs is the LSTM Cell cr.
computed as:
C₁ Σ fnn C₁a' in O R. n'EPo(m)
h2 = On O tanh (c,.)
(2)
with hidden state output h. and (n) denoting the parents of node n in a DAG G. In standard LSTMs, the input é,, forget fn. output gates o,, and the Cell update 2,, dependent external inputs E, and on the previous/parent hidden states (h,,)) (nee)))))))m))) which makes them non-linear. In contrast, for undirected graphs, DAG LSTMs can be applied by ordering the nodes first, thereby fixing edge directions and using DAG LSTM bi-directionally. Note that there is typically more than one ordering option, where different orders can lead to different outcomes.
3.2 Linear RNNs
M
Linear RNNs resemble the structure of the original LSTM. For linearization, they remove the previous hidden state dependency of the gates and cell update. Additionally, they commonly in- clude state expansion dimension query key dimension. This enables a form of associative memory with inner-product retrieval in relation to quadratic energy-basse based models [Hopfield] 1982, Schmidhuber] 1992]. To include all relevant dimensions, switch the Einstein sum convention notation (so dimen- sions are visible in the indices) and absorb the input gate in the key and the output gate in the query. In reverse, query and key can also be interpreted as extended input and output gates. Additional normalization, as in mLSTM [Beck et al., [2025b), can be interpreted as parallel exccution with similar inputs and reduced dimensionality.
T
Input/output Jo
gatt
Figure 3: Transition from input/- forget/output gating in between nodes (and input/output) in (linear) RNNs towards Source Transition/- Mark gating between edges and in- put/output (bottom/top) in pLSTMs. As for linear RNNs, the gates do not depend on states, but inputs only.
Ciro FfCt-1Ru + Ki Vev H1, Quk
(3)
(4)
For certain architectures, the scalar forget gate F, is extended to the key dimension as Fii in GLA
Yang et al. 2023) or matrix in DeltaNet as Cr -1' with FIRN = = 3,,, K [Schlag et al. 2021 Y anng et al. 2024] or Gated DeltaNet (Yang et al., [025). DeltaProduct (Siems et al. 2025 uses product of multiple of these Householder-matrices, The important part is the linearization, by removing the dependence on the previous hidden state. All of F., Kik- V., and Qu depend only on the input of the layer Fia for this time step. This structure enables a meeunkise-reeuuren form for
efficient exccution modern hardware [Beck et al. 2025a Yang et al. 2023 2024].

---

pLSTM
The Cell state receives an external input via the Source Sg-(e) at the incoming node in addition to a sum of the Cell states of incoming edges to the incoming node modified by the respective Transition Similar to how LSTMs are extended to general DAGs, this can be applied to linear RNNs. Here, we extend the setting of MDRNNs and DAG LSTMs enable guided directional propagation. Instead of being associated with node in the graph, the Cell state C. is now associated with an edge. The inputs and outputs hidden states are still associated with nodes. Using this reformulation, can define generalized linear DAG networks that resemble the LSTM structure (see Figures 3 and
Σ
The cc'
Ceko Kn- vo Sa sm
Ceku
(5)
d'EE((n((e)
H Q%
Σ (ER_(n)
McCCku + K Q
Th The hidden state network output H, at node n is a sum of all incoming edge Cell states Ceke modified by Mark M,," and projected by query Q Note that the Source, Transition, and Mark gates are influenced by the input at node n, and optionally by the edge features of their associated local edges. Functionally, the Source replaces the input gate, the Transition the forget gate, and the Mark the output gate of a traditional LSTM (Hochreiter and Schmidhuber, 1997 Gers 1999]. while acting here on the node-edge, edge-edge, and edge-node combinations. The direct D term represents skip directly from input to output a node and also integrates the Q K V terms. These equations define general parallelizable Linear Source Transition Mark (pLSTM) networks on DAGs. This S-T-M structure can be computed sequentially, going through the DAG topological node order.
4.1 Parallelization on DAGs
4.1.1 Naïve Parallelization Paths
Since all operations are linear, the iterative definition of Equation also be resolved sum over all ancestor nodes and a combination of Source, path Transition and Mark over all paths P(e',e) connecting nodes ní and n via first edge e' and last edge e:
H
Σ K
V
Σ St
Σ
I
enP
M oh +K D" Q
n'n
d.e
PEP(e
APEP
Gon
Gn'n
The middle STM part in Equation 7 can be precomputed into the resulting gating matrix 6"'n and
the state expansion and projection via Q and KT can be integrated externally, similar to classical (linear) self-attention /Vaswani al. 2017 Katharopoulos et al. 2020). the D" term represents the diagonal entries. Computing the gating matrix via all paths is infeasible due to their exponential number (see Appendix D.I).
4.1.2 Hierarchical Parallelization
Similar to an associative scan in ID, we can apply divide-and-comque approach for hierarchical parallelization. Recursively, we merge Source-Transition, Transition- Transition-Transitionn, and Transition-Ne Mark combinations to higher-order Source, Transition, and Mark objects that cover more than a single node. In Appendix B.2] we show how this procedure works for DAGs. For regular grids such as sequences in ID and images in 2D, we can derive scheme of einsum, concatenation, and padding operations, which are well supported by modern hardware and parallel computation frameworks. In Appendix B.3] we show the parallelization in 1D. In Appendix [B.] we show the parallelization scheme for 2D grids. To better convey the core idea of TTM parallelization, we include sketch for all three cases and pLSTM in general in Figure X of the appendix.
5

---

4.2 Long-Range Stability
For sequential data, it is known since the seminal LSTM work of Hochreiter and Schmidhuber [1997] that having forget gates larger or smaller than one in magnitude leads to aaplodin/vaanshiin activations and gradients. For general DAGs, this was not explored yet, with recent analyses on undirected GNNs [Arroyo et al. 2025]. In DAG LSTM, Tree LSTM, Grid LSTM, or MD LSTM, potentially exploding parts are limited only in the hidden state by non-linear, bounded activation functions (i.e.. tanh) [Hochreiter and Schmidhuber 1997]. The Cell states, however, can grow exponentially if the forget gates are limited to only one. See Appendix D.I for details.
Using the more general setting of Transitions Ted from edge e' to edge e, as introduced in this work, we can derive two modes of stability. First, a propagation mode (P-mode) that covers all paths in the DAG, but is effectively limited to a line as mode of propagation. Second, distribution mode (D-mode) that limits the Transitions to a subset, such that their resulting line-graph is reduced from a DAG to a multi-tree Jung 1978).
The line graph of DAG G = (N, is the DAG (E, E) formed by the edges E of the original DAG and the line edges E E' = ē= e2), E N E1 E (_(n) neg E (n),,, which connect two original edges if these are connected by a node in Ç [Whitney 1932] . Given the previous definitions of Cell states C, and Transitions T these can be interpreted states and edges on the nodes of the line graph S. The Transitions Tee' can be viewed entries of the adjacency matrix T of G. This way, we can apply the theory of power iterations and matrix norms. The application of PLSTM of Equation iS equivalent to the application of the power iteration and can be bounded by compatible matrix norms
Pg To
|c|
Σ p=0
T
s|
s
Pøls|
(8)
p=0
As 9 DAG, also its line graph & is a DAG and this one's adjacency matrix Tee nilpotent, for a certain power 323 s |N] Vp > Pg T = O. For simplicity, we use matrix-vector notation here for Cell states c, Transitions T, and Sources & without key-value extension.
The first stabilization option for the P-mode is to limit the norm of T as in Equation g A node-local option to achieve this is to limit the Transition entries The per node as:
Σ
Iha '
(EE
This limits the respective column e' in the line graph adjacency matrix T. and in turn limits its Li-norm ||T]|1 Since the L1 -norm is uub-muliplicatiiv [Horn and Johnson 1985], this can be applied to the matrix powers in Equation X and keeps the T.r-norms of the Cell states C. limited. Keeping the norm fixed at one exactly enables long-range propagation (e.g.. by division of the column entries by the norm). This stabilizes the gradient propagation with an Loo bound, as gradients propagated backwards and the applied transposed adjacency matrix T normooiied and therefore T..o bounded (the dual norm to the 7.1 norm).
For the second stabilization option, the D-mode, we do not limit matrix norm of T, but instead reduce T' or 9 from a DAG multitree [Jung] 1978]. J. i.e., directed graph, such that two nodes are only connected by single path. The original graph ς is not reduced to such a multitree by that. Since there is only one path in the line graph structure, there is only one path term in Equation 7 so there is no exponential explosion of path numbers, and limiting single T entries by one is sufficient.
4.3 PLSTM on regular grids
Although PLSTMs can be used and parallelized on general DAGs, the additional structure greatly helps with the parallelization for regular grids. All operations can be decomposed into generalized matrix operations in the form of views, einsums, concatenations, and paddings. We show this for the ID case of sequences in Appendix B.3

---

4.3.1 PLSTM in 2D images
On a 2D undirected grid, pLSTM can be applied in different ways since there are more than two options for DAGs covering the undirected graph. There are, however, four distinet DAG covers, for which the local structure translates to global structure: allowing only Transitions exclusively left or right and up or down, both in combination: ;; 2, We focus on the first, the right-down combination, for the description, but all four combinations should be used to cover all directional interactions. Note that here, nodes correspond to pixels/patches and edges to the respective connections between them, including the introduced direction. For the parallelization, a hierarchy of different levels of Source, Transition and Mark, covering larger chunks of multiple nodes. This is also depicted in Figure and As Transitions follow all edge combinations, there are four options again: 7, 4,
derive
and described in Appendix Section B.2 and B.4
The level zero Source, Transition, and Mark tensors are therefore: S Say. Tay. Tay Tay. my M M The general DAG is depicted in Figure 6 in the center right. As the in- and out-degrees of this DAG are now larger than one, there are more stabilization options, namely the P-mode and D-mode (Section 4.2)
P-mode
When we allow all Transition options 7, 4,, 2, 4, T
we are in the P-mode. In general, the
T
Transition matrix at position z, y looks like Tay
Try. The P-mode stabilization now limits
T2y
this matrix to absolute column sums 1. This limits the structure at criticality to [
8
1-a
1- 3
When setting 3 = a, we even arrive at a geometrical interpretation of this Transition: it directionally propagates the signal with maximal amplitude along the line defined by Az a (see Appendix D.2 By for the derivation). In Figure 6 bottom right, we depict the receptive field of a single non-zero Source with constant Transitions of this structure across the grid. For practical parameterization, we, therefore, fix this structure with the 'angle' of propagation and an additional decay factor ry:
T= T=
a (1 (1-a)
(10)
D-mode While the P-mode offers directional propagation, the receptive field of a node/patch/pixel limited to roughly a line, it cannot reach all other nodes equally. The D-mode can provide this by losing the notion of directionality. For the D-mode, either TVy or Try have to be set to zero globally. In this way, two edges can only be connected via a single path. In Figure the bottom right, the second option is shown, with an additional decay. At criticality, this node can still reach all other nodes to its bottom right with magnitude one, long-range propagation in two dimensions.
Directional Initialization and Combination Using the concept of multiple heads, as used in Vaswani et al.] [2017). criticality on initialization. Similarly, we can initialize different decay scales for both P- and D-mode for different heads. By using the P- and D-mode in alternating layers, leverage both their potential.
initialize our network to cover multiple directions in P-mode at
State-Trache Trackira Expansion Appendix B.5, we show how PLSTMM can be extended for state tracking via non-diagonal Transitions MMerrll et al. 2024 Siems et al. 2025 Grazzi al. 2025)
5
Experiments
We first showcase the theoretical advantages of PLSSTM on the synthetic arrow-pointing extrapolation task (see Section Then we highlight the benefits of PLSTM for two-dimensional input data on ImageNet- IK [Dngg et al. 2009 Russakovsky al. 2015), demonstrating scalability large-scale
datasets (see Sections 5.2 and 5.3). Finally, we illustrate how pLSTM scales to more than two input
dimensions on established graph benchmarks (see Section 5.4 5.4
5.1 Arrow Pointing Extrapolation
In this task, an image containing an arrow and a circle must be classified as to whether the arrow pointing to the circle (see also Figure 2). This is long-range directional task as it involves the relative
7

---

PLSTM
2DMamba
EfficientNet
VIL
Mamba2D
VIT
1.0
1.0
1.0
Val. Acc. (Ext.)
Train. Acc.
Val. Acc.
0.8
0.8
0.8
0.6
0.6
0.6
10
20
30
40
10 20
30
40
0
10 20
30
40
Training Steps (k)
Training Steps (k)
Training Steps (k)
Figure 4: Training curves for the Arrow Pointing Extrapolation task, averaged over 5 seeds with 90% CI. ViT and EfficientNet can quickly match the training set (left), and EfficientNet reaches the best validation performance of all models the samples of the same resolution seen during training. PLSTM performs best on the extrapolation set (right) by significant gap. While Mamba2D and 2DMamba should cover restricted modes of pLSTM, their learning shows high variance. ViL and Mamba2D do not extrapolate at all beyond random performance.
positioning of two objects, in conjunction with local information from one of them the direction of the arrow. To test for arrow pointing extrapolation capabilities, we generate balanced dataset of 100k arrow pointing images resolution 192 x 192, with positions of the arrow and circle randomly sampled. For validation, we generate 5120 images in the same resolution and at resolution 384 x 384 to test for extrapolation capabilities. For all models, including the VIT [Dosovitskiy et al. 2021 baseline, we use a bicubic interpolation of the positional embeddings (via jax. image........ .mmge.resize) for scaling to higher resolutions (see also [Dosovitskiy al., (02T)). For details the training, Appendix F.2 In Figure we show the learning curves for PLSTM and all baselines. We find that PLSTM performs well both the training and validation tasks. Importantly, PLSTM exhibits the strongest extrapolation abilities. In contrast, ViL performs poorly because it traverses the image patches sequentially along the scanline form and misses out on important directional information. Despite its global receptive field, ViT falls short on extrapolation to larger image resolutions, which we surmise is due the its positional encoding. Similarly, EfficientNet exhibits strong performance on training/validation tasks, but fails to extrapolate. Finally, pLSTM considerably outperforms Mamba2D (which covers the P-mode) and 2DMamba (which covers the D-mode), despite their multidimensional nature. Also note that, in practice, PLSTM is much faster than Mamba2D or 2DMamba on real hardware, despite missing kernels, and on par with ViL. (see(6) In Figure 5 we highlight the performance differences between D- and P-mode-only models on the arrow pointing task.
PLSTM
PLSTM (P-mode)
PLSTM (STM bias only)
PLSTM (no pos emb.)
PLSTM (D-mode)
VIT
1.0
1.0
Val. Acc. (Ext.)
Train. Acc.
Val. Acc.
0.8
0.8
0.8
0.6
0.6
0
10
20
30
40
10
20
30
40
10 20
30
40
Training Steps (k)
Training Steps (k)
Training Steps (k)
Figure Training curves for Arrow Pointing Extrapolation task, averaged over seeds with 90% CI on different model ablations. P-mode by itself performs worse, D-mode by itself is not as general in interpolation, but performs better on extrapolation compared to other models. pLSTM does not rely the positional embedding for strong performance.

---

Table 1: ImageNet-IKK pre-training accuracy. All models use a patch size of 16 16 with 224 x 224 resolution at maximum. Models with "+" in their "Epochs" column pre- pre-train on lower resolution, followed by fine-tuning on 224 x 224 resolution for some epochs. Values of reference models are taken from Alkin et al. [025 and the original work for 2DMamba, Mamba2 and EfficientNet [Tan] and Le, 2ha9] et al. 2024b Baty et al. 2024). Note that EfficientNet, as CNN-based baseline, was also trained on larger resolutions for the scaled-up versions. Due to the chunkwise-recurrent
option, PLSTM FLOPs could be optimized further (see Appendix B.7p.
Model
Epochs
#Params
FLOPS
IN-1K
EfficientNet-B0 Tan and Le
2019
5M
0.39G
77.1
DeïT-T Touvron et al. 2021
300
6M
1.3G
72.2
DeiT-III-T (reimpl.) Touvron al.
2022
800+20
6M
1.1G
75.4
VRWKV-T [Duan et al.
2024
300
6M
1.2G
75.1
Vim-T Zhu
2024
300
7M
1.5G
76.1
Vil-T Alkin et al. 2025
800+20
6M
1.3G
78.3
PLSTM-ViS-1
800+20
6M
1.4G
75.2
EfficientNet-B44 (an and Le
2019
19M
18G
82.9
DeïT-S Touvron al., 2021
300
22M
4.6G
79.8
DeiT-111-S (reimpl.)
Touvron et al.
2022
400+20
22M
4.6G
80.3
ConvNeXt-S (iso.)
Liu et
2022
300
22M
4.3G
79.7
VRWKV-S Duan et al.
2024
300
24M
4.6G
80.1
Vim-S Zhu et al. 20224
300
26M
5.3G
80.5
Mamba2D-T Baty al.
2024
300
27M
82.4
2DMamba-T T Zhang et al
2024a
30M
4.9G
82.8
ViL-S Alkin et al. 2025
400+20
23M
4.7G
81.5
PLSTTM- Vis-S
400+20
23M
4.9G
80.7
Efficien(Net-B6 Tan and Le
2019)
43M
19G
84.0
DeiT-B (Tovrron et al. 2021
300
86M
17.6G
81.8
DeiT-IlI-B (remmll.)
Touvron al.
2022
400+20
87M
16.8G
83.5
ConvNeXt-B (iso.)
Liu et al. 2022
300
87M
16.9G
82.0
VRWKV-B [Duan al.
2024
300
94M
18.2G
82.0
2DMamba-S ['Zang et
2024a
50M
8.8G
83.8
Vil.B [Alkin et 2025
400+5
89M
17.9G
82.4
DIS-B-VIs-B
400+20
89M
18.2 G
82.5
5.2 ImageNetlk
To test for real-world image model capabilities, we train PLSTM using the schedule of DeiT- III [Tourron et al.) [022), comparing our architecture to other vision models. PL.TT performs similar to other popular approaches. With the integration of local features, e.g., including additional convolutions such as in Vision LSTM [Alkin et al. 2025), we see room to narrow the gap to SOTA. For detailed discussion the results and relation previous work, see Appendix E.
5.3 ImageNet-Ik Ablation
We ablate our PLSTM and compare against ViT trained using the VIT-T model scale and simpler single pre-training schedule close to DeiT [Touvron et al. 202T] at 224 224 resolution (see Appendix I 1 for details. At this scale, pLSTM outperforms ViT by significant margin. Results are shown in Table 2

---

Table Ablation of PLSTM variants and VIT (DeiT) re-training on ImageNet-Ik.
Table 3: Long-Range Graph Benchmark Re- sults [Dwivedi et al. 2022
Model
ImageNet- 1k (top-1)
Model
Peptides -truct (MAE)
Peptides
-func
(Avg. Prec.) T
PLSTM
75.51
GAT Veličković et al. 2018
0.270
0,525
GCN Kipt and Welling 2017
0.264
0.534
PLSTM / (no posemb.)
75.22
GIN
[Xu al. 2018
0.328
0.593
PLSTM (P-mode only)
74.86
LSTM-GNN [Lang et al. 2016
0.274
0.502
PLSTM /
(D-mode only)
75.13
MPNN [Gilmer al.
0.260
0.557
PLSTM (STM bias only)
75.10
PLSTM
0.264
0.536
ViT
73.49
5.4
Graphs
Small molecules and proteins are typically depicted as an undirected graph. However, both exist as 3D structures in the real world, and GNNs trained such datasets might benefit from directional information propagation, as there is an underlying spatial relation. We test this hypothesis on popular small molecules and bioinformatics datasets from the TUDataset benchmark [Morris et al.. [2020] and the Long- Range Graph Benchmark [Dwivedi et al., [022]. The results in Tables and show that PLTM can compete with popular GNN architectures on those datasets. Similar to the computer vision experiments (Section 5.2). PLSTM alternates between the P-mode and the D)mode. Since there is no external notion of direction, we use node and edge features to adapt the Transitions.
Table 4: Test accuracy on small molecule and bioinformatics datasets as provided by TUDataset.
Model
MUTAG
NCTI
PROTEINS
PTC_FM
AVG
GAT Veličković et al. 2018
0.7822 + 0.09
0.7968 + 0.03
0.7215 0.03
0.6105 + 0.05
0.7277 + 0.06
GCN
Kipt Welling 2017
0.7234 0.08
0.7852
0.7395
0.03
0.6162 0.04
0.7161 : 0.05
GIN
Xu 2018
0.8251 0.10
0.8175
0.02
0.7350
0.04
0.6097 1 0.10
0.7468 i 0.08
LSTM-GNN
Liang ct al. 2016
0.7450 + 0.11
0.7951
0.02
0.7503
0,04
0.6076 + 0.04
0.7245 t 0.06
MPNN GGimerr al. 2017
0.7450 + 0.09
0,8012 + 0.02
0.7350 + 0.04
0.5786
+ 0.07
0.7149 : 0.06
PLSTM
0.8512 = 0.06
0.7324 2 0.03
0.7502 2 0.05
0.6133 = 0.08
0.7368 : 0.06
Conclusion
In this work, we introduce pLSTM, which unites the benefits of MDRNNs (Graves et al. 2007] and the recently introduced xLSTM (Beck et al., (025b). PLSTM overcomes the limitations of modern recurrent architectures when applied to multi-dimensional data, such as images and graphs. To achieve this, modify the gating structure of xLSTMM and introduce Source, Transition, and Mark gates. Then, we derive a parallelization scheme that enables processing data in multiple dimensions concurrently. SLS comes with P-mode and D-mode, which together enable large and auto-tunable effective receptive field. We demonstrate the theoretical advantages of pLSTM on the arrow-pointing task and highlight its ability to generalize to varying grid resolutions. Finally, we show the efficacy of PLSTM both on classical computer vision and graph benchmarks.
Limitations & Future Work PLSTM shows promising results across different domains, however, there is still room for improvement compared to highly optimized -omain-specific models. With the incorporation of domain-specific inductive biases. are positive that the results can be further improved. While PLSTMs should thcoretically enable long-range propagation of activations and gradients in multiple dimensions for recurrent, subquadratic models, the arrow pointing extrapolation task only tests this in a restricted way. Moreover, while PLSTM models generalize better than ViTs, which use position embeddings to encode spatial relations, there is still a gap to perfect extrapolation. long-range
This leaves room for improvements. both on the data side of harder multi-dimensione benchmarks and the architectural side to better generalize to that data. Given the flexibility of PLSTM to handle data with rich multi-dimensional structure, we anticipate that STMSMM can be successfully applied across a broader spectrum of challenging domains, including biology, chemistry, medical imaging, and other scientific domains.
10

---

Acknowledgments and Disclosure of Funding
We thank Benedikt Alkin, Pieter-Jan Hoedt, Phillip Lippe, Maximilian Beck, Kajetan Schweighofer, Erich Kobler, and Günter Klambauer for helpful discussions and feedback.
Korbinian Pöppel was supported by the Konrad Zuse School of Excellence in Learning and Intelligent Systems (ELIZA) through the DAAD program Konrad Zuse Schools of Excellence in Artificial Intelligence, sponsored by the Federal Ministry of Education and Research during his ELLIS visit at Uni Freiburg, where this work finished.
The ELLIS Unit Linz, the LIT Al Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects FWF AIRI FG 9-N (10.55776/FG9) Al4GreenHeeaingGrrid (FFG- 899943), 4tars4Waters (HORIZON-CL6-2222 10--2002CCLMMAA-01--000000 FWF Bilateral Artificial Intelligence (10.55776/COE12), We thank NXAI GmbH, Audi AG, Silicon Austria Labs (SAL), Merck Healthcare KGaA, GLS (Univ. Waterloo), TÜV Holding GmbH, Software Competence Center Hagenberg GmbH, dSPACE GmbH, TRUMPF SE + Co. KG.
We acknowledge EPr Joint Undertaking for awarding us access to Leonardo at CINECA, Italy, and Mare at BSC, Spain.
References
B. Alkin, M. Beck, K. Pöppel. S. Hochreiter, and J. Brandstetter. Vision-LSTM: xLSTM as Generic Vision Backhone. In The Thirteenth International Conference on Learning Representations, 2025.
URL httpgp///pppateeeeeeeeetttfffrro
A. Arroyo, A. Gravina, B. Gutteridge, F. Barbero, C. Gallicchio, X. Dong, M. Bronstein, and P. Van- dergheynst. On Vanishing Gradients, Ever-Smoothing, and Over-Squashing in GNNs: Bridg- ing Recurrent and Graph Learning, Feb. 2025. URL sattp://arxile org/ abs/ htre//bb/250022 2502 10818 arXiv:2502. 10818 [es].
S. B. Atitallah, C. B. Rabah, M. Driss, W. Boulila, and A. Koubaa. Exploring graph mamba: A comprehensive survey on state-space models for graph learning. arXiv preprint arXiv:2412.18322. 2024.
E. Baron, I. Zimerman, and L. Wolf. A 2-Dimensional State Space Layer for Spatial Inductive Bias. In The Twelfth International Conference on Learning Representations, 2024. URL https //openreview www.nneeet/oddeeeeeeoett BdBGkaammSmm
F. Baty, A. H. Díaz, C. Bridges, R. Davidson, S. Eckersley, and S. Hadfield. Mamba2D: A Natively Multi-Dimensional State-Space Model for Vision Tasks, Dec. 2024. URL http:://ax///i. org/ abs/ Abs/2412.em 16146, #rXiv:2412.16146 [cs].
M. Beck, K. Pöppel. P. Lippe, and S. Hochreiter. Tiled flash linear attention: More efficient linear mn and xlstm kernels. arXiv preprint 4rX1v:2503.14376, 2025a.
M. Beck, K. Pöppel, M. Spanring, Auer, O. Prudnikova, M. Kopp, Klambauer, Brandstetter, and S. Hochreiter. XLSTM: Extended long short-term memory. Advances in Neural Information Processing Systems, 37:107547-107603, 2025b.
A. Behrouz and F. Hashemi. Graph Mamba: Towards Learning on Graphs with State Space Models. CoRR, abs/2400.008678, 2024. httts:s/:///oi org/10 415.....////r// 2tt.//222.vv........ 2800208877
A. Behrouz, P. Zhong, and V. Mirrokni. Titans: Learning to memorize at test time. arXiv preprint N0Xv:22001.00633. 2024.
Y. Bengio and P. Frasconi. Credit assignment through time: Alternatives to backpropagation. Advances in neural information processing systems, 6, 1993.
Y. Bengio, Simard, and Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions neural networks, 5(2):157-166, 1994.

---

X. Chen, Z. Shi, X. Qiu, and X. Huang. Dag-based long short-term memory for neural word segmentation. arXiv preprint vrXvv11077002255 2017.
T. Dao and A. Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint arXiv:2405.21060, 2024.
T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with 1w-Arennsss, June 2022. (tLp:tt///are orgoorgg///. 2205 14135 2XXi::2055.14355 [cs].
J. Deng. W. Dong, R. Socher, Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 258.255. leee, 2009.
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly. J. Uszkoreit, and N. Houlsby. An Image is Worth 16x Words: Transformers for Image Recognition at Scale. In International Conference on Learning poot/forum?dd-iiooroa
Representations, 2021. URL https:
Y. Duan, W. Wang. Z. Chen, X. Zhu, L. Lu, T. Lu, Y. Qiao, H. Li, Dai, and W. Wang. Vision- rwkv: Efficient and scalable visual perception with rwkv-like architectures. arXiv preprint arXiv:2403.003388, 2024.
V. P. Dwivedi and X. Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699, 2020.
V. Dwivedi, A. Luu, T. Laurent, Y. Bengio, and X. Bresson. Graph neural networks with learnable structural and positional representations. arXiv preprint arXiv:21 2021.
V. P. Dwivedi, L. Rampášek, M. Galkin, A. Parviz, G. Wolf, A. T. Luu, and D. Beaini. Long range graph benchmark. Advances in Neural Information Processing Systems, 35:22326-22340, 2022.
F. Gers. Learning to forget: continual prediction with LSTM. In 9th International Conference Artificial Neural Networks: ICANN '99, volume 1999, pages 850-855, Edinburgh, UK, 1999.
IEE. ISBN 978-0-85296-721-8
doi: 10. 10.1049/cp: 19991218. URL https://digitall-11 1049/ 11w01/0p1999911111 19991218
theiet. org/ tll.istt.rrg//oottttt/coo...o.t. www.oooteeetttooooooooeoeeeeeeee
J. Gilmer, S. S. Schoenholz, F. Riley, O. Vinyals, and G. E. Dahl. Neural Message Passing for Quantum Chemistry. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference Machine Learning, volume 70, pages 1263-1272, 2017.
Graves, A.R.. Mohamed, and G. E. Hinton. Speech recognition with deep recurrent neural networks. In JEEE International Conference on Acoustics, Speech and Signal Processing, pages IEEE, 2013. doi: 10.
A. Graves, S. Fernández, and J. Schmidhuber. Multi-dimensional recurrent neural networks. In International Conference on Artificial Neural Networks (ICANN), volume 4668 of Lecture Notes in Computer Science, pages 6645-6649. Springer, Berlin, Heidelberg, 2007.
R. Grazzi, J. Siems, J. K. H. Franke, A. Zela, F. Hutter, and M. Pontil. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. In The Thirteenth International Conference on Learning Representations, 2025. URL httpt:tttts/:rrreettee
A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
S. Hochreiter. Untersuchungen zu dynamischen neuronalen Netzen. Diploma thesis, Institut für Informatik, Lehrstuhl Prof. Brauer, Technische Universität München, 1991.
S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Comput., 98):1735- 1997.
S. Hochreiter, Y. Bengio, P. Frasconi, J. Schmidhuber, al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001.
12

---

J. Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences, 79(8):255--258, Apr. 1982. ISSN 0027-8424. 1091-6490. doi: 2..1073/pnas.79.8 29.2. 79.8.255 2554
httpaa.orr/ooofffl///// pnas
R. A. Horn and C. R. Johnson. Matrix Analysis. Cambridge University Press, edition, Dec. 1985. ISBN 978-0-521-3886222 978-0-221-00588665 978--11111177-------11811-1-- doi: 10.1017/
1B09780511810817. URL http:://www..aooo.g...ggg//////t
97806118108817/yyp/bbboal
H. A. Jung. On a class of posets and the corresponding comparability graphs. Journal of Combi- natorial Theory, Series B, 24(2):125-1333 1978. ISSN 0095-8956. doi: https://doi.org/10..0000
0095-8956678/90003344 URL //ww://wwwww httpp::/www.scceeneeeee.eeeeeeee www.:sseeeeeeedee.ccoommm http:///eteeeeeeeeeeeeeeeettcc/p 0095895678900138
N. Kalchbrenner, I. Danihelka, and A. Graves. Grid long short-term memory. arXiv, 100700152, 2015.
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. '20. JMLR.org. 2020.
T. N. Kipf and M. Welling. Semi Semi-Supervised Classification with Graph Convolutional Networks. In International Conference Learning Representations, 2017.
Li, H. Singh, and A. Grover. Mamba-ND: Selective State Space Modeling for
Data. In A. Leonardis, E. Ricci, S. Roth, O. Russakovsky, T. Sattler, and G. Varol, editors, Computer Vision ECCV 2024, pages 75-92, Cham, 2025. Springer Nature Switzerland. ISBN 978-3-031-7341444
X. Liang, X. Shen, J. Feng, L. Lin, and S. Yan. Semantic object parsing with graph LSTM. CoRR, 2016. URL http::t:://axxi http://ww.orggaa///...... 07063
Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11976 1 1986, 2022.
W. Merrill, J. Petty, and A. Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.0887, 2024.
M. Minsky. Steps toward artificial intelligence. Proceedings of the IRE, 49(1):8-33)) 30, 1961. doi: 10.100/JRPROC..99118877..
C. Morris, N. M. Kriege, F. Bause, K. Kersting, P. Mutzel, and M. Neumann. TUDataset: A collection of benchmark datasets for learning with graphs. In JCML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020), 2020.
A. Orvieto, L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670-26698. PMLR, 2023.
Park. deit3-jax, 2024. URL Title: GitHub repository.
Publication
R. Pascanu. T. Mikolov, and Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318, Pmlr, 2013.
N. Peng, H. Poon, C. Quirk, K. Toutanova, and W.-t. Yih. Cross-sentence n-ary relation extraction with graph Istms. Transactions of the Association for Computational Linguistics, 5:101-115, 2017.
K. Pöppel, M. Beck, and S. Hochreiter. FlashRNN: 1/O-Aware Optimization of Traditional RNNs on modern hardware. In The Thirteenth International Conference on Learning Representations, 2025. URL httpps///operrevier httpppppaeeeeieeeeeeoooooorraaaa
13

---

L. Rampášek, M. Galkin, V. P. Dwivedi, A. T. Luu, G. Wolf, and D. Beaini. Recipe for general. powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35: 14501-14515, 2022.
O. Russakovsky, J. Deng. H. Su, Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei. ImageNet Large Scale Visual Recog- nition Challenge. International Journal of Computer Vision (IJCV), 115(3):211-255 2015. doi: 10.1007/ss 1208001--0-11..
F. Scarselli, M. Gori, A. Tsoi, M. Hagenbuchner, and G. Monfardini. The Graph Neural Network Model. IEEE Transactions on Neural Networks, 20(1):61-80, 2009. doi: 1000 2005605.
Y. Schiff, C.-H. Kao, Gokaslan, T. Dao, Gu, and V. Kuleshov. Caduceus: Bi-directional equivariant long-rang raage dna sequence modeling, arXiv preprint wrXiv:2403.0322.4 2024.
I. Schlag, K. Irie, and J. Schmidhuber. Linear transformers secretly fast weight programmers. In International conference on machine learning, pages 9355-9366. PMLR, 2021.
J. Schmidhuber. Learning to Control Fast-Weight Memories: An Alternative to Dynamic Recurrent Networks. Neural Computation, 4:131-139, 1992. URL https://api sem.aaicccemanticccooolaaa
org/CorpusIll /t.//Grrpus1I::666333477 16683347
Schmidhuber. Deep learning in neural networks: An overview. Neural networks, 61:85-117, 2015.
N. Schmidinger, L. Schneckenreiter, P. Seidl, J. Schimunck, P.-J. Hoedt, J. Brandstetter, A. Mayr, S. Luukkonen, S. Hochreiter, and G. Klambauer. Bio-xlstm: Generative modeling, representation and in-context learning of biological and chemical sequences. arXiv preprint 1r.iv:2411.04165 2024.
Schmied, T. Adler, V. Patil, M. Beck, Pöppel. Brandstetter, G. Klambauer, R. Pascanu, and S. Hochreiter. A large recurrent action model: xlstm enables fast inference for robotics tasks. arXiv preprint arXiv:2410.22391, 2024.
J. Siems, T. Carstensen, A. Zela, F. Hutter, M. Pontil, and R. Grazzi. Deltaproduct: Improving state-tracking in linear rnns via householder products. arXiv preprint arXiv:2502.10297 2025.
M. Stollenga, W. Byeon, M. Liwicki, and J. Schmidhuber. Parallel multi-dimeensonal Istm, with appli- cation to fast biomedical volumetric image segmentation. In Proceedings of the 28th International Conference on Neural Information Processing Systems (NIPS), pages 2998-3006, 2015.
Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang. Y. Dubois, X. Chen, X. Wang, S. Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.046200 2024.
K. S. Tai, R. Socher, and C. D. Manning. Improved semantic representations from tree-structured long short-term memory networks. arXiv preprint arXiv: 2015.
M. Tan and Q Le. Efffcciettnt: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pages 6105-6114, PMLR, 2019.
H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. International conference on machine learning, pages 10347-10357. PMLR, 2021.
H. Touvron, M. Cord, and H. Jégou. DeiT III: Revenge of the VIT. In Computer Vision ECCV 2022: 17th European Conference. Tel Aviv, Israel, October 23-27, 2022, Proceedings. Part XXIV. pages 516-533, Berlin, Heidelberg, 2022. Springer ISBN 978-3-031-20052-6, doi: 10. 111199788---11-200333 event-place: Tel Aviv, Israel.
A. Van Den Oord, N. Kalchbrenner, and K. Kavukcuoglu. Pixel recurrent neural networks. In International conference machine learning, pages 1747-1756. PMLR, 2016.
14

---

A. van den Oord, N. Kalchbrenner, and K. Kavukcuoglu. Pixel Recurrent Neural Networks. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd Internationall Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1747-1756. New York, New York, USA, June 2016. PMLR. URL https::/ppoceeedinggg .ariing......p.p...... press// v48/oord16. html
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio. Graph Attention Networks. In International Conference on Learning Representations, 2018.
C. Wang, O. Tsepa, Ma, and B. Wang. Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces, Feb. 2024a. StLLhttp://axxiv org/ abs/ /abbe/2242...o 00789 2.0002 [es].
C. Wang, W. Zheng, Y. Huang. Zhou, and Lu. V2M: Visual 2-Dimensional Mamba for Image Representation Learning. arXiv preprint w7Xi::241011038,, 2024b.
H. Whitney, Congruent Graphs and the Connectivity of Graphs. American Journal of Mathematics, 54(1):150, Jan. 1932. ISSN 00029327. doi: 10.2307/2371066. [RL]htts:://wwt jswtor org/ stable/
S. Williams, A. Waterman, and D. Patterson. Roofline: insightful visual performance model for multicrre architectures. Communications the ACM, 52(4))65-76, Apr. 2009. ISSN 0001-
0782. 1557-7317, doi: 10.1145/149876511498753 URL https::/d/ll acm ooc.moorr//oo//////////// 1145/ 1498765. 1498785
K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? arXiv preprint arXiv: 10t::000000000 0.8066 2018.
K. Xu, W. Hu, J. Leskovee, and S. Jegelka. How Powerful are Graph Neural Networks? In International Conference Learning Representations, 2019.
S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware- efficient training. arXiv preprint arXiv:2312.06633. 2023.
S. Yang, B. Wang. Y. Zhang, Y. Shen, and Y. Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406. 06484, 2024.
S. Yang, J. Kautz, and A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. In The Thirteenth International Conference Learning Representations, 2025. URL
J. Zhang, A. T. Nguyen. X. Han, V. Q.-. Trinh, H. Qin, D. Samaras, and M. S. Hosseini. 2dmamba: Efficient state space model for image representation with applications on giga-pixel whole slide image classification. arXiv preprint arXiv:2412.00667, 2024a.
J. Zhang, A. T. Nguyen, X. Han, V. Q.-H. Trinh, H. Qin, D. Samaras, and M. S. Hosseini. 2DMamba: Efficient State Space Model for Image Representation with Applications Giga-Pixel Whole Slide Image Classification, Dec. 2024b. URL httptt//:///titt httpppp//waaav.or///.//// aa////ba//22222 arXiv:2412.00678 [cs].
L. Zhu, B. Liao, Q. Zhang, X. Wang, and W. Liu. Vision mamba: Efficient visual representation learning with bidirectional state space model (preprint). arXiv preprint arXiv:2401.09417, 2024.
X. Zhu, P. Sobhani, and H. Guo. Dag-structure long short-term memory for semantic compositional- ity. In Proceedings of the 2016 conference of the north american chapter of the association for computational linguistics: Human language technologies, pages 917-926, 2016.
15

---

A Method Visualization
1D
S Source
Te'e Transition
M Mark
M
Sp
Sy]
Σ
Tax
T ay
T
Tyx
T Tyy
M
M
My
= Σ
Σ
Σ
Gating Matrix Entries
Σ
Σ
Stable Modes:
|
1
dinectional
P-mode
P-mode
fiels
D-mode
|
D-mode
receptive
line graph
Inc prosh mutteres
Input
mode
edge
Mutremed out
Transtion
edges
"subgraph
grid mode
grid odgo
grid cége (dow))
exemplary
Figure 6: Illustration PLSTM on general DAGs (left), and on 1D 2D grids (right): the top-left part, general DAG is visualized, with Sources S Transitions Tee: and Marks M populating mude-outgoing-edger
and scoomnngeedge-an associative structure of these lincar operators, they be combined; some examples shown in the center left. The end-to-end (node-to-node) result is the Gating matrix. the bottom left, the two stable modes are depicted: The P-mode, where the edge have to limited by one, and the D-mode, where the line graph over edges is reduced to multitree, requiring the remaining Transitions each to be limited in absolute to one. On the right, the application of this to regular ID and 2D grids is shown. For 2D, the P-mode results in a directional propagation, whereas the D-mode results in diffusive/ un-directed distribution (shown with decay here).
macoming-cdge edge-outgoing-edee..
pairs. Due to the
absolute Transitions of in-coming
16

---

B
Method Details
B.1
Loosely self-similar graphs
Graphs modeling hierarchical order or spatial extensions, as in meshes or grids do typically have low node degree and can be decomposed recursively into subgraphs of similar structure. We use the term loosely self-similar graphs for these. In particular, graph 9 := g( should be decomposable into M subgraphs (g/h-11) (N!(-1), ) ))E((... and boundary edges Bub-") c cN("-1)
N,1 (L(1-1
such that g(L) (U. (N -U,, Blb- U.. Ee D This decomposition should be
applicable recursively down single nodes, such that | = Va with a (ar-11--00) the multi-index of successive decompositions. Ideally, all decompositions should be balanced, I.v." | E {-1,0,1} Vl.a, for optimal parallclization which should serve as the definition of loosely self-similar. Choosing arbitrary subsets of nodes for subgraphs, one can always arrive at decomposition of any graph, but the edges across subgraphs might not be balanced then.
B.2 Hierarchical Parallelization
Given decomposition a (loosely self-similar) DAG, we can apply divide-and-conqque strategy to the expensive Transition sum/product combination. We want to calculate now the Gm matrix entry for two nodes n. E N; s n, such that, at level ++ 1, both are in the subgraph gl+1 :n, E N but going one level lower, they are in distinct subgraphs n' No E N # b. With incoming and outgoing boundary edges e' 1a e B Brosa we define the generalized Source, Transition, and Mark tensors at level l:
S L.ot
Σ
ΣE PEP(d.e)
T on
(11)
dEE-(n))
Thack ée
Σ
T en
(12)
PEP(d.e)
Madaab
Σ
Σ
T ex o
M
(13)
(EE.(n)
(PEP(e)e) Peer
Given these, the gating matrix entry can be calculated as:
Gm'n
Σ
Σ
S Jonta
Σ
Lose Lase hove m mmb
M daech
CE
e'EB'...
(e)P
Here, P is the set of "meta-path"s that connect two edges e', via subgraphs at level 1, in these meta-paths the subgraphs act as "nodes" and the boundary edges B"asa B'.. "edges". Note that, by Einstein sum convention we implicitly sum over multiple in-coming and out-going edges Lase
!aaee E-m! and e-m!! for generalized Transitions plaae in Gm'n and in the following. See Figure 6 for an illustration of this decomposition. In reverse, this means that higher-level Source, Transition, and Mark can be constructed recursively bottom-up from one level lower (a chosen such that n' e N'd or n. No )
141,
Σ
s A'oiaa
Σ
Tlacbb Lork m mm
Se
(14)
e'EBLot ota
P'ep bept
++1,a Le
Σ
I
I L Linch mmm mmm
(15)
(e, bep!
Maltlaa
Σ
Σ
Thais hais hai mn mm)
(16)
P'Epl eCBLaca
(e, bept
17

---

Note that using this hierarchy only up to certain level, one can optimize the parallelization to the used hardware. Whereas first, all the Source, Transition, Mark, and Gating matrix objects are constructed in parallel up a certain level, the recurrent mode of Equation 5 is applied with generalized Source, Transition and Mark matrices of level I, in the topological ordering of the subgraphs at this level. This resembles the structure of the chunk-wise parallel formulation for linear RNNs.
B.3 Parallelization of pLSTM 1D sequences
For sequence DAG, PLSTMs reduce to linear RNNs. Here, all nodes and subgraphs have a single incoming and outgoing (boundary) edge. For parallelization, the decomposition is done hierarchically in powers-of-wo splits. The S, T, M and G tensors are constructed follows:
S.n
st-1
( S (20+1))m)))
(17)
T- Taay ((0+
(18)
M M (22++1)m))) ((2)) M(20+10ml))
(19)
Gan'n (an)) 0]
(20)
1-1 (z)) M(20+1)mh)) Gaa-
with the external index (Source index (0..8//2 1} of level l), the internal index (node index {0..2' - 1} within generalized Source Sy,) Here, [], denotes concatenation along the axis i, and Ja and 1, denote these indices are used as axis for this concatenation one taking the lower half, one the upper half, e.g. for vectors a R.V: b, N/2.
6n
n N/2 n 2 N/2
[b(m]))) Amb)lo :=
(21)
Cn-N/2
Note that we use a causal setting here. For a non-causal structure, one would apply this bi- bidirectionally, effectively filling the top-right entries of the gating matrix with different Souree/Mark combination for the opposite order. Equivalently, both directions can be added in their outputs while computing the parallel parts twice.
B.4 Paraliclization of PLSTM in 2D - -images
Here, we now use u for internal indices along the x-axis, and v for internal indices along the y-axis, and a, b, c for index along the edge/boundary direction, while keeping x, for outer indices. These are needed for higher-level tensors, leading to a recursion shown below. The Source, Transition and Mark tensors level I have the following structure: TT. zgabt Th,> T M M with u, v, a, 20..''
SL. s'.- S Sxguva: Toyos
1+1,, Syawa
d-- [ (20) 1 (aa "anne oatp Oat Mow.se/e.......oooooooooo.eeba.e
(22)
[ [ (2x+ T2 1)(2))b The (20t + I.I (22 13 (2 (2
(a (22)
I/I (2 (20 1)++) 1obelo [S (20 (24 w
(23)
18

---

d+11I Szyucca
[ oe Mohloollalllllooo [
(24)
(e (20 (2 (22) (26)
s
Tan 1)6e T (2x 12y
(200 22
(24+ 10 ( (2x) (adl)) )o (2x+ 122 Me)))) (2x 122 Now
zyab
= [e (ee (ee (12))))))))))))))))) (ab))))
(25)
1.1 (20) (20) (2.+ 1120+))(01 + (anxew T 1)(2)) (2++1) 1)d(b)1))
Thes
The (2 (2x+ 1200 (4)))) ota
syah
= [[ (2w. (200 (6bb)0600
(26)
[T (2mneemah (22) (2.+ (22++1)((11))) + T ( ( 1222) (2++1) ()((b))))
The
T/e (20+)) (2x 0(ax))111))))
Th% syah
= [ (0=-0)m)0))))))) T2m (ee
(27)
1hes (22)eeMaaa Thes (2 (2x+ 1)22+1))(0))) (2otee (a (ew.
T/e
zyab
The (200)2+ (antevr [[ ht
(28)
This Th1 (22) (2n) yad (2x+ 1)22)+)))0))))) (2n) The T The
(22+1)(2)))) (22+1) 1)d())))
(100 (22+ 1opp))
d.-
= [IM [e (2ntte Maana wlor [1 M (2n
zyuva
(29)
The T Ma Me T
(2.) Ma
[I [[Mesrewwtutthhhhaa
M
[T [T M
(30)
d. Te The
[[owboo)))) M
(1at (twhhoo The
(2e (20 (
19

---

B.5 State-Tracking Extension
Recent research has shown that having scalar, positive Transitions T__ does not allow for state tracking capabilities in these linear RNNs [Merrill et al.] 2024] Grazzi et al. [025]. We can generalize the Transition scalars to Transition matrices here as well: Tehee Le i, E (I..Jr), Jr being the Transition state tracking dimension. Note that in this way, the Source and Mark matrix have to
include this extended dimension as well. We can define generalizations with additional state tracking dimensions Jg, JM:
Sa I.. Seki
k E
(31)
r/o
T The
i.j (L..Tr)
(32)
de
M MMltla jm
j E (1..Jr), (1.JJ)
(33)
Gm'n
Gkm
mn
(I....s),, (1..Ju)
(34)
This extension includes specific parameterizations such as defining the Transition matrix as Ti 1 + Bli,k, known from DeltaNet [Schlag et al. 2021]. For these, the chunk-wise parallelization of Yang et al. 2024 can be applied along multiple dimensions.
B.6 Stability in the State Tracking Extension
The stabilization of Section B.5 can be applied in the state tracking extensions as well. In the extended case, the absolute values of Transitions replaced by the spectral norm of the Transition matrix (Teest))s defined by its largest singular value. The stability can be ensured by limiting the sum of the largest singular values (instead of scalar entries) for the P-mode zero matrix entries, resulting in a line graph tree with maximally unit norm matrices for tree Transitions in the D-mode. There are several options to limit set the Transition matrices by/to one in magnitude. A straightforward option is to parametrize not the Transition entries Ti, directly, but to parametrize
its singular value decomposition, with orthogonal (or unitary) matrices T UVV" The singular values Σ can be limited in magnitude by a tanh (for multiplicative decay) or fixed to +1 (long- range propagation). In the case of multiple directions (see 2D grid case), in the P-mode, they are multiplied by an additional softmax- limited propagation factor ("angle") limiting the sum over the direction pairs Equation 9 The orthogonal matrices U, V can be parametrized by the product of Householder matrices (generated from vectors), or the exponential of the Lie-group generating group of special orthogonal matrices: the skew-symmetric matrices (directly parameterized), With these parametrizations, in turn depending on the network inputs (at nodes), state-tracking capabilities can be achieved while maintaining long-range stability [Merrill et al. 2024 Grazzi et al.) 2025).
B.7
Shunkwise-Paralle Form
Given the hierarchical structure of the parallelization shown in Sections B.2 B.3 and B.4 one can stop at certain level of it. At this level, the resulting objects will again form higher- leve. DAG or grid, which can now be processed sequentially, in the topological ordering of this DAG. This way, the quadratic complexity introduced by the parallelization only introduced up to the optimal level of the hardware's 's parallelization capabilities. While typically the FLOOP-optimll point is at lower parallelization, Beck et al. 20225a] showed that certain degree of parallelization is hardware optimal.
Nulti-Directional Form For the 2D grid case of images, we use PLSTM in all four direction combinations of the 2D DAG at once in parallel, at each node going down/right, down/left, up/right, up-left. By pre-computing all of them in full parallelization, we can add up their gating matrices, resulting in just single pass of the gated linear attention computation (see Equation 77. This fully parallelized form was used for the vision experiments.
20

---

C
Full Model Architecture
C.1 PLSTM for Vision
For the PLSTM-Vis, use the ViT (Dosovitskiy RMSNorm) structure as backbone architecture, replacing the Multi-Head Attention with PLSTM layer. Note that in addition query-, key-, value-, and out out- projjeecctsons. we have linear layers for Source, Transition, and Mark gates in the four direction combinations. In addition, we apply multi-head RMSNorm. Note that the Source, Mark, and Direct gates are regmoid-activatee and the Transition is tanh-activated. For details, we refer to the attached source code.
2021 with pre-LayerNorm (actually
D Theoretical Analysis
D.1 Exponential growth of path counts on 2D grid
Assume a 2D grid structure of a DAG, then the number of paths between two nodes can be calculated the following way. total, the number of Transitions to the right and Transitions to the bottom are fixed by the position offsets A₂ and Ay between the nodes. In total, the number of paths can
be counted by the number of orderings of string consisting of A₂ times and Ay times +.. This results in the combinatorial factor, which is exponential in the path length A., + =: A:
A + D,
(A, + Ay)! Stirling
(A2 + A2) (A + A A
Paths
4
A M
2
A A
A=1=
2T 2:1 B) (3 (1-11
(35)
Given Transitions forget gates potentially all at 1, this exponential number of paths leads to an exponentially growing magnitude of the Cell state C, of Equation For non-linear MD-RNN, like DAG-LSTM, this path count applies as well, assuming that all forget gates are fixed to one. Also in this case, the cell states accumulate all their ancestors' values via all paths.
D.2 Long-Range Decay for P-mode in 2D
As the P-mode offers directional propagation, here, we want to calculate the decay of one Source signal along the leading direction given all Transitions are equal within the grid and at criticality:
a T (1 (1 Mark Ma+ A.) + A Dy with path length A A + A, > 1, direction 3 = A
a This results in the following overall Transition from Source Say at x. y to
(1-
A
(1- A Ag a A.
Σ TATA =
(1 - a) (2-8)4
22(0+ 2y(0++..19+2.))
3A
Paths
(36)
Notice the binomial structure of the equation that fixes the total sum Σx, e(0..a) 1 Ky+a,) to 1, so the total activation is conserved along the diagonal. Differentiating this equation by 3, we can get the direction of largest propagation:
dull
0 1 0g log full
(1((11))4 ()4 + + Ey(==t. -0 (log (og r((1 -
+ Alog(a) - A log(1 a
(37)
Stirling
-0, (8A log( + (1-13)) log(1 B)) + Aoog(a) A Alog(1 a.)
a(11 5)
= =Alog
3(1-a)
21

---

This implies 3 = a for the direction of largest propagation. Now, inserting this direction into the Transition product:
A
a1 a))
(1-1a))
Stirling
VTT
A
a
ab
(1-
a
2
04
aa
V 1m
هامماا (1-0)
1
V
2Ta(1 a)
(38)
So, even in the critical case of the P-mode, the signal is decaying, but only power law O(A 1(3) of the path length.
D.3
Computational complexity analysis
For the (chunkwise) recurrent mode of operation, the state updates need to be taken into consideration. Let's ignore the and M multiplications for the moment. For each chunk (higher level node), there are F Fout state transition operations (T C) and Nodee updates of key-value outer products. The source terms can be absorbed into K V, and the pre-computation overhead neglected this analysis. This leads to complexity per chunk for the recurrent part (C' = + SKV):
Free
PLOP
Nein Nemut * Dr Dv + Nnodes Dr * Dv Noout
(39)
There is also the memory complexity of loaded memory parts, assuming that the loaded chunk data (K, V) can be kept in lower-level memory, iterating over edge combinations:
F MEM rec
Nain Neaut * DR Dy + Naodes * (Dx Dv)
(40)
The edge-combination memory overhead reduced by the amount of edge/cell states that can be cut down to (Nein + Neoa) Dx * Dv For the parallel computation (Q C + (G o Q K). V), get the following complexities, assuming that the gating and mark terms are pre-computed with no overhead:
FLOP
Nein * Nnodes * Dx Dy + Nnodes * (Dx + Dv)
(41)
par
F.EMM Nein * Dr Dy + Nondes (2 Dr + Dv)
(42)
par
Note that this can be further fused, if Dr. Dy are small enough. Actually, for the concrete reduction of the einsums, we rely on JAX+XLA, which might peform this suboptimally. Now, we add this up with certain overhead M per memory op (333 FLOPs/byte for a H100 on TensorCores), and add scaling for the Nein; Neout
Nnodes; a 1/2 in 2D, Nootll Nnodes
Nehunke
F EFF Neotal/ Nnoder [[2 N No N 2 * Dg Dy
(43)
.M (3 Nooess Dx Dv Nnodes (3* DR * Dv))
(44)
Removing constants and scales and calculating the derivative of the leading order terms (Dx * Dv). arrive at:
dF EFF 2(x * Nords + (2xx 1) * N 2a-22 + M 3(xx 1) N
(45)
which leads an optimal number of nodes per chunk Nnodes (derivative equals 0) for 2D 1/2):
dF EFF
N Nnodes + 3/2
nodes 3/2 M
= =
(46)
So, the optimum is:
22

---

(3/2 * M) N modes S (3 M)
(47)
For the sequential linear RNN case (a 0), it is:
N Nondes 2VM
(48)
Note that lower-order terms are neglected here, and FLOPs counting might vary in prefactors, but the overall scaling of the optimum remains the same. Also, interestingly, for graphs where the edges in between chunks scale as the number of contained nodes (or worse, a: 2 1), the derivative always positive and the optimum is the pure recurrent computation.
Modern GPUs need high arithmetic intensity [Williams al., 2009 Dao al. 2022 Pöppel et al. 2025] such that FLOPs become the bottleneck of computation. Because of this trade-off between FLOPs versus runtime + memory, typically more parallelization is beneficial for actual runtime up to certain Pareto limit (see also Beck et al. (2025a]). ). For larger images, full parallelization will not be the optimal case due to the quadratic FLOPs complexity in the image size.
E Experimental Results
jax based imple- mentation to enable efficient compilation of the computational graph. In particular, we re-use parts of Park [2024). Early experiments on using torch. compile on the torch implementation showed slowdown rather than speed-up of the model computations, which is why we use jax. Our source code, released with this work, has detached configuration that works for both jax and torch and should enable fast switch of frameworks for future changes. We use jax version 0.4.32 and CUDA 12.2.
Due to the complex structure of PLSTM in 2D and its parallelization,
E.1 Initialization
While for the ablations and arrow pointing extrapolation, we use zeros initialization for the weight terms of Source, Mark, Transition, Direct as well as Orientation (P mode angle) layers, on the DeiT-III style training on ImageNetlk, using a non-zero small normal init leads to significantly better results reported here. For ViT, we observe that LayerScale initialized with le-4 is important for ImageNet training, whereas on Arrow Pointing Extrapolation, important to initialize the LayerScale to reach the observed performance.
Table General pLSTM initialization settings
Parameter
Value
Source Bias Init
Mark Bias Init
-4
Direct Bias Init
-6
Transition Bias Init
1
Transition Scaling Factor
5 Headwise Range in [-2, 2] 1e-5
Orientation Bias Init
Multi-Head RMSNorm
Pooling
Comnr Patches
Mode
P+D (alternating)
E.2 Arrow Pointing
Here, train for 50 epochs with batch size 128, using learning rates [ le-4, 3e-4, 1e-3 ] and report the mean validation curves over five seeds at the best learning rate. For Vil., we include le-5 as the
https://jax.dev
https://pytorch.org
23

---

learning rate, as it fails to improve for higher learning rates. The validation datasets (standard and extrapolation) are generated from the same validation seed for all runs. We use linear-warmup + cosine-decay schedule with one warmup epoch starting from zero and ending at 0.001 x peak Ir. The models take about one hour of training on single NVIDIA H100-64GB.
E.2.1 Test results
Table 6: Test Results on Anrow Pointing Extrapolation (5 seeds, 90% confidence interval) and training time on one H100, * is just trained on one learning rate, without flash attention for fair comparison
Best LR
Test Acc.
Test Acc. (Ext.)
Training Time (h)
Model
PLSTM
0.0001
0.972 + 0.003
0.778 0.031
2.0
PLSTM / (no posemb.)
0.0001
0.975 + 0.003
0.769 + 0.018
2.0
PLSTM/ (P-mode only)
0.0003
0.978 + 0.005
0.746 0.027
2.0
PLSTM / (D-mode only)
0,0001
0.957 + 0.002
0.828 0.028
2.0
PLSTM/ / (STM bias only)
0.0001
0.975 + 0.003
0.784 0.020
2.0
VIT
0.0003
0.915 # 0.019
0.707 0.014
0.7"
ViL
3e-05
0.823 + 0.157
0.503 0.004
1.9
EfficientNet
0.001
0.998 + 0.001
4649
1.2
2DMamba
0 8
0.964 + 0.045
0.584 + 0.144
8.1
Mamba2I
0.0001 8
0.553 + 0.121
0.500 0.001
7.3
V2M
0.0001 8
0.948 + 0.009
0.557 + 0.036
26.7
V2M+Mamba2D
0.0001 R
0.496 + 0.005
0.500 0.001
17.9
E.3
ImageNet-1k
For training ImageNet-Ik,, we match the original training schedule of DeiT-IIII [Tourron al. 2022]. EfficientNets [Tan and Le Le]2019) 2019] as a convolution-basee baseline architecture are still the SOTA these scales, but is important to note that larger models were also trained at larger resolutions, whereas training resolution was not scaled with the models for all other models. EfficientNet, Mamba2D, and 2DMamba are non-isotropic in that their embedding dimensions are scaled up with depth. ViT, Vil., and pLSTM are isotropic, as is the reported isotropic ConvNeXt. For Xon isotropic models showed lower performance compared to non-isotropic ones. For PLSTMM a transition to non-isotropic model could therefore lead to performance gains as well.
All of the models pre-trained for up to 24 hours 4 nodes, with 4 NVIDIA H100-64GB GPUs each. ViT models are faster (about 30-50%), as our models do not yet utilize specific kernels. For counting FL.OPs, we use fvcore on the Py Tyrchh implementation. For arrow pointing extrapolation training, the Mamba2D and 2DMamba variants were about twice as slow as pLSTM despite utilizing custom kernels.
E.3.1 Ablation Settings
For the ablation studies, we use simplified training setting, without an additional fine-tuning stage, resembling a DeiT-T training over 400 epochs [ouvron et al. 202T]. In Table we provide summary of the used training hyperparameters.
E.4 Graph Benchmarks
For the graph version of PLSTM, we use torch, as the dynamic computation graph support is better for this framework. For simplicity, we also do not implement any parallelization of the graph computation, but use the recurrent form only.

---

Table 7: ImageNetlk DeiT-IIII style training parameters
Parameter
Value - Fine-Tuning"
Image Resolution
192 x 192 224 x 224
Training Epochs
800 (T), 400 (S), 400 (B) - 20
Hidden Dimension
192 (T). 384 (S), 768 (B)
Num Heads
(T), 6 (S), 12 (B)
DropPath Rate
0. (T), 0.05 (S), 0.1 (B)
LayerScale
Warmup Epochs
Peak Learning Rate
4e-3 (T), 4e-3 (S), 3e-3 (B) le-5 0.2
Weight Decay
Gradient Clip Norm
1.0
Optimizer
Lamb AdamW
Loss Type
Binary Cross Entropy 0.8
Cross Entropy
MixUp
CutMix
1.0
Label Smoothing
0.0 - 0.1
Global Batch Size
2048 512
ColorJitter
0.3 - 0.0
Auto-Aument
"3a" "randdm9--mstd.---innn
RandomErasing
0.
AugmentRepeats
3
TestCropRatio
1.0
RandomCrop
Irc
Datasets. Our experiments are conducted with 10-fold voos-validattio popular TUDatasets [Morris al., 2020]. For each fold, we use 1/10 of the respective training data as test set, 1/10 for validation and 5/10 for training. For every fold, we pick the test accuracy of the epoch with the best validation accuracy and report the average over all folds. For PLSTM, we encode the number of neighbors for each node with a standard positional encoding |Vaswani et al. 2017). Otherwise, no data transformation, augmentation, or normalization used.
Training procedure. We train all models for 100 epochs with AdamW, a learning rate of 0.001, batch size 64, and cosine decay learning rate schedule with warmup epochs.
Models.
The compared models are GAT (Velickoić et al., 2018 GCN Kipf and Welling 2017]
GIN [Xu et al.] 2019). LSTM GNN [Liang et al. 2016), MPNN GGimmer et al. | [017]. All GNNs consist of an encoder, decoder, and or 8 message passing layers. The best model configuration is selected based on the validation accuracy. pLSTM also consists of the same encoder and decoder, but has 2 layers that operate in D-mode and 2 layers that operate in P-mode. To achieve an approximately similar parameter count, we fix the hidden dimension pLSTM to 96, while the other GNNs have a hidden dimension of 128. All trained models have parameter count of < 300,000.
Graph Positional Embedding. Encodings are frequently used to boost performance on graph benchmarks. For example, encoding based eigenvectors and eigenvalues of the graph Laplacian (LapPE) (Dwivedi and Bresson| [220) has been frequently used to improve results on peptides-strus struct, while random-wall structural encoding (RWSE) performs well peptides-funt (Dwivedi et al.) 2021 |Ramášsk et al. 2022). Especially the RWSE could skew the results, as is very similar message passing (which compare against). However, in experiments, we do not include LapPE or RWSE for STT or any other method for a cleaner architectural comparison. In this setting, PLSTM performs comparable better than several baselines the peptide benchmarks. PLSTM is also fully compatible with such encodings, and integrating them could further improve performance. However, we intentionally did not include them in this study better isolate the effect of our proposed architecture incorporating the structural information.
25

---

Table 8: ImageNetlk DeiT style ablation training parameters
Parameter
Value
Image Resolution
224 x 224
Training Epochs
400
Hidden Dimension
192
Num Heads
3
DropPath Rate
0.
LayerScale
Warmup Epochs
5
Peak Learning Rate
le-3
Weight Decay
0.05
Gradient Clip Norm
1.0
Optimizer
Adam W
Loss Type
Cross Entropy
MixUp
0.8
CutMix
1.0
Label Smoothing
0.0 - 0.1
Global Batch Size
2048
ColorJitter
0.0
AutoAugment
frandm99mmstd.55
RandomErasing
0.25
AegeentRepeats
+ 1
RestCropRatio
1.0
RandomCrop
mrc
26

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper the
relevant information is Not Available.
Please provide short (1-2 sentence) justiffcaton right after your answer (even for NA).
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims of contribution are reflected by the respective related work sections, which includes to the best of our knowledge all related works, and our distinctive qualities compared to those. We believe that the abstract provides good summary of the paper's storyline.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We provide a respective paragraph as part of the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
27

---

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Our theoretical results are accompanied by the relevant assumptions and proven within the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: addition the detailed experiment setup description, we (will) provide the code including training configurations as part of our submission supplement.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
28

---

(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We (will) include the training code including training configurations as part of submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/mips.cc/ppute httpp//////ii/igggiid/////eeeee CCoeee/mddeSubeeeee setisabbmssiiiio wssionooiccyyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We specify the training details in the Experimental Section of the supplement and include the training code with configurations as part of our submission.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
29

---

7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: While for the large scale experiments on ImageNetlk, computational restric- tions limit us report just single run accuracies (as it is common for those), our smaller scale experiments all include eπror bars (and description) showing the statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide detailed training times and compute resources in the Experiments section of the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification: We use common, public and synthetic datasets that not contain PII. While we cannot rule out negative sociental impacts of this work in the distant future, we do not see mid-term impact. Our aim to benefit society by expanding the technical capabilities of machine learning systems.
30

---

Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction)
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We do not see significant societal impacts of our work in the short-term future, but improvements in understanding, capabilities and efficieny of general model architectures might well have some over long time scales.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
We do not think our models have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
31

---

12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Code used from other sources under Apache 2.0 License is properly credited both in the adapted code and the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paterssi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: While an experimentation codebase by definition experimental, we try to have core assets (model architecture, training pipeline) well tested and documented.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/codeemmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification: None of our experiments / research contains newly crowd-sourced data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
32

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our work does involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We did use LLMs for substantial, non-standard parts of this research.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.M for what should or should not be described.
33


Reviews:
[{"review_id": "YPHtfToYP9", "reviewer": "Reviewer_hoGY", "summary": "This paper extends recent sequence models based on linear RNNs and linear attention to graph structured data, particularly directed acyclic graphs (DAGs). In the introduced pLSTM framework hidden states are ascribed to edges and propagate through nodes depending on source (=input), transition (=forget), and mark (=output) gates. The linearity of the transition function allows to unroll this recurrent description (under topological ordering) into the quadratic form of linear attention with a connectivity-dependent decay/gating term. This enables both an analysis of potential parallelization and two vanishing/exploding signal stabilization techniques (P-Mode and D-Mode). Finally, the framework is instantiated and evaluated on images (2d-regular grid) and molecular graph tasks.", "questions": "- Why is it important to associate the states with edges instead of the nodes? From what I understand, the source and mark gates could also be subsumed into KV and Q respectively without significantly impacting Equations 5,6,7. This would also avoid extra notation on line-graphs in Sec. 4.2.\n- Would your analysis based on power iterations allow for other (potentially input-dependent) stabilization techniques beyond P-mode and D-mode? What would be the effect of limiting the norm of the row (i.e. node local input normalization)? Would row-wise softmax recover GATs? Would the exponential gating with output normalization of xLSTM also be stable?\n- Could you compare V2M and Mamba2D to pLSTM on the arrow task? Since the final model interleaves D-Mode and P-Mode layer, could you furthermore compare to a model with interleaved V2M and Mamba2D layers?\n- In recent gated linear attention models [2,3], an important aspect for parallelization is that the gating matrix $G^{nn’}$ is a rank-1 structure which allows to decompose an outer product and subsume the resulting vectors into the key and query terms. Significant speedups are realized by using the associative scan only once to precompute cumulative gating as a preprocessing step. Can you explain how this translates to parallelization on DAGs and how you can compute pLSTM on regular grids (in particular the 1D case) without cumulative sums/products?\n- Can you explain how your parallelization scheme differs from applying parallel kernels along each dimension independently, as is (presumably) done in the works that you discussed as most similar to yours (2DMamba, V2M, Mamba2D)?\n\n------------\n\n[2] Yang et al. (2023), *“Gated Linear Attention Transformers with Hardware-Efficient Training”*, https://arxiv.org/abs/2312.06635\n\n[3] Dao et al. (2024), *“Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”*, https://arxiv.org/abs/2405.21060", "limitations": "Beyond the limitations already discussed in the paper, a discussion of the potential for parallelization in practice would be interesting. This is particularly important, as the parallel algorithm for general Graphs is purely conceptual without practical evidence that speedups could be realized as opposed to the recurrent implementation.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The pLSTM framework provides a precise yet intuitive generalization of linear RNNs and linear attention to DAGs. It particularly avoids unnecessary complexities introduced by state-space model (SSM) perspective such as discretization.\n- Stability guarantees in forward and backward pass can then be derived through general conditions on the (input-dependent) values of the adjacency matrix. This understanding potentially enables design and analysis of various stabilization techniques beyond the parametrizations induced by SSM discretization.\n- Promising results for two instantiations of the framework in very distinct application domains highlight the potential of the approach.\n\n**Weaknesses**\n- A comparison to e.g. Mamba for graph data is missing (cf. survey [1]) and on 2D grids, the resulting stabilization techniques (P-mode and D-mode) turn out to be very minor modifications to existing works, as depicted in Figure 1. Therefore, for the resulting pLSTM models, it could be better discussed what the key novelties are. **Update:** *addressed in the rebuttal*\n- Along that line, the empirical evaluation does not provide a complete understanding how pLSTM models compare to related work. In particular, the baselines for Graph data are all more than 7 years old. On the 2D arrow task which tests spatial reasoning, the claimed conceptual advantage over works which are considered most related (2DMamba, V2M, Mamba2D) is not empirically verified. **Update:** *partially addressed in the rebuttal*\n\n**Remarks**\n- In Figure 3, it would be helpful to include a green triangle for the edge on the left and variable names $n’$, $e$, $n$, $z\\_n$, $h\\_n$, and $e’$, $n$, $e$, $K^nV^n$, $H^n$ to the graph elements to support understanding of Equation 1,2 and 5,6. The caption could be more descriptive.\n- There is an inconsistency in the notation of $T^n\\_{ee’}$. While Equation 5 uses $T\\_{\\\\text{outgoing, incoming}}^{\\\\text{node}}$, Equation 7 uses (mistakingly?) $T\\_{\\\\text{incoming, outgoing}}^{\\\\text{node}}$. That Sec. 4.2 switches the convention to using $e’$ as outgoing and $e$ as incoming edge complicates things further. I would suggest to maintain the original notation $T^n\\_{ee’}$ and explicitly mention *“from edge $e’$ to $e$”* in L200. In L213,  $T^n\\_{e’e}$ is (mistakenly?) used as the adjacency matrix $\\mathbf{T}$.\n- From Equation 5,6 to 7, the order of symbols is reversed which makes it harder to follow and possibly caused notational inconsistency above. Maintaining the original QKV form and assuming topological ordering would make it more clear that $G^{nn’}$ is a graph-causal gating matrix which is applied element-wise to the attention matrix QK.\n- The conceptual step from DAGs to (bidirectional) grids or general graphs is not very well explained. In particular the notation switches completely from $T\\_{\\\\text{edge}}^{\\\\text{node}}$ to a different grid 2D specific notation with higher-level tensors which is only introduced in the appendix. Finally it is unclear how a DAG pLSTM layer is applied to general graphs.\n- State Tracking Extension seems very tangential. In theory the diagonal transition of any linear RNN can always be trivially replaced by a (structured) dense transition matrix. In practice however, the difficulty lies in efficient parallelization long-range stability.\n\n**Conclusion**\nI believe that the concepts introduced in this paper will have a positive impact and therefore recommend it for acceptance. I would, however,  recommend to the authors to focus more on how their framework can contribute to novel stabilization and parallelisation schemes. \n\n----------\n\n[1] Atitallah et al. (2024), *“Exploring Graph Mamba: A Comprehensive Survey on State-Space Models for Graph Learning”*, https://arxiv.org/abs/2412.18322", "was_revised": true, "final_justification": "The authors addressed my main concerns: potential extensions of pLSTM for other stabilization modes and how their parallelization method relates to previous works.\n\nOverall I believe that the pLSTM graph mixing framework with its novel stabilization analysis and parallelization technique poses a significant and unique contribution for the field. Although the evaluation is still limited, I believe that the results are promising especially since confounding factors such as backbone and hyperparameters generally favour established architectures over emerging ones. I would, however, recommend to the authors to further shift the focus from introducing a new architecture (which evokes an expectation for sota results on benchmarks) to introducing a novel framework for stabilization/parallelization (where promising results on benchmarks suffice). Since the authors addressed most of my concerns regarding the framework perspective, I increase the score to 5.", "rebuttal": "Thank you for your detailed feedback and your questions. \n\n## Notation inconsistencies and visual improvement\nWe have integrated your helpful feedback on the descriptive Figure 3, and the inconsistency of the T matrix notation is fixed now, where of in-coming and out-going edge index was switched. Also we added the suggested notation phrase “from in-coming edge e’ to out-going edge e”. We also made the order of S, Q, K, V consistent across all equations.\n\n## DAGs to pLSTM\nRegarding the transition from a description of general DAGs to pLSTM on the 2D grids, we agree that the hierarchy used for parallelization could be better introduced in the grid case. The reference to Figure 5 should be emphasized as well, as the visualization greatly helps the understanding of the parallelization hierarchy. We have now revise our manuscript to improve the introduction of the grid case, moving parts of the additional notation to the appendix as it is not needed for understanding in the main paper, only for the full formulae in the appendix.\n\n## State-tracking extension\nThe state tracking extension is indeed tangential to the main contributions and is added for clarity on how to extend this framework further for DeltaNet and other architectures. Note that additional gating is crucial here (in P-mode), such that activations remain stable.\n\n## pLSTM vs. Mamba-based models\nWe have added a discussion of Mamba and Transformers applied to graph data for a broader conceptual comparison. Note that other works apply Mamba on graph data merely via randomly chosen paths in the graph. Regarding the 2D grid case (images), we acknowledge and cite the previous works on extending Mamba to 2D. pLSTM provides a unifying framework for these approaches and generalizes the applicability to DAGs. Also, pLSTM introduces a general principle of parallelization not explored in these works, as well as an extended understanding of stabilization of linear networks on DAGs (and other grids), including a geometrical interpretation thereof. \n\nRegarding additional experiments on graph data, we have results on Peptides-struct and Peptides-func. Also, on these larger-scale benchmarks, pLSTM performs well, despite no use of structural embeddings/encodings. For the results, see the answer to reviewer T3S3. \n\n## Association of recurrent states with edges\n\nAssociating states with edges rather than nodes is indeed a deliberate design choice, primarily motivated by the goal of enabling efficient parallelization and directional information flow.\n\nWhile traditional recurrent models (e.g., MPNNs) often associate states with nodes, this becomes limiting when trying to model directional transitions across complex structures like DAGs or grids in a parallelizable way. In our framework, information is propagated between larger regions. By tying states to edges, we naturally capture the directionality of transitions, crucial for enabling our directed propagation mode (P-mode). Without associating state with edges, transitions would either be direction-agnostic or require duplicating logic for each edge type (e.g., incoming vs. outgoing), leading to inefficiencies and loss of structure-specific inductive bias.\n\nRegarding the Source and Mark gates: while it is possible to fold them into standard KV/Q formulations, we chose to represent them explicitly for structural clarity and interpretability. Together with the Transition scalar, they form a complete gating mechanism and offer points of fast adaptation similar to learnable residual pathways. This decomposition helps separate concerns (source selection, direction marking, and propagation strength), making the model easier to analyze and potentially more flexible.\n\n## Other stabilization modes\nThank you for the insightful questions. Yes, our power iteration-based analysis provides a flexible framework that can accommodate various stabilization techniques beyond the proposed P-mode and D-mode.\n\nUsing row-wise normalization, as you suggest, would correspond to controlling the $ L_{\\inf} $ norm (affecting activations) instead of the $ L_{1} $ norm (affecting gradients), offering an alternative tradeoff in stability. This opens the door to input-dependent normalizations such as softmax, which would resemble GAT-like mechanisms. However, GAT also includes self-loops and nonlinearities, which are not present in our formulation. Our design omits these explicitly to preserve long-range propagation. #CHANGE#\nSimilarly, xLSTM-style exponential gating with output normalization is compatible with our formulation. Such normalization can be applied externally to the pLSTM core, for instance by adapting the input gates accordingly. Normalizer terms can be computed as running pLSTM with $ V_v = 1 $, stabilization of an exponential Source gate (<-> input gate) here can be achieved by subtracting the maximal pre-activation value from all others. This allows pLSTM to inherit desirable stability properties while maintaining modularity.\n\nIn general, the modular structure of pLSTM supports such augmentations and we believe further exploration of these stabilization variants is a promising direction for future work.\n\n## Arrow pointing results\nIn the response to reviewer T3S3, we include additional results on arrow pointing extrapolation for other models: Mamba2D, 2DMamba, V2M, ViL, V2M+Mamba2D, and EfficientNet. pLSTM outperforms all of these on the extrapolation test set due to its superior inductive bias on directional long-range interaction. We also found that a small update on weight initialization strongly improves ImageNet-1k training results for the larger scales, now outperforming the ViL baseline. \n\n## Structural difference between the difference scans / parallelization forms\nOur parallelization strategy differs significantly from those used in 2DMamba, V2M, and Mamba2D, both in efficiency and in how directional dependencies are handled.\nV2M performs two sequential 1D scans (along rows and columns), often followed by naive parallel fusion. This approach is inherently limited by the sequential nature of 1D scans and requires explicit materialization of many intermediate states, which leads to increased memory usage and reduced arithmetic intensity, suboptimal for GPU execution.\nMamba2D adopts a wavefront (diagonal) parallelization strategy. While this reduces the depth of sequential steps from O(P) to O(sqrt(P)), it still scales poorly with larger input sizes and lacks full parallelism, particularly when compared to scan-based or logarithmic-time alternatives.\n2DMamba is structurally similar to V2M, but fuses some operations in 2D according to the GPU structure, which might still be sub-optimal for linear RNNs regarding arithmetic intensity.\n\nIn contrast, pLSTM generalizes associative scan-style parallelism to DAGs and 2D grids by operating on the line graph representation. For regular grids (2D), this enables precomputing the gating matrix in O(½ log P) sequential steps, using tensorized operations like einsum, concatenation and padding.\nThe gated/weighted aggregation step, analogous to masked linear attention, is performed in a single pass with O(P^2) FLOPs. In practice, we can use chunkwise-recurrent formulations to reduce the FLOP complexity to O(P) complexity for large images.\nAlthough chunking was not used in our experiments (up to 384×384 images), we instead performed a full fusion over all four DAG directions (right-bottom, left-bottom, right-top, left-top), allowing us to consolidate four fully parallel traversals into one unified computation with complete gating precomputation.\n\n\n## Graph Parallelization Limitation\nThank you for pointing this out. We agree that the parallel algorithm for general DAGs remains conceptual at this stage and that practical implementation and benchmarking are important directions for future work.\n\nThat said, we emphasize that our current recurrent implementation for general DAGs is significantly slower than the fully parallelized baselines (e.g., CNNs, ViTs, Mamba2D). This clear performance gap strongly suggests that substantial speedups are achievable once parallelization is realized. Moreover, our method is structurally designed to enable parallel execution using line-graph representations and associative scan-style operations—tools already optimized in many modern compute frameworks.\n\n\n## Creation of the DAG from a regular graph\nA general undirected graph $G$ can be converted into a directed acyclic graph $\\mathrm{DAG}$ using the following procedure:\n\nInput: An undirected graph $G = (V, E)$, where $V$ is the set of nodes and $E \\subseteq \\{\\{u, v\\} \\mid u, v \\in V, u \\neq v\\}$ is the set of undirected edges.\n\nOutput: A directed acyclic graph $\\mathrm{DAG} = (V, E')$.\n\n* Assign a unique index to each node: $\\mathrm{idx}[v] \\gets$ index of $v$ in a fixed ordering of $V$.\n* Initialize the DAG with all nodes $v \\in V$.\n* Initialize the edge set: $E' \\gets \\emptyset$.\n* For each edge $\\{(u, v)\\} \\in E$:\n         If $\\mathrm{idx}[u] < \\mathrm{idx}[v], E' \\gets E' \\cup \\{(u, v)\\}$\n         Otherwise,  $E' \\gets E' \\cup \\{(v, u)\\}$.\n* Return the DAG: $\\mathrm{DAG} = (V, E')$.\n\n\nThe output is a directed acyclic graph $\\mathrm{DAG} = (V, E')$, which admits a valid topological ordering and can be used as input to a pLSTM or other DAG-based neural architectures. In our experiments, we use a random fixed ordering for the graph benchmarks. Note that this ordering is non-unique. Future work could address this limitation by using multiple random orderings instead of one.\nThe 2D grid admits a special ordering that is self-similar, which we use here.\n\n\nOur manuscript has improved substantially by addressing and incorporating your feedback. If any questions remain, we are happy to engage in further discussion. If you find our responses and additional results useful, we would appreciate it if you would consider changing your score."}, {"review_id": "r9ofspLw55", "reviewer": "Reviewer_T3S3", "summary": "The authors propose an innovative architecture that extends linear RNNs to Directed Acyclic Graphs (DAGs), enabling chunkwise-parallel processing. To address parallelization on DAGs, they introduce two distinct parallelization strategies: the Propagation Model (P Model), which covers all paths but faces efficiency constraints, and the Distribution Model (D Model), which retains only a subset of paths for enhanced parallelism. Experimental results on both 2D (image) and graph data demonstrate that this architecture surpasses Vision Transformers (ViT) and achieves performance comparable to popular state-of-the-art models.", "questions": "See above.\n\n1. Please explain about the efficiency of the architecture. What is the theoretical computational complexity for each parallelization policy? Why the experiments do not show the better efficiency?\n2. More results on arrow datasets and graph datasets.", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The architecture is novel. The idea to extend the linear RNN to non-regular DAGs data is interesting. By combining the two distinct model (P model and D model), different propagation paradigms can be constructed for structured data.\n2. The architecture promise well parallelizablity.\n\nWeakness:\n1. The most weakness is the not outstanding experiment results for both performance and the efficiency. The motivation of linear RNN is the efficiency. But the flops results shown in the experiments show similar efficiency compared to other models. And the experiments on graph data lacks the discussion about the efficiency. \n2. In the arrow pointing experiments, the only one baseline is the ViT, why other baseline models are not tested on this datasets?\n3. The architecture is proposed for DAGs, but the main experiments are on image datasets. I think more experiments on graph datasets could help the authors to proof their novelty and significance. At the same time, both the accuracy and the efficiency should be discussed.", "was_revised": true, "final_justification": "Because my main concern has been resolved, I will raise my score to 4.", "rebuttal": "We thank the reviewer for the constructive feedback. We are glad that you consider our architecture novel and the extension to non-regular DAGs interesting.\n\nWe have now extended the arrow-pointing extrapolation experiments to include additional baselines. As shown, pLSTM outperforms all other baselines by a clear margin on the extrapolation task, while remaining competitive in terms of training time. \n\n| Model          |  Test acc. (ext) - 5 runs | Training time (h) |\n|----------------|---------------------------|-------------------|\n| EfficientNet   |  0.649 ± 0.023            |       1.2         |\n| ViT            |   0.707 ± 0.014           |       0.4         |\n| ViL            |  0.503 ± 0.004            |       1.9         |\n| Mamba2D        |  0.500 ± 0.001            |       7.3         |\n| 2DMamba        |  0.584 ± 0.144            |       8.1         |\n| V2M            | 0.557 ± 0.036             |      26.7         |\n| V2M + Mamba2D  |  0.500 ± 0.001            |      17.9         |\n| pLSTM          |   **0.778 ± 0.031**       |       2.0         |\n\nMoreover, we are happy to provide additional results on ‘Peptides-struct’ and ‘Peptides-func’ as new graph datasets part of the Long Range Graph Benchmark:\n\n| Model                |     Peptides-struct (MAE)↓    |     Peptides-func (Avg. Prec.)↑   |\n|----------------------|-------------------------------|-----------------------------------|\n| GAT                          |  0.270       |  0.525     |\n| GCN                         |  0.264     |  0.534    |\n| GIN                          |  0.328        |  0.593   |\n| LSTM-GNN             |   0.274                                        |  0.502 |\n| MPNN                 |   0.260                                     |  0.557    |\n| pLSTM              |  0.264                                      |  0.536  |\n\npLSTM can achieve results similar to re-trained baseline models on these benchmarks, being the second best for the ‘Peptides-struct’. We encoded atoms as nodes and bonds as edges and did not use any additional encoding such as Laplacian encoding or RWSE. This has a negative effect on the results, however it allows for a more precise comparison without any confounding factors.\n\n\nAdditionally, we observed that a minor change to the weight initialization (increasing the initialization scale) improves pLSTM's performance on ImageNet-1k, particularly at larger model sizes. This change likely introduces a beneficial regularization effect via increased activation diversity. \nWe would like to emphasize that our primary goal is not to outperform current state-of-the-art architectures, but rather to introduce a novel and generalizable architectural mechanism. Our work opens the door for future improvements and integration into broader model designs. Note also that both Mamba2D and 2DMamba use non-isotropic backbones, which often yield stronger performance at scale, making the performance of isotropic pLSTM even more encouraging.\n\n| Model      | Test acc. |\n|-------------|-------------|\n| pLSTM-T |      75.2  |\n| pLSTM-S |       80.7 |\n| pLSTM-B |       82.5 | \n\n## Efficiency considerations \nThank you for the question. The parallelization scheme for DAGs was not used in our graph benchmark experiments due to the irregularity of the graph structures. Instead, we used a recurrent implementation, which has low FLOP complexity but higher runtime due to limited parallelism.\n\nFor structured data like 1D/2D grids, we do implement efficient parallel versions, which benefit from reduced memory overhead and fast batched operations. We believe the parallelization scheme has strong theoretical potential and plan to explore efficient implementations for general graphs in future work.\n\n## FLOPs considerations\nFor the (chunkwise) recurrent mode of operation, the state updates need to be taken into consideration. Let's ignore the S and M multiplications for the moment.\nFor each chunk (higher level node), there are $E_{in} * E_{out}$ state transition operations $(T * C)$ and $N_{nodes}$ updates of key-value outer products. The source terms can be absorbed into K / V, and the pre-computation overhead is neglected in this analysis. This leads to a complexity per chunk for the recurrent part ( $C' = \\sum T \\cdot C + S K V$ ):\n\n$ F_{rec}^{FLOP} = N_{ein} * N_{eout} * D_{K} * D_{V} + N_{nodes} * D_{K} * D_{V} * N_{eout}$ \n\nThere is also the memory complexity of loaded memory parts, assuming that the loaded chunk data (K, V) can be kept in lower-level memory, iterating over edge combinations:\n\n$F_{rec}^{MEM} = N_{ein} * N_{eout} * D_{K} * D_{V} + N_{nodes} * (D_{K} + D_{V})$\n\nThe edge-combination memory overhead is reduced by the amount of edge/cell states that can be cut down to $(N_{ein} + N_{eout}) * D_{K} *D_{V}$ .\nFor the parallel computation ($Q \\cdot C + (G \\odot Q \\cdot K) \\cdot V$ ), we get the following complexities, assuming that the gating and mark terms are pre-computed with no overhead, and all :\n\n$F_{par}^{FLOP} = N_{ein} * N_{nodes} * D_{K} * D_{V} + N_{nodes}^2 * (D_{K} + D_{V})$\n\n$F_{par}^{MEM} = N_{ein} * D_{K} * D_{V} + N_{nodes} (2 * D_{K}  + D_{V})$\n\nNote that this can be further fused if $D_{K}, D_{V}$ are small enough. Actually, for the concrete reduction of the einsums, we rely on JAX+XLA, which might perform this suboptimally.\nNow, if we add this up with a certain overhead $M$ per memory op (333 FLOPs/byte for a H100 on TensorCores), and add a scaling for the $N_{ein}, N_{eout} = N_{nodes}^{\\alpha} , \\alpha = 1/2$ in 2D, $N_{total} = N_{nodes} * N_{chunks}$:\n\n\n$F^{EFF} = N_{total} / N_{nodes} * ( (2 * N_{nodes}^{1+\\alpha} + N_{nodes}^{2 \\alpha} + N_{nodes}^{2}) * D_{K} * D_{V}  +  M * (3 * N_{nodes}^\\alpha * D_{K} * D_{V}   + N_{nodes} * ( 3 * D_{K} + 2 * D_{V} )))  $\n\nRemoving constants and scales and calculating the derivative of the leading order terms (D_K * D_V), we arrive at:\n\n$d\\tilde{F}^{EFF} = 2 \\alpha * N_{nodes}^{\\alpha-1} + (2 \\alpha - 1) * N_{nodes}^{2 \\alpha - 2} + 1  +  M * 3 (\\alpha - 1) * N_{nodes}^{\\alpha - 2} = 0$\n\nwhich leads to an optimal number of nodes per chunk $N_{nodes}$ (derivative equals 0) for 2D ($\\alpha = 1/2$):\n\n$d\\tilde{F}^{EFF} = N_{nodes} + N_{nodes}^{3/2}  -  3/2 * M = 0$\n\nSo the optimum is:\n\n$(3/2 * M )^{2/3} \\leq N_{nodes}^{2D*} \\leq (3 * M)^{2/3}$\n\nFor the sequential linear RNN case ($\\alpha = 0$), it is:\n\n$N^{1D*}_{nodes} = 2 \\sqrt{M}$\n\nNote that lower-order terms are neglected here, and FLOPs counting might vary in prefactors, but the overall scaling of the optimum remains the same.\nAlso, interestingly, for graphs where edges in between chunks scale as the number of contained nodes (or worse, $\\alpha \\geq 1$), the derivative is always positive and the optimum is the pure recurrent computation. \n\nModern GPUs need a high arithmetic intensity such that FLOPs become the bottleneck of computation. Because of this tradeoff between FLOPs vs. runtime+memory, typically more parallelization is beneficial for actual runtime up to a certain Pareto limit (see [TFLA]). For larger images, full parallelization will not be the optimal case due to the quadratic FLOPs complexity in the image size.\nThank you for raising this important point. We have added an extended version of this analysis to our updated manuscript. \nNote that we integrated several other models now into our codebase, and the training time for the arrow pointing experiments shows that this is not strongly correlated with FLOPs (B=128, R=192x192, 1 NVIDIA H100 GPU), but rather implementation efficiency.\n\nOur manuscript has improved considerably by incorporating your feedback: thank you. If you find our responses and additional results useful, we would appreciate it if you would consider changing your score."}, {"review_id": "7SGdNgbeLT", "reviewer": "Reviewer_hu73", "summary": "This paper demonstrates a gated linear recurrent network with associative cell, similar to MAMBA, designed for arbitrary DAGs.  Two modes of propagation are proposed: a D-mode that is similar to 2DMAMBA and has the same decay properties assuming the transition matrix is norm-constrained as described there, and a P-mode that propagates with decay exponential in the path length in an arbitrary learned direction provided that the L1-norm of each column of the transition matrix is less than 1.  Results are demonstrated for three applications: On ImageNet-1K it performs within 4% of SoTA, on several small-molecule graph sets it performs within 1% of SoTA, and on a new long-distance 2D inference task it outperforms ViT.", "questions": "Does 2DMAMBA also outperform pLSTM for the arrow-pointing task?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThe framework elegantly generalizes linear gated recurrent networks to arbitrary graphs. The proof of non-explosion is clear and elegant and might even imply a statement more general than what is claimed here.  \n\nWeaknesses:\n\nThere is a lot of recent research trying to generalize linear gated recurrent units to 2D.  Although this formalism is more elegant than others I've seen, its performance considerably lags 2DMAMBA.\n\nMinor issues/ possible typos:\n\nThe second term on RHS in both Eq. (6) and Eq. (7) has no factor with a \"v\" subscript, suggesting that this term is constant w.r.t. the v subscript on the LHS.  That seems wasteful and inconsistent.  The recurrent terms suggest that V_k^n should have been V_v^n, whereas the phrase \"skip directly from input to output\" suggests that Q_kK_kV_k should have been Q_vK_vV_v; I wonder if this is a typo, and if so, which correction is correct?\n\nIs Eq. (8) intended to be approximate?  It is true if |s| means any vector norm and |T| means the consistent matrix norm, but it is untrue for some other common interpretations of those symbols, e.g., I have often seen |T| used to denote the matrix determinant, which would make Eq. (8) untrue. It seems that Eq. (8) is only really used to define the P-mode in Eq. (9), since the D-mode uses a different derivation, but if you were to use double-bars instead of single-bars in Eq. (8), it could be interpreted as correctly licensing other propagation modes besides P-mode and D-mode (specifically, it seems to license any propagation mode in which ||T|| is the matrix norm consistent with ||s|| and ||T||<=1).", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and, in particular, the detailed look at the presented equations.\n\nIndeed, in the RHS terms of Eq. (6) and (7), the $V$ term should have subscript $V_{v}$ instead of $V_{k}$. Thank you for pointing out this typo. We have now fixed this in our updated version of our manuscript.\n\n## P-mode norm choice and notation\n\nRegarding the norms in Eq. 8, these are intended to denote any vector norm and a respective consistent matrix norm (we have updated to a double | notation). You are right that this allows for multiple options for the P-mode. Note, however, that there are only two matrix norms that can be computed locally without iterating over the graph globally: The $L_{1}$ and $L_{\\inf}$ norms. Note that these are also dual to each other, which provides both stability for the activations and gradients (see l.218ff). One could flip the norms for activations and gradients, leading to a slightly different P-mode indeed. But, for other norms, the global graph structure would need to be considered, which creates an undesirable additional overhead. Therefore, we restrict the norm to $L_{1}$ for our proposed P-mode.\n\n## Additional arrow pointing extrapolation results\n\nThank you for the question. We address 2DMamba's performance on the arrow-pointing task in our response to Reviewer T3S3, where we include additional results. Notably, 2DMamba underperforms pLSTM on this task, despite its stronger results on ImageNet.\n\n## ImageNet improvements\n\nWe also improved pLSTM's ImageNet results with modified initialization (see T3S3 response). Importantly, pLSTM is trained in an isotropic setting (no hidden dimension upscaling), whereas 2DMamba, V2M, and Mamba2D use non-isotropic configurations, which are known to boost performance (e.g., in the V2M paper). In the revised version, we will clearly mark this distinction.\n\nThank you again for your positive assessment of our work. If any questions remain or come up during the discussion period, we would be happy to engage in further discussion."}]

Author rebuttal:

Meta-review: This paper introduces pLSTM, a novel framework that generalizes linear RNNs to directed acyclic graphs (DAGs) via edge-focused gating and a hierarchical parallelization scheme, and provides a stability analysis yielding two operational modes (P-mode and D-mode).

During the discussion, the reviewers raised concerns about how the approach differs from Mamba-based approaches, the limited and somewhat outdated baselines (especially for graph data), and the lack of demonstrated efficiency gains given the motivation of linear RNNs. The authors responded thoroughly with new results (on ImageNet, APE, and peptide graph benchmarks), clarified novelty vs. prior work, corrected notational issues, and provided detailed complexity/runtime analyses. They also emphasized that pLSTM is a general framework rather than an optimized benchmark competitor.

Most reviewers agreed that the framework is conceptually strong and opens new directions, even if benchmark performance is not yet state of the art. Given the conceptual advances I believe asking for state-of-the-art performance is not necessary and recommend accepting the paper.
Decision: Accept (poster)