Add offline sparse supervision control

README.md

@@ -904,6 +904,62 @@ Recommended experimental procedure:
 Class-balanced/focal loss, a learned turn-query token, and LoRA are later
 ablations. They should not replace the frozen-Omni linear/MLP-head baseline.
 
+### Offline frozen states with transition-sparse supervision
+
+This is the clean sampling-only control.  It does not reconstruct cached AV
+tokens and does not run the Thinker during training.  It loads the same dense
+5 Hz frozen states used by the standard head baseline, then retains only the
+`target_index` values listed by the existing transition-sparse cache
+manifests.  The representation, head architecture, optimizer, 50 epochs, loss,
+targets, and final dense test evaluator are otherwise unchanged.
+
+Set `FROZEN_FEATURE_ROOT` to the exact feature directory used by the matching
+dense-supervision baseline.  Do not compare sparse 8-second states with a
+dense 20-second baseline (or vice versa):
+
+```bash
+FROZEN_FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_20s_decord
+TARGET_ROOT=/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame
+SPARSE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s
+
+CUDA_VISIBLE_DEVICES=3 \
+python train_turn_head.py \
+  --feature-root "$FROZEN_FEATURE_ROOT" \
+  --target-root "$TARGET_ROOT" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_7" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_11" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_10" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_6" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_9" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_4" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_8" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_5" \
+  --sparse-anchor-cache-dir "$SPARSE_ROOT/Session_12" \
+  --output checkpoints/easycom_frozen_states_sparse_supervision_50ep.pth \
+  --epochs 50 \
+  --batch-size 512 \
+  --expected-mode any \
+  --device cuda:0
+```
+
+Evaluation remains dense 5 Hz on the official test sessions:
+
+```bash
+CUDA_VISIBLE_DEVICES=3 \
+python infer_turn_head.py \
+  --feature-root "$FROZEN_FEATURE_ROOT" \
+  --target-root "$TARGET_ROOT" \
+  --checkpoint checkpoints/easycom_frozen_states_sparse_supervision_50ep.pth \
+  --output-dir results/easycom_frozen_states_sparse_supervision_50ep \
+  --expected-mode any \
+  --device cuda:0
+```
+
+Compare this result against a head trained for the same 50 epochs on every
+dense 5 Hz anchor from the same `FROZEN_FEATURE_ROOT`.  A large drop isolates
+the problem to sparse supervision; a close result means sparse supervision is
+not sufficient to explain the LoRA regression.
+
 ## Thinker-LoRA sampling and cache policy
 
 This is an implemented second-stage experiment, not yet a reported benchmark
@@ -1021,6 +1077,79 @@ process uses its `LOCAL_RANK`.  DDP intentionally rejects
 `--online-current-tail`; the online producer-consumer path remains a separate
 benchmark.
 
+### Frozen-Thinker sparse head-only control
+
+To determine whether the regression comes from transition-sparse sampling or
+from updating the Thinker representation, rerun the same sparse-anchor
+training with `--freeze-thinker`.  This still forwards every reconstructed
+8-second sequence through the original Thinker, but runs that forward under
+`no_grad` and updates only the turn head initialized from the frozen 8-second
+checkpoint.  It does not create an adapter directory.
+
+For a schedule-matched comparison with a one-epoch LoRA run at global batch
+16, use the same cache directories, head learning rate, warmup, seed, and
+approximately `ceil(num_examples / 16)` updates:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+OMP_NUM_THREADS=2 \
+MKL_NUM_THREADS=2 \
+OPENBLAS_NUM_THREADS=2 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+torchrun --standalone --nproc_per_node=4 train_omni_turn_lora.py \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_7 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_11 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_10 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_6 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_9 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_4 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_8 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_5 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_12 \
+  --freeze-thinker \
+  --head-checkpoint checkpoints/easycom_open_stream_2fps_8s_head_seed0.pth \
+  --output-dir checkpoints/easycom_sparse_head_only_from_frozen_8s_ddp4_epoch1 \
+  --max-steps 3390 \
+  --batch-size 2 \
+  --gradient-accumulation-steps 2 \
+  --num-workers 0 \
+  --block-cache-size 32 \
+  --head-learning-rate 1e-4 \
+  --warmup-steps 500 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --log-every 20
+```
+
+Evaluate the resulting head on the same dense test caches without an adapter:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2 \
+OMP_NUM_THREADS=1 \
+MKL_NUM_THREADS=1 \
+OPENBLAS_NUM_THREADS=1 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \
+torchrun --standalone --nproc_per_node=3 eval_omni_turn_lora.py \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_3 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_2 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_1 \
+  --head-checkpoint checkpoints/easycom_sparse_head_only_from_frozen_8s_ddp4_epoch1/turn_head.pt \
+  --output-dir results/easycom_sparse_head_only_dense_eval_epoch1 \
+  --batch-size 2 \
+  --block-cache-size 32 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --ddp-timeout-minutes 60
+```
+
+If this head-only control also falls well below the frozen `58.16` dense-cache
+control, sparse sampling/head optimization is sufficient to explain the
+regression.  If it stays near `58.16` while LoRA falls, the main failure is
+Thinker representation drift.
+
 To continue from a completed run, pass both its `adapter/` directory through
 `--adapter-checkpoint` and its `turn_head.pt` through `--head-checkpoint` while
 writing to a new output directory.  This is weights-only continuation: the

train_omni_turn_lora.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches.
+"""Train a frozen or LoRA Qwen2.5-Omni Thinker + 10x3 head from AV caches.
 
 This is a custom task runner intended to execute inside the ms-swift image.
 It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI
@@ -219,6 +219,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "scheduler are reinitialized"
         ),
     )
+    parser.add_argument(
+        "--freeze-thinker",
+        action="store_true",
+        help=(
+            "Do not create or update LoRA parameters; run the same cached "
+            "Thinker sequences under no_grad and train only the turn head"
+        ),
+    )
     parser.add_argument("--head-checkpoint", type=Path)
     parser.add_argument("--max-examples", type=int, default=0)
     parser.add_argument("--max-steps", type=int, default=200)
@@ -900,6 +908,14 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                 "Warning: --adapter-checkpoint was provided without "
                 "--head-checkpoint; the turn head will be reinitialized."
             )
+    if args.freeze_thinker and args.adapter_checkpoint is not None:
+        raise ValueError(
+            "--freeze-thinker cannot be combined with --adapter-checkpoint"
+        )
+    if args.freeze_thinker and args.gradient_checkpointing:
+        raise ValueError(
+            "--gradient-checkpointing is unnecessary with --freeze-thinker"
+        )
     if (
         args.audio_right_padding_samples is not None
         and args.audio_right_padding_samples < 0
@@ -1001,6 +1017,17 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     )
     if not target_modules:
         raise ValueError("--lora-target-modules is empty")
+    if args.freeze_thinker:
+        base_model.requires_grad_(False)
+        language_model = base_model.model
+        language_model.config.use_cache = False
+        language_model.eval()
+        if run.is_main:
+            print(
+                "Frozen-Thinker control: no LoRA adapter; only the turn head "
+                "is trainable"
+            )
+    else:
         language_model = _prepare_language_lora(
             base_model,
             adapter_checkpoint=args.adapter_checkpoint,
@@ -1080,7 +1107,11 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
 
     raw_language_model = language_model
     raw_head = head
-    _synchronize_trainable_parameters((raw_language_model, raw_head), run)
+    synchronization_modules = (
+        (raw_head,) if args.freeze_thinker else (raw_language_model, raw_head)
+    )
+    _synchronize_trainable_parameters(synchronization_modules, run)
+    if not args.freeze_thinker:
         language_model = _wrap_ddp(raw_language_model, run)
     head = _wrap_ddp(raw_head, run)
     # Model initialization is synchronized above; rank-specific RNG streams
@@ -1092,8 +1123,15 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         for parameter in raw_language_model.parameters()
         if parameter.requires_grad
     ]
-    if not lora_parameters:
+    if not args.freeze_thinker and not lora_parameters:
         raise RuntimeError("LoRA preparation produced no trainable parameters")
+    if args.freeze_thinker:
+        optimizer = torch.optim.AdamW(
+            [{"params": raw_head.parameters(), "lr": args.head_learning_rate}],
+            weight_decay=args.weight_decay,
+        )
+        head_group_index = 0
+    else:
         optimizer = torch.optim.AdamW(
             [
                 {"params": lora_parameters, "lr": args.learning_rate},
@@ -1101,6 +1139,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             ],
             weight_decay=args.weight_decay,
         )
+        head_group_index = 1
     scheduler = torch.optim.lr_scheduler.LambdaLR(
         optimizer,
         lr_lambda=lambda step: _lr_lambda(
@@ -1122,6 +1161,9 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         print(f"Initial cached loss: {initial_loss:.6f}")
     if materialize_batch is not None:
         materialize_batch.reset_timing()
+    if args.freeze_thinker:
+        language_model.eval()
+    else:
         language_model.train()
     head.train()
     optimizer.zero_grad(set_to_none=True)
@@ -1144,7 +1186,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     history: list[dict[str, float | int]] = []
     progress = tqdm(
         total=args.max_steps,
-        desc="Thinker LoRA",
+        desc=("Frozen Thinker head" if args.freeze_thinker else "Thinker LoRA"),
         unit="update",
         disable=args.no_progress or not run.is_main,
     )
@@ -1177,10 +1219,21 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             )
             with ExitStack() as stack:
                 if run.enabled and not synchronize_gradients:
-                    assert isinstance(language_model, DDP)
                     assert isinstance(head, DDP)
+                    if not args.freeze_thinker:
+                        assert isinstance(language_model, DDP)
                         stack.enter_context(language_model.no_sync())
                     stack.enter_context(head.no_sync())
+                if args.freeze_thinker:
+                    with torch.no_grad():
+                        states = forward_last_state(
+                            language_model,
+                            inputs_embeds=inputs_embeds,
+                            position_ids=position_ids,
+                            attention_mask=attention_mask,
+                            readout_indices=readout_indices,
+                        )
+                else:
                     states = forward_last_state(
                         language_model,
                         inputs_embeds=inputs_embeds,
@@ -1193,7 +1246,11 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                 (loss / args.gradient_accumulation_steps).backward()
             accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps
 
-        trainable_parameters = [*lora_parameters, *raw_head.parameters()]
+        trainable_parameters = (
+            list(raw_head.parameters())
+            if args.freeze_thinker
+            else [*lora_parameters, *raw_head.parameters()]
+        )
         grad_norm = torch.nn.utils.clip_grad_norm_(
             trainable_parameters, args.max_grad_norm
         )
@@ -1201,25 +1258,33 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
         reported_loss = _distributed_average(accumulated_loss, run)
+        lora_lr = (
+            0.0 if args.freeze_thinker else float(optimizer.param_groups[0]["lr"])
+        )
+        head_lr = float(optimizer.param_groups[head_group_index]["lr"])
         record = {
             "step": update_step,
             "loss": reported_loss,
             "grad_norm": float(grad_norm),
-            "lora_lr": float(optimizer.param_groups[0]["lr"]),
+            "lora_lr": lora_lr,
-            "head_lr": float(optimizer.param_groups[1]["lr"]),
+            "head_lr": head_lr,
         }
         if run.is_main:
             history.append(record)
             progress.update(1)
             progress.set_postfix(
                 loss=f"{reported_loss:.4f}",
-                lr=f"{optimizer.param_groups[0]['lr']:.2g}",
+                lr=f"{head_lr:.2g}" if args.freeze_thinker else f"{lora_lr:.2g}",
             )
         if run.is_main and (update_step == 1 or update_step % args.log_every == 0):
             print(
                 f"step={update_step} loss={reported_loss:.6f} "
                 f"grad_norm={float(grad_norm):.4f} "
-                f"lora_lr={optimizer.param_groups[0]['lr']:.3g}",
+                + (
+                    f"head_lr={head_lr:.3g}"
+                    if args.freeze_thinker
+                    else f"lora_lr={lora_lr:.3g} head_lr={head_lr:.3g}"
+                ),
                 flush=True,
             )
     progress.close()
@@ -1264,6 +1329,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
 
     adapter_dir = output_dir / "adapter"
     if run.is_main:
+        if not args.freeze_thinker:
             raw_language_model.save_pretrained(adapter_dir)
         torch.save(
             {
@@ -1288,6 +1354,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         "schema_version": 1,
         "status": "pilot" if args.max_examples else "training",
         "model_id": args.model_id,
+        "model_variant": (
+            "frozen_base_thinker_head_only"
+            if args.freeze_thinker
+            else "thinker_lora"
+        ),
+        "thinker_frozen": bool(args.freeze_thinker),
         "initial_adapter_checkpoint": (
             str(args.adapter_checkpoint.expanduser().resolve())
             if args.adapter_checkpoint is not None
@@ -1351,6 +1423,9 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                     f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
                     f"producer batches={online_timing['producer_batches']}"
                 )
+        if args.freeze_thinker:
+            print("Frozen Thinker: no LoRA adapter was created")
+        else:
             print(f"Saved LoRA adapter to {adapter_dir}")
         print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
         print(

train_turn_head.py

@@ -72,6 +72,18 @@ def parse_args() -> argparse.Namespace:
             "They are automatically removed from --sessions. Never use 1/2/3."
         ),
     )
+    parser.add_argument(
+        "--sparse-anchor-cache-dir",
+        type=Path,
+        action="append",
+        default=[],
+        help=(
+            "Optional transition-sparse AV cache directory containing "
+            "manifest.json. Repeat once per training session. Only its anchor "
+            "target_index values select examples; states still come from "
+            "--feature-root. Development examples always remain dense."
+        ),
+    )
     parser.add_argument("--epochs", type=int, default=50)
     parser.add_argument("--batch-size", type=int, default=512)
     parser.add_argument("--learning-rate", type=float, default=7e-5)
@@ -138,6 +150,7 @@ def _load_examples(
     sessions: list[str],
     expected_mode: str | None,
     dataset: str,
+    selected_indices: dict[str, set[int]] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
@@ -151,7 +164,21 @@ def _load_examples(
             expected_mode=expected_mode,
             dataset=dataset,
         )
-        states, targets, _ = data.future_examples(NUM_HORIZONS)
+        states, targets, future_indices = data.future_examples(NUM_HORIZONS)
+        dense_example_count = len(states)
+        if selected_indices is not None:
+            if name not in selected_indices:
+                raise ValueError(f"No sparse anchor manifest was supplied for {name}")
+            wanted = selected_indices[name]
+            current_indices = future_indices[:, 0] - 1
+            keep = torch.tensor(
+                [int(index) in wanted for index in current_indices.tolist()],
+                dtype=torch.bool,
+            )
+            states = states[keep]
+            targets = targets[keep]
+            if not len(states):
+                raise ValueError(f"Sparse anchor selection kept no examples for {name}")
         if hidden_size is None:
             hidden_size = data.hidden_size
         elif data.hidden_size != hidden_size:
@@ -166,12 +193,19 @@ def _load_examples(
                 "session": name,
                 "num_states": len(data.states),
                 "num_training_examples": len(states),
+                "num_dense_training_examples": dense_example_count,
+                "supervision_policy": (
+                    "transition-sparse-manifest"
+                    if selected_indices is not None
+                    else "dense-5hz"
+                ),
                 **data.feature_metadata,
             }
         )
         print(
             f"{name}: states={len(data.states)}, "
-            f"valid_10_horizon_examples={len(states)}"
+            f"valid_10_horizon_examples={dense_example_count}, "
+            f"selected_examples={len(states)}"
         )
     if hidden_size is None:
         raise ValueError("No sessions selected")
@@ -183,6 +217,38 @@ def _load_examples(
     )
 
 
+def _load_sparse_anchor_indices(
+    cache_dirs: list[Path],
+) -> dict[str, set[int]] | None:
+    if not cache_dirs:
+        return None
+    result: dict[str, set[int]] = {}
+    for cache_dir in cache_dirs:
+        manifest_path = cache_dir.expanduser().resolve() / "manifest.json"
+        if not manifest_path.is_file():
+            raise FileNotFoundError(manifest_path)
+        manifest = json.loads(manifest_path.read_text())
+        session = normalize_sample_name(str(manifest.get("session", "")), "easycom")
+        policy = manifest.get("anchor_selection", {}).get("policy")
+        if policy != "transition-sparse":
+            raise ValueError(
+                f"Expected transition-sparse manifest, got {policy!r}: "
+                f"{manifest_path}"
+            )
+        if session in result:
+            raise ValueError(f"Duplicate sparse anchor manifest for {session}")
+        indices = {
+            int(anchor["target_index"])
+            for anchor in manifest.get("anchors", [])
+            if bool(anchor.get("has_future_target", True))
+        }
+        if not indices:
+            raise ValueError(f"No valid sparse anchors in {manifest_path}")
+        result[session] = indices
+        print(f"sparse anchors {session}: {len(indices)} from {manifest_path}")
+    return result
+
+
 def _scheduler(
     optimizer: torch.optim.Optimizer,
     *,
@@ -297,6 +363,15 @@ def _checkpoint(
             "scheduler": "linear_warmup_times_cosine",
             "seed": args.seed,
             "expected_mode": args.expected_mode,
+            "supervision_policy": (
+                "transition-sparse-manifest"
+                if args.sparse_anchor_cache_dir
+                else "dense-5hz"
+            ),
+            "sparse_anchor_cache_dirs": [
+                str(path.expanduser().resolve())
+                for path in args.sparse_anchor_cache_dir
+            ],
         },
         "history": history,
     }
@@ -308,6 +383,11 @@ def main() -> None:
     _seed_everything(args.seed)
     device = _resolve_device(args.device)
     expected_mode = None if args.expected_mode == "any" else args.expected_mode
+    sparse_anchor_indices = _load_sparse_anchor_indices(
+        args.sparse_anchor_cache_dir
+    )
+    if sparse_anchor_indices is not None and args.dataset != "easycom":
+        raise ValueError("--sparse-anchor-cache-dir currently supports EasyCom only")
 
     if args.dataset == "ego4d":
         assert args.manifest is not None
@@ -342,6 +422,14 @@ def main() -> None:
     validate_disjoint_sessions(train_sessions, dev_sessions)
     if not train_sessions:
         raise ValueError("No training sessions remain after removing dev sessions")
+    if sparse_anchor_indices is not None:
+        missing = set(train_sessions) - set(sparse_anchor_indices)
+        extra = set(sparse_anchor_indices) - set(train_sessions)
+        if missing or extra:
+            raise ValueError(
+                "Sparse anchor manifests must match training sessions exactly; "
+                f"missing={sorted(missing)}, extra={sorted(extra)}"
+            )
     leaked = official_test & (set(train_sessions) | set(dev_sessions))
     if leaked:
         raise ValueError(
@@ -358,6 +446,7 @@ def main() -> None:
         train_sessions,
         expected_mode,
         args.dataset,
+        sparse_anchor_indices,
     )
     print(f"training examples: {len(states)}, feature_dim={hidden_size}")
     dev_states: torch.Tensor | None = None
@@ -369,6 +458,7 @@ def main() -> None:
             dev_sessions,
             expected_mode,
             args.dataset,
+            None,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(