pLSTM: Redes de Transição de Fonte Linear Paralelizáveis

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

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

cs.AI

Resumo

Arquiteturas recorrentes modernas, como xLSTM e Mamba, recentemente desafiaram o Transformer na modelagem de linguagem. No entanto, sua estrutura limita sua aplicabilidade a sequências ou exige o processamento de estruturas de dados multidimensionais, como imagens ou grafos moleculares, em uma ordem sequencial predefinida. Em contraste, Redes Neurais Recorrentes Multidimensionais (MDRNNs) são bem adequadas para dados com uma estrutura de nível superior, como grades 2D, árvores e grafos acíclicos direcionados (DAGs). Neste trabalho, estendemos a noção de multidimensionalidade para RNNs lineares. Introduzimos redes Lineares de Transição de Fonte Marcada paralelizáveis (pLSTMs) usando portas de Fonte, Transição e Marca que atuam no grafo de linha de um DAG geral. Isso permite a paralelização em analogia a varreduras associativas paralelas e à forma recorrente por blocos de RNNs lineares sequenciais, mas para DAGs. Para grades regulares (1D e 2D), como imagens, esse esquema pode ser implementado de forma eficiente usando operações einsum, concatenações e preenchimento em tempo logarítmico. As pLSTMs abordam o problema de ativação/gradiente que desaparece/explode para longas distâncias em DAGs por meio de dois modos distintos: um modo de propagação direcionada (modo P) e um modo de distribuição difusiva (modo D). Para demonstrar as capacidades de longo alcance da pLSTM, introduzimos a extrapolação de apontamento de seta como uma tarefa sintética de visão computacional que contém informações direcionais de longa distância. Demonstramos que as pLSTMs generalizam bem para tamanhos de imagem maiores, enquanto os Transformers têm dificuldade em extrapolar. Em benchmarks estabelecidos de grafos moleculares e visão computacional, as pLSTMs também mostram um desempenho forte. Código e conjuntos de dados estão disponíveis em: 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.