{"id":15600959,"url":"https://github.com/lucidrains/block-recurrent-transformer-pytorch","last_synced_at":"2025-04-05T06:03:57.178Z","repository":{"id":65963538,"uuid":"598360852","full_name":"lucidrains/block-recurrent-transformer-pytorch","owner":"lucidrains","description":"Implementation of Block Recurrent Transformer - Pytorch","archived":false,"fork":false,"pushed_at":"2024-08-20T17:00:12.000Z","size":35891,"stargazers_count":218,"open_issues_count":0,"forks_count":20,"subscribers_count":7,"default_branch":"main","last_synced_at":"2025-03-29T05:03:11.425Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","deep-learning","long-context-attention","long-context-transformers","memory","recurrence"],"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-02-07T00:09:55.000Z","updated_at":"2025-02-11T22:49:09.000Z","dependencies_parsed_at":"2024-08-20T19:07:06.669Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/block-recurrent-transformer-pytorch","commit_stats":{"total_commits":37,"total_committers":1,"mean_commits":37.0,"dds":0.0,"last_synced_commit":"bfa31c3078a24d7e244a0f859212582cba791fc3"},"previous_names":[],"tags_count":33,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fblock-recurrent-transformer-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fblock-recurrent-transformer-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fblock-recurrent-transformer-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fblock-recurrent-transformer-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/block-recurrent-transformer-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247294516,"owners_count":20915340,"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-attention","long-context-transformers","memory","recurrence"],"created_at":"2024-10-03T02:10:23.599Z","updated_at":"2025-04-05T06:03:57.163Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./block-recurrent-transformer.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Block Recurrent Transformer - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2203.07852\"\u003eBlock Recurrent Transformer\u003c/a\u003e - Pytorch. The highlight of the paper is its reported ability to remember something up to 60k tokens ago.\n\nThis design is SOTA for recurrent transformers line of research, afaict.\n\nIt will also include \u003ca href=\"https://arxiv.org/abs/2205.14135\"\u003eflash attention\u003c/a\u003e as well as routed memories of up to 250k tokens using ideas from \u003ca href=\"https://github.com/lucidrains/CoLT5-attention\"\u003ethis paper\u003c/a\u003e\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStability.ai\u003c/a\u003e for the generous sponsorship to work and open source cutting edge artificial intelligence research\n\n## Install\n\n```bash\n$ pip install block-recurrent-transformer-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom block_recurrent_transformer_pytorch import BlockRecurrentTransformer\n\nmodel = BlockRecurrentTransformer(\n    num_tokens = 20000,             # vocab size\n    dim = 512,                      # model dimensions\n    depth = 6,                      # depth\n    dim_head = 64,                  # attention head dimensions\n    heads = 8,                      # number of attention heads\n    max_seq_len = 1024,             # the total receptive field of the transformer, in the paper this was 2 * block size\n    block_width = 512,              # block size - total receptive field is max_seq_len, 2 * block size in paper. the block furthest forwards becomes the new cached xl memories, which is a block size of 1 (please open an issue if i am wrong)\n    num_state_vectors = 512,        # number of state vectors, i believe this was a single block size in the paper, but can be any amount\n    recurrent_layers = (4,),        # where to place the recurrent layer(s) for states with fixed simple gating\n    use_compressed_mem = False,     # whether to use compressed memories of a single block width, from https://arxiv.org/abs/1911.05507\n    compressed_mem_factor = 4,      # compression factor of compressed memories\n    use_flash_attn = True           # use flash attention, if on pytorch 2.0\n)\n\nseq = torch.randint(0, 2000, (1, 1024))\n\nout, mems1, states1 = model(seq)\nout, mems2, states2 = model(seq, xl_memories = mems1, states = states1)\nout, mems3, states3 = model(seq, xl_memories = mems2, states = states2)\n```\n\n## Test on Enwik8\n\nFirst `pip install -r requirements.txt`, then\n\n```bash\n$ python train.py\n```\n\n## Todo\n\n- [x] use dynamic positional bias\n- [x] add enhanced recurrence\n- [x] setup local attention blocks, as in the paper\n- [x] wrapper transformer class for training\n- [x] take care of generation with recurrence in `RecurrentTrainWrapper`\n- [x] add ability to dropout to entire memories and states during each segment step during trainng\n- [x] test full system on enwik8 locally and ablate states and memories and see effects first  hand\n- [x] make sure attention allow for single head key / values too\n- [x] run a few experiments of fixed gating in regular transformers - does not work\n- [x] integrate \u003ca href=\"https://github.com/hazyresearch/flash-attention\"\u003eflash attention\u003c/a\u003e\n- [x] cache attention mask + rotary embeddings\n- [x] add \u003ca href=\"https://github.com/lucidrains/compressive-transformer-pytorch\"\u003ecompressed memories\u003c/a\u003e\n\n- [ ] revisit \u003ca href=\"https://github.com/lucidrains/memformer\"\u003ememformer\u003c/a\u003e\n- [ ] try routing long distance memories of up to 250k using coordinate descent (Wright et al.)\n\n## Citations\n\n```bibtex\n@article{Hutchins2022BlockRecurrentT,\n    title   = {Block-Recurrent Transformers},\n    author  = {DeLesley S. Hutchins and Imanol Schlag and Yuhuai Wu and Ethan Dyer and Behnam Neyshabur},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2203.07852}\n}\n```\n\n```bibtex\n@article{Shazeer2019FastTD,\n    title   = {Fast Transformer Decoding: One Write-Head is All You Need},\n    author  = {Noam M. Shazeer},\n    journal = {ArXiv},\n    year    = {2019},\n    volume  = {abs/1911.02150}\n}\n```\n\n```bibtex\n@inproceedings{Sun2022ALT,\n    title     = {A Length-Extrapolatable Transformer},\n    author    = {Yutao Sun and Li Dong and Barun Patra and Shuming Ma and Shaohan Huang and Alon Benhaim and Vishrav Chaudhary and Xia Song and Furu Wei},\n    year      = {2022}\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@inproceedings{Ainslie2023CoLT5FL,\n    title   = {CoLT5: Faster Long-Range Transformers with Conditional Computation},\n    author  = {Joshua Ainslie and Tao Lei and Michiel de Jong and Santiago Ontan'on and Siddhartha Brahma and Yury Zemlyanskiy and David Uthus and Mandy Guo and James Lee-Thorp and Yi Tay and Yun-Hsuan Sung and Sumit Sanghai},\n    year    = {2023}\n}\n```\n\n*Memory is Attention through Time* - Alex Graves\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fblock-recurrent-transformer-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fblock-recurrent-transformer-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fblock-recurrent-transformer-pytorch/lists"}