1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_reg.py",
4  "start_line": 246,
5  "end_line": 273,
6  "content": "def compute_regularization(model, inputs, outputs, targets, config):\n    \"\"\"Adaptive Architecture-Aware Regularization (AAR).\n\n    Combines three complementary penalties:\n    1. Feature Diversity Penalty: Penalizes high correlation between feature maps\n       to encourage diverse feature learning.\n    2. Weight Smoothness Penalty: Encourages smooth convolutional filters by\n       penalizing large differences between adjacent weights (local smoothness).\n    3. Adaptive scaling based on layer depth and training progress.\n\n    Strength schedule: Warm-up for first 50 epochs, then decay with cosine.\n    \"\"\"\n    device = outputs.device\n    epoch = config['epoch']\n    total_epochs = config['total_epochs']\n    \n    # Base regularization strength\n    lambda_base = 1e-3\n    \n    # Warm-up then cosine decay schedule\n    warmup_epochs = 50\n    if epoch < warmup_epochs:\n        progress = epoch / warmup_epochs\n        lambda_scale = lambda_base * progress\n    else:\n        # Cosine decay after warmup\n        progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)\n        lambda_scale = lambda_base * (0.5 + 0.5 * math.cos(math.pi * progress))\n    \n    # Collect regularization terms\n    feature_div_penalty = torch.tensor(0.0, device=device)\n    weight_smooth_penalty = torch.tensor(0.0, device=device)\n    \n    # Hook to collect feature maps from conv layers\n    feature_maps = []\n    hooks = []\n    \n    # Register forward hooks on conv layers\n    for name, module in model.named_modules():\n        if isinstance(module, nn.Conv2d):\n            def hook(module, input, output, fm_list=feature_maps):\n                fm_list.append(output)\n            hooks.append(module.register_forward_hook(hook))\n    \n    # Forward pass to collect feature maps (in eval mode to avoid BN stats corruption)\n    was_training = model.training\n    model.eval()\n    with torch.no_grad():\n        _ = model(inputs)\n    if was_training:\n        model.train()\n    \n    # Remove hooks\n    for h in hooks:\n        h.remove()\n    \n    # 1. Feature Diversity Penalty: penalize high correlation between feature maps\n    for feat in feature_maps:\n        # feat shape: [B, C, H, W]\n        B, C, H, W = feat.shape\n        if C < 2 or H*W < 2:\n            continue\n        \n        # Compute mean across spatial dimensions -> [B, C]\n        feat_mean = feat.mean(dim=(2, 3))\n        # Center features\n        feat_centered = feat_mean - feat_mean.mean(dim=0, keepdim=True)\n        \n        # Compute correlation matrix across channels\n        norm = torch.norm(feat_centered, dim=0, keepdim=True)\n        norm = torch.clamp(norm, min=1e-8)\n        feat_normalized = feat_centered / norm\n        correlation = torch.mm(feat_normalized.t(), feat_normalized) / B\n        \n        # Penalize off-diagonal correlations (ignore diagonal)\n        mask = torch.eye(C, device=device)\n        off_diag_corr = correlation * (1 - mask)\n        # Average squared correlation\n        feature_div_penalty = feature_div_penalty + (off_diag_corr ** 2).mean()\n    \n    # 2. Weight Smoothness Penalty for convolutional filters\n    for name, param in model.named_parameters():\n        if 'conv' in name and 'weight' in name and param.dim() == 4:\n            # param shape: [out_channels, in_channels, kH, kW]\n            out_c, in_c, kH, kW = param.shape\n            \n            # Horizontal smoothness: penalize large differences between adjacent columns\n            if kW > 1:\n                diff_h = param[:, :, :, 1:] - param[:, :, :, :-1]\n                weight_smooth_penalty = weight_smooth_penalty + (diff_h ** 2).mean()\n            \n            # Vertical smoothness: penalize large differences between adjacent rows\n            if kH > 1:\n                diff_v = param[:, :, 1:, :] - param[:, :, :-1, :]\n                weight_smooth_penalty = weight_smooth_penalty + (diff_v ** 2).mean()\n    \n    # 3. Adaptive layer scaling: deeper layers get stronger regularization\n    # Count total conv layers and assign depth-based weight\n    conv_layers = [name for name, module in model.named_modules() if isinstance(module, nn.Conv2d)]\n    total_conv = len(conv_layers)\n    \n    if total_conv > 0:\n        # Feature diversity penalty scaled by depth factor (linear increase)\n        depth_factor = 1.0  # baseline, could be made depth-dependent but simpler for now\n    \n    # Combine penalties\n    total_penalty = lambda_scale * (feature_div_penalty + 0.1 * weight_smooth_penalty)\n    \n    return total_penalty"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_reg.py",
4  "start_line": 246,
5  "end_line": 354,
6  "content": "def compute_regularization(model, inputs, outputs, targets, config):\n    \"\"\"Adaptive Architecture-Aware Regularization (AAR) v2.\n\n    Combines:\n    1. Feature Decorrelation Penalty: Penalizes high correlation between\n       feature maps to encourage diverse feature learning.\n    2. Weight Smoothness Penalty: Encourages smooth convolutional filters.\n    3. Adaptive scaling with warm-up and cosine decay.\n    \n    Key improvements:\n    - Average penalties over layers to avoid accumulation\n    - Reduced base strength\n    - Skip small feature maps\n    - More efficient computation\n    \"\"\"\n    device = outputs.device\n    epoch = config['epoch']\n    total_epochs = config['total_epochs']\n    \n    # Base regularization strength (reduced)\n    lambda_base = 5e-5\n    \n    # Warm-up then cosine decay schedule\n    warmup_epochs = 30\n    if epoch < warmup_epochs:\n        progress = epoch / warmup_epochs\n        lambda_scale = lambda_base * progress\n    else:\n        # Cosine decay after warmup\n        progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)\n        lambda_scale = lambda_base * (0.5 + 0.5 * math.cos(math.pi * progress))\n    \n    # If lambda_scale is negligible, skip computation\n    if lambda_scale < 1e-8:\n        return torch.tensor(0.0, device=device)\n    \n    # Collect regularization terms\n    feature_penalties = []\n    weight_penalties = []\n    \n    # Hook to collect feature maps from conv layers\n    feature_maps = []\n    hooks = []\n    \n    # Register forward hooks on conv layers\n    for name, module in model.named_modules():\n        if isinstance(module, nn.Conv2d):\n            def hook(module, input, output, fm_list=feature_maps):\n                fm_list.append(output)\n            hooks.append(module.register_forward_hook(hook))\n    \n    # Forward pass to collect feature maps (in eval mode to avoid BN stats corruption)\n    was_training = model.training\n    model.eval()\n    with torch.no_grad():\n        _ = model(inputs)\n    if was_training:\n        model.train()\n    \n    # Remove hooks\n    for h in hooks:\n        h.remove()\n    \n    # 1. Feature Decorrelation Penalty\n    for feat in feature_maps:\n        # feat shape: [B, C, H, W]\n        B, C, H, W = feat.shape\n        if C < 2 or H*W < 4:\n            continue\n        \n        # Compute spatial mean -> [B, C]\n        feat_mean = feat.mean(dim=(2, 3))\n        # Center features\n        feat_centered = feat_mean - feat_mean.mean(dim=0, keepdim=True)\n        \n        # Compute correlation matrix across channels\n        norm = torch.norm(feat_centered, dim=0, keepdim=True)\n        norm = torch.clamp(norm, min=1e-8)\n        feat_normalized = feat_centered / norm\n        correlation = torch.mm(feat_normalized.t(), feat_normalized) / B\n        \n        # Penalize off-diagonal correlations\n        mask = torch.eye(C, device=device)\n        off_diag_corr = correlation * (1 - mask)\n        # Mean squared correlation (averaged over off-diagonal elements)\n        penalty = (off_diag_corr ** 2).sum() / (C * (C - 1))\n        feature_penalties.append(penalty)\n    \n    # 2. Weight Smoothness Penalty for convolutional filters\n    for name, param in model.named_parameters():\n        if 'conv' in name and 'weight' in name and param.dim() == 4:\n            out_c, in_c, kH, kW = param.shape\n            penalty = 0.0\n            \n            # Horizontal smoothness\n            if kW > 1:\n                diff_h = param[:, :, :, 1:] - param[:, :, :, :-1]\n                penalty = penalty + (diff_h ** 2).mean()\n            \n            # Vertical smoothness\n            if kH > 1:\n                diff_v = param[:, :, 1:, :] - param[:, :, :-1, :]\n                penalty = penalty + (diff_v ** 2).mean()\n            \n            if penalty > 0:\n                weight_penalties.append(penalty)\n    \n    # Average over layers\n    avg_feature_penalty = torch.tensor(0.0, device=device)\n    if feature_penalties:\n        avg_feature_penalty = torch.stack(feature_penalties).mean()\n    \n    avg_weight_penalty = torch.tensor(0.0, device=device)\n    if weight_penalties:\n        avg_weight_penalty = torch.stack(weight_penalties).mean()\n    \n    # Combine penalties with weighting\n    total_penalty = lambda_scale * (avg_feature_penalty + 0.05 * avg_weight_penalty)\n    \n    return total_penalty"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_reg.py",
4  "start_line": 246,
5  "end_line": 365,
6  "content": "def compute_regularization(model, inputs, outputs, targets, config):\n    \"\"\"Weight Smoothness Regularization (WSR) - simplified test.\n    \n    Only applies weight smoothness penalty to convolutional filters.\n    Very small coefficient to ensure training works.\n    \"\"\"\n    device = outputs.device\n    epoch = config['epoch']\n    total_epochs = config['total_epochs']\n    \n    # Base regularization strength (very small)\n    lambda_base = 1e-6\n    \n    # Warm-up for first 10 epochs\n    warmup_epochs = 10\n    if epoch < warmup_epochs:\n        progress = epoch / warmup_epochs\n        lambda_scale = lambda_base * progress\n    else:\n        lambda_scale = lambda_base\n    \n    if lambda_scale < 1e-9:\n        return torch.tensor(0.0, device=device)\n    \n    weight_penalties = []\n    \n    for name, param in model.named_parameters():\n        if 'conv' in name and 'weight' in name and param.dim() == 4:\n            out_c, in_c, kH, kW = param.shape\n            penalty = 0.0\n            \n            # Horizontal smoothness\n            if kW > 1:\n                diff_h = param[:, :, :, 1:] - param[:, :, :, :-1]\n                penalty = penalty + (diff_h ** 2).mean()\n            \n            # Vertical smoothness\n            if kH > 1:\n                diff_v = param[:, :, 1:, :] - param[:, :, :-1, :]\n                penalty = penalty + (diff_v ** 2).mean()\n            \n            if penalty > 0:\n                weight_penalties.append(penalty)\n    \n    if not weight_penalties:\n        return torch.tensor(0.0, device=device)\n    \n    avg_penalty = torch.stack(weight_penalties).mean()\n    return lambda_scale * avg_penalty"
7}

1{
2  "n": 3
3}