|
|
import math
|
|
|
|
|
|
import numpy as np
|
|
|
import torch
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
loss_fn = torch.nn.L1Loss()
|
|
|
|
|
|
|
|
|
def ad_loss(
|
|
|
q_list, ks_list, vs_list, self_out_list, scale=1, source_mask=None, target_mask=None
|
|
|
):
|
|
|
loss = 0
|
|
|
attn_mask = None
|
|
|
for q, ks, vs, self_out in zip(q_list, ks_list, vs_list, self_out_list):
|
|
|
if source_mask is not None and target_mask is not None:
|
|
|
w = h = int(np.sqrt(q.shape[2]))
|
|
|
mask_1 = torch.flatten(F.interpolate(source_mask, size=(h, w)))
|
|
|
mask_2 = torch.flatten(F.interpolate(target_mask, size=(h, w)))
|
|
|
attn_mask = mask_1.unsqueeze(0) == mask_2.unsqueeze(1)
|
|
|
attn_mask=attn_mask.to(q.device)
|
|
|
|
|
|
target_out = F.scaled_dot_product_attention(
|
|
|
q * scale,
|
|
|
torch.cat(torch.chunk(ks, ks.shape[0]), 2).repeat(q.shape[0], 1, 1, 1),
|
|
|
torch.cat(torch.chunk(vs, vs.shape[0]), 2).repeat(q.shape[0], 1, 1, 1),
|
|
|
attn_mask=attn_mask
|
|
|
)
|
|
|
loss += loss_fn(self_out, target_out.detach())
|
|
|
return loss
|
|
|
|
|
|
|
|
|
|
|
|
def q_loss(q_list, qc_list):
|
|
|
loss = 0
|
|
|
for q, qc in zip(q_list, qc_list):
|
|
|
loss += loss_fn(q, qc.detach())
|
|
|
return loss
|
|
|
|
|
|
|
|
|
def qk_loss(q_list, k_list, qc_list, kc_list):
|
|
|
loss = 0
|
|
|
for q, k, qc, kc in zip(q_list, k_list, qc_list, kc_list):
|
|
|
scale_factor = 1 / math.sqrt(q.size(-1))
|
|
|
self_map = torch.softmax(q @ k.transpose(-2, -1) * scale_factor, dim=-1)
|
|
|
target_map = torch.softmax(qc @ kc.transpose(-2, -1) * scale_factor, dim=-1)
|
|
|
loss += loss_fn(self_map, target_map.detach())
|
|
|
return loss
|
|
|
|
|
|
|
|
|
def qkv_loss(q_list, k_list, vc_list, c_out_list):
|
|
|
loss = 0
|
|
|
for q, k, vc, target_out in zip(q_list, k_list, vc_list, c_out_list):
|
|
|
self_out = F.scaled_dot_product_attention(q, k, vc)
|
|
|
loss += loss_fn(self_out, target_out.detach())
|
|
|
return loss
|
|
|
|
