Add overlap-aware turn prediction objectives

README.md

@@ -954,6 +954,21 @@ without normalizing the target. The previous L2-normalized objective belongs
 to MiniROAD and is available only as the explicit
 `--loss miniroad-normalized` ablation.
 
+Two label-structure ablations are also available without changing the saved
+three-column target files:
+
+- `--loss independent-bce` keeps a `10 x 3` head, trains the three labels with
+  independent sigmoid/BCE, and evaluates the sigmoid marginals directly;
+- `--loss four-state-ce` changes the head to `10 x 4` and maps each target to
+  `silence`, `wearer_only`, `others_only`, or `overlap`. At evaluation time it
+  returns the original three marginal scores as `P(silence)`,
+  `P(wearer_only)+P(overlap)`, and `P(others_only)+P(overlap)`.
+
+The four-state marginals are intentionally not renormalized. Keep features,
+split, fusion, optimizer, schedule, and seeds fixed when comparing these losses.
+For checkpoint selection, reserve one official training session such as
+Session 12 with `--dev-sessions 12`; Sessions 1--3 remain test-only.
+
 Recommended experimental procedure:
 
 1. Freeze Qwen and extract all training-session states offline.

easycom_head.py

@@ -30,9 +30,22 @@ CLASS_NAMES = (
     "camera_wearer_speaking",
     "other_people_speaking",
 )
+FOUR_STATE_NAMES = (
+    "silence",
+    "wearer_only",
+    "others_only",
+    "overlap",
+)
 STEP_SECONDS = 0.2
 NUM_HORIZONS = 10
 NUM_CLASSES = 3
+NUM_FOUR_STATES = len(FOUR_STATE_NAMES)
+LOSS_NAMES = (
+    "lstr-mce",
+    "miniroad-normalized",
+    "independent-bce",
+    "four-state-ce",
+)
 DATASET_NAMES = ("easycom", "ego4d")
 LAYER_FUSION_MODES = (
     "single",
@@ -44,7 +57,7 @@ DatasetName = Literal["easycom", "ego4d"]
 
 
 class TurnPredictionHead(nn.Module):
-    """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
+    """Map one Omni state to configurable ``[horizon, output]`` logits."""
 
     def __init__(
         self,
@@ -150,18 +163,76 @@ class TurnPredictionHead(nn.Module):
         return self.layer_logits.softmax(dim=0)
 
 
-def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
+def _validate_loss_inputs(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    num_output_classes: int = NUM_CLASSES,
+) -> None:
     if logits.shape != targets.shape:
         raise ValueError(
             f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
             f"{tuple(targets.shape)}"
         )
-    if logits.shape[-1] != NUM_CLASSES:
+    if logits.shape[-1] != num_output_classes:
         raise ValueError(
-            f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
+            f"Expected {num_output_classes} classes, got {logits.shape[-1]}"
         )
 
 
+def _validated_binary_activity_targets(targets: torch.Tensor) -> torch.Tensor:
+    """Validate the four legal EasyCom/EgoSpeak three-label activity rows."""
+
+    if targets.ndim < 1 or targets.shape[-1] != NUM_CLASSES:
+        raise ValueError(
+            f"Expected activity targets [...,{NUM_CLASSES}], got "
+            f"{tuple(targets.shape)}"
+        )
+    if not torch.isfinite(targets).all():
+        raise ValueError("Activity targets contain NaN or Inf")
+    rounded = targets.round()
+    if not torch.allclose(targets, rounded, atol=1e-6, rtol=0.0):
+        raise ValueError("Activity targets must contain only binary 0/1 values")
+    expected_background = (rounded[..., 1:].sum(dim=-1) == 0).to(rounded.dtype)
+    if not torch.equal(rounded[..., 0], expected_background):
+        raise ValueError(
+            "Activity targets must be one of [1,0,0], [0,1,0], "
+            "[0,0,1], or [0,1,1]"
+        )
+    return rounded
+
+
+def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor:
+    """Map three-label activity targets to silence/wearer/others/overlap."""
+
+    targets = _validated_binary_activity_targets(targets)
+    return (targets[..., 1] + 2 * targets[..., 2]).long()
+
+
+def output_classes_for_loss(loss_name: str) -> int:
+    """Return the number of head logits required by one training objective."""
+
+    if loss_name not in LOSS_NAMES:
+        raise ValueError(f"Unknown loss {loss_name!r}; expected one of {LOSS_NAMES}")
+    return NUM_FOUR_STATES if loss_name == "four-state-ce" else NUM_CLASSES
+
+
+def output_space_for_loss(loss_name: str) -> str:
+    """Return a stable checkpoint description of the head output space."""
+
+    output_classes_for_loss(loss_name)
+    if loss_name == "four-state-ce":
+        return "four-state-softmax"
+    if loss_name == "independent-bce":
+        return "three-label-sigmoid"
+    return "three-class-softmax"
+
+
+def head_output_names_for_loss(loss_name: str) -> tuple[str, ...]:
+    output_classes_for_loss(loss_name)
+    return FOUR_STATE_NAMES if loss_name == "four-state-ce" else CLASS_NAMES
+
+
 def egospeak_multihorizon_loss(
     logits: torch.Tensor,
     targets: torch.Tensor,
@@ -211,6 +282,99 @@ def miniroad_multihorizon_loss(
     raise ValueError(f"Unknown reduction: {reduction}")
 
 
+def independent_bce_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Independent sigmoid/BCE loss for the three marginal activity labels."""
+
+    _validate_loss_inputs(logits, targets)
+    targets = targets.to(
+        device=logits.device,
+        dtype=logits.dtype,
+    )
+    element_loss = nn.functional.binary_cross_entropy_with_logits(
+        logits,
+        targets,
+        reduction="none",
+    )
+    if reduction == "mean":
+        return element_loss.mean()
+    if reduction == "sum":
+        return element_loss.sum()
+    if reduction == "none":
+        return element_loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def four_state_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Four-state CE for silence, wearer-only, others-only, and overlap."""
+
+    if logits.ndim < 1 or logits.shape[-1] != NUM_FOUR_STATES:
+        raise ValueError(
+            f"Expected four-state logits [...,{NUM_FOUR_STATES}], got "
+            f"{tuple(logits.shape)}"
+        )
+    if logits.shape[:-1] != targets.shape[:-1] or targets.shape[-1] != NUM_CLASSES:
+        raise ValueError(
+            "Four-state logits/activity target shape mismatch: "
+            f"{tuple(logits.shape)} vs {tuple(targets.shape)}"
+        )
+    state_targets = (targets[..., 1] + 2 * targets[..., 2]).long().to(
+        logits.device
+    )
+    loss = nn.functional.cross_entropy(
+        logits.reshape(-1, NUM_FOUR_STATES),
+        state_targets.reshape(-1),
+        reduction="none",
+    ).reshape(logits.shape[:-1])
+    if reduction == "mean":
+        return loss.mean()
+    if reduction == "sum":
+        return loss.sum()
+    if reduction == "none":
+        return loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def marginal_probabilities_from_logits(
+    logits: torch.Tensor,
+    loss_name: str,
+) -> torch.Tensor:
+    """Convert any supported head output to three EgoSpeak marginals.
+
+    Four-state marginals deliberately are not renormalized: wearer and others
+    may both be near one when the overlap state is likely.
+    """
+
+    expected_classes = output_classes_for_loss(loss_name)
+    if logits.ndim < 1 or logits.shape[-1] != expected_classes:
+        raise ValueError(
+            f"Loss {loss_name!r} expects {expected_classes} logits, got "
+            f"{tuple(logits.shape)}"
+        )
+    if loss_name == "independent-bce":
+        return logits.sigmoid()
+    probabilities = logits.softmax(dim=-1)
+    if loss_name != "four-state-ce":
+        return probabilities
+    return torch.stack(
+        (
+            probabilities[..., 0],
+            probabilities[..., 1] + probabilities[..., 3],
+            probabilities[..., 2] + probabilities[..., 3],
+        ),
+        dim=-1,
+    )
+
+
 def multihorizon_cross_entropy(
     logits: torch.Tensor,
     targets: torch.Tensor,
@@ -380,6 +544,7 @@ def load_session_data(
             f"{target_path}"
         )
     targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
+    _validated_binary_activity_targets(targets)
     if len(label_indices) and int(label_indices[-1]) >= len(targets):
         raise ValueError(
             f"Last label index {int(label_indices[-1])} exceeds target length "
@@ -455,6 +620,30 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     return config
 
 
+def checkpoint_loss_name(payload: dict[str, Any]) -> str:
+    """Read the objective from a checkpoint, defaulting legacy files to MCE."""
+
+    training_config = payload.get("training_config", {})
+    if not isinstance(training_config, dict):
+        raise ValueError("Checkpoint training_config must be a mapping")
+    loss_name = str(training_config.get("loss", "lstr-mce"))
+    expected_classes = output_classes_for_loss(loss_name)
+    actual_classes = int(checkpoint_head_config(payload)["num_classes"])
+    if actual_classes != expected_classes:
+        raise ValueError(
+            f"Checkpoint loss {loss_name!r} requires {expected_classes} head "
+            f"outputs, but head_config has {actual_classes}"
+        )
+    recorded_space = payload.get("output_space")
+    expected_space = output_space_for_loss(loss_name)
+    if recorded_space is not None and recorded_space != expected_space:
+        raise ValueError(
+            f"Checkpoint output_space is {recorded_space!r}, expected "
+            f"{expected_space!r} for loss {loss_name!r}"
+        )
+    return loss_name
+
+
 def validate_disjoint_sessions(
     first: Sequence[str],
     second: Sequence[str],

infer_turn_head.py

@@ -21,9 +21,13 @@ from easycom_head import (
     TurnPredictionHead,
     average_precision_metrics,
     checkpoint_head_config,
+    checkpoint_loss_name,
     load_ego4d_manifest_split,
     load_session_data,
+    marginal_probabilities_from_logits,
     normalize_sample_name,
+    output_classes_for_loss,
+    output_space_for_loss,
 )
 
 
@@ -86,9 +90,12 @@ def _load_head(
         raise ValueError(
             f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}"
         )
-    if int(config["num_classes"]) != NUM_CLASSES:
+    loss_name = checkpoint_loss_name(payload)
+    expected_classes = output_classes_for_loss(loss_name)
+    if int(config["num_classes"]) != expected_classes:
         raise ValueError(
-            f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}"
+            f"Checkpoint has {config['num_classes']} outputs, expected "
+            f"{expected_classes} for loss {loss_name!r}"
         )
     head = TurnPredictionHead(**config)
     head.load_state_dict(state_dict)
@@ -104,6 +111,7 @@ def _predict_states(
     *,
     description: str,
     show_progress: bool,
+    loss_name: str,
 ) -> np.ndarray:
     chunks: list[torch.Tensor] = []
     starts = range(0, len(states), batch_size)
@@ -120,7 +128,9 @@ def _predict_states(
             dtype=torch.float32,
             non_blocking=True,
         )
-        chunks.append(head(batch).softmax(dim=-1).cpu())
+        chunks.append(
+            marginal_probabilities_from_logits(head(batch), loss_name).cpu()
+        )
     return torch.cat(chunks, dim=0).numpy()
 
 
@@ -258,6 +268,7 @@ def main() -> None:
             f"{args.dataset!r}"
         )
     head_config = checkpoint_head_config(checkpoint_payload)
+    loss_name = checkpoint_loss_name(checkpoint_payload)
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
     prediction_dir = args.output_dir / "pred"
@@ -269,6 +280,9 @@ def main() -> None:
     session_targets: dict[str, np.ndarray] = {}
 
     print(f"device: {device}")
+    print(
+        f"loss: {loss_name}; output_space: {output_space_for_loss(loss_name)}"
+    )
     print(f"sessions: {', '.join(sessions)}")
     for name in sessions:
         data = load_session_data(
@@ -295,6 +309,7 @@ def main() -> None:
             args.batch_size,
             description=f"infer {name}",
             show_progress=not args.no_progress,
+            loss_name=loss_name,
         )
         predictions, valid = _align_to_egospeak_timeline(
             probabilities,
@@ -323,6 +338,8 @@ def main() -> None:
                 else None
             ),
             "classes": list(CLASS_NAMES),
+            "loss": loss_name,
+            "output_space": output_space_for_loss(loss_name),
             "prediction_layout": "[target_time, class, future_horizon]",
             "checkpoint": str(args.checkpoint.resolve()),
             "checkpoint_epoch": checkpoint_payload.get("epoch"),

streaming_turn_model.py

@@ -44,7 +44,9 @@ import torch
 from easycom_head import (
     TurnPredictionHead,
     checkpoint_head_config,
+    checkpoint_loss_name,
     egospeak_multihorizon_loss,
+    marginal_probabilities_from_logits,
     miniroad_multihorizon_loss,
     multihorizon_cross_entropy,
 )
@@ -2141,10 +2143,11 @@ def _load_head(
     checkpoint: Path,
     hidden_size: int,
     device: torch.device,
-) -> TurnPredictionHead:
+) -> tuple[TurnPredictionHead, str]:
     payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
     if isinstance(payload, dict) and "head" in payload:
         config = checkpoint_head_config(payload)
+        loss_name = checkpoint_loss_name(payload)
         if int(config["hidden_size"]) != hidden_size:
             raise ValueError(
                 f"Head hidden size {config['hidden_size']} != model {hidden_size}"
@@ -2154,8 +2157,9 @@ def _load_head(
     else:
         head = TurnPredictionHead(hidden_size=hidden_size)
         state_dict = payload
+        loss_name = "lstr-mce"
     head.load_state_dict(state_dict)
-    return head.to(device).eval()
+    return head.to(device).eval(), loss_name
 
 
 def parse_args() -> argparse.Namespace:
@@ -2393,11 +2397,19 @@ def main() -> None:
         "verification": result.verification,
     }
     if args.head_checkpoint is not None:
-        head = _load_head(args.head_checkpoint, _hidden_size(model), model.device)
+        head, head_loss_name = _load_head(
+            args.head_checkpoint,
+            _hidden_size(model),
+            model.device,
+        )
         with torch.inference_mode():
             logits = head(result.states.to(model.device, dtype=model.dtype)).float().cpu()
         payload["logits"] = logits
-        payload["probabilities"] = logits.softmax(dim=-1)
+        payload["probabilities"] = marginal_probabilities_from_logits(
+            logits,
+            head_loss_name,
+        )
+        payload["head_loss"] = head_loss_name
 
     if args.num_steps is not None and len(result.times) != args.num_steps:
         raise RuntimeError(

streaming_turn_model_directly.py

@@ -1540,7 +1540,7 @@ def main(
         "verification": result.verification,
     }
     if args.head_checkpoint is not None:
-        head = core._load_head(
+        head, head_loss_name = core._load_head(
             args.head_checkpoint, core._hidden_size(model), model.device
         )
         with torch.inference_mode():
@@ -1548,7 +1548,11 @@ def main(
                 result.states.to(model.device, dtype=model.dtype)
             ).float().cpu()
         payload["logits"] = logits
-        payload["probabilities"] = logits.softmax(dim=-1)
+        payload["probabilities"] = core.marginal_probabilities_from_logits(
+            logits,
+            head_loss_name,
+        )
+        payload["head_loss"] = head_loss_name
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     torch.save(payload, args.output)

test_easycom_head_objectives.py

@@ -0,0 +1,109 @@
+"""Focused tests for EasyCom loss/output-space ablations."""
+
+from __future__ import annotations
+
+import math
+import unittest
+
+import torch
+
+from easycom_head import (
+    TurnPredictionHead,
+    activity_targets_to_four_states,
+    checkpoint_loss_name,
+    four_state_multihorizon_loss,
+    independent_bce_multihorizon_loss,
+    marginal_probabilities_from_logits,
+    output_classes_for_loss,
+)
+
+
+class ObjectiveTests(unittest.TestCase):
+    def setUp(self) -> None:
+        self.targets = torch.tensor(
+            [
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+                [0.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0],
+            ]
+        )
+
+    def test_activity_target_mapping(self) -> None:
+        actual = activity_targets_to_four_states(self.targets)
+        self.assertTrue(torch.equal(actual, torch.tensor([0, 1, 2, 3])))
+
+    def test_invalid_activity_target_is_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "must be one of"):
+            activity_targets_to_four_states(torch.tensor([[0.0, 0.0, 0.0]]))
+
+    def test_independent_bce_uses_standard_mean(self) -> None:
+        logits = torch.zeros_like(self.targets)
+        loss = independent_bce_multihorizon_loss(logits, self.targets)
+        self.assertAlmostEqual(float(loss), math.log(2.0), places=6)
+
+    def test_four_state_loss_accepts_three_label_targets(self) -> None:
+        logits = torch.zeros(4, 4)
+        loss = four_state_multihorizon_loss(logits, self.targets)
+        self.assertAlmostEqual(float(loss), math.log(4.0), places=6)
+
+    def test_four_state_probabilities_are_marginalized(self) -> None:
+        state_probabilities = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
+        logits = state_probabilities.log()
+        actual = marginal_probabilities_from_logits(logits, "four-state-ce")
+        expected = torch.tensor([[0.1, 0.6, 0.7]])
+        self.assertTrue(torch.allclose(actual, expected, atol=1e-6))
+        self.assertGreater(float(actual.sum()), 1.0)
+
+    def test_bce_probabilities_use_sigmoid(self) -> None:
+        actual = marginal_probabilities_from_logits(
+            torch.zeros(2, 3),
+            "independent-bce",
+        )
+        self.assertTrue(torch.equal(actual, torch.full((2, 3), 0.5)))
+
+    def test_four_state_head_shape(self) -> None:
+        head = TurnPredictionHead(
+            hidden_size=8,
+            num_horizons=2,
+            num_classes=output_classes_for_loss("four-state-ce"),
+            dropout=0.0,
+        )
+        self.assertEqual(tuple(head(torch.randn(3, 8)).shape), (3, 2, 4))
+
+    def test_all_new_objectives_backpropagate_to_marginals(self) -> None:
+        states = torch.randn(4, 8)
+        future_targets = self.targets[:, None, :].repeat(1, 2, 1)
+        for loss_name, loss_function in (
+            ("independent-bce", independent_bce_multihorizon_loss),
+            ("four-state-ce", four_state_multihorizon_loss),
+        ):
+            with self.subTest(loss=loss_name):
+                head = TurnPredictionHead(
+                    hidden_size=8,
+                    num_horizons=2,
+                    num_classes=output_classes_for_loss(loss_name),
+                    dropout=0.0,
+                )
+                logits = head(states)
+                loss = loss_function(logits, future_targets)
+                loss.backward()
+                scores = marginal_probabilities_from_logits(logits, loss_name)
+                self.assertEqual(tuple(scores.shape), (4, 2, 3))
+                self.assertTrue(torch.isfinite(scores).all())
+                self.assertIsNotNone(head.net[-1].weight.grad)
+
+    def test_legacy_checkpoint_defaults_to_lstr(self) -> None:
+        self.assertEqual(checkpoint_loss_name({}), "lstr-mce")
+
+    def test_checkpoint_rejects_loss_output_mismatch(self) -> None:
+        payload = {
+            "training_config": {"loss": "four-state-ce"},
+            "head_config": {"num_classes": 3},
+        }
+        with self.assertRaisesRegex(ValueError, "requires 4"):
+            checkpoint_loss_name(payload)
+
+
+if __name__ == "__main__":
+    unittest.main()

train_turn_head.py

@@ -24,17 +24,23 @@ from easycom_head import (
     CLASS_NAMES,
     DATASET_NAMES,
     LAYER_FUSION_MODES,
-    NUM_CLASSES,
+    LOSS_NAMES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
     STEP_SECONDS,
     TurnPredictionHead,
     average_precision_metrics,
     egospeak_multihorizon_loss,
+    four_state_multihorizon_loss,
+    head_output_names_for_loss,
+    independent_bce_multihorizon_loss,
     load_ego4d_manifest_split,
     load_session_data,
+    marginal_probabilities_from_logits,
     miniroad_multihorizon_loss,
     normalize_sample_name,
+    output_classes_for_loss,
+    output_space_for_loss,
     validate_disjoint_sessions,
 )
 
@@ -111,8 +117,13 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--loss",
-        choices=("lstr-mce", "miniroad-normalized"),
+        choices=LOSS_NAMES,
         default="lstr-mce",
+        help=(
+            "lstr-mce/miniroad-normalized use three competing softmax logits; "
+            "independent-bce uses three sigmoid marginals; four-state-ce uses "
+            "silence/wearer-only/others-only/overlap softmax logits."
+        ),
     )
     parser.add_argument("--device", default="cuda:0")
     parser.add_argument("--num-workers", type=int, default=0)
@@ -314,6 +325,7 @@ def _evaluate_valid_horizons(
     targets: torch.Tensor,
     device: torch.device,
     batch_size: int,
+    loss_name: str,
 ) -> dict[str, Any]:
     head.eval()
     chunks: list[torch.Tensor] = []
@@ -323,7 +335,9 @@ def _evaluate_valid_horizons(
             dtype=torch.float32,
             non_blocking=True,
         )
-        chunks.append(head(x).softmax(dim=-1).cpu())
+        chunks.append(
+            marginal_probabilities_from_logits(head(x), loss_name).cpu()
+        )
     probabilities = torch.cat(chunks, dim=0).numpy()
     target_array = targets.numpy()
 
@@ -358,12 +372,12 @@ def _checkpoint(
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
-        "format_version": 1,
+        "format_version": 2,
         "head": head.state_dict(),
         "head_config": {
             "hidden_size": hidden_size,
             "num_horizons": NUM_HORIZONS,
-            "num_classes": NUM_CLASSES,
+            "num_classes": head.num_classes,
             "dropout": args.dropout,
             "num_input_layers": num_input_layers,
             "layer_fusion": args.layer_fusion,
@@ -377,6 +391,8 @@ def _checkpoint(
             ],
         },
         "classes": list(CLASS_NAMES),
+        "head_output_names": list(head_output_names_for_loss(args.loss)),
+        "output_space": output_space_for_loss(args.loss),
         "dataset": args.dataset,
         "manifest": (
             str(args.manifest.expanduser().resolve())
@@ -548,10 +564,11 @@ def main() -> None:
         persistent_workers=args.num_workers > 0,
         generator=generator,
     )
+    num_output_classes = output_classes_for_loss(args.loss)
     head = TurnPredictionHead(
         hidden_size=hidden_size,
         num_horizons=NUM_HORIZONS,
-        num_classes=NUM_CLASSES,
+        num_classes=num_output_classes,
         dropout=args.dropout,
         num_input_layers=num_input_layers,
         learnable_layer_mixture=args.learnable_layer_mixture,
@@ -566,7 +583,8 @@ def main() -> None:
         else (head.layer_logits.numel() if head.layer_logits is not None else 0)
     )
     print(
-        f"layer_fusion={args.layer_fusion}, "
+        f"loss={args.loss}, output_space={output_space_for_loss(args.loss)}, "
+        f"head_outputs={num_output_classes}, layer_fusion={args.layer_fusion}, "
         f"fusion_parameters={fusion_parameters:,}, "
         f"head_trainable_parameters={trainable_parameters:,}"
     )
@@ -601,11 +619,12 @@ def main() -> None:
         warmup_epochs=args.warmup_epochs,
         warmup_factor=args.warmup_factor,
     )
-    loss_function = (
+    loss_function = {
-        egospeak_multihorizon_loss
+        "lstr-mce": egospeak_multihorizon_loss,
-        if args.loss == "lstr-mce"
+        "miniroad-normalized": miniroad_multihorizon_loss,
-        else miniroad_multihorizon_loss
+        "independent-bce": independent_bce_multihorizon_loss,
-    )
+        "four-state-ce": four_state_multihorizon_loss,
+    }[args.loss]
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     history: list[dict[str, Any]] = []
@@ -675,6 +694,7 @@ def main() -> None:
                 dev_targets,
                 device,
                 args.batch_size,
+                args.loss,
             )
             dev_map = float(dev_metrics["mean_map"])
             record["dev_valid_mean_map"] = dev_map