Skip to content
/ plstm Public

Main implementation in flax.nnx, flax.linen and torch of parallellizable Linear Source Transition Mark networks

License

Notifications You must be signed in to change notification settings

ml-jku/plstm

Repository files navigation

pLSTM: parallelizable Linear Source Transition Mark Networks - Core

arXiv License: MIT

Korbinian Pöppel1,2, Richard Freinschlag1, Thomas Schmied1, Wei Lin1,, Sepp Hochreiter1,2

1ELLIS Unit Linz and LIT AI Lab, Institute for Machine Learning, Johannes Kepler University Linz, Austria
2NXAI GmbH, Linz, Austria

This repository contains the pLSTM (parallelizable Linear Source Transition Mark networks) core implementations in flax.nnx, flax.linen and torch. pLSTMs inherit ideas from Multi-Dimensional RNNs Graves et al. 2007 and linear RNNs. With the linearity, and changing the gating structure to the Source, Transition and Mark gates, we introduce a multi-dimensional parallel associative scan, on general directed acyclic graphs (DAGs) for parallelization.

pLSTMs also solve the vanishing/exploding gradient/activation problem on DAGs, similar to how the LSTM tackled them for RNNs on sequences.

Configuration

All layers within pLSTM can be configured using the config classes in plstm.config composed by way of compoconf library.

Framework Implementations

pLSTM offers implementations across multiple popular deep learning frameworks:

  • nnx
  • linen
  • torch

Graph Implementation

Please note that plstm_graph is currently only implemented in torch.

References

MD-RNNs:

linear RNNs (among lots of others):

State-Tracking:

License

MIT License

Citation

If you use this dataset in your research, please cite:

@misc{poppel_plstm_2025,
	title = {{pLSTM}: parallelizable {Linear} {Source} {Transition} {Mark} networks},
	shorttitle = {{pLSTM}},
	url = {http://arxiv.org/abs/2506.11997},
	doi = {10.48550/arXiv.2506.11997},
	urldate = {2025-06-16},
	publisher = {arXiv},
	author = {Pöppel, Korbinian and Freinschlag, Richard and Schmied, Thomas and Lin, Wei and Hochreiter, Sepp},
	month = jun,
	year = {2025},
	note = {arXiv:2506.11997 [cs]},
	keywords = {Computer Science - Machine Learning, Statistics - Machine Learning},
}

About

Main implementation in flax.nnx, flax.linen and torch of parallellizable Linear Source Transition Mark networks

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published