pLSTM: parallelizable Linear Source Transition Mark networks

要約

XLSTMやMambaなどの最新の再発アーキテクチャは、最近、言語モデリングの変圧器に挑戦しました。
ただし、その構造は、シーケンスのみへの適用性を制約するか、画像や分子グラフなどの多次元データ構造を事前に定義された順序で処理する必要があります。
対照的に、多次元RNN（MDRNNS）は、2Dグリッド、ツリー、監督された非環式グラフ（DAG）など、より高いレベル構造を持つデータに適しています。
この作業では、多次元の概念を線形RNNに拡張します。
一般的なDAGの線グラフに作用するソース、遷移、およびマークゲートを使用して、並列化可能な線形ソース遷移ネットワーク（PLSTMS）を導入します。
これにより、並行した連想スキャンと連続的な線形RNNの塊状回復型と類似した並列化が可能になりますが、DAGの場合。
画像などの通常のグリッド（1Dおよび2D）の場合、このスキームは、名誉操作、連結、および対数時間のパディングを使用して効率的に実装できます。
PLSTMSは、2つの異なるモードを介してDAGの長距離の消失/爆発のアクティベーション/勾配問題に取り組みます：指向伝播モード（Pモード）と拡散分布モード（Dモード）。
PLSTMの長距離機能を紹介するために、長距離の方向情報を含む合成コンピュータービジョンタスクとして矢印を指している外挿を導入します。
私たちは、PLSTMがより大きな画像サイズによく一般化するのに対し、トランスは外挿に苦労することを実証します。
確立された分子グラフとコンピュータービジョンベンチマークでは、PLSTMも強力なパフォーマンスを示しています。
コードとデータセットは、https：//github.com/ml-jku/plstm_experimentsで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google