{"id":13752933,"url":"https://github.com/lucidrains/recurrent-memory-transformer-pytorch","last_synced_at":"2025-05-15T09:07:15.875Z","repository":{"id":154140790,"uuid":"631755045","full_name":"lucidrains/recurrent-memory-transformer-pytorch","owner":"lucidrains","description":"Implementation of Recurrent Memory Transformer, Neurips 2022 paper, in Pytorch","archived":false,"fork":false,"pushed_at":"2025-01-06T16:55:15.000Z","size":35956,"stargazers_count":407,"open_issues_count":8,"forks_count":16,"subscribers_count":13,"default_branch":"main","last_synced_at":"2025-04-10T06:39:13.760Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","deep-learning","long-context","memory","recurrence","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","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":"2023-04-24T01:35:47.000Z","updated_at":"2025-04-08T20:30:09.000Z","dependencies_parsed_at":null,"dependency_job_id":"0114c1ef-985d-480c-ac8e-97b94e69049c","html_url":"https://github.com/lucidrains/recurrent-memory-transformer-pytorch","commit_stats":{"total_commits":68,"total_committers":2,"mean_commits":34.0,"dds":"0.014705882352941124","last_synced_commit":"5027c9e457dc6b098875387b39e7d9e4dc43d7ee"},"previous_names":[],"tags_count":44,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Frecurrent-memory-transformer-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Frecurrent-memory-transformer-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Frecurrent-memory-transformer-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Frecurrent-memory-transformer-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/recurrent-memory-transformer-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254310515,"owners_count":22049469,"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":["artificial-intelligence","attention-mechanisms","deep-learning","long-context","memory","recurrence","transformers"],"created_at":"2024-08-03T09:01:12.813Z","updated_at":"2025-05-15T09:07:10.867Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Transformer库与优化","Python"],"sub_categories":[],"readme":"\u003cimg src=\"./rmt.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Recurrent Memory Transformer - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2207.06881\"\u003eRecurrent Memory Transformer\u003c/a\u003e \u003ca href=\"https://openreview.net/forum?id=Uynr3iPhksa\"\u003e(openreview)\u003c/a\u003e in Pytorch. They had \u003ca href=\"https://arxiv.org/abs/2304.11062\"\u003ea short follow up paper\u003c/a\u003e recently that demonstrated it was able to copy information across 1 million tokens at the very least.\n\nThere is no doubt in my mind that RMT would make a stronger RL agent than \u003ca href=\"https://sites.google.com/view/adaptive-agent/\"\u003eAdA\u003c/a\u003e, which is just a Transformer-XL - Update: \u003ca href=\"https://sites.google.com/view/rate-model/\"\u003eRecurrent Action Transformer with Memory (RATE)\u003c/a\u003e \n\n\u003ca href=\"https://www.youtube.com/watch?v=4Cclp6yPDuw\"\u003eYannic Kilcher paper review\u003c/a\u003e\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStability\u003c/a\u003e and \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for their generous sponsorships to work on and open source cutting edge artificial intelligence research\n\n## Install\n\n```bash\n$ pip install recurrent-memory-transformer-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom recurrent_memory_transformer_pytorch import RecurrentMemoryTransformer\n\nmodel = RecurrentMemoryTransformer(\n    num_tokens = 20000,               # number of tokens\n    num_memory_tokens = 128,          # number of memory tokens, this will determine the bottleneck for information being passed to the future\n    dim = 512,                        # model dimensions\n    depth = 6,                        # transformer depth\n    causal = True,                    # autoregressive or not\n    dim_head = 64,                    # dimension per head\n    heads = 8,                        # heads\n    seq_len = 1024,                   # sequence length of a segment\n    use_flash_attn = True             # whether to use flash attention\n)\n\nx = torch.randint(0, 256, (1, 1024))\n\nlogits1, mem1, _ = model(x)        # (1, 1024, 20000), (1, 128, 512), None\nlogits2, mem2, _ = model(x, mem1)  # (1, 1024, 20000), (1, 128, 512), None\nlogits3, mem3, _ = model(x, mem2)  # (1, 1024, 20000), (1, 128, 512), None\n\n# and so on ...\n\n```\n\nWith XL memories\n\n```python\nimport torch\nfrom recurrent_memory_transformer_pytorch import RecurrentMemoryTransformer\n\nmodel = RecurrentMemoryTransformer(\n    num_tokens = 20000,\n    num_memory_tokens = 128,\n    dim = 512,\n    depth = 6,\n    causal = True,\n    dim_head = 64,\n    heads = 8,\n    seq_len = 1024,\n    use_flash_attn = True,\n    use_xl_memories = True,    # set this to True\n    xl_mem_len = 512           # can be shorter than the seq len - i think just having a bit of the past will prevent much of the RMT memories  memorizing the immediate preceding text\n)\n\nx = torch.randint(0, 256, (1, 1024))\n\nlogits1, mem1, xl_mem1 = model(x)                               # (1, 1024, 20000), (1, 128, 512), [(2, 1, 512, 512)]\nlogits2, mem2, xl_mem2 = model(x, mem1, xl_memories = xl_mem1)  # (1, 1024, 20000), (1, 128, 512), [(2, 1, 512, 512)]\nlogits3, mem3, xl_mem3 = model(x, mem2, xl_memories = xl_mem2)  # (1, 1024, 20000), (1, 128, 512), [(2, 1, 512, 512)]\n\n# and so on ...\n```\n\nTrain on an absurdly long sequence\n\n```python\nimport torch\nfrom recurrent_memory_transformer_pytorch import (\n    RecurrentMemoryTransformer,\n    RecurrentMemoryTransformerWrapper\n)\n\nmodel = RecurrentMemoryTransformer(\n    num_tokens = 256,\n    num_memory_tokens = 128,\n    dim = 512,\n    depth = 6,\n    seq_len = 1024,\n    use_flash_attn = True,\n    causal = True\n)\n\nmodel = RecurrentMemoryTransformerWrapper(model).cuda()\n\nseq = torch.randint(0, 256, (4, 65536)).cuda()   # absurdly long sequence, in reality, they curriculum learned this starting with 1 segment to about 7-8 segments\n\nloss = model(seq, memory_replay_backprop = True) # memory efficient training from memformer paper\n\n```\n\n## Todo\n\n- [ ] move the memory replay backprop into a torch.function, test out bidirectional, then test on a real problem\n\n- [x] get rotary embeddings working properly with xl memories\n- [x] add xl memories, detached\n- [x] offer a way to turn off rotary embeddings, absolute positional embeddings, and add token shift\n- [x] make memories being causally masked an option\n- [x] add the memory replay backprop technique from memformer paper\n- [x] relative positional encoding\n\n## Alternatives\n\n- \u003ca href=\"https://github.com/lucidrains/block-recurrent-transformer-pytorch\"\u003eBlock Recurrent Transformer\u003c/a\u003e\n\n- \u003ca href=\"https://github.com/lucidrains/memformer\"\u003eMemformer\u003c/a\u003e\n\n## Citations\n\n```bibtex\n@inproceedings{bulatov2022recurrent,\n  title     = {Recurrent Memory Transformer},\n  author    = {Aydar Bulatov and Yuri Kuratov and Mikhail Burtsev},\n  booktitle = {Advances in Neural Information Processing Systems},\n  editor    = {Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho},\n  year      = {2022},\n  url       = {https://openreview.net/forum?id=Uynr3iPhksa}\n}\n```\n\n```bibtex\n@misc{bulatov2023scaling,\n  title     = {Scaling Transformer to 1M tokens and beyond with RMT},\n  author    = {Aydar Bulatov and Yuri Kuratov and Mikhail S. Burtsev},\n  year      = {2023},\n  eprint    = {2304.11062},\n  archivePrefix = {arXiv},\n  primaryClass = {cs.CL}\n}\n```\n\n```bibtex\n@inproceedings{dao2022flashattention,\n  title     = {Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},\n  author    = {Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\\'e}, Christopher},\n  booktitle = {Advances in Neural Information Processing Systems},\n  year      = {2022}\n}\n```\n\n```bibtex\n@misc{shazeer2020glu,\n    title   = {GLU Variants Improve Transformer},\n    author  = {Noam Shazeer},\n    year    = {2020},\n    url     = {https://arxiv.org/abs/2002.05202}\n}\n```\n\n```bibtex\n@misc{su2021roformer,\n    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding},\n    author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},\n    year    = {2021},\n    eprint  = {2104.09864},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n```\n\n```bibtex\n@inproceedings{Wu2020MemformerAM,\n    title   = {Memformer: A Memory-Augmented Transformer for Sequence Modeling},\n    author  = {Qingyang Wu and Zhenzhong Lan and Kun Qian and Jing Gu and Alborz Geramifard and Zhou Yu},\n    booktitle = {AACL/IJCNLP},\n    year    = {2020}\n}\n```\n\n```bibtex\n@software{peng_bo_2021_5196578,\n    author    = {PENG Bo},\n    title     = {BlinkDL/RWKV-LM: 0.01},\n    month     = {aug},\n    year      = {2021},\n    publisher = {Zenodo},\n    version   = {0.01},\n    doi       = {10.5281/zenodo.5196578},\n    url       = {https://doi.org/10.5281/zenodo.5196578}\n}\n```\n\n```bibtex\n@misc{ding2021cogview,\n    title   = {CogView: Mastering Text-to-Image Generation via Transformers},\n    author  = {Ming Ding and Zhuoyi Yang and Wenyi Hong and Wendi Zheng and Chang Zhou and Da Yin and Junyang Lin and Xu Zou and Zhou Shao and Hongxia Yang and Jie Tang},\n    year    = {2021},\n    eprint  = {2105.13290},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n\n```bibtex\n@software{Dayma_DALLE_Mini_2021,\n    author  = {Dayma, Boris and Patil, Suraj and Cuenca, Pedro and Saifullah, Khalid and Abraham, Tanishq and Lê Khắc, Phúc and Melas, Luke and Ghosh, Ritobrata},\n    doi     = {10.5281/zenodo.5146400},\n    license = {Apache-2.0},\n    month   = {jul},\n    title   = {{DALL·E Mini}},\n    url     = {https://github.com/borisdayma/dalle-mini},\n    version = {v0.1-alpha},\n    year    = {2021}}\n```\n\n```bibtex\n@inproceedings{anonymous2022normformer,\n    title   = {NormFormer: Improved Transformer Pretraining with Extra Normalization},\n    author  = {Anonymous},\n    booktitle = {Submitted to The Tenth International Conference on Learning Representations },\n    year    = {2022},\n    url     = {https://openreview.net/forum?id=GMYWzWztDx5},\n    note    = {under review}\n}\n```\n\n```bibtex\n@misc{ding2021erniedoc,\n    title   = {ERNIE-Doc: A Retrospective Long-Document Modeling Transformer},\n    author  = {Siyu Ding and Junyuan Shang and Shuohuan Wang and Yu Sun and Hao Tian and Hua Wu and Haifeng Wang},\n    year    = {2021},\n    eprint  = {2012.15688},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n```\n\n```bibtex\n@article{Zhu2024HyperConnections,\n    title   = {Hyper-Connections},\n    author  = {Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou},\n    journal = {ArXiv},\n    year    = {2024},\n    volume  = {abs/2409.19606},\n    url     = {https://api.semanticscholar.org/CorpusID:272987528}\n}\n```\n\n```bibtex\n@inproceedings{Zhou2024ValueRL,\n    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},\n    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},\n    year    = {2024},\n    url     = {https://api.semanticscholar.org/CorpusID:273532030}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Frecurrent-memory-transformer-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Frecurrent-memory-transformer-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Frecurrent-memory-transformer-pytorch/lists"}