Add higher-capacity last-four-layer fusion

train_turn_head.py

@@ -23,6 +23,7 @@ from tqdm.auto import tqdm
 from easycom_head import (
     CLASS_NAMES,
     DATASET_NAMES,
+    LAYER_FUSION_MODES,
     NUM_CLASSES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
@@ -95,8 +96,17 @@ def parse_args() -> argparse.Namespace:
         "--learnable-layer-mixture",
         action="store_true",
         help=(
-            "Learn softmax weights over stacked [T,L,D] query features before "
-            "the anticipation head."
+            "Legacy alias for --layer-fusion scalar-mixture."
+        ),
+    )
+    parser.add_argument(
+        "--layer-fusion",
+        choices=tuple(mode for mode in LAYER_FUSION_MODES if mode != "single"),
+        default=None,
+        help=(
+            "Fusion for stacked [T,L,D] query features. fixed-mean has no "
+            "fusion parameters; scalar-mixture learns L global weights; "
+            "concat-mlp flattens L*D and learns a full projection to D."
         ),
     )
     parser.add_argument(
@@ -135,6 +145,14 @@ def _validate_args(args: argparse.Namespace) -> None:
         raise ValueError("--num-workers cannot be negative")
     if args.dataset == "ego4d" and args.manifest is None:
         raise ValueError("--manifest is required when --dataset ego4d")
+    if (
+        args.learnable_layer_mixture
+        and args.layer_fusion not in {None, "scalar-mixture"}
+    ):
+        raise ValueError(
+            "--learnable-layer-mixture is an alias for "
+            "--layer-fusion scalar-mixture and cannot select another fusion"
+        )
 
 
 def _resolve_device(name: str) -> torch.device:
@@ -348,6 +366,7 @@ def _checkpoint(
             "num_classes": NUM_CLASSES,
             "dropout": args.dropout,
             "num_input_layers": num_input_layers,
+            "layer_fusion": args.layer_fusion,
             "learnable_layer_mixture": args.learnable_layer_mixture,
         },
         "alignment": {
@@ -383,7 +402,8 @@ def _checkpoint(
             "scheduler": "linear_warmup_times_cosine",
             "seed": args.seed,
             "expected_mode": args.expected_mode,
-            "learnable_layer_mixture": args.learnable_layer_mixture,
+            "layer_fusion": args.layer_fusion,
+            "learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
             "supervision_policy": (
                 "transition-sparse-manifest"
                 if args.sparse_anchor_cache_dir
@@ -401,6 +421,11 @@ def _checkpoint(
 def main() -> None:
     args = parse_args()
     _validate_args(args)
+    if args.layer_fusion is None:
+        args.layer_fusion = (
+            "scalar-mixture" if args.learnable_layer_mixture else "single"
+        )
+    args.learnable_layer_mixture = args.layer_fusion == "scalar-mixture"
     _seed_everything(args.seed)
     device = _resolve_device(args.device)
     expected_mode = None if args.expected_mode == "any" else args.expected_mode
@@ -473,14 +498,15 @@ def main() -> None:
         f"training examples: {len(states)}, feature_dim={hidden_size}, "
         f"input_layers={num_input_layers}"
     )
-    if num_input_layers > 1 and not args.learnable_layer_mixture:
+    if num_input_layers > 1 and args.layer_fusion == "single":
         raise ValueError(
             f"Features contain {num_input_layers} stacked layers; pass "
-            "--learnable-layer-mixture"
+            "--layer-fusion fixed-mean, scalar-mixture, or concat-mlp"
         )
-    if num_input_layers == 1 and args.learnable_layer_mixture:
+    if num_input_layers == 1 and args.layer_fusion != "single":
         raise ValueError(
-            "--learnable-layer-mixture requires stacked [T,L,D] features"
+            f"--layer-fusion {args.layer_fusion} requires stacked [T,L,D] "
+            "features"
         )
     dev_states: torch.Tensor | None = None
     dev_targets: torch.Tensor | None = None
@@ -529,7 +555,21 @@ def main() -> None:
         dropout=args.dropout,
         num_input_layers=num_input_layers,
         learnable_layer_mixture=args.learnable_layer_mixture,
+        layer_fusion=args.layer_fusion,
     ).to(device)
+    trainable_parameters = sum(
+        parameter.numel() for parameter in head.parameters() if parameter.requires_grad
+    )
+    fusion_parameters = (
+        sum(parameter.numel() for parameter in head.layer_fusion_net.parameters())
+        if head.layer_fusion_net is not None
+        else (head.layer_logits.numel() if head.layer_logits is not None else 0)
+    )
+    print(
+        f"layer_fusion={args.layer_fusion}, "
+        f"fusion_parameters={fusion_parameters:,}, "
+        f"head_trainable_parameters={trainable_parameters:,}"
+    )
     if head.layer_logits is None:
         optimizer_parameters: Any = head.parameters()
     else:
@@ -698,7 +738,7 @@ def main() -> None:
     summary = {key: value for key, value in payload.items() if key != "head"}
     summary_path.write_text(json.dumps(summary, indent=2) + "\n")
     print(f"training summary: {summary_path}")
-    if args.learnable_layer_mixture:
+    if payload["head_config"].get("layer_fusion") == "scalar-mixture":
         saved_logits = payload["head"]["layer_logits"].float()
         saved_weights = saved_logits.softmax(dim=0).tolist()
         print(