pLSTM: Parallelisierbare Lineare Quellübergangs-Markierungsnetzwerke

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

Autoren: Korbinian Pöppel, Richard Freinschlag, Thomas Schmied, Wei Lin, Sepp Hochreiter

cs.AI

Zusammenfassung

Moderne rekurrente Architekturen wie xLSTM und Mamba haben kürzlich den Transformer in der Sprachmodellierung herausgefordert. Ihre Struktur beschränkt jedoch ihre Anwendbarkeit auf Sequenzen oder erfordert die Verarbeitung mehrdimensionaler Datenstrukturen, wie Bilder oder molekulare Graphen, in einer vordefinierten sequenziellen Reihenfolge. Im Gegensatz dazu sind mehrdimensionale RNNs (MDRNNs) gut geeignet für Daten mit einer höheren Struktur, wie 2D-Gitter, Bäume und gerichtete azyklische Graphen (DAGs). In dieser Arbeit erweitern wir den Begriff der Mehrdimensionalität auf lineare RNNs. Wir führen parallelisierbare Linear Source Transition Mark-Netzwerke (pLSTMs) ein, die Source-, Transition- und Mark-Gates verwenden, die auf dem Linien-Graphen eines allgemeinen DAGs wirken. Dies ermöglicht eine Parallelisierung analog zu parallelen assoziativen Scans und der chunkweise-rekurrenten Form sequenzieller linearer RNNs, jedoch für DAGs. Für regelmäßige Gitter (1D und 2D), wie Bilder, kann dieses Schema effizient mit Einsum-Operationen, Verkettungen und Padding in logarithmischer Zeit implementiert werden. pLSTMs adressieren das Problem verschwindender/explodierender Aktivierungen/Gradienten für lange Distanzen in DAGs durch zwei verschiedene Modi: einen gerichteten Propagationsmodus (P-Modus) und einen diffusiven Verteilungsmodus (D-Modus). Um die Langstreckenfähigkeiten von pLSTM zu demonstrieren, führen wir die Pfeilzeiger-Extrapolation als synthetische Computer-Vision-Aufgabe ein, die langstreckige Richtungsinformationen enthält. Wir zeigen, dass pLSTMs gut auf größere Bildgrößen verallgemeinern, während Transformer Schwierigkeiten haben, zu extrapolieren. Auf etablierten molekularen Graphen- und Computer-Vision-Benchmarks zeigen pLSTMs ebenfalls starke Leistungen. Code und Datensätze sind verfügbar unter: https://github.com/ml-jku/plstm_experiments.

English

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