first

README.md

@@ -0,0 +1,115 @@
+---
+language:
+- en
+tags:
+- summarization
+datasets:
+- multi_news
+metrics:
+- rouge
+model-index:
+- name: ccdv/lsg-bart-base-4096-multinews
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+**This model relies on a custom modeling file, you need to add trust_remote_code=True**\
+**See [\#13467](https://github.com/huggingface/transformers/pull/13467)**
+
+```python
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+
+tokenizer = AutoTokenizer.from_pretrained("ccdv/lsg-bart-base-4096-multinews", trust_remote_code=True)
+model = AutoModelForSeq2SeqLM.from_pretrained("ccdv/lsg-bart-base-4096-multinews", trust_remote_code=True)
+
+text = "Replace by what you want."
+pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=0)
+generated_text = pipe(
+  text, 
+  truncation=True, 
+  max_length=64, 
+  no_repeat_ngram_size=7,
+  num_beams=2,
+  early_stopping=True
+  )
+```
+
+# ccdv/lsg-bart-base-4096-multinews
+
+This model is a fine-tuned version of [ccdv/lsg-bart-base-4096](https://huggingface.co/ccdv/lsg-bart-base-4096) on the multi_news default dataset. \
+It achieves the following results on the test set:
+
+| Length | Sparse Type  | Block Size | Sparsity | Connexions | R1    | R2    | RL    | RLsum |
+|:------ |:------------ |:---------- |:-------- | :--------- |:----- |:----- |:----- |:----- |
+| 4096   | Local        | 256        | 0        | 768        | 47.10 | 18.94 | 25.22 | 43.13 |
+| 4096   | Local        | 128        | 0        | 384        | 46.73 | 18.79 | 25.13 | 42.76 |
+| 4096   | Pooling      | 128        | 4        | 644        | 46.83 | 18.87 | 25.23 | 42.86 |
+| 4096   | Stride       | 128        | 4        | 644        | 46.83 | 18.68 | 24.98 | 42.88 |
+| 4096   | Block Stride | 128        | 4        | 644        | 46.83 | 18.72 | 25.06 | 42.88 |
+| 4096   | Norm         | 128        | 4        | 644        | 46.74 | 18.60 | 24.93 | 42.79 |
+| 4096   | LSH          | 128        | 4        | 644        | 46.74 | 18.82 | 25.19 | 42.77 |
+
+With smaller block size (lower ressources):
+
+| Length | Sparse Type  | Block Size | Sparsity | Connexions | R1    | R2    | RL    | RLsum |
+|:------ |:------------ |:---------- |:-------- | :--------- |:----- |:----- |:----- |:----- |
+| 4096   | Pooling      | 32         | 4        | 160        | 44.77 | 17.31 | 24.16 | 40.86 |
+| 4096   | Stride       | 32         | 4        | 160        | 45.29 | 17.81 | 24.45 | 41.40 |
+| 4096   | Block Stride | 32         | 4        | 160        | 45.39 | 17.86 | 24.51 | 41.43 |
+| 4096   | Norm         | 32         | 4        | 160        | 44.65 | 17.25 | 24.09 | 40.76 |
+| 4096   | LSH          | 32         | 4        | 160        | 44.44 | 17.20 | 24.00 | 40.57 |
+
+## Model description
+The model relies on Local-Sparse-Global attention to handle long sequences:
+![attn](attn.png)
+
+The model has about ~145 millions parameters (6 encoder layers - 6 decoder layers). \
+The model is warm started from BART-base, converted to handle long sequences (encoder only) and fine tuned. \
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 8e-05
+- train_batch_size: 8
+- seed: 42
+- gradient_accumulation_steps: 4
+- total_train_batch_size: 32
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- lr_scheduler_type: linear
+- lr_scheduler_warmup_ratio: 0.1
+- num_epochs: 12.0
+
+### Generate hyperparameters
+
+The following hyperparameters were used during generation:
+- dataset_name: multi_news
+- dataset_config_name: default
+- eval_batch_size: 8
+- eval_samples: 5622
+- early_stopping: True
+- ignore_pad_token_for_loss: True
+- length_penalty: 2.0
+- max_length: 320
+- min_length: 32
+- num_beams: 5
+- no_repeat_ngram_size: None
+- seed: 123
+
+### Framework versions
+
+- Transformers 4.18.0
+- Pytorch 1.10.1+cu102
+- Datasets 2.1.0
+- Tokenizers 0.11.6

config.json

@@ -0,0 +1,96 @@
+{
+  "_name_or_path": "models/ccdv/lsg-bart-base-4096-multinews",
+  "activation_dropout": 0.1,
+  "activation_function": "gelu",
+  "adaptive": true,
+  "add_bias_logits": false,
+  "add_final_layer_norm": false,
+  "architectures": [
+    "LSGBartForConditionalGeneration"
+  ],
+  "attention_dropout": 0.1,
+  "auto_map": {
+    "AutoConfig": "modeling_lsg_bart.LSGBartConfig",
+    "AutoModel": "modeling_lsg_bart.LSGBartModel",
+    "AutoModelForCausalLM": "modeling_lsg_bart.LSGBartForCausalLM",
+    "AutoModelForQuestionAnswering": "modeling_lsg_bart.LSGBartForQuestionAnswering",
+    "AutoModelForSeq2SeqLM": "modeling_lsg_bart.LSGBartForConditionalGeneration",
+    "AutoModelForSequenceClassification": "modeling_lsg_bart.LSGBartForSequenceClassification"
+  },
+  "base_model_prefix": "lsg",
+  "block_size": 256,
+  "bos_token_id": 0,
+  "classif_dropout": 0.1,
+  "classifier_dropout": 0.0,
+  "d_model": 768,
+  "decoder_attention_heads": 12,
+  "decoder_ffn_dim": 3072,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 6,
+  "decoder_start_token_id": 2,
+  "dropout": 0.1,
+  "early_stopping": true,
+  "encoder_attention_heads": 12,
+  "encoder_ffn_dim": 3072,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 6,
+  "eos_token_id": 2,
+  "forced_bos_token_id": 0,
+  "forced_eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2"
+  },
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2
+  },
+  "length_penalty": 2.0,
+  "lsh_num_pre_rounds": 1,
+  "max_length": 320,
+  "max_position_embeddings": 4096,
+  "min_length": 32,
+  "model_type": "bart",
+  "no_repeat_ngram_size": null,
+  "normalize_before": false,
+  "normalize_embedding": true,
+  "num_beams": 5,
+  "num_global_tokens": 1,
+  "num_hidden_layers": 6,
+  "pad_token_id": 1,
+  "pass_global_tokens_to_decoder": true,
+  "pool_with_global": true,
+  "scale_embedding": false,
+  "sparse_block_size": 0,
+  "sparsity_factor": 2,
+  "sparsity_type": "none",
+  "task_specific_params": {
+    "summarization": {
+      "length_penalty": 1.0,
+      "max_length": 128,
+      "min_length": 12,
+      "num_beams": 4
+    },
+    "summarization_cnn": {
+      "length_penalty": 2.0,
+      "max_length": 142,
+      "min_length": 56,
+      "num_beams": 4
+    },
+    "summarization_xsum": {
+      "length_penalty": 1.0,
+      "max_length": 62,
+      "min_length": 11,
+      "num_beams": 6
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.18.0",
+  "use_cache": true,
+  "vocab_size": 50265
+}

modeling_lsg_bart.py

@@ -0,0 +1,1759 @@
+from logging import warn
+import torch
+from transformers.models.bart.modeling_bart import *
+from transformers.models.bart.modeling_bart import _expand_mask
+import torch.nn as nn
+from torch.nn import BCEWithLogitsLoss
+import sys
+
+AUTO_MAP = {
+        "AutoModel": "modeling_lsg_bart.LSGBartModel",
+        "AutoModelForCausalLM": "modeling_lsg_bart.LSGBartForCausalLM",
+        "AutoModelForQuestionAnswering": "modeling_lsg_bart.LSGBartForQuestionAnswering",
+        "AutoModelForSequenceClassification": "modeling_lsg_bart.LSGBartForSequenceClassification",
+        "AutoModelForSeq2SeqLM": "modeling_lsg_bart.LSGBartForConditionalGeneration"
+    }
+
+class LSGBartConfig(BartConfig):
+    """
+    This class overrides :class:`~transformers.RobertaConfig`. Please check the superclass for the appropriate
+    documentation alongside usage examples.
+    """
+
+    base_model_prefix = "lsg"
+    model_type = "bart"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
+
+    def __init__(
+        self,
+        adaptive=True,
+        base_model_prefix="lsg",
+        block_size=128,
+        lsh_num_pre_rounds=1,
+        num_global_tokens=1,
+        pass_global_tokens_to_decoder=True,
+        pool_with_global=True,
+        sparse_block_size=128,
+        sparsity_factor=2,
+        sparsity_type="norm",
+        **kwargs
+        ):
+        """Constructs LSGConfig."""
+        super().__init__(**kwargs)
+        
+        self.adaptive = adaptive
+        self.auto_map = AUTO_MAP
+        self.base_model_prefix = base_model_prefix
+        self.block_size = block_size
+        self.lsh_num_pre_rounds = lsh_num_pre_rounds
+        self.num_global_tokens = num_global_tokens
+        self.pass_global_tokens_to_decoder = pass_global_tokens_to_decoder
+        self.pool_with_global = pool_with_global
+        self.sparse_block_size = sparse_block_size
+        self.sparsity_factor = sparsity_factor
+        self.sparsity_type = sparsity_type
+
+        if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride"]:
+            logger.warning(
+                "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride'], setting sparsity_type=None, computation will skip sparse attention")
+            self.sparsity_type = None
+
+        if self.sparsity_type == "stride":
+            if self.sparsity_factor > self.encoder_attention_heads:
+                logger.warning(
+                "[WARNING CONFIG]: sparsity_factor > encoder_attention_heads is not recommended for stride sparsity"
+            )
+        
+        if self.num_global_tokens < 1:
+            logger.warning(
+                "[WARNING CONFIG]: num_global_tokens < 1 is not compatible, setting num_global_tokens=1"
+            )
+            self.num_global_tokens = 1
+        elif self.num_global_tokens > 512:
+            logger.warning(
+                "[WARNING CONFIG]: num_global_tokens > 512 is not compatible, setting num_global_tokens=512"
+            )
+            self.num_global_tokens = 512
+        
+        if self.sparsity_factor > 0:
+            assert self.block_size % self.sparsity_factor == 0, "[ERROR CONFIG]: block_size must be divisible by sparsity_factor"
+            assert self.block_size//self.sparsity_factor >= 1, "[ERROR CONFIG]: make sure block_size >= sparsity_factor"
+            
+        
+def shift_tokens_right(input_ids, pad_token_id, decoder_start_token_id):
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("self.model.config.pad_token_id has to be defined.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids
+
+
+def _make_causal_mask(input_ids_shape, dtype, past_key_values_length=0):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), float("-inf"))
+    mask_cond = torch.arange(mask.size(-1))
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+def _expand_mask(mask, dtype, tgt_len=None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)
+
+
+class BaseSelfAttention(nn.Module):
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        is_decoder=False,
+        bias=True,
+        ):
+
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.scaling = self.head_dim ** -0.5
+        self.is_decoder = is_decoder
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (
+            self.num_heads,
+            self.head_dim,
+        )
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def reshape_output(self, context_layer):
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.embed_dim,)
+        return context_layer.view(*new_context_layer_shape)
+
+    def project_QKV(self, hidden_states):
+
+        query_layer = self.transpose_for_scores(self.q_proj(hidden_states))
+        key_layer = self.transpose_for_scores(self.k_proj(hidden_states))
+        value_layer = self.transpose_for_scores(self.v_proj(hidden_states))
+        return query_layer, key_layer, value_layer
+
+    
+class BaseAttentionProduct(nn.Module):
+
+    def __init__(self, config):
+        """
+        Compute attention: softmax(Q @ K.T) @ V
+        """
+        super().__init__()
+        self.dropout = nn.Dropout(config.attention_dropout)
+
+    def forward(self, query_layer, key_layer, value_layer, attention_mask=None):
+        
+        d = query_layer.shape[-1]
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)
+
+        del query_layer
+        del key_layer
+
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
+            attention_scores = attention_scores + attention_mask
+            del attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        context_layer = self.dropout(attention_probs) @ value_layer
+
+        return context_layer
+
+
+class LSGAttentionProduct(nn.Module):
+
+    def __init__(self, config, block_size=None, sparse_block_size=None, sparsity_factor=4):
+        """
+        Compute block or overlapping blocks attention products
+        """
+        super().__init__()
+ 
+        self.block_size = block_size
+        self.sparse_block_size = sparse_block_size
+        self.sparsity_factor = sparsity_factor
+
+        if self.block_size is None:
+            self.block_size = config.block_size
+
+        if self.sparse_block_size is None:
+            self.sparse_block_size = config.sparse_block_size
+
+        # Shape of blocks
+        self.local_shapes = (self.block_size*3, self.block_size)
+        if self.sparse_block_size and self.sparsity_factor > 0:
+            self.sparse_shapes = (self.sparse_block_size*3, self.block_size//self.sparsity_factor)
+
+        self.attention = BaseAttentionProduct(config)
+        
+    def build_lsg_inputs(self, hidden_states, sparse_hidden_states, global_hidden_states, is_attn_mask=False):
+        
+        # Build local tokens
+        local_hidden_states = self.reshape_to_local_block(hidden_states, is_attn_mask)
+        del hidden_states
+
+        # Build sparse tokens
+        if sparse_hidden_states is not None:
+            sparse_hidden_states = self.reshape_to_sparse_block(sparse_hidden_states, is_attn_mask)
+        
+        return self.cat_global_sparse_local_tokens(global_hidden_states, sparse_hidden_states, local_hidden_states)
+
+    def forward(
+        self, 
+        query_layer, 
+        key_layer, 
+        value_layer, 
+        attention_mask=None, 
+        sparse_key=None,
+        sparse_value=None, 
+        sparse_mask=None, 
+        global_key=None, 
+        global_value=None, 
+        global_mask=None
+        ):
+
+        # Input batch, heads, length, hidden_size
+        n, h, t, d = query_layer.size()
+        n_blocks = t // self.block_size
+        assert t % self.block_size == 0
+
+        key_layer = self.build_lsg_inputs(
+            key_layer, 
+            sparse_key, 
+            global_key
+            )
+        del sparse_key
+        del global_key
+
+        value_layer = self.build_lsg_inputs(
+            value_layer, 
+            sparse_value, 
+            global_value
+            )
+        del sparse_value
+        del global_value
+
+        attention_mask = self.build_lsg_inputs(
+            attention_mask, 
+            sparse_mask, 
+            global_mask.transpose(-1, -2), 
+            is_attn_mask=True
+            ).transpose(-1, -2)
+        del sparse_mask
+        del global_mask
+
+        # expect (..., t, d) shape
+        # Compute attention
+        context_layer = self.attention(
+                query_layer=self.chunk(query_layer, n_blocks), 
+                key_layer=key_layer,
+                value_layer=value_layer,
+                attention_mask=attention_mask
+                )
+                
+        return context_layer.reshape(n, h, -1, d)
+    
+    def reshape_to_local_block(self, hidden_states, is_attn_mask=False):
+        
+        size, step = self.local_shapes
+        s = (size - step) // 2
+
+        # Pad before block reshaping
+        if is_attn_mask:
+            pad_value = -10000  
+            hidden_states = hidden_states.transpose(-1, -2)
+        else: 
+            pad_value = 0
+
+        hidden_states = torch.nn.functional.pad(
+            hidden_states.transpose(-1, -2), 
+            pad=(s, s),
+            value=pad_value
+            ).transpose(-1, -2)
+
+        # Make blocks
+        hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
+
+        return hidden_states
+
+    def reshape_to_sparse_block(self, hidden_states, is_attn_mask=False):
+        
+        size, step = self.sparse_shapes
+
+        # In case of odd case
+        odd_offset = (step % 2)
+
+        # n, h, t, d*2 + 1
+        size = size*2 
+        s = (size - step) // 2 + odd_offset
+
+        # Pad before block reshaping
+        if is_attn_mask:
+            pad_value = -10000  
+            hidden_states = hidden_states.transpose(-1, -2)
+        else: 
+            pad_value = 0
+
+        hidden_states = torch.nn.functional.pad(
+            hidden_states.transpose(-1, -2), 
+            pad=(s, s),
+            value=pad_value
+            ).transpose(-1, -2)
+
+        # Make blocks
+        hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
+
+        # Fix case where block_size == sparsify_factor
+        if odd_offset: 
+            hidden_states = hidden_states[..., :-1, :, :]
+
+        # Indexes for selection
+        u = (size - self.block_size * 3 // self.sparsity_factor) // 2 + odd_offset
+        s = self.sparse_block_size
+
+        u_ = u + odd_offset
+        return torch.cat([hidden_states[..., u-s:u, :], hidden_states[..., -u_:-u_+s, :]], dim=-2)
+
+    def cat_global_sparse_local_tokens(self, x_global, x_sparse=None, x_local=None, dim=-2):
+
+        n, h, b, t, d = x_local.size()
+        x_global = x_global.unsqueeze(-3).expand(-1, -1, b, -1, -1)
+        if x_sparse is not None:
+            return torch.cat([x_global, x_sparse, x_local], dim=dim)
+        return torch.cat([x_global, x_local], dim=dim)
+
+    def chunk(self, x, n_blocks):
+
+        t, d = x.size()[-2:]
+        return x.reshape(*x.size()[:-2], n_blocks, -1, d)
+
+
+class LSGBartEncoderAttention(BaseSelfAttention):
+    '''
+    Compute local attention with overlapping blocs
+    Use global attention for tokens with highest norm
+    '''
+    def __init__(
+        self, 
+        config, 
+        embed_dim,
+        num_heads,
+        dropout
+        ):
+
+        super().__init__(embed_dim, num_heads, dropout)
+
+        self.block_size = config.block_size
+        self.sparse_block_size = config.sparse_block_size
+        self.num_global_tokens = config.num_global_tokens
+        self.sparsity_factor = config.sparsity_factor
+
+        self.attention = LSGAttentionProduct(
+            config, 
+            block_size=config.block_size, 
+            sparse_block_size=config.sparse_block_size, 
+            sparsity_factor=self.sparsity_factor, 
+            )
+
+        self.full_attention = BaseAttentionProduct(config)
+
+        sparse_functions = {
+            "norm": self.get_sparse_tokens_with_norm, 
+            "pooling": self.get_sparse_tokens_with_pooling,
+            "lsh": self.get_sparse_tokens_with_lsh,
+            "stride": self.get_sparse_tokens_with_stride,
+            }
+        
+        self.sparsity_type = config.sparsity_type
+        self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda 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):
+        
+        if self.sparsity_factor == 1:
+            return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
+
+        with torch.no_grad():
+
+            block_size = min(self.block_size, self.sparse_block_size)
+            key_norm = keys.detach().norm(dim=-1, keepdim=True)
+            key_norm = key_norm * ~mask.transpose(-1, -2).bool()
+            key_norm = self.chunk(key_norm, block_size)
+
+            n, h, b, t, d = key_norm.size()
+            
+            idx = key_norm.argsort(dim=-2) 
+            del key_norm
+            idx += (torch.arange(b, device=keys.device)*t).reshape(1, 1, b, 1, 1)
+
+            split = (t - block_size // self.sparsity_factor, block_size // self.sparsity_factor)
+            sparse_idx = idx.split(split, -2)[-1].reshape(n, h, -1, 1)
+        
+        d = keys.size()[-1]
+        keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
+        values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
+        mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
+
+        return keys, values, mask
+
+    def get_sparse_tokens_with_pooling(self, keys, values, mask):
+        
+        if self.sparsity_factor == 1:
+            return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
+
+        keys = self.chunk(keys, self.sparsity_factor)
+        values = self.chunk(values, self.sparsity_factor)
+
+        n, h, b, t, d = keys.size()
+        mask = mask.reshape(n, 1, b, 1, t)
+        mask = ~mask.transpose(-1, -2).bool()
+
+        keys = keys * mask
+        values = values * mask
+
+        mask = mask.sum(dim=-2)
+        keys = keys.sum(dim=-2) / (mask + 1e-6)
+        values = values.sum(dim=-2) / (mask + 1e-6)
+
+        mask = - (1. - mask.clamp(0, 1)) * 1e4
+        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):
+
+        if self.sparsity_factor == 1:
+            return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
+
+        n, h, t, d = keys.size()
+