{"id":15600985,"url":"https://github.com/lucidrains/agent-attention-pytorch","last_synced_at":"2025-05-11T00:56:15.880Z","repository":{"id":213154252,"uuid":"733217742","full_name":"lucidrains/agent-attention-pytorch","owner":"lucidrains","description":"Implementation of Agent Attention in Pytorch","archived":false,"fork":false,"pushed_at":"2024-07-10T16:34:37.000Z","size":528,"stargazers_count":89,"open_issues_count":1,"forks_count":5,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-05-09T21:35:26.170Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","deep-learning","linear-attention"],"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-12-18T20:38:16.000Z","updated_at":"2025-04-21T03:01:39.000Z","dependencies_parsed_at":"2023-12-18T23:36:23.785Z","dependency_job_id":"9713c5c5-a108-41dd-82f5-d0c9519c899b","html_url":"https://github.com/lucidrains/agent-attention-pytorch","commit_stats":{"total_commits":20,"total_committers":1,"mean_commits":20.0,"dds":0.0,"last_synced_commit":"091f6d447d5005cd8c8e16843688b5cc2a9b8cd2"},"previous_names":["lucidrains/agent-attention-pytorch"],"tags_count":13,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fagent-attention-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fagent-attention-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fagent-attention-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fagent-attention-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/agent-attention-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":253329007,"owners_count":21891562,"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","linear-attention"],"created_at":"2024-10-03T02:11:01.227Z","updated_at":"2025-05-09T21:36:29.674Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./agent-attention.png\" width=\"350px\"\u003e\u003c/img\u003e\n\n## Agent Attention - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2312.08874\"\u003eAgent Attention\u003c/a\u003e in Pytorch.\n\nThis work seems to be an elegant simplification of \u003ca href=\"https://github.com/lucidrains/isab-pytorch\"\u003e`ISAB`\u003c/a\u003e architecture from the \u003ca href=\"https://arxiv.org/abs/1810.00825\"\u003eSet Transformers\u003c/a\u003e paper (requires only one attention block rather than two). While ISAB works, I have found it to be a bit unstable, thus wondering if the simplification in this work resolves that issue.\n\nThis repository will add support for variable sequence lengths (masking) and post-softmax talking heads.\n\n## Appreciation\n\n- \u003ca href=\"https://a16z.com/supporting-the-open-source-ai-community/\"\u003eA16Z Open Source AI Grant Program\u003c/a\u003e and \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for the generous sponsorships, as well as my other sponsors, for affording me the independence to open source current artificial intelligence research\n\n## Install\n\n```bash\n$ pip install agent-attention-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom agent_attention_pytorch import AgentSelfAttention\n\nattn = AgentSelfAttention(\n    dim = 512,\n    num_agent_tokens = 256,       # number of \"agent\" tokens\n    dim_head = 64,                # attention head dimension\n    heads = 8                     # number of heads\n)\n\nx = torch.randn(3, 65536, 512)\nmask = torch.ones(3, 65536).bool()\n\nout = attn(x, mask = mask)\n\nassert out.shape == x.shape\n```\n\nFor a full fledged linear transformer based on agent tokens, just import `AgentTransformer`\n\n```python\nimport torch\nfrom agent_attention_pytorch import AgentTransformer\n\ntransformer = AgentTransformer(\n    dim = 512,\n    depth = 6,\n    num_agent_tokens = 128,\n    dim_head = 64,\n    heads = 8\n)\n\nx = torch.randn(3, 65536, 512)\nmask = torch.ones(3, 65536).bool()\n\nout, agent_tokens = transformer(x, mask = mask, return_agent_tokens = True)\n\n# (3, 65536, 512), (3, 128, 512)\nassert out.shape == x.shape\n```\n\n## Citations\n\n```bibtex\n@inproceedings{Han2023AgentAO,\n    title   = {Agent Attention: On the Integration of Softmax and Linear Attention},\n    author  = {Dongchen Han and Tianzhu Ye and Yizeng Han and Zhuofan Xia and Shiji Song and Gao Huang},\n    year    = {2023},\n    url     = {https://api.semanticscholar.org/CorpusID:266210414}\n}\n```\n\n```bibtex\n@misc{shazeer2020talkingheads,\n    title   = {Talking-Heads Attention}, \n    author  = {Noam Shazeer and Zhenzhong Lan and Youlong Cheng and Nan Ding and Le Hou},\n    year    = {2020},\n    eprint  = {2003.02436},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n\n```bibtex\n@article{Bondarenko2023QuantizableTR,\n    title   = {Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing},\n    author  = {Yelysei Bondarenko and Markus Nagel and Tijmen Blankevoort},\n    journal = {ArXiv},\n    year    = {2023},\n    volume  = {abs/2306.12929},\n    url     = {https://api.semanticscholar.org/CorpusID:259224568}\n}\n```\n\n```bibtex\n@article{Wang2022FoundationT,\n    title   = {Foundation Transformers},\n    author  = {Hongyu Wang and Shuming Ma and Shaohan Huang and Li Dong and Wenhui Wang and Zhiliang Peng and Yu Wu and Payal Bajaj and Saksham Singhal and Alon Benhaim and Barun Patra and Zhun Liu and Vishrav Chaudhary and Xia Song and Furu Wei},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2210.06423},\n    url     = {https://api.semanticscholar.org/CorpusID:252846241}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fagent-attention-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fagent-attention-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fagent-attention-pytorch/lists"}