Add higher-capacity last-four-layer fusion

easycom_head.py

@@ -34,6 +34,12 @@ STEP_SECONDS = 0.2
 NUM_HORIZONS = 10
 NUM_CLASSES = 3
 DATASET_NAMES = ("easycom", "ego4d")
+LAYER_FUSION_MODES = (
+    "single",
+    "fixed-mean",
+    "scalar-mixture",
+    "concat-mlp",
+)
 DatasetName = Literal["easycom", "ego4d"]
 
 
@@ -48,30 +54,55 @@ class TurnPredictionHead(nn.Module):
         dropout: float = 0.1,
         num_input_layers: int = 1,
         learnable_layer_mixture: bool = False,
+        layer_fusion: str | None = None,
     ) -> None:
         super().__init__()
         if num_input_layers <= 0:
             raise ValueError("num_input_layers must be positive")
-        if learnable_layer_mixture and num_input_layers < 2:
-            raise ValueError(
-                "learnable_layer_mixture requires at least two input layers"
+        if layer_fusion is None:
+            layer_fusion = (
+                "scalar-mixture" if learnable_layer_mixture else "single"
             )
-        if not learnable_layer_mixture and num_input_layers != 1:
+        if layer_fusion not in LAYER_FUSION_MODES:
             raise ValueError(
-                "Multiple input layers require learnable_layer_mixture=True"
+                f"Unknown layer_fusion={layer_fusion!r}; expected one of "
+                f"{LAYER_FUSION_MODES}"
+            )
+        if learnable_layer_mixture and layer_fusion != "scalar-mixture":
+            raise ValueError(
+                "learnable_layer_mixture=True is only compatible with "
+                "layer_fusion='scalar-mixture'"
+            )
+        if layer_fusion != "single" and num_input_layers < 2:
+            raise ValueError(
+                f"layer_fusion={layer_fusion!r} requires at least two input layers"
+            )
+        if layer_fusion == "single" and num_input_layers != 1:
+            raise ValueError(
+                "Multiple input layers require a non-single layer_fusion"
             )
         self.hidden_size = hidden_size
         self.num_horizons = num_horizons
         self.num_classes = num_classes
         self.dropout = dropout
         self.num_input_layers = num_input_layers
-        self.learnable_layer_mixture = learnable_layer_mixture
-        if learnable_layer_mixture:
+        self.layer_fusion = layer_fusion
+        self.learnable_layer_mixture = layer_fusion == "scalar-mixture"
+        if self.learnable_layer_mixture:
             # Zero initialization exactly reproduces a uniform layer mean at
             # step zero while allowing training to select useful query layers.
             self.layer_logits = nn.Parameter(torch.zeros(num_input_layers))
         else:
             self.register_parameter("layer_logits", None)
+        if layer_fusion == "concat-mlp":
+            fusion_input_size = num_input_layers * hidden_size
+            self.layer_fusion_net: nn.Module | None = nn.Sequential(
+                nn.LayerNorm(fusion_input_size),
+                nn.Linear(fusion_input_size, hidden_size),
+                nn.GELU(),
+            )
+        else:
+            self.layer_fusion_net = None
         self.net = nn.Sequential(
             nn.LayerNorm(hidden_size),
             nn.Linear(hidden_size, hidden_size // 2),
@@ -85,13 +116,16 @@ class TurnPredictionHead(nn.Module):
             raise ValueError(
                 f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
             )
-        if self.learnable_layer_mixture:
+        if self.layer_fusion != "single":
             if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
                 raise ValueError(
-                    "Layer-mixture input must end in "
+                    "Multi-layer input must end in "
                     f"[{self.num_input_layers},{self.hidden_size}], got "
                     f"{tuple(states.shape)}"
                 )
+        if self.layer_fusion == "fixed-mean":
+            states = states.mean(dim=-2)
+        elif self.layer_fusion == "scalar-mixture":
             assert self.layer_logits is not None
             weights = self.layer_logits.softmax(dim=0).to(states.dtype)
             view_shape = (1,) * (states.ndim - 2) + (
@@ -99,6 +133,10 @@ class TurnPredictionHead(nn.Module):
                 1,
             )
             states = (states * weights.view(view_shape)).sum(dim=-2)
+        elif self.layer_fusion == "concat-mlp":
+            assert self.layer_fusion_net is not None
+            states = states.flatten(start_dim=-2)
+            states = self.layer_fusion_net(states)
         logits = self.net(states)
         return logits.view(
             *states.shape[:-1], self.num_horizons, self.num_classes
@@ -406,6 +444,14 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("dropout", 0.1)
     config.setdefault("num_input_layers", 1)
     config.setdefault("learnable_layer_mixture", False)
+    config.setdefault(
+        "layer_fusion",
+        (
+            "scalar-mixture"
+            if config["learnable_layer_mixture"]
+            else "single"
+        ),
+    )
     return config