same feature map for q and k

facebookresearch · Jan 18, 2022 · f3b613c · f3b613c
1 parent 60e94e5
commit f3b613c
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/xformers/components/attention/favor.py b/xformers/components/attention/favor.py
@@ -98,8 +98,7 @@ def __init__(
             "normalize_inputs": self.normalize_inputs,
         }
 
-        self.feature_map_query: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
-        self.feature_map_key: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
+        self.feature_map: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
 
     @staticmethod
     def _maybe_promote(x: torch.Tensor) -> torch.Tensor:
@@ -135,8 +134,8 @@ def forward(
     ):
 
         # Project key and queries onto the feature map space
-        k_prime = self.feature_map_key(k)
-        q_prime = self.feature_map_query(q)
+        k_prime = self.feature_map(k)
+        q_prime = self.feature_map(q)
 
         with autocast(enabled=False):
             # The softmax kernel approximation for Favor will easily overflow