pLSTM: Redes de Transición de Fuente Lineal Paralelizables

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

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

cs.AI

Resumen

Las arquitecturas recurrentes modernas, como xLSTM y Mamba, han desafiado recientemente al Transformer en el modelado del lenguaje. Sin embargo, su estructura limita su aplicabilidad a secuencias o requiere procesar estructuras de datos multidimensionales, como imágenes o grafos moleculares, en un orden secuencial predefinido. En contraste, las Redes Neuronales Recurrentes Multidimensionales (MDRNNs) son adecuadas para datos con un nivel de estructura más alto, como cuadrículas 2D, árboles y grafos acíclicos dirigidos (DAGs). En este trabajo, extendemos la noción de multidimensionalidad a las RNNs lineales. Introducimos las redes de Transición de Fuente Lineal paralelizables (pLSTMs) utilizando puertas de Fuente, Transición y Marca que actúan sobre el grafo lineal de un DAG general. Esto permite la paralelización en analogía con los escaneos asociativos paralelos y la forma recurrente por bloques de las RNNs lineales secuenciales, pero para DAGs. Para cuadrículas regulares (1D y 2D), como imágenes, este esquema puede implementarse eficientemente utilizando operaciones einsum, concatenaciones y relleno en tiempo logarítmico. Las pLSTMs abordan el problema de la desaparición/explosión de activaciones/gradientes para distancias largas en DAGs mediante dos modos distintos: un modo de propagación dirigida (modo P) y un modo de distribución difusiva (modo D). Para demostrar las capacidades de largo alcance de las pLSTMs, introducimos la extrapolación de apuntado de flechas como una tarea sintética de visión por computadora que contiene información direccional de larga distancia. Demostramos que las pLSTMs generalizan bien a tamaños de imagen más grandes, mientras que los Transformers tienen dificultades para extrapolar. En benchmarks establecidos de grafos moleculares y visión por computadora, las pLSTMs también muestran un rendimiento sólido. El código y los conjuntos de datos están disponibles en: 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.