Open King4819 opened 3 months ago
Sorry for the late reply, I didn’t check my GitHub in time. For dTPS, we use gumbel-softmax to help train learnable token scoring head ,which follows dynamicViT. So, in the training stage, the tokens can not be removed actually but with the policy mask , we can remove the pruned tokens' influence on the cost matrix in dTPS modules and the following attention operation
` class TPS(M.Module):
`
if self.variant == 'dTPS' and self.training:
cost_matrix = cal_cosine_similarity( reserved, None, mask_eye=-100) ...
else: cost_matrix = cal_cosine_similarity( pruned, reserved, mask_eye=None) ,,,
seems like in dtps training stage, it calculate the cosine similarity between reserved token set itself instead of reserved token set and pruned token set, which is different from inference stage I'm wondering how does it work ? Thanks!!!