models/mae.py

from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random

from models.encoders import MultiEncoder
from models.decoders import Decoder
from models.fusion import FusionUpsampleBlock
from models.transformer_modules import RevIN
from configs.global_cfg import TARGET_RESAMPLE



class MultiMaskedAutoencoderConv1d(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """
    def __init__(self, sig_len=6000, patch_size=10,
                 embed_dim=[32,64,128,256], depth=[2,2,4], specific_depth=[2,2,2], num_heads=8,
                 decoder_embed_dim=256, decoder_depth=6, decoder_num_heads=8,
                 mlp_ratio=4, norm_layer=nn.LayerNorm,
                 modalities=["EEG","EOG"], num_experts=2, num_experts_per=1,
                 use_shared=True, use_moe=True, 
                 fusion_method='concat', fusion_depth=2, fusion_heads=8, 
                 use_revin=False,
                 mask_ratio=0.5,
                 contrastive_training=False, koleo_training=False, pos_enc='sinusoidal'):

        super().__init__()         
        self.patch_size = patch_size        
        self.modalities = modalities        
        self.num_experts = num_experts
        self.num_experts_per = num_experts_per        
        self.embed_dim = embed_dim
        self.depth = depth
        self.specific_depth = specific_depth
        self.mask_ratio = mask_ratio
        self.encoder = MultiEncoder(sig_len=sig_len, patch_size=patch_size,
                 embed_dim=embed_dim, depth=depth, num_heads=num_heads, specific_depth=specific_depth,
                 mlp_ratio=mlp_ratio, norm_layer=norm_layer,
                 modalities=self.modalities, num_experts=num_experts, num_experts_per=num_experts_per,
                 mask_ratio=mask_ratio, pos_enc=pos_enc, use_shared=use_shared, use_moe=use_moe,)

        self.num_patches = self.encoder.num_patches                
        
        self.use_revin = use_revin
        if self.use_revin:
            self.revin = nn.ModuleDict()
            for m in self.modalities:
                if m != 'SaO2':
                    self.revin[m] = RevIN(num_features=1, affine=False)
                    
        self.use_shared = use_shared
        self.use_moe = use_moe
        
        self.contrastive_training = contrastive_training
        self.koleo_training = koleo_training

        # contrastive loss hyperparameters (InfoNCE style)
        self.contrastive_temp = 0.07
        
        self.fusion = nn.ModuleDict()                
        self.decoders = nn.ModuleDict()
        self.upsampler = nn.ModuleDict()
        

        self.fuser = FusionUpsampleBlock(dim_list=embed_dim[2:], 
                                         fusion_method=fusion_method, 
                                         fusion_depth=fusion_depth, 
                                         fusion_heads=fusion_heads,
                                         mlp_ratio=mlp_ratio, 
                                         norm_layer=norm_layer, 
                                         pos_enc=pos_enc, 
                                         use_shared=use_shared, 
                                         use_upsample=True, 
                                         modalities=self.modalities)
        
        for m in self.modalities:
            decoder_pred_dim = self.patch_size if m != 'SaO2' else 1
            self.decoders[m] = Decoder(num_patches=self.num_patches, encoder_embed_dim=embed_dim[2], decoder_embed_dim=decoder_embed_dim, 
                    decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
                    mlp_ratio=mlp_ratio, norm_layer=norm_layer, decoder_pred_dim=decoder_pred_dim, pos_enc=pos_enc)                        
        
        
        
        self.apply(self._init_weights)

    def koleo_loss(self, x, ys=None, eps=1e-8):
        """
        Compute KoLeo (Kozachenko-Leonenko) regularization loss.
        Encourages uniform distribution of representations on the hypersphere.
        
        Args:
            x: Tensor of shape (N, D) where N is batch size and D is feature dimension
            ys: Optional list of tensors to also consider for distance computation
            eps: Small constant for numerical stability
        
        Returns:
            KoLeo loss (scalar)
        """
        # Normalize representations to unit hypersphere
        x = F.normalize(x, p=2, dim=-1)
        
        # Compute pairwise distances
        # Using efficient computation: ||a - b||^2 = ||a||^2 + ||b||^2 - 2<a,b>
        # Since normalized, ||a||^2 = ||b||^2 = 1, so ||a - b||^2 = 2 - 2<a,b>
        gram_matrix = torch.matmul(x, x.t())  # (N, N)
        sq_distances = 2 - 2 * gram_matrix  # (N, N)
        sq_distances = torch.clamp(sq_distances, min=eps, max=4.0)  # numerical stability, max distance on unit sphere is 2
        
        # Set diagonal to large value (exclude self-distance)
        sq_distances = sq_distances + torch.eye(x.shape[0], device=x.device) * 1e10

        if ys is not None:
            for y in ys:
                y = F.normalize(y, p=2, dim=-1)
                y_distance = (x - y).pow(2).sum(dim=-1, keepdim=True)  # (N, 1)
                y_distance = torch.clamp(y_distance, min=eps, max=4.0)
                sq_distances = torch.cat([sq_distances, y_distance], dim=1)  # (N, N + len(ys))
        
        # Find minimum distance for each sample (nearest neighbor)
        min_sq_distances, _ = torch.min(sq_distances, dim=1)  # (N,)
        
        # KoLeo loss is negative log of minimum distances
        # We want to maximize minimum distances -> minimize negative log
        distances = torch.sqrt(min_sq_distances)
        
        # Additional clamping on distances before log
        distances = torch.clamp(distances, min=eps)
        
        koleo_loss = -torch.log(distances).mean()
        
        return koleo_loss

    def forward(self, inputs, mask_dict=None, ids_keep_dict=None):
        # embed patches
        # m1, m2 = self.modalities[0], self.modalities[1]
        modalities_to_use = list(inputs.keys())
        # random.shuffle(modalities_to_use)
        
        m1, m2 = modalities_to_use[0], modalities_to_use[1]
        if mask_dict is None or ids_keep_dict is None:
            mask_dict, ids_keep_dict = self.random_masking(inputs[m1], self.mask_ratio)       
        
        # ids_shuffle, mask, ids_keep, ids_restore = self.random_masking(inputs[m1], mask_ratio)
        
        N = inputs[m1].shape[0]
        d = self.embed_dim[2]
        inputs_ = dict()
        if self.use_revin:
            for i, m in enumerate(modalities_to_use):
                # mask_m = mask_dict[m].reshape(N, -1, 1)                
                # mask_raw_m = torch.tile(mask_m, (1,  1, self.patch_size)).reshape(N, -1, 1)
                if m == 'SaO2':                    
                    inputs_[m] = inputs[m][:, ::TARGET_RESAMPLE, :] # N,L* TARGET_RESAMPLE,1 -> N,L,1
                else:
                    inputs_[m] = self.revin[m](inputs[m], None, mode='norm')
        else:
            inputs_ = inputs

        preds = dict()
        latents_shared, latents_specific, gating_scores = self.encoder(inputs_, mask_dict, ids_keep_dict, modalities_to_use)
        
        latents_fused, upsampled = self.fuser(latents_shared, latents_specific)
        
        if self.mask_ratio == 0:
            losses = dict()
            for m in modalities_to_use:                                     
                losses[m] = 0.0                          
            
        else:
            for i, m in enumerate(modalities_to_use):                     
                ids_keep_m = ids_keep_dict[m] # (N, L * large portion)        
                preds[m] = self.decoders[m](upsampled[i], ids_keep_m)     
            # Compute reconstruction losses
            losses = self.forward_loss(inputs_, preds, mask_dict, ids_keep_dict, modalities_to_use)

        # Contrastive learning between latents_shared and latents_specific
        # Positives: (latents_shared[m1], latents_shared[m2]) for the same timestep (index)
        # Negatives: (latents_shared[m1], latents_specific[m1]) for the same timestep
        contrastive_loss = 0.0
        koleo_loss = 0.0
        num_contrast_pairs = 0
        
        ctrst_grp = 30 # (seconds)
        ctrst_grp = (ctrst_grp * (100 / self.patch_size)) * (1 - self.mask_ratio) # 0.5 is mask ratio
        # How many patches form a patch group for contrastive learning?
        ctrst_grp = int(ctrst_grp // (5 * 2 ** (len(self.depth) - 2))) # The first downsampling is 5x, then 2x for each level
    
        
        if self.contrastive_training:      

            z_shared, z_specific = [], []
            for i, m in enumerate(modalities_to_use):
                shared_m = latents_shared[-1][i]
                N, L, D = shared_m.shape
                Lm = shared_m.reshape(-1, ctrst_grp*D)
                # Lm = self.pooler1[m](Lm)
                Lm = F.normalize(Lm, dim=-1)                
                z_shared.append(Lm)  # (N, L, D')

                specific_m = latents_specific[-1][i]
                # N, L, D = specific_m.shape
                # specific_m = specific_m.reshape(N, L, 4, D).mean(dim=2)
                Lm = specific_m.reshape(-1, ctrst_grp*D)
                # Lm = self.pooler2[m](Lm)                
                Lm = F.normalize(Lm, dim=-1)                
                z_specific.append(Lm)  # (N, L, D')

            
            for i, m1 in enumerate(modalities_to_use):                  
                shared_specific = []  
                z_anchor = z_shared[i]
                
                
                sims_neg1 = z_anchor @ z_anchor.permute(1,0) # (N*L, N*L)                
                diag_idx = torch.arange(sims_neg1.size(1), device=sims_neg1.device)
                sims_neg1[diag_idx, diag_idx] = float('-inf')
                
                
                z_pos = [z for k, z in enumerate(z_shared) if k != i]
                z_pos = torch.mean(torch.stack(z_pos, dim=0), dim=0)  # Average over other modalities
                
                sims_pos = torch.sum(z_anchor * z_pos, dim=-1, keepdim=True)  # (N*L, 1)            
                
                if self.use_shared:
                    zs = z_specific[i]
                    sims_neg2 = torch.sum(z_anchor * zs, dim=-1, keepdim=True)
                    shared_specific.append(sims_neg2)
                    sims_neg2 = torch.cat(shared_specific, dim=-1) # (N*L, num_modalities)              

                    logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) / (self.contrastive_temp + 1e-8)                     
                else:
                    logits = torch.cat([sims_pos, sims_neg1], dim=-1) / (self.contrastive_temp + 1e-8)
                # logits = torch.cat([sims_pos, sims_neg1, sims_neg2], dim=-1) * logit_scale 
                labels = torch.zeros(logits.shape[0], dtype=torch.long, device=logits.device)                    
                loss_c = F.cross_entropy(logits, labels)
                contrastive_loss = contrastive_loss + loss_c
                num_contrast_pairs += 1
            

            if num_contrast_pairs > 0:
                contrastive_loss = contrastive_loss / float(num_contrast_pairs)
        
        if self.koleo_training: 
            # Compute KoLeo regularization on shared representations
            # This encourages uniform distribution of shared features on the hypersphere
            # for i in range(len(modalities_to_use)):
            #     shared_flat = z_shared[i]
                          
               
            #     specific_flats = [z_specific[i]]
            #     # combined_flat = torch.cat([shared_flat, specific_m1_flat, specific_m2_flat], dim=0)
                        
            #     koleo_loss = koleo_loss + self.koleo_loss(shared_flat, ys=specific_flats)
            koleo_loss = self.koleo_loss(torch.cat(z_shared, dim=0), ys=torch.cat(z_specific, dim=0))
         

        
        aux_loss = self.load_balancing_loss(gating_scores)
        # return losses, aux_loss_lg, aux_loss_sm, preds, mask_dict, ids_keep_dict
        return losses, aux_loss, contrastive_loss, koleo_loss, preds, mask_dict, ids_keep_dict

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0) 

    def patchify(self, sig):
        """
        sig: (N, L, 1)
        patchfied: (N, L, patch_size)
        """
        # p = self.patch_size
        B, L, _ = sig.shape

        num_patches = L // self.patch_size

        x = sig.reshape(shape=(B, num_patches, self.patch_size))

        return x
    
    def random_masking(self, x, mask_ratio=0.5):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """

        mask_dict, ids_keep_dict = dict(), dict()
        
        # If no mask is used. Keep all patches.
        if mask_ratio == 0:
            for m in self.modalities:
                mask_dict[m] = torch.zeros(x.shape[0], x.shape[1] // self.patch_size, device=x.device, dtype=x.dtype)
                ids_keep_dict[m] = torch.arange(x.shape[1] // self.patch_size, device=x.device).unsqueeze(0).repeat([x.shape[0], 1])
            return mask_dict, ids_keep_dict

        N = x.shape[0]        
        L = x.shape[1] // self.patch_size # number of large patches        

        
        len_keep = int(L * (1 - mask_ratio))  # Number of keep large patches

        # noise = torch.rand(N, L_lg, device=x.device)  # noise in [0, 1]
        # ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove     
        
        ids_ordered = torch.arange(L, device=x.device).unsqueeze(0).repeat([N, 1])
        # ids_ordered = ids_ordered.view([N, L //2, 2])
        rand_ids = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
        rand_ids2 = torch.randint(0, 2, [N, len_keep, 1], device=x.device)
        ids_keep_m1 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids, dim=-1).view([N, -1])
        ids_keep_m2 = torch.gather(ids_ordered.view([N, len_keep, -1]), index=rand_ids2, dim=-1).view([N, -1])
        
        num_patches_slow = (200 // self.patch_size)
        len_keep_slow = int((L // num_patches_slow) * (1 - mask_ratio))  # Number of keep large patches for slow modalities
        
        rand_ids = torch.randint(0, 2, [N, len_keep_slow, 1], device=x.device)
        ids_keep_slow = torch.gather(ids_ordered.view([N, len_keep_slow, -1, num_patches_slow]), 
                                index=rand_ids.unsqueeze(-1).repeat([1,1,1,num_patches_slow]), 
                                dim=2).view([N, -1])
        
        
        # ids_keep_m1 = ids[:, 0::2]  ## Large keep area, modality 1
        # ids_keep_m2 = ids[:, 1::2]  ## Large keep area, modality 2
        
         # Mask construction for large modeling area
        mask_m1 = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, len_keep], device=x.device, dtype=x.dtype)
        mask_m1.scatter_(dim=1, index=ids_keep_m1, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
 
        mask_m2 = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
        mask_m2.scatter_(dim=1, index=ids_keep_m2, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
        
        mask_slow = torch.ones([N, L], device=x.device, dtype=x.dtype)
        src = torch.zeros([N, L - len_keep], device=x.device, dtype=x.dtype)
        mask_slow.scatter_(dim=1, index=ids_keep_slow, src=src)  # 0 is keep, 1 is remove, Large mask + small modeling area masked
        # print(rand_ids)
        # print(ids_keep_slow)
        

        for i, m in enumerate(self.modalities):
            idx = (i % 2) + 1            
            if m in ['RESP','SaO2']:
            # if False: 
                ids_keep_dict[m] = ids_keep_slow
                mask_dict[m] = mask_slow  # Both use the same mask
            else:
                ids_keep_dict[m] = ids_keep_m1 if idx == 1 else ids_keep_m2            
                mask_dict[m] = mask_m1 if idx == 1 else mask_m2            


        return mask_dict, ids_keep_dict

    def forward_loss(self, inputs, preds, mask_dict, ids_keep_dict, modalities_to_use):
        """
        inputs: m * [N, L]
        pred: [N, L, p*p*3]
        mask: [N, L], 0 is keep, 1 is remove, 
        """
        losses = dict()        
        # lg_ids, sm_ids = ids_keep_dict['ids_lg'], ids_keep_dict['ids_sm']
        # for m in inputs.keys():
        for i, m in enumerate(modalities_to_use):            
            
            sig, pred, mask= inputs[m], preds[m], mask_dict[m]
            
            N, L, _ = sig.shape

            if m == 'SaO2':
                repeat_factor = 100 // self.patch_size
                target = sig.repeat_interleave(repeat_factor, dim=1)
                pred = pred.view(target.shape)            
                
                loss = (pred - target) ** 2               
                loss = (loss.mean(dim=-1) * mask).sum() / mask.sum()            

            else:            
                target = self.patchify(sig) # (N, L, patch_size)            
                pred = pred.view(target.shape)            
                
                sig_smooth = F.avg_pool1d(sig.view(N, -1, L), 
                                kernel_size=11, 
                                stride=1, 
                                padding=5, 
                                count_include_pad=False).view(N, L, -1)
                
                target_smooth = self.patchify(sig_smooth)
                loss_smooth = (pred - target_smooth) ** 2               
                loss_smooth = (loss_smooth.mean(dim=-1) * mask).sum() / mask.sum()            
                
                loss = loss_smooth
            
            
            losses[m] = loss
        return losses
    
    def load_balancing_loss(self, gating_scores):
        """
        Compute the load balancing loss for the gating scores.
        ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py#L591
        """
        num_moe = len(gating_scores)
        if num_moe > 0:
            overall_loss = 0.0
            for i in range(num_moe):
                gating_scores_i = gating_scores[i]  # stop gradient
                # gating_scores = torch.stack(gating_scores, dim=0)  # [num_MoEblocks, B, num_experts]
                routing_weights = gating_scores_i.softmax(dim=-1)  # [B, L, num_experts]
                _, selected_experts = torch.topk(routing_weights, self.num_experts_per, dim=-1) # Index of selected experts
                expert_mask = torch.nn.functional.one_hot(selected_experts, self.num_experts) # [B, L, num_experts_per, num_experts]

                # Compute the average number of data points routed to each experts
                # tokens_per_expert = torch.mean(expert_mask.float(), dim=1) # [B, num_experts]

                # # Compute the average probability of routing to these experts
                # router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts]

                # overall_loss += torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(1))
                tokens_per_expert = torch.mean(expert_mask.float(), dim=(0,1,2)) # [num_experts]
                router_prob_per_expert = torch.mean(routing_weights, dim=(0,1)) # [num_experts]
                overall_loss += torch.mean(tokens_per_expert * router_prob_per_expert)
            
            return overall_loss / num_moe
        else:
            return 0.0