pLSTM: parallelizable Linear Source Transition Mark networks

现代循环架构，例如 xLSTM 和 Mamba，在语言建模方面最近对 Transformer 构成了挑战。然而，它们的结构限制了其仅适用于序列数据，或者需要以预定义的顺序处理多维数据结构，如图像或分子图。相比之下，多维循环神经网络（MDRNNs）非常适合具有更高层次结构的数据，例如二维网格、树和有向无环图（DAGs）。在本工作中，我们将多维性的概念扩展到了线性 RNN。我们引入了可并行化的线性源转换标记网络（pLSTMs），该网络使用源门、转换门和标记门作用于一般 DAG 的线图。这使得可以类似并行关联扫描以及顺序线性 RNN 的分块递归形式的方式对 DAG 进行并行化。对于规则网格（如一维和二维图像），该方案可以通过 einsum 操作、拼接和填充在对数时间内高效实现。pLSTMs 通过两种不同的模式解决了 DAG 中长距离的激活/梯度消失或爆炸问题：一种是有向传播模式（P-模式），另一种是扩散分布模式（D-模式）。为了展示 pLSTM 的长程能力，我们引入了一种合成计算机视觉任务——箭头指向外推，其中包含长距离的方向信息。我们证明了 pLSTMs 在更大的图像尺寸上具有良好的泛化能力，而 Transformer 在外推时表现困难。在已建立的分子图和计算机视觉基准测试中，pLSTMs 同样表现出强大的性能。代码和数据集可在以下链接获取：https://github.com/ml-jku/plstm_experiments。

论文试图解决长距离依赖问题在多维和非线性数据结构（如图像、分子图等）中的建模挑战。现有的RNN变体或Transformer在处理这些结构时，通常受限于序列化输入或难以扩展到复杂的数据形式。这是一个重要的问题，但并非全新的研究方向，而是对现有方法的改进与扩展。

论文提出了一种新的架构pLSTM（parallelizable Linear Source Transition Mark networks），通过Source、Transition和Mark门操作，在DAG（有向无环图）的线图上实现并行化计算。这种方法不仅适用于1D/2D网格数据（如图像），还能高效处理树形结构和DAG。相比传统的RNN或Transformer，pLSTM引入了P-mode（定向传播模式）和D-mode（扩散分布模式）来缓解长距离激活/梯度消失或爆炸问题。

1. 提出了针对DAG的并行化RNN架构pLSTM，支持高效处理多维数据结构；2. 在图像箭头指向预测任务中展示了其优秀的长距离信息捕捉能力，而Transformer在此类任务中表现较差；3. 在分子图和计算机视觉基准测试中表现出强大的性能；4. 提供了开源代码和数据集（https://github.com/ml-jku/plstm_experiments），便于复现和进一步研究；5. 值得继续探索的方向包括如何将pLSTM应用于更复杂的动态图结构以及更大规模的真实世界任务。

近期相关工作包括：1. xLSTM和Mamba等现代RNN架构，尝试挑战Transformer在语言建模中的主导地位；2. Multi-Dimensional RNNs (MDRNNs)，专门设计用于处理具有高阶结构的数据；3. Graph Transformers，例如Graphormer和GAT，专注于图结构数据的建模；4. 其他结合RNN和图结构的研究，如Graph RNNs（GRNNs）。这些方法为pLSTM的设计提供了理论基础和技术参考。