Update lavila/models/prompt_tuning.py

lavila/models/prompt_tuning.py

@@ -58,7 +58,7 @@ class PromptPoolLearner(nn.Module):
 
             if istrain:
                 inv_freq = self.id_table.sum() / self.id_table.float()
-                weights = torch.softmax((1 + similarity) / 2 + 0.5 * (1 - gamma) * inv_freq / inv_freq.sum(), dim=1)
+                weights = (similarity + 1) / 2 * gamma + (1 - gamma) * torch.softmax(inv_freq, dim=-1)
                 idx = torch.multinomial(weights, k, replacement=False)
             else:
                 idx = torch.argsort(similarity, dim=-1, descending=True)[:, :k]
@@ -74,8 +74,8 @@ class PromptPoolLearner(nn.Module):
         out['prompts'] = prompts
         sel_sim = similarity[torch.arange(BZ).view(-1, 1), idx]
         sel_key = keys[idx.flatten(), ...].view(BZ, k, self.prompt_dim)
-        diff = F.mse_loss((sel_sim.unsqueeze(1) @ sel_key).squeeze(), query.detach(), reduction='sum') / BZ
-        ksim = torch.mean(torch.matmul(keys, keys.t()) - torch.eye(self.size).to(keys.device))
+        diff = F.mse_loss((sel_sim.unsqueeze(1) @ sel_key).squeeze(1), query.detach(), reduction='sum') / BZ
+        ksim = torch.sum(torch.abs(torch.matmul(keys, keys.t()) - torch.eye(self.size).to(keys.device))) / BZ
         out['ps_loss'] = diff + ksim
 
         return out