""" Multimodal PC Fault Detection - Model Architecture v2 ====================================================== Changes from v1: - Auxiliary unimodal classification heads (force each branch to independently classify) - Asymmetric loss weighting: λ_visual=1.5 (boost weak), λ_audio=0.5 (dampen dominant) - OGM-GE (On-the-fly Gradient Modulation + Generalization Enhancement) support - Forward returns per-branch logits + embeddings for OGM-GE gradient modulation Two-branch architecture: - Visual: ViT-B/16 pretrained on ImageNet-21k - Audio: AST pretrained on AudioSet - Fusion: Late fusion (concat / weighted sum / attention) Supports LoRA, full fine-tuning, and linear probe modes. References: - OGM-GE: Peng et al., "Balanced Multimodal Learning via On-the-fly Gradient Modulation", CVPR 2022 (arXiv: 2203.15332) """ import torch import torch.nn as nn import torch.nn.functional as F from typing import Dict, Optional, Literal from transformers import ViTModel, ASTModel, ViTImageProcessor, ASTFeatureExtractor from peft import LoraConfig, get_peft_model from config import ModelConfig, LoRAConfig, FAULT_CLASSES # =========================================================================== # Branch Modules (unchanged from v1) # =========================================================================== class VisualBranch(nn.Module): def __init__(self, config, lora_config=None, finetune_method="lora"): super().__init__() self.vit = ViTModel.from_pretrained(config.vit_model_name) if finetune_method == "lora" and lora_config and lora_config.enabled: peft_config = LoraConfig( r=lora_config.r, lora_alpha=lora_config.lora_alpha, target_modules=lora_config.vit_target_modules, lora_dropout=lora_config.lora_dropout, bias=lora_config.bias) self.vit = get_peft_model(self.vit, peft_config) self.vit.print_trainable_parameters() elif finetune_method == "linear_probe": for param in self.vit.parameters(): param.requires_grad = False def forward(self, pixel_values): return self.vit(pixel_values=pixel_values).last_hidden_state[:, 0, :] class AudioBranch(nn.Module): def __init__(self, config, lora_config=None, finetune_method="lora"): super().__init__() self.ast = ASTModel.from_pretrained(config.ast_model_name) if finetune_method == "lora" and lora_config and lora_config.enabled: peft_config = LoraConfig( r=lora_config.r, lora_alpha=lora_config.lora_alpha, target_modules=lora_config.ast_target_modules, lora_dropout=lora_config.lora_dropout, bias=lora_config.bias) self.ast = get_peft_model(self.ast, peft_config) self.ast.print_trainable_parameters() elif finetune_method == "linear_probe": for param in self.ast.parameters(): param.requires_grad = False def forward(self, input_values): return self.ast(input_values=input_values).last_hidden_state[:, 0, :] # =========================================================================== # Fusion Module (unchanged from v1) # =========================================================================== class LateFusion(nn.Module): def __init__(self, config): super().__init__() self.fusion_type = config.fusion_type if config.fusion_type == "concat": self.visual_proj = nn.Linear(config.vit_embed_dim, config.fusion_dim) self.audio_proj = nn.Linear(config.ast_embed_dim, config.fusion_dim) self.classifier = nn.Sequential( nn.LayerNorm(config.fusion_dim * 2), nn.Dropout(config.fusion_dropout), nn.Linear(config.fusion_dim * 2, config.fusion_dim), nn.GELU(), nn.Dropout(config.fusion_dropout), nn.Linear(config.fusion_dim, config.num_classes)) elif config.fusion_type == "weighted_sum": self.visual_head = nn.Linear(config.vit_embed_dim, config.num_classes) self.audio_head = nn.Linear(config.ast_embed_dim, config.num_classes) self.fusion_weights = nn.Parameter(torch.tensor([0.5, 0.5])) elif config.fusion_type == "attention": self.visual_proj = nn.Linear(config.vit_embed_dim, config.fusion_dim) self.audio_proj = nn.Linear(config.ast_embed_dim, config.fusion_dim) self.cross_attn = nn.MultiheadAttention( embed_dim=config.fusion_dim, num_heads=8, dropout=config.fusion_dropout, batch_first=True) self.classifier = nn.Sequential( nn.LayerNorm(config.fusion_dim), nn.Dropout(config.fusion_dropout), nn.Linear(config.fusion_dim, config.num_classes)) def forward(self, visual_emb, audio_emb, modality_mask=None): if modality_mask: visual_emb = visual_emb * modality_mask.get("visual", 1.0) audio_emb = audio_emb * modality_mask.get("audio", 1.0) if self.fusion_type == "concat": fused = torch.cat([self.visual_proj(visual_emb), self.audio_proj(audio_emb)], dim=-1) return self.classifier(fused) elif self.fusion_type == "weighted_sum": w = torch.softmax(self.fusion_weights, dim=0) return w[0] * self.visual_head(visual_emb) + w[1] * self.audio_head(audio_emb) elif self.fusion_type == "attention": tokens = torch.cat([ self.visual_proj(visual_emb).unsqueeze(1), self.audio_proj(audio_emb).unsqueeze(1)], dim=1) return self.classifier(self.cross_attn(tokens, tokens, tokens)[0].mean(dim=1)) # =========================================================================== # OGM-GE: On-the-fly Gradient Modulation with Generalization Enhancement # =========================================================================== class OGMGEModulator: """ Implements OGM-GE from Peng et al., CVPR 2022. After loss.backward(), this computes per-modality confidence ratios and modulates encoder gradients to suppress the dominant modality and boost the weaker one. Gaussian noise is added to suppressed gradients for generalization enhancement. Usage in training loop: loss.backward() coeff_v, coeff_a, stats = ogm.compute_modulation_coefficients( visual_logits, audio_logits, labels) ogm.apply_gradient_modulation(model, coeff_v, coeff_a) optimizer.step() """ def __init__(self, alpha=0.3, noise_sigma=0.1): """ Args: alpha: Modulation strength. Higher = more aggressive suppression of dominant modality. Paper uses 0.3-0.5. noise_sigma: Std of Gaussian noise added to suppressed modality's gradients (Generalization Enhancement). Paper uses 0.1. """ self.alpha = alpha self.noise_sigma = noise_sigma @torch.no_grad() def compute_modulation_coefficients(self, visual_logits, audio_logits, labels): """ Compute OGM-GE modulation coefficients based on per-modality confidence. For each modality, we compute the average softmax probability of the correct class (confidence). The modality with higher confidence is considered dominant and gets its gradients scaled down. Args: visual_logits: (B, C) logits from the auxiliary visual head audio_logits: (B, C) logits from the auxiliary audio head labels: (B,) ground truth class indices Returns: coeff_visual: gradient scaling factor for visual encoder coeff_audio: gradient scaling factor for audio encoder stats: dict with debugging info """ # Softmax probabilities v_probs = F.softmax(visual_logits, dim=-1) a_probs = F.softmax(audio_logits, dim=-1) # Confidence = avg probability assigned to correct class batch_indices = torch.arange(labels.size(0), device=labels.device) v_conf = v_probs[batch_indices, labels].mean().item() a_conf = a_probs[batch_indices, labels].mean().item() # Confidence ratio: how much better one modality is than the other # ratio > 1 means visual is dominant, < 1 means audio is dominant eps = 1e-8 ratio = (v_conf + eps) / (a_conf + eps) # Modulation: scale down the dominant modality's gradients # If ratio > 1 (visual dominant): coeff_v < 1, coeff_a = 1 # If ratio < 1 (audio dominant): coeff_v = 1, coeff_a < 1 if ratio > 1.0: # Visual is dominant — suppress visual, keep audio coeff_visual = 1.0 - self.alpha * torch.tanh(torch.tensor(ratio - 1.0)).item() coeff_audio = 1.0 else: # Audio is dominant — suppress audio, keep visual coeff_visual = 1.0 coeff_audio = 1.0 - self.alpha * torch.tanh(torch.tensor(1.0 / ratio - 1.0)).item() stats = { "visual_conf": v_conf, "audio_conf": a_conf, "ratio": ratio, "coeff_visual": coeff_visual, "coeff_audio": coeff_audio, } return coeff_visual, coeff_audio, stats def apply_gradient_modulation(self, model, coeff_visual, coeff_audio): """ Scale gradients of encoder parameters. Only affects the visual_branch and audio_branch encoder weights — NOT the fusion head or auxiliary heads. For the suppressed modality (coeff < 1), also adds Gaussian noise to gradients (Generalization Enhancement from the paper). """ for name, param in model.named_parameters(): if param.grad is None: continue if "visual_branch" in name: param.grad.data.mul_(coeff_visual) # GE: add noise to suppressed modality if coeff_visual < 1.0 and self.noise_sigma > 0: noise = torch.randn_like(param.grad.data) * self.noise_sigma * param.grad.data.abs().mean() param.grad.data.add_(noise) elif "audio_branch" in name: param.grad.data.mul_(coeff_audio) if coeff_audio < 1.0 and self.noise_sigma > 0: noise = torch.randn_like(param.grad.data) * self.noise_sigma * param.grad.data.abs().mean() param.grad.data.add_(noise) # =========================================================================== # Main Model v2 — with auxiliary heads and OGM-GE support # =========================================================================== class MultimodalPCFaultDetector(nn.Module): """ v2 changes: - Auxiliary classification heads on each branch (visual_head, audio_head) - Forward returns per-branch logits for OGM-GE gradient modulation - Loss = loss_fusion + λ_v * loss_visual + λ_a * loss_audio - Asymmetric λ weights: λ_visual=1.5 (boost weak), λ_audio=0.5 (dampen dominant) """ def __init__(self, model_config, lora_config=None, finetune_method="lora", mode="multimodal", use_ogm=True, lambda_visual=1.5, lambda_audio=0.5): super().__init__() self.mode = mode self.modality_dropout_p = model_config.modality_dropout_p self.use_ogm = use_ogm self.lambda_visual = lambda_visual self.lambda_audio = lambda_audio # --- Branches --- self.visual_branch = ( VisualBranch(model_config, lora_config, finetune_method) if mode in ("multimodal", "visual_only") else None) self.audio_branch = ( AudioBranch(model_config, lora_config, finetune_method) if mode in ("multimodal", "audio_only") else None) # --- Fusion / classifier --- if mode == "multimodal": self.fusion = LateFusion(model_config) # NEW: Auxiliary unimodal classification heads # These force each branch to independently learn discriminative features self.visual_head = nn.Sequential( nn.LayerNorm(model_config.vit_embed_dim), nn.Dropout(0.2), nn.Linear(model_config.vit_embed_dim, model_config.num_classes)) self.audio_head = nn.Sequential( nn.LayerNorm(model_config.ast_embed_dim), nn.Dropout(0.2), nn.Linear(model_config.ast_embed_dim, model_config.num_classes)) else: embed_dim = (model_config.vit_embed_dim if mode == "visual_only" else model_config.ast_embed_dim) self.classifier = nn.Sequential( nn.LayerNorm(embed_dim), nn.Dropout(model_config.fusion_dropout), nn.Linear(embed_dim, model_config.fusion_dim), nn.GELU(), nn.Dropout(model_config.fusion_dropout), nn.Linear(model_config.fusion_dim, model_config.num_classes)) self.loss_fn = nn.CrossEntropyLoss() # --- Print parameter counts --- total = sum(p.numel() for p in self.parameters()) trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) print(f"[Model v2] Mode={mode}, Total={total:,}, Trainable={trainable:,} " f"({100*trainable/total:.2f}%)") if mode == "multimodal": print(f"[Model v2] OGM-GE={'ON' if use_ogm else 'OFF'}, " f"λ_visual={lambda_visual}, λ_audio={lambda_audio}") def forward(self, pixel_values=None, audio_values=None, labels=None): if self.mode == "multimodal": v_emb = self.visual_branch(pixel_values) a_emb = self.audio_branch(audio_values) # Modality dropout (training only) mask = None if self.training and self.modality_dropout_p > 0: mask = { "visual": 0.0 if torch.rand(1).item() < self.modality_dropout_p else 1.0, "audio": 0.0 if torch.rand(1).item() < self.modality_dropout_p else 1.0, } # Ensure at least one modality is active if mask["visual"] == 0.0 and mask["audio"] == 0.0: mask["visual" if torch.rand(1).item() < 0.5 else "audio"] = 1.0 # Fusion logits logits = self.fusion(v_emb, a_emb, mask) # Auxiliary unimodal logits (always computed, needed for OGM-GE) visual_logits = self.visual_head(v_emb) audio_logits = self.audio_head(a_emb) outputs = { "logits": logits, "visual_logits": visual_logits, "audio_logits": audio_logits, "visual_emb": v_emb, "audio_emb": a_emb, } if labels is not None: loss_fusion = self.loss_fn(logits, labels) loss_visual = self.loss_fn(visual_logits, labels) loss_audio = self.loss_fn(audio_logits, labels) # Total loss with asymmetric weighting loss = (loss_fusion + self.lambda_visual * loss_visual + self.lambda_audio * loss_audio) outputs["loss"] = loss outputs["loss_fusion"] = loss_fusion.item() outputs["loss_visual"] = loss_visual.item() outputs["loss_audio"] = loss_audio.item() elif self.mode == "visual_only": logits = self.classifier(self.visual_branch(pixel_values)) outputs = {"logits": logits} if labels is not None: outputs["loss"] = self.loss_fn(logits, labels) else: # audio_only logits = self.classifier(self.audio_branch(audio_values)) outputs = {"logits": logits} if labels is not None: outputs["loss"] = self.loss_fn(logits, labels) return outputs # =========================================================================== # Factory functions # =========================================================================== def create_model(model_config, lora_config, mode="multimodal", finetune_method="lora", use_ogm=True, lambda_visual=1.5, lambda_audio=0.5): """Create model with v2 anti-collapse features.""" return MultimodalPCFaultDetector( model_config, lora_config, finetune_method, mode, use_ogm=use_ogm, lambda_visual=lambda_visual, lambda_audio=lambda_audio) def get_processors(model_config): """Load ViT image processor and AST feature extractor.""" return ( ViTImageProcessor.from_pretrained(model_config.vit_model_name), ASTFeatureExtractor.from_pretrained(model_config.ast_model_name))