{"id":13644168,"url":"https://github.com/lucidrains/memorizing-transformers-pytorch","last_synced_at":"2025-05-16T10:06:27.335Z","repository":{"id":37620439,"uuid":"472525693","full_name":"lucidrains/memorizing-transformers-pytorch","owner":"lucidrains","description":"Implementation of Memorizing Transformers (ICLR 2022), attention net augmented with indexing and retrieval of memories using approximate nearest neighbors, in Pytorch","archived":false,"fork":false,"pushed_at":"2023-07-17T00:08:50.000Z","size":35839,"stargazers_count":632,"open_issues_count":10,"forks_count":46,"subscribers_count":10,"default_branch":"main","last_synced_at":"2025-05-11T01:35:49.454Z","etag":null,"topics":["approximate-nearest-neighbors","artificial-intelligence","attention-mechanism","deep-learning","memory","retrieval","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":"2022-03-21T21:56:45.000Z","updated_at":"2025-04-16T10:00:26.000Z","dependencies_parsed_at":"2024-08-02T01:28:07.140Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/memorizing-transformers-pytorch","commit_stats":{"total_commits":113,"total_committers":2,"mean_commits":56.5,"dds":"0.017699115044247815","last_synced_commit":"25db92330cc52e806d7066e53103ee08a69e609d"},"previous_names":[],"tags_count":67,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemorizing-transformers-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemorizing-transformers-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemorizing-transformers-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemorizing-transformers-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/memorizing-transformers-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254509475,"owners_count":22082891,"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":["approximate-nearest-neighbors","artificial-intelligence","attention-mechanism","deep-learning","memory","retrieval","transformers"],"created_at":"2024-08-02T01:01:58.472Z","updated_at":"2025-05-16T10:06:27.306Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Reimplementations","Python"],"sub_categories":[],"readme":"\u003cimg src=\"./diagram.png\" width=\"500px\"\u003e\u003c/img\u003e\n\n## Memorizing Transformers - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2203.08913\"\u003eMemorizing Transformers\u003c/a\u003e (ICLR 2022), attention net augmented with indexing and retrieval of memories using approximate nearest neighbors, in Pytorch\n\nThis repository deviates from the paper slightly, using a hybrid attention across attention logits local and distant (rather than the sigmoid gate setup). It also uses cosine similarity attention (with learned temperature) for the KNN attention layer.\n\n## Install\n\n```bash\n$ pip install memorizing-transformers-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom memorizing_transformers_pytorch import MemorizingTransformer\n\nmodel = MemorizingTransformer(\n    num_tokens = 20000,                 # number of tokens\n    dim = 512,                          # dimension\n    dim_head = 64,                      # dimension per attention head\n    depth = 8,                          # number of layers\n    memorizing_layers = (4, 5),         # which layers to have ANN memories\n    max_knn_memories = 64000,           # maximum ANN memories to keep (once it hits this capacity, it will be reset for now, due to limitations in faiss' ability to remove entries)\n    num_retrieved_memories = 32,        # number of ANN memories to retrieve\n    clear_memories_on_sos_token_id = 1, # clear passed in ANN memories automatically for batch indices which contain this specified SOS token id - otherwise, you can also manually iterate through the ANN memories and clear the indices before the next iteration\n)\n\ndata = torch.randint(0, 20000, (2, 1024)) # mock data\n\nknn_memories = model.create_knn_memories(batch_size = 2) # create collection of KNN memories with the correct batch size (2 in example)\n\nlogits = model(data, knn_memories = knn_memories) # (1, 1024, 20000)\n```\n\nYou can make the KNN memories read-only by setting `add_knn_memory` on forward to `False`\n\nex.\n\n```python\nlogits = model(data, knn_memories = knn_memories, add_knn_memory = False) # knn memories will not be updated\n```\n\nWith Transformer-XL memories (only the memories that will be discarded will be added to the KNN memory)\n\n```python\nimport torch\nfrom memorizing_transformers_pytorch import MemorizingTransformer\n\nmodel = MemorizingTransformer(\n    num_tokens = 20000,\n    dim = 512,\n    depth = 8,\n    memorizing_layers = (4, 5),\n    max_knn_memories = 64000,\n    num_retrieved_memories = 32,\n    clear_memories_on_sos_token_id = 1,\n    xl_memory_layers = (2, 3, 4, 5),      # xl memory layers - (https://arxiv.org/abs/2007.03356 shows you do not need XL memory on all layers, just the latter ones) - if a KNNAttention layer ends up using XL memories, only the XL memories that will be discarded will be added to long term memory\n    xl_max_memories = 512,                # number of xl memories to keep\n    shift_knn_memories_down = 1,          # let a layer look at the KNN memories this number of layers above\n    shift_xl_memories_down = 1,           # let a layer look at the XL memories this number of layers above, shown to enhance receptive field in ernie-doc paper\n)\n\ndata = torch.randint(0, 20000, (2, 1024)) # mock data\n\nxl_memories = None\n\nwith model.knn_memories_context(batch_size = 2) as knn_memories:\n    logits1, xl_memories = model(data, knn_memories = knn_memories, xl_memories = xl_memories)\n    logits2, xl_memories = model(data, knn_memories = knn_memories, xl_memories = xl_memories)\n    logits3, xl_memories = model(data, knn_memories = knn_memories, xl_memories = xl_memories)\n\n    # ... and so on\n```\n\n## KNN Memory\n\nThis repository contains a wrapper around Faiss that can automatically store and retrieve key / values\n\n```python\nimport torch\nfrom memorizing_transformers_pytorch import KNNMemory\n\nmemory = KNNMemory(\n    dim = 64,                   # dimension of key / values\n    max_memories = 64000,       # maximum number of memories to keep (will throw out the oldest memories for now if it overfills)\n    num_indices = 2             # this should be equivalent to batch dimension, as each batch keeps track of its own memories, expiring when it sees a new document\n)\n\nmemory.add(torch.randn(2, 512, 2, 64))  # (batch, seq, key | value, feature dim)\nmemory.add(torch.randn(2, 512, 2, 64))\n\nmemory.clear([0]) # clear batch 0, if it saw an \u003csos\u003e\n\nmemory.add(torch.randn(2, 512, 2, 64))\nmemory.add(torch.randn(2, 512, 2, 64))\n\nkey_values, mask = memory.search(torch.randn(2, 512, 64), topk = 32)\n```\n\n## Training\n\nEnwik8 training\n\n```bash\n$ python train.py\n```\n\n## Todo\n\n- [x] switch to ivfhnsw and just remember all memories\n- [x] enwik8 demo\n- [x] validation for enwik8\n- [x] solve gradient accumulation problem by offering some way to scope reads and writes to knn memories with another indices array\n- [ ] setup text generation with memories\n- [ ] figure out how to deal with memories efficiently once capacity has been hit\n- [ ] try to speed up reading and writing to knn memories collection with multiprocessing\n\n## Citations\n\n```bibtex\n@article{wu2022memorizing,\n  title   = {Memorizing transformers},\n  author  = {Wu, Yuhuai and Rabe, Markus N and Hutchins, DeLesley and Szegedy, Christian},\n  journal = {arXiv preprint arXiv:2203.08913},\n  year    = {2022}\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@Article{AlphaFold2021,\n  author  = {Jumper, John and Evans, Richard and Pritzel, Alexander and Green, Tim and Figurnov, Michael and Ronneberger, Olaf and Tunyasuvunakool, Kathryn and Bates, Russ and {\\v{Z}}{\\'\\i}dek, Augustin and Potapenko, Anna and Bridgland, Alex and Meyer, Clemens and Kohl, Simon A A and Ballard, Andrew J and Cowie, Andrew and Romera-Paredes, Bernardino and Nikolov, Stanislav and Jain, Rishub and Adler, Jonas and Back, Trevor and Petersen, Stig and Reiman, David and Clancy, Ellen and Zielinski, Michal and Steinegger, Martin and Pacholska, Michalina and Berghammer, Tamas and Bodenstein, Sebastian and Silver, David and Vinyals, Oriol and Senior, Andrew W and Kavukcuoglu, Koray and Kohli, Pushmeet and Hassabis, Demis},\n  journal = {Nature},\n  title   = {Highly accurate protein structure prediction with {AlphaFold}},\n  year    = {2021},\n  doi     = {10.1038/s41586-021-03819-2},\n  note    = {(Accelerated article preview)},\n}\n```\n\n```bibtex\n@inproceedings{Rae2020DoTN,\n  title   = {Do Transformers Need Deep Long-Range Memory?},\n  author  = {Jack W. Rae and Ali Razavi},\n  booktitle = {ACL},\n  year    = {2020}\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@misc{henry2020querykey,\n    title   = {Query-Key Normalization for Transformers},\n    author  = {Alex Henry and Prudhvi Raj Dachapally and Shubham Pawar and Yuxuan Chen},\n    year    = {2020},\n    eprint  = {2010.04245},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n```\n\n*Memory is Attention through Time* - Alex Graves\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmemorizing-transformers-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fmemorizing-transformers-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmemorizing-transformers-pytorch/lists"}