loss.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class OrthogonalSVDLoss(nn.Module):
    def __init__(self, lambda_1=0.1, lambda_2=0.1, r=8):
        super().__init__()
        self.lambda_1 = lambda_1
        self.lambda_2 = lambda_2
        self.r = r
        
    def forward(self, U_gen_dict, log_Sigma_gen_dict, V_gen_dict, 
                U_star_dict, log_Sigma_star_dict, V_star_dict):
        """
        Computes the geometrically constrained loss for Zero-Shot Semantic HyperLoRA:
        L_total = L_MSE + L_ortho
        """
        loss_mse = 0.0
        loss_ortho = 0.0
        
        for name in U_gen_dict.keys():
            U_gen = U_gen_dict[name]
            log_Sigma_gen = log_Sigma_gen_dict[name]
            # V_gen mathematically acts as V^T per the spec formulation (r, d_in).
            V_gen = V_gen_dict[name]
            
            U_star = U_star_dict[name].to(U_gen.device)
            log_Sigma_star = log_Sigma_star_dict[name].to(log_Sigma_gen.device)
            V_star_T = V_star_dict[name].to(V_gen.device) 
            
            # 1. Reconstruction Loss (MSE)
            loss_mse += F.mse_loss(U_gen, U_star)
            loss_mse += F.mse_loss(log_Sigma_gen, log_Sigma_star)
            loss_mse += F.mse_loss(V_gen, V_star_T)
            
            # 2. Orthogonal Regularization (Stiefel Manifold Constraint)
            # U_gen is (B, d_out, r). U_gen^T U_gen should be I_r
            B = U_gen.size(0)
            I_r = torch.eye(self.r, device=U_gen.device).unsqueeze(0).expand(B, -1, -1)
            
            U_ortho = torch.bmm(U_gen.transpose(1, 2), U_gen)
            loss_ortho += self.lambda_1 * F.mse_loss(U_ortho, I_r)
            
            # V_gen is (B, r, d_in), representing V^T. Thus V_gen V_gen^T should be I_r.
            V_ortho = torch.bmm(V_gen, V_gen.transpose(1, 2))
            loss_ortho += self.lambda_2 * F.mse_loss(V_ortho, I_r)
            
        total_loss = loss_mse + loss_ortho
        return total_loss, loss_mse, loss_ortho