Add learnable last-four-layer query fusion

easycom_head.py

@@ -46,12 +46,32 @@ class TurnPredictionHead(nn.Module):
         num_horizons: int = NUM_HORIZONS,
         num_classes: int = NUM_CLASSES,
         dropout: float = 0.1,
+        num_input_layers: int = 1,
+        learnable_layer_mixture: bool = False,
     ) -> 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 not learnable_layer_mixture and num_input_layers != 1:
+            raise ValueError(
+                "Multiple input layers require learnable_layer_mixture=True"
+            )
         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:
+            # 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)
         self.net = nn.Sequential(
             nn.LayerNorm(hidden_size),
             nn.Linear(hidden_size, hidden_size // 2),
@@ -61,11 +81,36 @@ class TurnPredictionHead(nn.Module):
         )
 
     def forward(self, states: torch.Tensor) -> torch.Tensor:
+        if states.shape[-1] != self.hidden_size:
+            raise ValueError(
+                f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
+            )
+        if self.learnable_layer_mixture:
+            if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
+                raise ValueError(
+                    "Layer-mixture input must end in "
+                    f"[{self.num_input_layers},{self.hidden_size}], got "
+                    f"{tuple(states.shape)}"
+                )
+            assert self.layer_logits is not None
+            weights = self.layer_logits.softmax(dim=0).to(states.dtype)
+            view_shape = (1,) * (states.ndim - 2) + (
+                self.num_input_layers,
+                1,
+            )
+            states = (states * weights.view(view_shape)).sum(dim=-2)
         logits = self.net(states)
         return logits.view(
             *states.shape[:-1], self.num_horizons, self.num_classes
         )
 
+    def layer_mixture_weights(self) -> torch.Tensor | None:
+        """Return normalized oldest-to-newest layer weights, if enabled."""
+
+        if self.layer_logits is None:
+            return None
+        return self.layer_logits.softmax(dim=0)
+
 
 def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
     if logits.shape != targets.shape:
@@ -219,6 +264,10 @@ class SessionData:
     def hidden_size(self) -> int:
         return int(self.states.shape[-1])
 
+    @property
+    def num_input_layers(self) -> int:
+        return 1 if self.states.ndim == 2 else int(self.states.shape[-2])
+
     def future_examples(
         self,
         num_horizons: int = NUM_HORIZONS,
@@ -256,8 +305,12 @@ def load_session_data(
     if not isinstance(payload, dict) or "states" not in payload:
         raise ValueError(f"Invalid feature payload: {feature_path}")
     states = payload["states"]
-    if not isinstance(states, torch.Tensor) or states.ndim != 2:
-        raise ValueError(f"states must be [T,D] in {feature_path}")
+    if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}:
+        raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}")
+    if states.ndim == 3 and states.shape[1] < 2:
+        raise ValueError(
+            f"Stacked states must contain at least two layers: {feature_path}"
+        )
     if not torch.is_floating_point(states):
         raise ValueError(f"states must be floating point in {feature_path}")
     if not torch.isfinite(states).all():
@@ -351,6 +404,8 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("num_horizons", NUM_HORIZONS)
     config.setdefault("num_classes", NUM_CLASSES)
     config.setdefault("dropout", 0.1)
+    config.setdefault("num_input_layers", 1)
+    config.setdefault("learnable_layer_mixture", False)
     return config
 
 
@@ -360,4 +415,4 @@ def validate_disjoint_sessions(
 ) -> None:
     overlap = set(first) & set(second)
     if overlap:
-        raise ValueError(f"Session sets overlap: {sorted(overlap)}")
+        raise ValueError(f"Session sets overlap: {sorted(overlap)}")