From edd37dc36a296815fe99be05015186c385864bd6 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 24 Aug 2026 11:45:11 +0800 Subject: [PATCH] Add offline sparse supervision control --- README.md | 56 +++++++++++++++++++++++++++ train_turn_head.py | 96 ++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 149 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 8c935ab..a6f4d27 100644 --- a/README.md +++ b/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 diff --git a/train_turn_head.py b/train_turn_head.py index 92c2e36..b84c47c 100644 --- a/train_turn_head.py +++ b/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( @@ -527,4 +617,4 @@ def main() -> None: if __name__ == "__main__": - main() \ No newline at end of file + main()