pLSTM：可并行化线性源转换标记网络

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