arXiv:2506.11997cs.LGstat.ML2025-06NeurIPS被引 2

提出可并行计算的图结构序列模型,突破传统RNN长距离依赖瓶颈。

pLSTM: parallelizable Linear Source Transition Mark networks

  • 基于有向无环图的线性RNN架构,通过源/转移/标记门实现并行化
  • 在图像和分子图上优于Transformer,支持大尺寸图像外推
  • 适用于图像、分子图等高维结构数据,适合需要长程依赖的场景

现代循环架构如xLSTM和Mamba已挑战Transformer在语言建模中的地位,但其结构限制了对非序列数据(如图像、分子图)的应用,通常需预定义顺序处理。相比之下,多维RNN(MDRNN)更适合具有更高结构的数据,如2D网格、树和有向无环图(DAG)。本文提出可并行化的线性源转移马尔可夫网络(pLSTM),利用源、转移和标记门作用于一般DAG的线图,实现类似并行关联扫描和分块递归形式的并行化。对于规则网格(1D和2D,如图像),可通过einsum操作、拼接和填充在对数时间内高效实现。pLSTM通过两种模式解决DAG中长距离的激活/梯度消失/爆炸问题:定向传播模式(P-mode)与扩散分布模式(D-mode)。为展示pLSTM的长程能力,引入箭头指向外推这一合成计算机视觉任务,包含长距离方向信息。实验表明,pLSTM能良好外推至更大图像尺寸,而Transformer表现不佳。在分子图和计算机视觉基准上,pLSTM也表现出色。代码与数据集见:https://github.com/ml-jku/plstm_experiments。

原文摘要 · Abstract (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.

序列模型图神经网络并行计算长程依赖

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。