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 # Since normalized, ||a||^2 = ||b||^2 = 1, so ||a - b||^2 = 2 - 2 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