a = torch.arange(dim, device=idx.device)

ndim = idx.ndim

dims = idx.shape

n_idx = dims[-1]

dims = dims[:-1] + (-1, )

for i in range(1, ndim):

a = a.unsqueeze(0)

a = a.expand(*dims)

masked = torch.scatter(a, -1, idx, 0)

compl, _ = torch.sort(masked, dim=-1, descending=False)

compl = compl.permute(-1, *tuple(range(ndim - 1)))

compl = compl[n_idx:].permute(*(tuple(range(1, ndim)) + (0,)))

attn_np = attn.to(dtype=torch.float32).cpu().numpy().flatten()

Q1 = np.percentile(attn_np, 25)

Q3 = np.percentile(attn_np, 75)

IQR = Q3 - Q1

# lower_bound = Q1 - 1.5 * IQR

upper_bound = Q3 + 1.5 * IQR

outlier_indices = np.where((attn_np > upper_bound))[0]

ratio = len(outlier_indices) / len(attn_np)

def __init__(self, vision_tower, args, delay_load=False):

super().__init__()

self.is_loaded = False

self.vision_tower_name = vision_tower

self.select_layer = args.mm_vision_select_layer # default: -2

self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')

self.total_tokens = 0

if not delay_load:

self.load_model()

else:

self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)

def load_model(self):

self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)

self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)

self.vision_tower.requires_grad_(False)

self.is_loaded = True

def feature_select(self, image_forward_outs):

image_features = image_forward_outs.hidden_states[self.select_layer] # penultimate layer output

if self.select_feature == 'patch':

image_features = image_features[:, 1:]

elif self.select_feature == 'cls_patch':

image_features = image_features

else:

raise ValueError(f'Unexpected select feature: {self.select_feature}')

return image_features

def token_prune_merge_advanced(self, images, if_adaptive=True, reduction_ratio = 1/8):

'''

version 10/03/2024 using the key*key matrix to calculate the cosine similarity

'''

# token_indix_list = []

# token_indix_dict = {}

#set hooks for extracting desired layer's k and q

hook_handle_k = self.vision_tower.vision_model.encoder.layers[23].self_attn.k_proj.register_forward_hook(hook_k)

hook_handle_q = self.vision_tower.vision_model.encoder.layers[23].self_attn.q_proj.register_forward_hook(hook_q)

#forward pass

image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)

cls_token_last_layer =image_forward_outs.hidden_states[self.select_layer][:, 0:1]

image_features = self.feature_select(image_forward_outs).to(images.dtype)

B, N, C = image_features.shape

#extract desired layer's k and q and remove hooks; calculate attention

desired_layer_k = outputs["desired_k"]

desired_layer_q = outputs["desired_q"]

hook_handle_k.remove()

hook_handle_q.remove()

attn = (desired_layer_q @ desired_layer_k.transpose(-2, -1)) * C ** -0.5

attn = F.softmax(attn, dim=-1)

cls_attn = attn[:, 0, 1:]

if if_adaptive:

reduction_ratio = outlier_dectection(cls_attn)#*3.5

_, idx = torch.topk(cls_attn, int(N*reduction_ratio), dim=1, largest=True) # [B, left_tokens] , sorted=True

index = idx.unsqueeze(-1).expand(-1, -1, C) # [B, left_tokens, C]

Key_wo_cls = desired_layer_k[:, 1:] # [B, N-1, C]

x_others = torch.gather(image_features, dim=1, index=index) # [B, left_tokens, C]

x_others_attn = torch.gather(cls_attn, dim=1, index=idx)

Key_others = torch.gather(Key_wo_cls, dim=1, index=index) # [B, left_tokens, C]

compl = complement_idx(idx, N) # [B, N-1-left_tokens]

non_topk = torch.gather(image_features, dim=1, index=compl.unsqueeze(-1).expand(-1, -1, C)) # [B, N-1-left_tokens, C]

non_topk_Key = torch.gather(Key_wo_cls, dim=1, index=compl.unsqueeze(-1).expand(-1, -1, C))

non_topk_attn = torch.gather(cls_attn, dim=1, index=compl) # [B, N-1-left_tokens]

Key_others_norm = F.normalize(Key_others, p=2, dim=-1)

non_topk_Key_norm = F.normalize(non_topk_Key, p=2, dim=-1)

# cos_sim = torch.bmm(Key_others_norm, non_topk_Key_norm.transpose(1, 2)) # [B, left_tokens, N-1-left_tokens]

# _, cluster_indices = torch.topk(cos_sim, k=4, dim=2, largest=True)

B, left_tokens, C = x_others.size()

updated_x_others = torch.zeros_like(x_others)

for b in range(B):

for i in range(left_tokens):

key_others_norm = Key_others_norm[b,i,:].unsqueeze(0).unsqueeze(0)

before_i_Key = Key_others_norm[b, :i, :].unsqueeze(0)

after_i_Key = Key_others_norm[b, i+1:, :].unsqueeze(0)

before_i_x_others = x_others[b, :i, :].unsqueeze(0)

after_i_x_others = x_others[b, i+1:, :].unsqueeze(0)

rest_x_others = torch.cat([before_i_x_others, after_i_x_others, non_topk[b,:,:].unsqueeze(0)], dim=1)

before_i_x_others_attn = x_others_attn[b, :i].unsqueeze(0)

after_i_x_others_attn = x_others_attn[b, i+1:].unsqueeze(0)

rest_x_others_attn = torch.cat([before_i_x_others_attn, after_i_x_others_attn, non_topk_attn[b,:].unsqueeze(0)], dim=1)

rest_Keys = torch.cat([before_i_Key, after_i_Key, non_topk_Key_norm[b,:,:].unsqueeze(0)], dim=1)

cos_sim_matrix = torch.bmm(key_others_norm, rest_Keys.transpose(1, 2))

_, cluster_indices = torch.topk(cos_sim_matrix, k=int(32), dim=2, largest=True)

cluster_tokens = rest_x_others[:,cluster_indices.squeeze(),:]

weights = rest_x_others_attn[:,cluster_indices.squeeze()].unsqueeze(-1)

# update cluster centers

weighted_avg = torch.sum(cluster_tokens * weights, dim=1) #/ torch.sum(weights)

updated_center = weighted_avg + x_others[b, i, :]

updated_x_others[b, i, :] = updated_center

extra_one_token = torch.sum(non_topk * non_topk_attn.unsqueeze(-1), dim=1, keepdim=True) # [B, 1, C]

updated_x_others = torch.cat([updated_x_others, extra_one_token],dim=1)

image_features = updated_x_others

return image_features

def token_prune_merge_advanced_plus(self, images, if_adaptive=True, reduction_ratio = 1/8):

'''

version 24/03/2024 using the spacially smapled tokens to supplement the pruned tokens

'''

# token_indix_list = []

# token_indix_dict = {}

#set hooks for extracting desired layer's k and q

hook_handle_k = self.vision_tower.vision_model.encoder.layers[23].self_attn.k_proj.register_forward_hook(hook_k)

hook_handle_q = self.vision_tower.vision_model.encoder.layers[23].self_attn.q_proj.register_forward_hook(hook_q)

#forward pass

image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)

cls_token_last_layer =image_forward_outs.hidden_states[self.select_layer][:, 0:1]

image_features = self.feature_select(image_forward_outs).to(images.dtype)

B, N, C = image_features.shape

#extract desired layer's k and q and remove hooks; calculate attention

desired_layer_k = outputs["desired_k"]

desired_layer_q = outputs["desired_q"]

hook_handle_k.remove()

hook_handle_q.remove()

attn = (desired_layer_q @ desired_layer_k.transpose(-2, -1)) * C ** -0.5

attn = F.softmax(attn, dim=-1)

cls_attn = attn[:, 0, 1:]

if if_adaptive:

reduction_ratio = outlier_dectection(cls_attn)#*3.5

_, idx = torch.topk(cls_attn, int(N*reduction_ratio), dim=1, largest=True) # [B, left_tokens] , sorted=True

# # # print("idx: ", idx)

if if_adaptive:

step_length = int(1/reduction_ratio)

arithmetic_sequence = torch.arange(0, 575, int(step_length/3)).to(device=self.device)

original_tensor_1d = idx.flatten().to(device=self.device)

filtered_sequence = torch.tensor([x for x in arithmetic_sequence if x not in original_tensor_1d]).to(device=self.device)

concatenated_tensor = torch.cat((idx, filtered_sequence.unsqueeze(0)), dim=1)

idx = concatenated_tensor

# # print("idx_new: ", idx)

else:

# # this is for training

step_length = int(1/reduction_ratio)

new_idx = torch.zeros((idx.size(0), idx.size(1)*2), dtype=torch.long).to(device=self.device)

for i in range(idx.size(0)):

arithmetic_sequence = torch.arange(int(step_length/2), 575, int(step_length)).to(device=self.device)

original_tensor_1d = idx[i].flatten().to(device=self.device)

filtered_sequence = arithmetic_sequence

# filtered_sequence = torch.tensor([x for x in arithmetic_sequence if x not in original_tensor_1d]).to(device=self.device)

concatenated_tensor = torch.cat((original_tensor_1d, filtered_sequence), dim=0)

new_idx[i] = concatenated_tensor

idx = new_idx

index = idx.unsqueeze(-1).expand(-1, -1, C) # [B, left_tokens, C]

Key_wo_cls = desired_layer_k[:, 1:] # [B, N-1, C]

x_others = torch.gather(image_features, dim=1, index=index) # [B, left_tokens, C]

x_others_attn = torch.gather(cls_attn, dim=1, index=idx)

Key_others = torch.gather(Key_wo_cls, dim=1, index=index) # [B, left_tokens, C]

compl = complement_idx(idx, N) # [B, N-1-left_tokens]

non_topk = torch.gather(image_features, dim=1, index=compl.unsqueeze(-1).expand(-1, -1, C)) # [B, N-1-left_tokens, C]

non_topk_Key = torch.gather(Key_wo_cls, dim=1, index=compl.unsqueeze(-1).expand(-1, -1, C))

non_topk_attn = torch.gather(cls_attn, dim=1, index=compl) # [B, N-1-left_tokens]

Key_others_norm = F.normalize(Key_others, p=2, dim=-1)

non_topk_Key_norm = F.normalize(non_topk_Key, p=2, dim=-1)

# cos_sim = torch.bmm(Key_others_norm, non_topk_Key_norm.transpose(1, 2)) # [B, left_tokens, N-1-left_tokens]

# _, cluster_indices = torch.topk(cos_sim, k=4, dim=2, largest=True)

B, left_tokens, C = x_others.size()

updated_x_others = torch.zeros_like(x_others)

for b in range(B):

for i in range(left_tokens):

key_others_norm = Key_others_norm[b,i,:].unsqueeze(0).unsqueeze(0)

before_i_Key = Key_others_norm[b, :i, :].unsqueeze(0)

after_i_Key = Key_others_norm[b, i+1:, :].unsqueeze(0)

before_i_x_others = x_others[b, :i, :].unsqueeze(0)

after_i_x_others = x_others[b, i+1:, :].unsqueeze(0)

rest_x_others = torch.cat([before_i_x_others, after_i_x_others, non_topk[b,:,:].unsqueeze(0)], dim=1)

before_i_x_others_attn = x_others_attn[b, :i].unsqueeze(0)

after_i_x_others_attn = x_others_attn[b, i+1:].unsqueeze(0)

rest_x_others_attn = torch.cat([before_i_x_others_attn, after_i_x_others_attn, non_topk_attn[b,:].unsqueeze(0)], dim=1)

rest_Keys = torch.cat([before_i_Key, after_i_Key, non_topk_Key_norm[b,:,:].unsqueeze(0)], dim=1)

cos_sim_matrix = torch.bmm(key_others_norm, rest_Keys.transpose(1, 2))

_, cluster_indices = torch.topk(cos_sim_matrix, k=int(32), dim=2, largest=True)

cluster_tokens = rest_x_others[:,cluster_indices.squeeze(),:]

weights = rest_x_others_attn[:,cluster_indices.squeeze()].unsqueeze(-1)

# update cluster centers

weighted_avg = torch.sum(cluster_tokens * weights, dim=1) #/ torch.sum(weights)

updated_center = x_others[b, i, :] + weighted_avg

updated_x_others[b, i, :] = updated_center

extra_one_token = torch.sum(non_topk * non_topk_attn.unsqueeze(-1), dim=1, keepdim=True) # [B, 1, C]

updated_x_others = torch.cat([updated_x_others, extra_one_token],dim=1)

image_features = updated_x_others

return image_features

@torch.no_grad()

def forward(self, images):

if type(images) is list:

image_features = []

for image in images:

image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)

image_feature = self.feature_select(image_forward_out).to(image.dtype)

image_features.append(image_feature)

else:

# image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)

# image_features = self.feature_select(image_forward_outs).to(images.dtype)

# image_features = self.token_prune_merge_advanced(images, if_adaptive=True, reduction_ratio=1/8)

image_features = self.token_prune_merge_advanced_plus(images, if_adaptive=True, reduction_ratio=1/8)

return image_features

@property

def dummy_feature(self):

return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

@property

def dtype(self):

return self.vision_tower.dtype

@property

def device(self):

return self.vision_tower.device

@property

def config(self):

if self.is_loaded:

return self.vision_tower.config

else:

return self.cfg_only

@property

def hidden_size(self):

return self.config.hidden_size

@property

def num_patches(self):