From 1167a3408adc49898aa4b00c6235adbfceabf9d1 Mon Sep 17 00:00:00 2001
From: Elisei Rykov <e.rykov@tbank.ru>
Date: Tue, 15 Oct 2024 10:38:57 +0300
Subject: [PATCH] in progress

---
 .../llava_with_textual_attention.json         |  81 +++++++
 .../train/multimodal/attention_pooling.json   |   6 +-
 .../attention_pooling_with_n_heads.json       | 205 ++++++++++++++++++
 .../train/multimodal/llava_with_replica.json  | 205 ++++++++++++++++++
 .../multimodal/projectors/__init__.py         |   1 +
 .../projectors/attention_pooling.py           | 132 ++++++++++-
 turbo_alignment/settings/modality.py          |   1 +
 7 files changed, 623 insertions(+), 8 deletions(-)
 create mode 100755 configs/exp/multimodal/llava_with_textual_attention.json
 create mode 100644 configs/exp/train/multimodal/attention_pooling_with_n_heads.json
 create mode 100644 configs/exp/train/multimodal/llava_with_replica.json

diff --git a/configs/exp/multimodal/llava_with_textual_attention.json b/configs/exp/multimodal/llava_with_textual_attention.json
new file mode 100755
index 0000000..71b023b
--- /dev/null
+++ b/configs/exp/multimodal/llava_with_textual_attention.json
@@ -0,0 +1,81 @@
+{
+    "inference_settings": [
+      {
+        "metric_settings": [],
+        "model_settings": {
+          "model_path": "/from_s3/model",
+          "projections_path": "/from_s3/projections/modality_adapters.pt",
+          "n_modality_embeddings": 256,
+          "model_type": "causal",
+          "transformers_settings": {},
+          "adapter_path": "/from_s3/adapter"
+        },
+        "modality_encoder_settings_mapping": {
+            "image": {
+                "modality_encoder_type": "clip",
+                "is_pickle": true,
+                "encoder_path": "openai/clip-vit-large-patch14"
+            },
+            "audio": null
+        },
+        "modality_projector_mapping": {
+            "image": "llava",
+            "audio": null
+        },
+        "tokenizer_settings": {
+          "use_fast": false,
+          "tokenizer_path": "/from_s3/tokenizer"
+        },
+        "generation_settings": [
+          {
+            "transformers_settings": {
+              "num_beams": 1,
+              "max_new_tokens": 128
+            },
+            "custom_settings": {}
+          }
+        ],
+        "use_vllm": false,
+        "batch": 1,
+        "micro_batch": 1
+      }
+    ],
+    "dataset_settings": {
+        "sources": [
+            {
+                "name": "rullava",
+                "records_path": "/app/data/multimodal/rullava/val_chat.jsonl",
+                "num_samples": 50
+            }
+        ],
+      "prompt_template": {
+        "role_tag_mapping": {
+            "bot": "<bot>",
+            "user": "<user>",
+            "system": "<system>"
+        },
+        "prefix_template": "<RS>{role}",
+        "suffix_template": "</RS>"
+    },
+    "dataset_type": "multimodal",
+    "max_tokens_count": 2000,
+    "n_modality_embeddings": 256,
+    "start_modality_token": "<MS>",
+    "end_modality_token": "</MS>",
+    "only_answer_loss": true,
+    "modality_token_mapping": {
+        "image": "<img>",
+        "audio": "<audio>"
+    },
+    "modality_reader_settings_mapping": {
+        "image": {
+          "reader_type": "pickle"
+        },
+        "audio": null
+    },
+    "truncate_top": false,
+    "random_cut": true
+    },
+    "save_path": "inference_output"
+  }
+
diff --git a/configs/exp/train/multimodal/attention_pooling.json b/configs/exp/train/multimodal/attention_pooling.json
index 453561c..34d30e2 100644
--- a/configs/exp/train/multimodal/attention_pooling.json
+++ b/configs/exp/train/multimodal/attention_pooling.json
@@ -27,7 +27,7 @@
                 },
                 "audio": null
         },
-        "n_modality_embeddings": 64,
+        "n_modality_embeddings": 128,
         "start_modality_token": "<MS>",
         "end_modality_token": "</MS>",
         "dataset_type": "multimodal",
@@ -63,7 +63,7 @@
                 },
                 "audio": null
             },
-        "n_modality_embeddings": 64,
+        "n_modality_embeddings": 128,
         "start_modality_token": "<MS>",
         "end_modality_token": "</MS>",
         "dataset_type": "multimodal",
@@ -182,7 +182,7 @@
             },
             "dataset_type": "multimodal",
             "max_tokens_count": 2000,
-            "n_modality_embeddings": 64,
+            "n_modality_embeddings": 128,
             "start_modality_token": "<MS>",
             "end_modality_token": "</MS>",
             "only_answer_loss": true,
diff --git a/configs/exp/train/multimodal/attention_pooling_with_n_heads.json b/configs/exp/train/multimodal/attention_pooling_with_n_heads.json
new file mode 100644
index 0000000..13c387c
--- /dev/null
+++ b/configs/exp/train/multimodal/attention_pooling_with_n_heads.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/train_chat.jsonl",
+                "sample_rate": 1.0
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 448,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                "num_samples": 5000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 448,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 400,
+        "save_steps": 400,
+        "per_device_train_batch_size": 2,
+        "per_device_eval_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 16,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "top_k_attention_pooling_with_n_heads",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                    "num_samples": 50
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 448,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/configs/exp/train/multimodal/llava_with_replica.json b/configs/exp/train/multimodal/llava_with_replica.json
new file mode 100644
index 0000000..a1e5a8c
--- /dev/null
+++ b/configs/exp/train/multimodal/llava_with_replica.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/joined_train_chat_v2.jsonl",
+                "sample_rate": 1.0
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 1,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/joined_test_chat_v2.jsonl",
+                "num_samples": 5000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 1,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 250,
+        "save_steps": 250,
+        "per_device_train_batch_size": 2,
+        "per_device_eval_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 16,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "llava_with_replica",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/joined_test_chat_v2.jsonl",
+                    "num_samples": 50
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 1,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/turbo_alignment/modeling/multimodal/projectors/__init__.py b/turbo_alignment/modeling/multimodal/projectors/__init__.py
index 53cb930..9566242 100644
--- a/turbo_alignment/modeling/multimodal/projectors/__init__.py
+++ b/turbo_alignment/modeling/multimodal/projectors/__init__.py
@@ -1,6 +1,7 @@
 from turbo_alignment.modeling.multimodal.projectors.attention_pooling import (
     AttentionPoolingMultiModalProjector,
     TopKAttentionPoolingMultiModalProjector,
+    TopKAttentionPoolingWithNHeadsMultiModalProjector
 )
 from turbo_alignment.modeling.multimodal.projectors.c_abstractor import CAbstractor
 from turbo_alignment.modeling.multimodal.projectors.llava import (
diff --git a/turbo_alignment/modeling/multimodal/projectors/attention_pooling.py b/turbo_alignment/modeling/multimodal/projectors/attention_pooling.py
index 22faffc..7763136 100644
--- a/turbo_alignment/modeling/multimodal/projectors/attention_pooling.py
+++ b/turbo_alignment/modeling/multimodal/projectors/attention_pooling.py
@@ -4,8 +4,80 @@
     MultiModalProjectorRegistry,
 )
 from turbo_alignment.settings.modality import ModalityProjectorType
+import torch.nn.functional as F
+import math
+import numpy as np
 
 
+def get_abs_pos(abs_pos, tgt_size):
+    # abs_pos: L, C
+    # tgt_size: M
+    # return: M, C
+    src_size = int(math.sqrt(abs_pos.size(0)))
+    tgt_size = int(math.sqrt(tgt_size))
+    dtype = abs_pos.dtype
+
+    if src_size != tgt_size:
+        return F.interpolate(
+            abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
+            size=(tgt_size, tgt_size),
+            mode="bicubic",
+            align_corners=False,
+        ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)
+    else:
+        return abs_pos
+
+# https://huggingface.co/Qwen/Qwen-VL-Chat/blob/main/visual.py
+# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
 @MultiModalProjectorRegistry.register(ModalityProjectorType.ATTENTION_POOLING)
 class AttentionPoolingMultiModalProjector(torch.nn.Module):
     def __init__(self, encoder_hidden_size: int, text_hidden_size: int, n_modality_embs: int):
@@ -32,21 +104,71 @@ def __init__(self, encoder_hidden_size: int, text_hidden_size: int, n_modality_e
         self.n_modality_embs = n_modality_embs
         self.linear_projection = torch.nn.Linear(encoder_hidden_size, text_hidden_size)
         self.attention_scores = torch.nn.Linear(text_hidden_size, 1)
+        self.top_k = 128
+        self.pos_embed = torch.nn.Parameter(
+            torch.from_numpy(get_2d_sincos_pos_embed(text_hidden_size, 24)).float()
+        ).requires_grad_(False)
 
     def forward(self, image_features: torch.Tensor) -> torch.Tensor:
+        pos_embed = get_abs_pos(self.pos_embed, image_features.size(1))
+
         projected_features = self.linear_projection(
             image_features
         )  # map each image patch to the language model dimension
+        
+        projected_features = projected_features + pos_embed
+        
         attention_scores = torch.softmax(
             self.attention_scores(projected_features), 1
         )  # calculate learnable attention scores for each patch
         top_indices = torch.topk(
-            attention_scores.squeeze(-1), k=self.n_modality_embs, dim=1
+            attention_scores.squeeze(-1), k=attention_scores.shape[1], dim=1
+        ).indices  # select indices top N patches according to attention scores
+
+        projected_features[:, top_indices[:, self.top_k:].squeeze(0)] = 0 # set zero for unselected tokens
+        projected_features = projected_features[(projected_features != 0).any(dim=-1)] # remove zero vectors
+
+        return projected_features.unsqueeze(0)
+
+
+@MultiModalProjectorRegistry.register(ModalityProjectorType.TOP_K_ATTENTION_POOLING_WITH_N_HEADS)
+class TopKAttentionPoolingWithNHeadsMultiModalProjector(torch.nn.Module):
+    def __init__(self, encoder_hidden_size: int, text_hidden_size: int, n_modality_embs: int):
+        super().__init__()
+        self.encoder_hidden_size = encoder_hidden_size
+        self.text_hidden_size = text_hidden_size
+        self.n_modality_embs = n_modality_embs
+        self.linear_projection = torch.nn.Linear(encoder_hidden_size, text_hidden_size)
+        self.num_heads = 1
+        self.attention_scores = torch.nn.Linear(text_hidden_size, self.num_heads)
+        self.top_k = n_modality_embs
+        # self.pos_embed = torch.nn.Parameter(
+        #     torch.from_numpy(get_2d_sincos_pos_embed(text_hidden_size, 15)).float()
+        # ).requires_grad_(False)
+
+    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
+        # pos_embed = get_abs_pos(self.pos_embed, image_features.size(1))
+
+        projected_features = self.linear_projection(
+            image_features
+        )  # map each image patch to the language model dimension
+
+        # projected_features = projected_features + pos_embed
+
+        scores = self.attention_scores(projected_features)
+        attention_scores = torch.softmax(
+            scores, 1
+        )  # calculate learnable attention scores for each patch
+        attention_scores = torch.max(attention_scores, -1).values
+        # attention_scores = torch.mean(attention_scores, -1)
+        top_indices = torch.topk(
+            attention_scores.squeeze(-1), k=attention_scores.shape[1], dim=1
         ).indices  # select indices top N patches according to attention scores
-        top_k_hidden_states = torch.gather(
-            projected_features, index=top_indices.unsqueeze(-1).expand(-1, -1, projected_features.size(-1)), dim=1
-        )  # select top patches
-        return top_k_hidden_states
+
+        projected_features[:, top_indices[:, self.top_k:].squeeze(0)] = 0 # set zero for unselected tokens
+        projected_features = projected_features[(projected_features != 0).any(dim=-1)] # remove zero vectors
+
+        return projected_features.unsqueeze(0)
 
 
 @MultiModalProjectorRegistry.register(ModalityProjectorType.THRESHOLD_SELECTOR)
diff --git a/turbo_alignment/settings/modality.py b/turbo_alignment/settings/modality.py
index 230845d..3fb1f55 100755
--- a/turbo_alignment/settings/modality.py
+++ b/turbo_alignment/settings/modality.py
@@ -41,3 +41,4 @@ class ModalityProjectorType(str, Enum):
     ATTENTION_POOLING = 'attention_pooling'
     TOP_K_ATTENTION_POOLING = 'top_k_attention_pooling'
     THRESHOLD_SELECTOR = 'threshold_selector'
+    TOP_K_ATTENTION_POOLING_WITH_N_HEADS = 'top_k_attention_pooling_with_n_heads'
\ No newline at end of file