pLSTM: parallelizable Linear Source Transition Mark networks

Thank you for your review and, in particular, the detailed look at the presented equations.

Indeed, 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.

P-mode norm choice and notation

Regarding 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.

Additional arrow pointing extrapolation results

Thank 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.

ImageNet improvements

We 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.

Thank 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.

As the discussion period draws to a close, we wanted to sincerely thank you for your time and thoughtful feedback. We've done our best to address your concerns in the rebuttal and hope that our responses have clarified the key points of the paper. If there are any remaining questions or uncertainties, we would be happy to elaborate further while the discussion is still open. Should you feel that your concerns have been addressed, we would greatly appreciate your consideration in reflecting that in your final evaluation. Thank you again for participating in the review process. We truly value your input.

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.

We 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.

Moreover, we are happy to provide additional results on ‘Peptides-struct’ and ‘Peptides-func’ as new graph datasets part of the Long Range Graph Benchmark:

pLSTM 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.

Additionally, 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. We 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.

Efficiency considerations

Thank 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.

For 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.

FLOPs considerations

For 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. For 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$ ):

$F_{rec}^{FLOP} = N_{ein} * N_{eout} * D_{K} * D_{V} + N_{nodes} * D_{K} * D_{V} * N_{eout}$

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_{rec}^{MEM} = N_{ein} * N_{eout} * D_{K} * D_{V} + N_{nodes} * (D_{K} + D_{V})$

The 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}$ . For 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 :

$F_{par}^{FLOP} = N_{ein} * N_{nodes} * D_{K} * D_{V} + N_{nodes}^2 * (D_{K} + D_{V})$

$F_{par}^{MEM} = N_{ein} * D_{K} * D_{V} + N_{nodes} (2 * D_{K} + D_{V})$

Note 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. Now, 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}$:

$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} )))$

Removing constants and scales and calculating the derivative of the leading order terms (D_K * D_V), we arrive at:

$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$

which leads to an optimal number of nodes per chunk $N_{nodes}$ (derivative equals 0) for 2D ($\alpha = 1/2$):

$d\tilde{F}^{EFF} = N_{nodes} + N_{nodes}^{3/2} - 3/2 * M = 0$

So the optimum is:

$(3/2 * M )^{2/3} \leq N_{nodes}^{2D*} \leq (3 * M)^{2/3}$

For the sequential linear RNN case ($\alpha = 0$), it is:

$N^{1D*}_{nodes} = 2 \sqrt{M}$

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 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.

Modern 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. Thank you for raising this important point. We have added an extended version of this analysis to our updated manuscript. Note 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.

Our 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.

Thank you for your thoughtful follow-up. To support our claim:

Yes, pLSTM is fully compatible with standard structural encodings, such as Laplacian positional encodings (LapPE) [1] and Random Walk Structural Encodings (RWSE) [2, 3].

1. For LapPE and basic RWSE, the node features are enhanced by precomputed positional encodings (e.g., Laplacian eigenvectors or random walk probabilities), which are simply added to or concatenated with the input embeddings. This is directly compatible with pLSTM and requires no architectural changes.

2. For more expressive RWSE variants [3], which introduce additional edge features or virtual edges (e.g., between distant nodes in the same structure), pLSTM can naturally operate on the extended graph by incorporating these edges and features in its source and transition computations. This is because the gating functions in pLSTM are flexible and can integrate arbitrary edge features.

In summary, incorporating these encodings into pLSTM is straightforward and opens up clear directions for improving performance. We did not include them in our current experiments to isolate the architectural contribution of pLSTM itself, but we agree that adding such encodings would likely lead to stronger results.

Thank you again for the engaging discussion and your openness to revising the score.

[1] A Generalization of Transformer Networks to Graphs, Dwivedi et al., 2021

[2] Graph Neural Networks with Learnable Structural and Positional Representations, Dwivedi et al., 2022

[3] Recipe for a General, Powerful, Scalable Graph Transformer, Rampášek et al., 2022

Thank you for your detailed feedback and your questions.

Notation inconsistencies and visual improvement

We 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.

DAGs to pLSTM

Regarding 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.

State-tracking extension

The 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.

pLSTM vs. Mamba-based models

We 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.

Regarding 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.

Association of recurrent states with edges

Associating 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.

While 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.

Regarding 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.

Other stabilization modes

Thank 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.

Using 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# Similarly, 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.

In 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.

Arrow pointing results

In 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.

Structural difference between the difference scans / parallelization forms

Our parallelization strategy differs significantly from those used in 2DMamba, V2M, and Mamba2D, both in efficiency and in how directional dependencies are handled. V2M 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. Mamba2D 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. 2DMamba 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.

In 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. The 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. Although 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.

Graph Parallelization Limitation

Thank 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.

That 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.

Creation of the DAG from a regular graph

A general undirected graph $G$ can be converted into a directed acyclic graph $\mathrm{DAG}$ using the following procedure:

Input: 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.

Output: A directed acyclic graph $\mathrm{DAG} = (V, E')$.

• Assign a unique index to each node: $\mathrm{idx}[v] \gets$ index of $v$ in a fixed ordering of $V$.
• Initialize the DAG with all nodes $v \in V$.
• Initialize the edge set: $E' \gets \emptyset$.
• For each edge $\{(u, v)\} \in E$: If $\mathrm{idx}[u] < \mathrm{idx}[v], E' \gets E' \cup \{(u, v)\}$ Otherwise, $E' \gets E' \cup \{(v, u)\}$.
• Return the DAG: $\mathrm{DAG} = (V, E')$.

The 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. The 2D grid admits a special ordering that is self-similar, which we use here.

Our 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.

I thank the authors for their detailed response. I particularly appreciate their discussion of my main concerns: potential extensions of pLSTM for other stabilization modes and how their parallelization method relates to previous works. Beyond that I would like to offer a few remarks.

• Association of recurrent states with edges: I was hoping that the notation could be simplified further by directly working with the line graph representation, but this design choice is ultimately up to the authors.

• pLSTM vs. Mamba-based models: Thank you for adding a discussion regarding Mamba and Transformers on graph data. While I’m not very familiar with the literature, I found for example [1] which seems to apply the LRU in a graph setting without randomly chosen paths in the graph.

• Limited evaluation spatial reasoning: Thank you for comparing pLSTM to more baselines, particularly “V2M + Mamba2D” on the arrow task; this addresses my concern regarding limited evaluation of spatial reasoning.

• Limited Evaluation graph data: The results on Graph data, however, are not particularly insightful since the used baselines are quite outdated. For example, the baselines used in [1] are more recent and also seem to outperform pLSTM on first glance although the experimental setup like model size etc might differ.

Overall 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 will increase the score to 5.

[1] Ding et al. (2024), “Recurrent Distance Filtering for Graph Representation Learning”, https://arxiv.org/abs/2312.01538