num routed key / value sets feature can also work for self attention

lucidrains · Apr 19, 2023 · 47724ac · 47724ac
1 parent 83677b0
commit 47724ac
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/colt5_attention/transformer_block.py b/colt5_attention/transformer_block.py
@@ -486,6 +486,7 @@ def __init__(
         *,
         num_heavy_tokens_q,
         num_heavy_tokens_kv,
+        num_routed_kv = 1,
         light_dim_head = 64,
         light_heads = 8,
         light_window_size = 128,        # each token would see ~ 64 tokens either way to left or right
@@ -525,6 +526,7 @@ def __init__(
 
         self.kv_router = router_klass(
             dim = dim,
+            num_routing_tokens = num_routed_kv,
             straight_through = router_straight_through,
             **router_kwargs
         )
@@ -726,6 +728,7 @@ def __init__(
         *,
         num_heavy_attn_tokens_q,
         num_heavy_attn_tokens_kv,
+        num_routed_kv = 1,
         num_heavy_ff_tokens,
         light_dim_head = 64,
         light_heads = 8,
@@ -758,6 +761,7 @@ def __init__(
             heavy_heads = heavy_heads,
             num_heavy_tokens_q = num_heavy_attn_tokens_q,
             num_heavy_tokens_kv = num_heavy_attn_tokens_kv,
+            num_routed_kv = num_routed_kv,
             router_straight_through = router_straight_through,
             router_type = router_type,
             router_kwargs = router_kwargs

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'CoLT5-attention',
   packages = find_packages(),
-  version = '0.2.2',
+  version = '0.2.3',
   license='MIT',
   description = 'Conditionally Routed Attention',
   long_description_content_type = 'text/markdown',