pLSTM: reti di transizione di sorgente lineare parallelizzabili

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

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

cs.AI

Abstract

Le moderne architetture ricorrenti, come xLSTM e Mamba, hanno recentemente sfidato il Transformer nel campo del language modeling. Tuttavia, la loro struttura limita la loro applicabilità a sequenze o richiede l'elaborazione di strutture dati multidimensionali, come immagini o grafi molecolari, in un ordine sequenziale predefinito. Al contrario, le RNN multidimensionali (MDRNN) sono particolarmente adatte per dati con una struttura più complessa, come griglie 2D, alberi e grafi aciclici diretti (DAG). In questo lavoro, estendiamo il concetto di multidimensionalità alle RNN lineari. Introduciamo le reti Linear Source Transition Mark parallelizzabili (pLSTM) utilizzando porte Source, Transition e Mark che agiscono sul grafo lineare di un DAG generico. Ciò consente la parallelizzazione in analogia agli scan associativi paralleli e alla forma chunkwise-recurrent delle RNN lineari sequenziali, ma per i DAG. Per griglie regolari (1D e 2D), come le immagini, questo schema può essere implementato in modo efficiente utilizzando operazioni einsum, concatenazioni e padding in tempo logaritmico. Le pLSTM affrontano il problema della scomparsa/esplosione delle attivazioni/gradienti per distanze lunghe nei DAG attraverso due modalità distinte: una modalità di propagazione diretta (P-mode) e una modalità di distribuzione diffusa (D-mode). Per dimostrare le capacità a lungo raggio delle pLSTM, introduciamo l'estrapolazione della direzione delle frecce come un compito sintetico di computer vision che contiene informazioni direzionali a lunga distanza. Dimostriamo che le pLSTM si generalizzano bene a dimensioni di immagine più grandi, mentre i Transformer faticano a estrapolare. Su benchmark consolidati di grafi molecolari e computer vision, le pLSTM mostrano anche prestazioni solide. Codice e dataset sono disponibili su: 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.