bos_token + readme

modeling_lsg_camembert.py

@@ -53,16 +53,16 @@ class LSGCamembertConfig(CamembertConfig):
         self.sparsity_factor = sparsity_factor
         self.sparsity_type = sparsity_type
 
-        if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride", "block_stride"]:
+        if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride", "block_stride", "bos_pooling"]:
             logger.warning(
-                "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride'], \
+                "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride', 'bos_pooling'], \
                     setting sparsity_type=None, computation will skip sparse attention")
             self.sparsity_type = None
 
         if self.sparsity_type in ["stride", "block_stride"]:
-            if self.sparsity_factor > self.encoder_attention_heads:
+            if self.sparsity_factor > self.num_attention_heads:
                 logger.warning(
-                "[WARNING CONFIG]: sparsity_factor > encoder_attention_heads is not recommended for stride/block_stride sparsity"
+                "[WARNING CONFIG]: sparsity_factor > num_attention_heads is not recommended for stride/block_stride sparsity"
             )
         
         if self.num_global_tokens < 1:
@@ -497,15 +497,16 @@ class LSGSelfAttention(BaseSelfAttention):
             "lsh": self.get_sparse_tokens_with_lsh,
             "stride": self.get_sparse_tokens_with_stride,
             "block_stride": self.get_sparse_tokens_with_block_stride,
+            "bos_pooling": self.get_sparse_tokens_with_bos_pooling
             }
         
         self.sparsity_type = config.sparsity_type
-        self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda x, y, z: (None, None, None))
+        self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda w, x, y, z: (None, None, None))
             
         if config.sparsity_type == "lsh":
             self.lsh_num_pre_rounds = config.lsh_num_pre_rounds
 
-    def get_sparse_tokens_with_norm(self, keys, values, mask):
+    def get_sparse_tokens_with_norm(self, queries, keys, values, mask):
         
         if self.sparsity_factor == 1:
             return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
@@ -533,7 +534,7 @@ class LSGSelfAttention(BaseSelfAttention):
 
         return keys, values, mask
 
-    def get_sparse_tokens_with_pooling(self, keys, values, mask):
+    def get_sparse_tokens_with_pooling(self, queries, keys, values, mask):
         
         if self.sparsity_factor == 1:
             return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
@@ -556,7 +557,7 @@ class LSGSelfAttention(BaseSelfAttention):
         mask *= torch.finfo(mask.dtype).min
         return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
 
-    def get_sparse_tokens_with_stride(self, keys, values, mask):
+    def get_sparse_tokens_with_stride(self, queries, keys, values, mask):
 
         if self.sparsity_factor == 1:
             return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
@@ -572,7 +573,7 @@ class LSGSelfAttention(BaseSelfAttention):
 
         return keys, values, mask
 
-    def get_sparse_tokens_with_block_stride(self, keys, values, mask):
+    def get_sparse_tokens_with_block_stride(self, queries, keys, values, mask):
 
         if self.sparsity_factor == 1:
             return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
@@ -592,11 +593,14 @@ class LSGSelfAttention(BaseSelfAttention):
 
         return keys, values, mask
         
-    def get_sparse_tokens_with_lsh(self, keys, values, mask):
+    def get_sparse_tokens_with_lsh(self, queries, keys, values, mask):
         
         if self.sparsity_factor == 1:
             return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
 
+        if self.sparsity_factor == self.sparse_block_size:
+            return self.get_sparse_tokens_with_bos_pooling(queries, keys, values, mask)
+        
         block_size = min(self.block_size, self.sparse_block_size)
         keys = self.chunk(keys, block_size)
         values = self.chunk(values, block_size)
@@ -644,6 +648,29 @@ class LSGSelfAttention(BaseSelfAttention):
 
         return keys[..., :output_size, :], values[..., :output_size, :], mask[..., :output_size, :]
 
+    def get_sparse_tokens_with_bos_pooling(self, queries, keys, values, mask):
+        
+        if self.sparsity_factor == 1:
+            return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
+
+        queries = queries.unsqueeze(-3)
+        mask = self.chunk(mask.transpose(-1, -2), self.sparsity_factor).transpose(-1, -2)
+        keys = self.chunk(keys, self.sparsity_factor)
+        values = self.chunk(values, self.sparsity_factor)
+        
+        n, h, b, t, d = keys.size()
+        scores = (queries[..., :1, :] @ keys.transpose(-1, -2)) / math.sqrt(d)
+        if mask is not None:
+            scores = scores + mask
+
+        scores = torch.softmax(scores, dim=-1)
+        keys = scores @ keys 
+        values = scores @ values
+        mask = mask.mean(dim=-1)
+        mask[mask != torch.finfo(mask.dtype).min] = 0
+        
+        return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
+    
     def forward(
         self,
         hidden_states,
@@ -765,7 +792,7 @@ class LSGSelfAttention(BaseSelfAttention):
         # Get sparse idx
         sparse_key, sparse_value, sparse_mask = (None, None, None)
         if self.sparse_block_size and self.sparsity_factor > 0:
-            sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(key_layer, value_layer, attention_mask)
+            sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(query_layer, key_layer, value_layer, attention_mask)
         
         # Expand masks on heads
         attention_mask = attention_mask.expand(-1, h, -1, -1)
@@ -838,7 +865,7 @@ class LSGSelfAttention(BaseSelfAttention):
         sparse_key, sparse_value, sparse_mask = (None, None, None)
 
         if self.sparse_block_size and self.sparsity_factor > 0:
-            sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(key_layer, value_layer, attention_mask)
+            sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(query_layer, key_layer, value_layer, attention_mask)
         
         # Expand masks on heads
         attention_mask = attention_mask.expand(-1, h, -1, -1)