hypernet.py

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

class HyperNet(nn.Module):
    def __init__(self, target_modules_shapes, hidden_dim=2048, r=8):
        """
        target_modules_shapes (dict): Mapping from target layer string to a tuple (d_out, d_in).
        Example: {'q_proj': (4096, 4096), 'v_proj': (1024, 4096)}
        """
        super().__init__()
        self.r = r
        self.hidden_dim = hidden_dim
        self.target_modules_shapes = target_modules_shapes
        
        # Input Layer map the z vector to the hidden dimensionality
        self.input_layer = nn.Sequential(
            nn.Linear(768, hidden_dim),
            nn.ReLU()
        )
        
        self.head_u = nn.ModuleDict()
        self.head_sigma = nn.ModuleDict()
        self.head_v = nn.ModuleDict()
        
        for name, (d_out, d_in) in target_modules_shapes.items():
            safe_name = name.replace(".", "_")
            
            # Phase 3 Output Heads per target matrix
            self.head_u[safe_name] = nn.Linear(hidden_dim, d_out * r)
            self.head_sigma[safe_name] = nn.Linear(hidden_dim, r)
            self.head_v[safe_name] = nn.Linear(hidden_dim, r * d_in)
            
    def forward(self, z):
        """
        Args:
            z (Tensor): A continuous vector of shape (batch_size, 768)
        Returns:
            U_gen_dict (dict): Maps layer name -> Tensor of shape (batch_size, d_out, r)
            log_Sigma_gen_dict (dict): Maps layer name -> Tensor of shape (batch_size, r)
            V_gen_dict (dict): Maps layer name -> Tensor of shape (batch_size, r, d_in) 
                               representing the V^T right singular vectors matrix.
        """
        h = self.input_layer(z)
        
        U_gen_dict = {}
        log_Sigma_gen_dict = {}
        V_gen_dict = {}
        
        for name, (d_out, d_in) in self.target_modules_shapes.items():
            safe_name = name.replace(".", "_")
            
            u = self.head_u[safe_name](h)
            sigma = self.head_sigma[safe_name](h)
            v = self.head_v[safe_name](h)
            
            U_gen_dict[name] = u.view(-1, d_out, self.r)
            log_Sigma_gen_dict[name] = sigma
            V_gen_dict[name] = v.view(-1, self.r, d_in)
            
        return U_gen_dict, log_Sigma_gen_dict, V_gen_dict