bos_token + readme

README.md

@@ -70,26 +70,31 @@ model = AutoModel.from_pretrained("ccdv/lsg-distilbert-base-uncased-4096",
 
 ## Sparse selection type
 
-There are 5 different sparse selection patterns. The best type is task dependent. \
+There are 6 different sparse selection patterns. The best type is task dependent. \
+If `sparse_block_size=0` or `sparsity_type="none"`, only local attention is considered. \
 Note that for sequences with length < 2*block_size, the type has no effect.
-
-* sparsity_type="norm", select highest norm tokens
+* `sparsity_type="bos_pooling"` (new)
+    * weighted average pooling using the BOS token 
+    * Works best in general, especially with a rather large sparsity_factor (8, 16, 32)
+    * Additional parameters:
+        * None
+* `sparsity_type="norm"`, select highest norm tokens
     * Works best for a small sparsity_factor (2 to 4)
     * Additional parameters:
         * None
-* sparsity_type="pooling", use average pooling to merge tokens
+* `sparsity_type="pooling"`, use average pooling to merge tokens
     * Works best for a small sparsity_factor (2 to 4)
     * Additional parameters:
         * None
-* sparsity_type="lsh", use the LSH algorithm to cluster similar tokens
+* `sparsity_type="lsh"`, use the LSH algorithm to cluster similar tokens
     * Works best for a large sparsity_factor (4+)
     * LSH relies on random projections, thus inference may differ slightly with different seeds
     * Additional parameters:
         * lsg_num_pre_rounds=1, pre merge tokens n times before computing centroids
-* sparsity_type="stride", use a striding mecanism per head
+* `sparsity_type="stride"`, use a striding mecanism per head
     * Each head will use different tokens strided by sparsify_factor
     * Not recommended if sparsify_factor > num_heads
-* sparsity_type="block_stride", use a striding mecanism per head
+* `sparsity_type="block_stride"`, use a striding mecanism per head
     * Each head will use block of tokens strided by sparsify_factor
     * Not recommended if sparsify_factor > num_heads
 

modeling_lsg_distilbert.py

@@ -50,16 +50,16 @@ class LSGDistilBertConfig(DistilBertConfig):
         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.n_heads:
                 logger.warning(
-                "[WARNING CONFIG]: sparsity_factor > encoder_attention_heads is not recommended for stride/block_stride sparsity"
+                "[WARNING CONFIG]: sparsity_factor > n_heads is not recommended for stride/block_stride sparsity"
             )
         
         if self.num_global_tokens < 1:
@@ -477,15 +477,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)
@@ -513,7 +514,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)
@@ -536,7 +537,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)
@@ -552,7 +553,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)
@@ -572,11 +573,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)
@@ -624,6 +628,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,
         query,
@@ -695,7 +722,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)