pLSTM: paralleliseerbare Lineaire Bron Transitie Mark-netwerken

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

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

cs.AI

Samenvatting

Moderne recurrente architecturen, zoals xLSTM en Mamba, hebben onlangs de Transformer uitgedaagd in taalmodelering. Hun structuur beperkt echter hun toepasbaarheid tot sequenties of vereist het verwerken van multidimensionale gegevensstructuren, zoals afbeeldingen of moleculaire grafieken, in een vooraf gedefinieerde volgorde. Daarentegen zijn Multidimensionale RNN's (MDRNN's) goed geschikt voor gegevens met een hoger niveau van structuur, zoals 2D-roosters, bomen en gerichte acyclische grafieken (DAG's). In dit werk breiden we het concept van multidimensionaliteit uit naar lineaire RNN's. We introduceren paralleliseerbare Lineaire Bron Transitie Mark-netwerken (pLSTM's) met Bron-, Transitie- en Mark-poorten die inwerken op de lijngrafiek van een algemene DAG. Dit maakt parallelisatie mogelijk, analoog aan parallelle associatieve scans en de chunkwise-recurrente vorm van sequentiële lineaire RNN's, maar voor DAG's. Voor regelmatige roosters (1D en 2D), zoals afbeeldingen, kan dit schema efficiënt worden geïmplementeerd met behulp van einsum-bewerkingen, concatenaties en padding in logaritmische tijd. pLSTM's pakken het probleem van verdwijnende/exploderende activatie/gradiënt voor lange afstanden in DAG's aan via twee verschillende modi: een gerichte propagatiemodus (P-modus) en een diffuse distributiemodus (D-modus). Om de langeafstandscapaciteiten van pLSTM te demonstreren, introduceren we pijlaanwijzende extrapolatie als een synthetische computervisietaak die langeafstandsrichtingsinformatie bevat. We tonen aan dat pLSTM's goed generaliseren naar grotere afbeeldingsformaten, terwijl Transformers moeite hebben met extrapoleren. Op gevestigde benchmarks voor moleculaire grafieken en computervisie laten pLSTM's ook sterke prestaties zien. Code en datasets zijn beschikbaar op: 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.