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に挑戦を投げかけています。しかし、これらの構造はシーケンスのみへの適用を制約するか、画像や分子グラフなどの多次元データ構造を事前に定義された順序で処理する必要があります。対照的に、多次元RNN（MDRNN）は、2Dグリッド、ツリー、有向非巡回グラフ（DAG）などのより高次元の構造を持つデータに適しています。本研究では、線形RNNに多次元性の概念を拡張します。一般DAGの線グラフ上で動作するSource、Transition、Markゲートを用いた並列化可能な線形Source Transition Markネットワーク（pLSTM）を導入します。これにより、並列連想スキャンやシーケンシャル線形RNNのチャンクワイズリカレント形式に類似した並列化がDAGに対して可能になります。画像のような規則的なグリッド（1Dおよび2D）では、このスキームはeinsum操作、連結、およびパディングを用いて対数時間で効率的に実装できます。pLSTMは、DAG内の長距離における活性化/勾配の消失/爆発問題を、2つの異なるモード、すなわち指向性伝播モード（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.