pLSTM : réseaux de marquage de transition de source linéaire parallélisables

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

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

cs.AI

Résumé

Les architectures récurrentes modernes, telles que xLSTM et Mamba, ont récemment remis en question le Transformer dans le domaine de la modélisation du langage. Cependant, leur structure limite leur applicabilité aux séquences uniquement ou nécessite le traitement de structures de données multidimensionnelles, telles que les images ou les graphes moléculaires, selon un ordre séquentiel prédéfini. En revanche, les réseaux de neurones récurrents multidimensionnels (MDRNNs) sont bien adaptés aux données présentant une structure de niveau supérieur, comme les grilles 2D, les arbres et les graphes orientés acycliques (DAGs). Dans ce travail, nous étendons la notion de multidimensionnalité aux RNNs linéaires. Nous introduisons les réseaux de transition de source linéaire parallélisables (pLSTMs) utilisant des portes de Source, de Transition et de Marque qui agissent sur le graphe linéaire d'un DAG général. Cela permet une parallélisation analogue aux scans associatifs parallèles et à la forme récurrente par blocs des RNNs linéaires séquentiels, mais pour les DAGs. Pour les grilles régulières (1D et 2D), comme les images, ce schéma peut être implémenté efficacement à l'aide d'opérations einsum, de concaténations et de remplissages en temps logarithmique. Les pLSTMs abordent le problème de la disparition/explosion des activations/gradients sur de longues distances dans les DAGs via deux modes distincts : un mode de propagation dirigée (mode P) et un mode de distribution diffuse (mode D). Pour mettre en avant les capacités à longue portée des pLSTMs, nous introduisons l'extrapolation de pointage de flèche comme une tâche synthétique de vision par ordinateur contenant des informations directionnelles à longue distance. Nous démontrons que les pLSTMs généralisent bien à des tailles d'images plus grandes, alors que les Transformers peinent à extrapoler. Sur des benchmarks établis en vision par ordinateur et en graphes moléculaires, les pLSTMs montrent également de solides performances. Le code et les jeux de données sont disponibles à l'adresse suivante : 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.