{"id":13444620,"url":"https://github.com/tobiaskatsch/GatedLinearRNN","last_synced_at":"2025-03-20T19:30:37.524Z","repository":{"id":221636131,"uuid":"754915630","full_name":"tobiaskatsch/GatedLinearRNN","owner":"tobiaskatsch","description":null,"archived":false,"fork":false,"pushed_at":"2024-02-27T04:52:42.000Z","size":1281,"stargazers_count":27,"open_issues_count":0,"forks_count":1,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-03-16T02:41:11.301Z","etag":null,"topics":["deep-learning","gateloop"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2311.01927","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/tobiaskatsch.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE.txt","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-02-09T02:28:15.000Z","updated_at":"2025-02-12T04:46:46.000Z","dependencies_parsed_at":"2024-02-27T05:36:02.150Z","dependency_job_id":null,"html_url":"https://github.com/tobiaskatsch/GatedLinearRNN","commit_stats":null,"previous_names":["tobiaskatsch/gateloop","tobiaskatsch/lineargatedrnn","tobiaskatsch/gatedlinearrnn"],"tags_count":4,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tobiaskatsch%2FGatedLinearRNN","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tobiaskatsch%2FGatedLinearRNN/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tobiaskatsch%2FGatedLinearRNN/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tobiaskatsch%2FGatedLinearRNN/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/tobiaskatsch","download_url":"https://codeload.github.com/tobiaskatsch/GatedLinearRNN/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":244676454,"owners_count":20491828,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","gateloop"],"created_at":"2024-07-31T04:00:32.331Z","updated_at":"2025-03-20T19:30:36.632Z","avatar_url":"https://github.com/tobiaskatsch.png","language":"Python","funding_links":[],"categories":["ICLR 2024 submissions"],"sub_categories":[],"readme":"# GatedLinearRNN\n\n\u003cdiv style=\"display: flex; justify-content: space-around; align-items: center;\"\u003e\n    \u003cimg src=\"assets/gated_linear_rnn.png\" alt=\"GatedLinearRNN\" width=\"300px\"/\u003e\n    \u003cimg src=\"assets/eq.png\" alt=\"Equation\" width=\"400px\"/\u003e\n\u003c/div\u003e\n\n\u003e **GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling**\\\n\u003e Tobias Katsch*\\\n\u003e Paper: https://arxiv.org/abs/2311.01927\n\u003e\n\n## About\n\nLinear Gated RNNs (Mamba, GateLoop, HGRN) form a novel class of sequence models which generalize and generalize linear recurrent \nmodels such as S4, S5, LRU and RetNet, by employing data-controlled state transitions. \nWhile having a low cost linear complexity inference mode, they can be trained extremely \nefficient in parallel with logarithmic complexity making use of the highly optimized JAX \nassociative scan implementation. This repository implements a practical gated linear recurrent model with default \nchoices for input-, hidden- and gate activations and provides a drop-in replacement for causal multi-head-attention \nand a linear gated RNN language model architecture. Furthermore, linear gated RNNs can be used to train true recurrent \nmodels (GRU, LSTM) extremely fast by first training using associative scans and switching to a true recurrent mode \n(by enabling recurrent weights) for finetuning.\n\n## Installation\n\n- `pip install flax-gated-linear-rnn`:\n\nOther requirements:\n- JAX 0.4.20+\n- FLAX 0.8.0+\n\n## Usage\nWe provide 2 main modules:\n- ### [gated_linear_rnn.py](flax_gated_linear_rnn/gated_linear_rnn.py)\n  A causal time mixing sequence model which can be used as a drop-in replacement for causal multi-head-attention.\n  Usage:\n  ```\n  import jax\n  import jax.numpy as jnp\n  from flax import linen as nn\n  from flax_gated_linear_rnn import GatedLinearRNN\n  \n  batch_size, sequence_length, input_dim, d_h, d_model = 2, 64, 16, 32, 18\n  key = jax.random.PRNGKey(0)\n  x = jax.random.normal(key, (batch_size, sequence_length, input_dim))\n  \n  model = GatedLinearRNN(\n      d_model=d_model,\n      d_h=d_h,\n      input_activation=nn.tanh,\n      hidden_activation=nn.tanh,\n      gate_activation=nn.sigmoid,\n      use_true_recurrence=False,\n      use_tied_gates=True,\n  )\n  \n  params = model.init(jax.random.PRNGKey(1), x)\n  carry, y = model.apply(params, x, carry=None)\n  assert y.shape == (batch_size, sequence_length, d_model)\n  assert carry.shape == (batch_size, d_h)\n  ```\n  `carry` holds the hidden model state, which can be used for fast linear complexity autoregressive inference.\n  #### Two Stage Training\n  - **Associative Recurrent Mode:** (`use_true_recurrence=False`) Extremely efficient training through associative scan. This disables the recurrent weights, allowing for much faster training compared to Transformer, GRU \u0026 LSTM.\n  - **True Recurrent Mode:** (`use_true_recurrence=True`) Can be used to train a more expressive model from a Linear Recurrent Model checkpoint. This variant introduces additional parameters such that the inputs and gates also depend on previous hidden states similar to GRU \u0026 LSTM. Due to the true recurrent nature, this mode cannot be parallelized and thus is less efficient. We recommend this for finetuning from an linear recurrent checkpoint.\n\n  #### Gate Tying\n  - **Disjoint Input \u0026 Forget gate** (`use_tied_gates=False`) Applies seperate projections for input- \u0026 forget gates\n  - **Tied Input \u0026 Forget gate** (`use_tied_gates=True`) Ties the input and forget gate through the relation `forget_gate = 1-input_gate`.\n\n\n- ## [gated_linear_rnn_lm.py](flax_gated_linear_rnn/language_models/gated_linear_rnn_lm.py)\n  A GatedLinearRNN-based language model.\n  ```\n  import jax\n  import jax.numpy as jnp\n  from flax import linen as nn\n  from flax_gated_linear_rnn import GatedLinearRNNLM\n  \n  # Model parameters\n  n_layer = 4\n  d_model = 256\n  d_channel_mixing = 64\n  eps = 1e-6\n  channel_mixing_dropout = 0.1\n  time_mixing_dropout = 0.1\n  input_vocab_size = 10000\n  output_vocab_size = 10000\n  max_seq_length = 512\n  embedding_dropout = 0.1\n  use_word_embedding = True\n  positional_encoding_mode = 'none'  # 'none', 'learned', 'sinusoidal'\n  d_h = 256\n  \n  model = GatedLinearRNNLM(\n    n_layer=n_layer,\n    d_model=d_model,\n    d_channel_mixing=d_channel_mixing,\n    eps=eps,\n    channel_mixing_dropout=channel_mixing_dropout,\n    time_mixing_dropout=time_mixing_dropout,\n    input_vocab_size=input_vocab_size,\n    output_vocab_size=output_vocab_size,\n    max_seq_length=max_seq_length,\n    embedding_dropout=embedding_dropout,\n    use_word_embedding=use_word_embedding,\n    positional_encoding_mode=positional_encoding_mode,\n    d_h=d_h,\n    use_head=True,\n    input_activation=nn.tanh,\n    hidden_activation=nn.tanh,\n    gate_activation=nn.sigmoid,\n    use_true_recurrence=False,\n    use_tied_gates=True,\n  )\n  \n  # Sample input\n  batch_size = 32\n  x = jax.random.randint(jax.random.PRNGKey(0), (batch_size, max_seq_length), 0, input_vocab_size)\n  \n  # Initialize and apply model\n  params = model.init(jax.random.PRNGKey(2), x, training=False)\n  carry, y = model.apply(params, x, training=True, carry=None, rngs={'dropout': jax.random.PRNGKey(0)})\n  assert y.shape == (batch_size, max_seq_length, output_vocab_size)\n  assert carry.shape == (batch_size, n_layer, d_h)\n  ```\n  `carry` holds the hidden GatedLinearRNN model states for all layers which can be used for fast linear complexity autoregressive inference.\n\n## Synthetic speech generation examples\n\nhttps://tobiaskatsch.github.io/GatedLinearRNN/\n\n\u003caudio controls\u003e\n  \u003csource src=\"https://tobiaskatsch.github.io/GatedLinearRNN/assets/generated_1.wav\" type=\"audio/wav\"\u003e\n\u003c/audio\u003e\n\n\u003caudio controls\u003e\n  \u003csource src=\"https://tobiaskatsch.github.io/GatedLinearRNN/assets/generated_2.wav\" type=\"audio/wav\"\u003e\n\u003c/audio\u003e\n\n\u003caudio controls\u003e\n  \u003csource src=\"https://tobiaskatsch.github.io/GatedLinearRNN/assets/generated_3.wav\" type=\"audio/wav\"\u003e\n\u003c/audio\u003e\n\n\u003caudio controls\u003e\n  \u003csource src=\"https://tobiaskatsch.github.io/GatedLinearRNN/assets/generated_4.wav\" type=\"audio/wav\"\u003e\n\u003c/audio\u003e\n\n\n\n## Citation\n\nIf you use this codebase, please cite:\n```\n@misc{katsch2024GateLoop,\n      title={GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling}, \n      author={Tobias Katsch},\n      year={2024},\n      eprint={2311.01927},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftobiaskatsch%2FGatedLinearRNN","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftobiaskatsch%2FGatedLinearRNN","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftobiaskatsch%2FGatedLinearRNN/lists"}