{"id":15601026,"url":"https://github.com/lucidrains/colt5-attention","last_synced_at":"2025-04-09T05:09:45.832Z","repository":{"id":144597746,"uuid":"616582355","full_name":"lucidrains/CoLT5-attention","owner":"lucidrains","description":"Implementation of the conditionally routed attention in the CoLT5 architecture, in Pytorch","archived":false,"fork":false,"pushed_at":"2024-09-06T14:55:40.000Z","size":191,"stargazers_count":226,"open_issues_count":3,"forks_count":13,"subscribers_count":7,"default_branch":"main","last_synced_at":"2025-04-02T04:02:55.560Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","deep-learning","efficient-attention","routing"],"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-03-20T17:06:54.000Z","updated_at":"2025-02-24T09:00:04.000Z","dependencies_parsed_at":null,"dependency_job_id":"cbfe0e89-82d5-49ae-9329-d4d8dc269392","html_url":"https://github.com/lucidrains/CoLT5-attention","commit_stats":{"total_commits":104,"total_committers":1,"mean_commits":104.0,"dds":0.0,"last_synced_commit":"f2a8da09644615e827fc49883d2652c5c1c22792"},"previous_names":[],"tags_count":72,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCoLT5-attention","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCoLT5-attention/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCoLT5-attention/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCoLT5-attention/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/CoLT5-attention/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247980837,"owners_count":21027808,"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","efficient-attention","routing"],"created_at":"2024-10-03T02:12:30.646Z","updated_at":"2025-04-09T05:09:45.816Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./colt5.png\" width=\"400px\"\u003e\u003c/img\u003e\n\n## CoLT5 Attention - Pytorch\n\nImplementation of the conditionally routed efficient attention in the proposed \u003ca href=\"https://arxiv.org/abs/2303.09752\"\u003eCoLT5\u003c/a\u003e architecture, in Pytorch.\n\nThey used coordinate descent from \u003ca href=\"https://arxiv.org/abs/2211.01267\"\u003ethis paper\u003c/a\u003e (main algorithm originally from \u003ca href=\"https://arxiv.org/abs/1502.04759\"\u003eWright et al\u003c/a\u003e) to route a subset of tokens for 'heavier' branches of the feedforward and attention blocks.\n\nUpdate: unsure of how the routing normalized scores for the key-values are used. Did some improvising there, \u003ca href=\"https://github.com/lucidrains/CoLT5-attention/blob/main/colt5_attention/transformer_block.py#L86\"\u003escaling the projected values\u003c/a\u003e, but if you think you know the answer, please open an issue\n\nUpdate 2: seems to work well with the improvisation above\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStability.ai\u003c/a\u003e for the generous sponsorship to work on cutting edge artificial intelligence research\n\n- \u003ca href=\"https://github.com/arogozhnikov/einops\"\u003eeinops\u003c/a\u003e for making my life easy\n\n- \u003ca href=\"https://github.com/openai/triton\"\u003eTriton\u003c/a\u003e for allowing me to speed up coordinate descent with a fused implementation in just 2 days, sparing me from having to write a thousand lines of CUDA code\n\n## Install\n\n```bash\n$ pip install colt5-attention\n```\n\n## Usage\n\n```python\nimport torch\n\nfrom colt5_attention import (\n    ConditionalRoutedFeedForward,\n    ConditionalRoutedAttention,\n    ConditionalRoutedTransformerBlock\n)\n\n# mock input, say it is 32768 length\n\ntokens = torch.randn(2, 32768, 512)\nmask = torch.ones(2, 32768).bool()  # can handle variable lengthed sequences\n\n# feedforward\n\nff = ConditionalRoutedFeedForward(\n    dim = 512,\n    light_ff_mult = 0.5,      # hidden dimension ratio of light branch\n    heavy_ff_mult = 4,        # hidden dimension ratio of heavy branch\n    num_heavy_tokens = 1024   # heavy branch receives only 1024 routed tokens of 32768\n)\n\nff_out = ff(tokens, mask = mask)  # (2, 32768, 512) - light and heavy branch summed\n\n# attention\n\nattn = ConditionalRoutedAttention(\n    dim = 512,\n    light_dim_head = 64,       # attention head dimension of light branch\n    light_heads = 8,           # number of attention heads for light branch\n    light_window_size = 128,   # local attention receptive field for light\n    heavy_dim_head = 64,       # attention head dimension of heavy branch\n    heavy_heads = 8,           # number of attention heads for heavy branch\n    num_heavy_tokens_q = 1024, # heavy branch receives only 1024 routed tokens of 32768\n    num_heavy_tokens_kv = 1024 # heavy branch receives only 1024 routed tokens of 32768\n)\n\nattn_out = attn(tokens, mask = mask) # (2, 32768, 512) - light and heavy branch summed\n\n# both attention and feedforward with residual\n# the complete transformer block\n# a stack of these would constitute the encoder of CoLT5\n\nblock = ConditionalRoutedTransformerBlock(\n    dim = 512,\n    light_dim_head = 64,\n    light_heads = 8,\n    light_window_size = 128,\n    heavy_dim_head = 64,\n    heavy_heads = 8,\n    light_ff_mult = 0.5,\n    heavy_ff_mult = 4,\n    num_heavy_ff_tokens = 1024,\n    num_heavy_attn_tokens_q = 1024,\n    num_heavy_attn_tokens_kv = 1024\n)\n\nblock_out = block(tokens, mask = mask) # (2, 32768, 512)\n```\n\nAlso included a variation of the conditionally routed attention for cross attention, to be tried with long context memories in a transformer-xl\n\n```python\nimport torch\nfrom colt5_attention import ConditionalRoutedCrossAttention\n\n# mock input, let us say it is a transformer of 1024 length attending to 1 million context past memories\n\ntokens = torch.randn(1, 1024, 512).cuda()\ntokens_mask = torch.ones(1, 1024).bool().cuda()\n\nmemories = torch.randn(1, 1_048_576, 512).cuda()\nmemories_mask = torch.ones(1, 1_048_576).bool().cuda()\n\n# conditionally routed cross attention\n\ncross_attn = ConditionalRoutedCrossAttention(\n    dim = 512,\n    dim_head = 64,\n    heads = 8,\n    num_tokens_q = 512,         # only 512 routed from 1024\n    num_tokens_kv = 1024,       # only 1024 routed from 1 million\n    kv_routing_tokens = 2,      # say you want 2 routing tokens to route different sets of key / values to the queries. 4 attention heads will be allocated to each routed set in this example (8 / 2)\n    use_triton = True,          # use cuda kernel\n    route_block_size = 131072   # route in blocks of 131072\n).cuda()\n\ncross_attn_out = cross_attn(\n    tokens,\n    context = memories,\n    mask = tokens_mask,\n    context_mask = memories_mask\n)\n\ncross_attn_out.shape # (1, 1024, 512) - same as tokens\n```\n\nThis repository also has an improvised version for autoregressive attention. The way this was achieved was by viewing the sequence in windows. Each window can only attend to windows of key / values into the past. The local attention of the light branch covers the intra-window attention.\n\nThe coordinate descent is made viable through a CUDA kernel written in \u003ca href=\"https://github.com/openai/triton\"\u003eTriton\u003c/a\u003e. Finally, to get autoregressive generation to work well, I had to make sure for the unrouted tokens (for queries), outputs a learned output embedding rather than just zeros.\n\nCurrently I am seeing occasional differences between the gradients (as high as 1e-1 for a very small fraction of elements) once the number of iterations exceed 20. However, enwik8 seems to train well and I can see the effects of the routing. Training is surprisingly stable too\n\nex.\n\n```python\nimport torch\nfrom colt5_attention import ConditionalRoutedAutoregressiveAttention\n\n# mock input, say it is 8192 length\n\ntokens = torch.randn(2, 8192, 512).cuda()\n\n# attention\n\nattn = ConditionalRoutedAutoregressiveAttention(\n    dim = 512,\n    light_dim_head = 64,          # attention head dimension of light branch\n    light_heads = 8,              # number of attention heads for light branch\n    light_window_size = 128,      # local attention receptive field for light\n    heavy_window_size = 128,      # the windowing for the routed heavy attention, by default, will be equal to the light window size. be aware if this is any greater than the light window size, there may be tokens that would be missed by attention\n    heavy_dim_head = 64,          # attention head dimension of heavy branch\n    heavy_heads = 8,              # number of attention heads for heavy branch\n    num_heavy_tokens_q = 32,      # heavy branch receives only 32 out of 128 of the windowed queries (1024 query tokens total)\n    num_heavy_tokens_kv = 1024,   # heavy branch receives only 1024 routed tokens for key-values\n    num_routed_kv = 2,            # one can split the attention heads so that groups of heads attend to different sets of key - values (2 routing tokens in this case)\n    use_triton = True,            # will need to use Triton for this to be viable, otherwise it is too slow and memory efficient with the number of iterations\n    use_flash_attn = True         # use flash attention in heavy branch\n).cuda()\n\nattn_out = attn(tokens) + tokens # (2, 8192, 512) - output of attention with residual (prenorm is included)\n```\n\nFinally, this repository contains a version for image feature maps. Typically a lot of research papers cannot do attention on image feature maps with dimensions greater than 32 by 32. This routed attention will use a local window patch for the light branch, and routed attention for the heavy\n\nex.\n\n```python\nimport torch\nfrom colt5_attention import ConditionalRoutedImageAttention\n\nattn = ConditionalRoutedImageAttention(\n    dim = 32,\n    light_dim_head = 64,       # attention head dimension of light branch\n    light_heads = 8,           # number of attention heads for light branch\n    light_window_size = 32,    # height and width of local window attention on the image feature map\n    channel_first = True,      # whether to accept images with channel first than last\n    heavy_dim_head = 64,       # attention head dimension of heavy branch\n    heavy_heads = 8,           # number of attention heads for heavy branch\n    num_heavy_tokens_q = 1024, # heavy branch receives only 1024 routed tokens of 65536\n    num_heavy_tokens_kv = 1024 # heavy branch receives only 1024 routed tokens of 65536\n).cuda()\n\nfmap = torch.randn(1, 32, 256, 256).cuda() # image feature map is too large for attention, given 256 ^ 2  == 65536 tokens\n\nout = attn(fmap)\n```\n\n\u003ca href=\"https://arxiv.org/abs/2205.01580\"\u003eSimple ViT\u003c/a\u003e using coordinate descent routed attention and feedforward\n\n```python\nimport torch\nfrom colt5_attention.vit import ConditionalRoutedViT\n\nvit = ConditionalRoutedViT(\n    image_size = 256,                # image size\n    patch_size = 32,                 # patch size\n    num_classes = 1000,              # number of output classes\n    dim = 1024,                      # feature dimension\n    depth = 6,                       # depth\n    attn_num_heavy_tokens_q = 16,    # number of routed queries for heavy attention\n    attn_num_heavy_tokens_kv = 16,   # number of routed key/values for heavy attention\n    attn_heavy_dim_head = 64,        # dimension per attention head for heavy\n    attn_heavy_heads = 8,            # number of attention heads for heavy\n    attn_light_window_size = 4,      # the local windowed attention for light branch\n    attn_light_dim_head = 32,        # dimension per head for local light attention\n    attn_light_heads = 4,            # number of attention heads for local windowed attention\n    ff_num_heavy_tokens = 16,        # number of tokens routed for heavy feedforward\n    ff_heavy_mult = 4,               # the expansion factor of the heavy feedforward branch\n    ff_light_mult = 2                # expansion factor of the light feedforward branch\n)\n\nimages = torch.randn(1, 3, 256, 256)\n\nlogits = vit(images) # (1, 1000)\n```\n\n## Differentiable Topk\n\nUse a small wrapper around coordinate descent for differentiable `topk`\n\n```python\nimport torch\nfrom colt5_attention import topk\n\nx = torch.randn(1024, 512)\n\nvalues, indices, coor_descent_values, gates = topk(x, k = 10, fused = True)\n\n# you can either use the topk indices + gates, or use the values directly (values have already been multiplied with the gates within the function)\n```\n\n## Todo\n\n- [x] add the coordinate descent method as another router\n- [x] allow for multi-headed routing (multiple routing tokens), only for key-values\n- [x] add an autoregressive version of the conditionally routed attention\n- [x] test out the autoregressive version and verify that more routed key / value tokens lead to better results - it works\n- [x] make flash attention compatible\n- [x] create a variant of CoLT5 for high resolution feature maps (image attention) - then try out for diffusion\n- [x] fused coordinate descent kernel using triton\n    - [x] forwards        \n    - [x] backwards\n    - [x] benchmark triton vs plain pytorch coor_descent - 50 iterations with 4 segments - 18.5x faster for forward (7.23 vs 0.39), 7.2x faster for backwards (5.77 vs 0.80)\n    - [x] fall back on plain coordinate descent for cpu\n    - [x] handle edge case for when a row is completely masked out for triton, or simply enforce it never to be so\n    - [x] fix masking in coordinate descent\n    - [x] simplified some logic within the triton kernel and the problem went away. probably some tiny quirk with the compiler\n    - [x] maximum block size in triton allowed is 131k, make sure at least quarter of million sequence length can be reached. to get around this initially, one can fold a million token sequence into ~9 131k and uniformly route. offer uniform routing scheme within router itself\n    - [x] remove sinkhorn and cumulative softmax approaches and cleanup; neither can work as well as coordinate descent\n    - [x] allow for saving intermediates every number of iterations - trading memory for recompute efficiency during backwards\n    - [x] in-place write to checkpointed a and b tensor for potentially savings on forward when recompute segments is high\n\n## Citations\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```bibtex\n@article{Tillet2019TritonAI,\n    title   = {Triton: an intermediate language and compiler for tiled neural network computations},\n    author  = {Philippe Tillet and H. Kung and D. Cox},\n    journal = {Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages},\n    year    = {2019}\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@article{Lei2023ConditionalAP,\n    title   = {Conditional Adapters: Parameter-efficient Transfer Learning with Fast Inference},\n    author  = {Tao Lei and Junwen Bai and Siddhartha Brahma and Joshua Ainslie and Kenton Lee and Yanqi Zhou and Nan Du and Vincent Zhao and Yuexin Wu and Bo Li and Yu Zhang and Ming-Wei Chang},\n    journal = {ArXiv},\n    year    = {2023},\n    volume  = {abs/2304.04947}\n}\n```\n\n```bibtex\n@article{Beyer2022BetterPV,\n    title   = {Better plain ViT baselines for ImageNet-1k},\n    author  = {Lucas Beyer and Xiaohua Zhai and Alexander Kolesnikov},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2205.01580}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcolt5-attention","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fcolt5-attention","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcolt5-attention/lists"}