Global CLS attention

configuration_bert.py

@@ -129,6 +129,7 @@ class JinaBertConfig(PretrainedConfig):
         feed_forward_type="original",
         emb_pooler=None,
         attn_implementation='torch',
+        cls_bias=None,
         **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, **kwargs)
@@ -151,6 +152,7 @@ class JinaBertConfig(PretrainedConfig):
         self.feed_forward_type = feed_forward_type
         self.emb_pooler = emb_pooler
         self.attn_implementation = attn_implementation
+        self.cls_bias = cls_bias
 
 class JinaBertOnnxConfig(OnnxConfig):
     @property

modeling_bert.py

@@ -701,12 +701,12 @@ class JinaBertEncoder(nn.Module):
         self.num_attention_heads = config.num_attention_heads
         self.register_buffer(
             "alibi",
-            self.rebuild_alibi_tensor(size=config.max_position_embeddings),
+            self.rebuild_alibi_tensor(size=config.max_position_embeddings, cls_bias=config.cls_bias),
             persistent=False,
         )
 
     def rebuild_alibi_tensor(
-        self, size: int, device: Optional[Union[torch.device, str]] = None
+        self, size: int, device: Optional[Union[torch.device, str]] = None, cls_bias=None
     ):
         # Alibi
         # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
@@ -747,6 +747,10 @@ class JinaBertEncoder(nn.Module):
         alibi = alibi.unsqueeze(0)
         assert alibi.shape == torch.Size([1, n_heads, size, size])
 
+        if cls_bias is not None:
+            alibi[:, :, 0, :] = cls_bias
+            alibi[:, :, :, 0] = cls_bias
+
         self._current_alibi_size = size
         return alibi
 
@@ -778,7 +782,8 @@ class JinaBertEncoder(nn.Module):
             )
             self.register_buffer(
                 "alibi",
-                self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device).to(
+                self.rebuild_alibi_tensor(size=seqlen, cls_bias=self.config.cls_bias,
+                                          device=hidden_states.device).to(
                     hidden_states.dtype
                 ),
                 persistent=False,