pLSTM: параллелизуемые сети с линейными переходами источников

pLSTM: parallelizable Linear Source Transition Mark networks

June 13, 2025

Авторы: Korbinian Pöppel, Richard Freinschlag, Thomas Schmied, Wei Lin, Sepp Hochreiter

cs.AI

Аннотация

Современные рекуррентные архитектуры, такие как xLSTM и Mamba, недавно бросили вызов Transformer в области языкового моделирования. Однако их структура ограничивает их применимость только к последовательностям или требует обработки многомерных структур данных, таких как изображения или молекулярные графы, в заранее определенном последовательном порядке. В отличие от них, многомерные рекуррентные нейронные сети (MDRNN) хорошо подходят для данных с более сложной структурой, например, двумерных сеток, деревьев и направленных ациклических графов (DAG). В данной работе мы расширяем понятие многомерности на линейные рекуррентные сети. Мы представляем параллелизуемые сети с линейными переходами и маркировкой (pLSTM), использующие ворота Источника, Перехода и Маркировки, которые действуют на линейный граф общего DAG. Это позволяет реализовать параллелизацию по аналогии с параллельными ассоциативными сканированиями и чанково-рекуррентной формой последовательных линейных RNN, но для DAG. Для регулярных сеток (1D и 2D), таких как изображения, эта схема может быть эффективно реализована с использованием операций einsum, конкатенаций и заполнения за логарифмическое время. pLSTM решают проблему исчезновения/взрыва активаций/градиентов на больших расстояниях в DAG с помощью двух различных режимов: режима направленного распространения (P-режим) и режима диффузного распределения (D-режим). Чтобы продемонстрировать способности pLSTM к работе с длинными дистанциями, мы вводим задачу экстраполяции направления стрелки как синтетическую задачу компьютерного зрения, содержащую информацию о направлении на больших расстояниях. Мы показываем, что pLSTM хорошо обобщаются на изображения большего размера, в то время как Transformer испытывают трудности с экстраполяцией. На устоявшихся бенчмарках для молекулярных графов и задач компьютерного зрения pLSTM также демонстрируют высокую производительность. Код и наборы данных доступны по адресу: 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.