Add files using upload-large-folder tool

analyze_d2_csv.py

@@ -0,0 +1,239 @@
+import argparse
+import csv
+import math
+from collections import defaultdict
+
+import numpy as np
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Analyze D2 frame-level CSV.")
+    parser.add_argument("--csv", required=True, help="Path to d2_llm_space.py or d2_basic.py CSV output.")
+    parser.add_argument("--beta", type=float, default=1.0)
+    parser.add_argument("--failure_iou", type=float, default=0.5)
+    parser.add_argument("--bottom_frac", type=float, default=0.2)
+    parser.add_argument("--pr_points", type=int, default=10)
+    return parser.parse_args()
+
+
+def read_rows(path, beta):
+    rows = []
+    with open(path, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            row_beta = float(row["beta"])
+            if abs(row_beta - beta) > 1e-8:
+                continue
+            q_col = "h_type" if "h_type" in row else "q_type"
+            rows.append(
+                {
+                    "sample_idx": int(row["sample_idx"]),
+                    "frame": int(row["frame"]),
+                    "anchor_type": row[q_col],
+                    "s_pred": float(row["s_pred"]),
+                    "s_gt": float(row["s_gt"]),
+                    "frame_iou": float(row["frame_iou"]),
+                    "iou_pred": float(row["iou_pred"]),
+                    "pred_area": float(row["pred_area"]),
+                    "gt_area": float(row["gt_area"]),
+                }
+            )
+    if not rows:
+        raise RuntimeError(f"No rows found for beta={beta} in {path}")
+    return rows
+
+
+def corr(x, y):
+    x = np.asarray(x, dtype=np.float64)
+    y = np.asarray(y, dtype=np.float64)
+    if len(x) < 2 or np.std(x) < 1e-12 or np.std(y) < 1e-12:
+        return float("nan")
+    return float(np.corrcoef(x, y)[0, 1])
+
+
+def residualize(y, controls):
+    y = np.asarray(y, dtype=np.float64)
+    cols = [np.ones(len(y), dtype=np.float64)]
+    for control in controls:
+        cols.append(np.asarray(control, dtype=np.float64))
+    x = np.stack(cols, axis=1)
+    coef, *_ = np.linalg.lstsq(x, y, rcond=None)
+    return y - x @ coef
+
+
+def r2_score(y, y_pred):
+    y = np.asarray(y, dtype=np.float64)
+    y_pred = np.asarray(y_pred, dtype=np.float64)
+    ss_res = np.sum((y - y_pred) ** 2)
+    ss_tot = np.sum((y - y.mean()) ** 2)
+    if ss_tot < 1e-12:
+        return float("nan")
+    return float(1.0 - ss_res / ss_tot)
+
+
+def linear_r2(y, features):
+    y = np.asarray(y, dtype=np.float64)
+    cols = [np.ones(len(y), dtype=np.float64)]
+    for feature in features:
+        cols.append(np.asarray(feature, dtype=np.float64))
+    x = np.stack(cols, axis=1)
+    coef, *_ = np.linalg.lstsq(x, y, rcond=None)
+    return r2_score(y, x @ coef)
+
+
+def real_rows(rows):
+    return [r for r in rows if r["anchor_type"] == "real"]
+
+
+def bottom_failure_enrichment(rows, failure_iou, bottom_frac):
+    rr = real_rows(rows)
+    n = len(rr)
+    k = max(1, int(round(n * bottom_frac)))
+    sorted_rows = sorted(rr, key=lambda r: r["s_pred"])
+    bottom = sorted_rows[:k]
+    baseline_rate = np.mean([r["frame_iou"] < failure_iou for r in rr])
+    bottom_rate = np.mean([r["frame_iou"] < failure_iou for r in bottom])
+    total_failures = sum(r["frame_iou"] < failure_iou for r in rr)
+    covered_failures = sum(r["frame_iou"] < failure_iou for r in bottom)
+    recall = covered_failures / max(total_failures, 1)
+    enrichment = bottom_rate / max(baseline_rate, 1e-12)
+    return {
+        "n": n,
+        "k": k,
+        "baseline_failure_rate": baseline_rate,
+        "bottom_failure_rate": bottom_rate,
+        "bottom_failure_recall": recall,
+        "enrichment": enrichment,
+        "total_failures": total_failures,
+    }
+
+
+def pr_curve(rows, failure_iou, points):
+    rr = sorted(real_rows(rows), key=lambda r: r["s_pred"])
+    total_failures = sum(r["frame_iou"] < failure_iou for r in rr)
+    out = []
+    for frac in np.linspace(0.05, 1.0, points):
+        k = max(1, int(round(len(rr) * frac)))
+        selected = rr[:k]
+        failures = sum(r["frame_iou"] < failure_iou for r in selected)
+        precision = failures / k
+        recall = failures / max(total_failures, 1)
+        out.append((frac, precision, recall))
+    return out
+
+
+def margin_rows(rows):
+    grouped = defaultdict(dict)
+    for r in rows:
+        key = (r["sample_idx"], r["frame"])
+        grouped[key][r["anchor_type"]] = r
+
+    out = []
+    for key, group in grouped.items():
+        if "real" not in group:
+            continue
+        controls = [group[name]["s_pred"] for name in ("shuffled", "wrong_ref") if name in group]
+        if not controls:
+            continue
+        real = group["real"]
+        item = dict(real)
+        item["s_margin"] = real["s_pred"] - max(controls)
+        out.append(item)
+    return out
+
+
+def bottom_failure_enrichment_for_score(rows, score_key, failure_iou, bottom_frac):
+    n = len(rows)
+    k = max(1, int(round(n * bottom_frac)))
+    sorted_rows = sorted(rows, key=lambda r: r[score_key])
+    bottom = sorted_rows[:k]
+    baseline_rate = np.mean([r["frame_iou"] < failure_iou for r in rows])
+    bottom_rate = np.mean([r["frame_iou"] < failure_iou for r in bottom])
+    total_failures = sum(r["frame_iou"] < failure_iou for r in rows)
+    covered_failures = sum(r["frame_iou"] < failure_iou for r in bottom)
+    return {
+        "n": n,
+        "k": k,
+        "baseline_failure_rate": baseline_rate,
+        "bottom_failure_rate": bottom_rate,
+        "bottom_failure_recall": covered_failures / max(total_failures, 1),
+        "enrichment": bottom_rate / max(baseline_rate, 1e-12),
+    }
+
+
+def main():
+    args = parse_args()
+    rows = read_rows(args.csv, args.beta)
+    rr = real_rows(rows)
+
+    print(f"CSV: {args.csv}")
+    print(f"beta: {args.beta}")
+    print(f"real frames: {len(rr)}")
+    print(f"failure definition: frame_iou < {args.failure_iou}")
+
+    print("\nReal s_pred Correlations")
+    print(f"corr(s_pred, frame_iou): {corr([r['s_pred'] for r in rr], [r['frame_iou'] for r in rr]):+.4f}")
+    print(f"corr(s_pred, iou_pred):  {corr([r['s_pred'] for r in rr], [r['iou_pred'] for r in rr]):+.4f}")
+    print(f"corr(s_pred, pred_area): {corr([r['s_pred'] for r in rr], [r['pred_area'] for r in rr]):+.4f}")
+
+    s_pred_values = [r["s_pred"] for r in rr]
+    frame_iou_values = [r["frame_iou"] for r in rr]
+    iou_pred_values = [r["iou_pred"] for r in rr]
+    pred_area_values = [r["pred_area"] for r in rr]
+    gt_area_values = [r["gt_area"] for r in rr]
+    partial_iou_pred = corr(
+        residualize(s_pred_values, [iou_pred_values]),
+        residualize(frame_iou_values, [iou_pred_values]),
+    )
+    partial_iou_area = corr(
+        residualize(s_pred_values, [iou_pred_values, pred_area_values]),
+        residualize(frame_iou_values, [iou_pred_values, pred_area_values]),
+    )
+    partial_iou_area_gt = corr(
+        residualize(s_pred_values, [iou_pred_values, pred_area_values, gt_area_values]),
+        residualize(frame_iou_values, [iou_pred_values, pred_area_values, gt_area_values]),
+    )
+    r2_iou_pred = linear_r2(frame_iou_values, [iou_pred_values])
+    r2_iou_pred_s = linear_r2(frame_iou_values, [iou_pred_values, s_pred_values])
+    r2_iou_pred_area = linear_r2(frame_iou_values, [iou_pred_values, pred_area_values])
+    r2_iou_pred_area_s = linear_r2(frame_iou_values, [iou_pred_values, pred_area_values, s_pred_values])
+
+    print("\nPartial Correlation / Residual Gain")
+    print(f"partial corr(s_pred, frame_iou | iou_pred):                 {partial_iou_pred:+.4f}")
+    print(f"partial corr(s_pred, frame_iou | iou_pred,pred_area):       {partial_iou_area:+.4f}")
+    print(f"partial corr(s_pred, frame_iou | iou_pred,pred_area,gt_area): {partial_iou_area_gt:+.4f}")
+    print(f"R2 frame_iou ~ iou_pred:                       {r2_iou_pred:.4f}")
+    print(f"R2 frame_iou ~ iou_pred + s_pred:              {r2_iou_pred_s:.4f} (gain {r2_iou_pred_s - r2_iou_pred:+.4f})")
+    print(f"R2 frame_iou ~ iou_pred + pred_area:           {r2_iou_pred_area:.4f}")
+    print(f"R2 frame_iou ~ iou_pred + pred_area + s_pred:  {r2_iou_pred_area_s:.4f} (gain {r2_iou_pred_area_s - r2_iou_pred_area:+.4f})")
+
+    stats = bottom_failure_enrichment(rows, args.failure_iou, args.bottom_frac)
+    print("\nBottom-k Failure Enrichment")
+    print(f"bottom_frac: {args.bottom_frac:.2f} ({stats['k']}/{stats['n']} frames)")
+    print(f"total failures: {stats['total_failures']}")
+    print(f"random/baseline failure rate: {stats['baseline_failure_rate']:.4f}")
+    print(f"bottom-s_pred failure rate:   {stats['bottom_failure_rate']:.4f}")
+    print(f"bottom-s_pred failure recall: {stats['bottom_failure_recall']:.4f}")
+    print(f"enrichment:                  {stats['enrichment']:.2f}x")
+
+    print("\nPR Curve Summary")
+    print("selected_frac | precision | recall")
+    for frac, precision, recall in pr_curve(rows, args.failure_iou, args.pr_points):
+        print(f"{frac:.2f} | {precision:.4f} | {recall:.4f}")
+
+    mr = margin_rows(rows)
+    if mr:
+        print("\nOffline Margin-D2")
+        print(f"margin frames: {len(mr)}")
+        print(f"corr(s_margin, frame_iou): {corr([r['s_margin'] for r in mr], [r['frame_iou'] for r in mr]):+.4f}")
+        print(f"corr(s_margin, pred_area): {corr([r['s_margin'] for r in mr], [r['pred_area'] for r in mr]):+.4f}")
+        mstats = bottom_failure_enrichment_for_score(mr, "s_margin", args.failure_iou, args.bottom_frac)
+        print(f"bottom-s_margin failure rate:   {mstats['bottom_failure_rate']:.4f}")
+        print(f"bottom-s_margin failure recall: {mstats['bottom_failure_recall']:.4f}")
+        print(f"margin enrichment:              {mstats['enrichment']:.2f}x")
+    else:
+        print("\nOffline Margin-D2 skipped: shuffled/wrong_ref controls not available.")
+
+
+if __name__ == "__main__":
+    main()

cache_q_features.py

@@ -0,0 +1,125 @@
+import json
+import os
+from functools import partial
+from itertools import islice
+
+import torch
+import transformers
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from configs import args
+from datasets import REFAVS
+from decoder_invariance_check import build_model, set_seed
+from load_model import collate_fn, dict_to_cuda
+
+
+def _jsonable_size(size):
+    if isinstance(size, torch.Tensor):
+        return [int(x) for x in size.detach().cpu().tolist()]
+    return [int(x) for x in size]
+
+
+def main():
+    set_seed(42)
+    torch.set_grad_enabled(False)
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+
+    dataset = REFAVS(args.cache_split, args, tokenizer, input_type="refer")
+    loader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
+
+    split_root = os.path.join(args.cache_root, args.cache_split)
+    os.makedirs(split_root, exist_ok=True)
+    index_path = os.path.join(split_root, "index.jsonl")
+    if os.path.exists(index_path) and not args.overwrite_cache:
+        raise FileExistsError(
+            f"{index_path} already exists. Pass --overwrite_cache to rebuild it."
+        )
+
+    limit = args.max_eval_rows if args.max_eval_rows > 0 else len(dataset)
+    print(f"cache split={args.cache_split} | samples={min(limit, len(dataset))}")
+    print(f"cache root: {split_root}")
+
+    model = build_model(tokenizer, seg_token_idx)
+    model.eval()
+
+    rows = []
+    for sample_idx, batch in enumerate(
+        tqdm(islice(loader, limit), total=min(limit, len(dataset)), desc=f"Caching {args.cache_split}")
+    ):
+        batch = dict_to_cuda(batch)
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            output = model.forward(
+                images=batch["images"],
+                images_clip=batch["images_clip"],
+                audio_features=batch["audio_feats"],
+                image_features=batch["image_feats"],
+                input_ids=batch["input_ids"],
+                labels=batch["labels"],
+                attention_masks=batch["attention_masks"],
+                masks_list=batch["masks"],
+                resize_list=batch["resizes"],
+                orgsize_list=batch["orgsizes"],
+                conversation_list=batch["convs"],
+                refs_num=batch["refs_num"],
+                fids=batch["fids"],
+                vids=batch["vids"],
+                contrast=args.ct_weight,
+                ref_ids=batch["ref_ids"],
+                inference=True,
+            )
+
+        cache_name = f"{sample_idx:06d}.pt"
+        cache_path = os.path.join(split_root, cache_name)
+        item = {
+            "sample_idx": sample_idx,
+            "vid": batch["vids"][0],
+            "refs": batch["refs"][0],
+            "fids": [int(x) for x in batch["fids"][0]],
+            "resize": _jsonable_size(batch["resizes"][0]),
+            "orgsize": _jsonable_size(batch["orgsizes"][0]),
+            "q": output["seg_embeddings"][0].detach().cpu().float(),
+        }
+        torch.save(item, cache_path)
+        rows.append(
+            {
+                "sample_idx": sample_idx,
+                "path": cache_name,
+                "vid": item["vid"],
+                "refs": item["refs"],
+                "fids": item["fids"],
+                "resize": item["resize"],
+                "orgsize": item["orgsize"],
+                "num_seg": int(item["q"].shape[0]),
+            }
+        )
+
+    if not rows:
+        raise RuntimeError("No samples were cached.")
+
+    with open(index_path, "w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+
+    print(f"cached samples: {len(rows)}")
+    print(f"saved index: {index_path}")
+
+
+if __name__ == "__main__":
+    main()

d2_basic.py

@@ -0,0 +1,340 @@
+import csv
+import math
+import os
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import transformers
+from torch.utils.data import DataLoader
+
+from configs import args
+from datasets import REFAVS
+from decoder_invariance_check import build_model, set_seed
+from load_model import collate_fn, dict_to_cuda
+
+
+def make_loader(tokenizer):
+    dataset = REFAVS(args.eval_split, args, tokenizer, input_type="refer")
+    return DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
+
+
+def build_tokenizer():
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+    return tokenizer, seg_token_idx
+
+
+def get_q(model, batch):
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        output = model.forward(
+            images=batch["images"],
+            images_clip=batch["images_clip"],
+            audio_features=batch["audio_feats"],
+            image_features=batch["image_feats"],
+            input_ids=batch["input_ids"],
+            labels=batch["labels"],
+            attention_masks=batch["attention_masks"],
+            masks_list=batch["masks"],
+            resize_list=batch["resizes"],
+            orgsize_list=batch["orgsizes"],
+            conversation_list=batch["convs"],
+            refs_num=batch["refs_num"],
+            fids=batch["fids"],
+            vids=batch["vids"],
+            contrast=args.ct_weight,
+            ref_ids=batch["ref_ids"],
+            inference=True,
+        )
+    return output["seg_embeddings"][0][0].float()
+
+
+def decode_low_res(model, batch, q):
+    visual_model = model.get_model().visual_model
+    sparse, dense = visual_model.prompt_encoder(
+        points=None,
+        boxes=None,
+        masks=None,
+        text_embeds=q.view(1, 1, -1).to(next(visual_model.parameters()).dtype),
+    )
+    sparse = sparse.to(q.dtype)
+    dense = dense.to(q.dtype)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        low_res_masks, iou_predictions = visual_model.mask_decoder(
+            image_embeddings=batch["image_feats"][0],
+            image_pe=visual_model.prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse,
+            dense_prompt_embeddings=dense,
+            multimask_output=False,
+        )
+    return low_res_masks.float(), iou_predictions.float().squeeze(-1)
+
+
+def masks_to_64(mask_logits_or_binary):
+    if mask_logits_or_binary.ndim == 3:
+        mask_logits_or_binary = mask_logits_or_binary.unsqueeze(1)
+    return F.interpolate(
+        mask_logits_or_binary.float(),
+        size=(64, 64),
+        mode="bilinear",
+        align_corners=False,
+    ).clamp(0.0, 1.0)
+
+
+def d2_scores(image_embeddings, mask64, q, beta):
+    feats = image_embeddings.float()
+    if mask64.shape[0] != feats.shape[0]:
+        raise ValueError(f"Mask/frame mismatch: {mask64.shape} vs {feats.shape}")
+
+    q = F.normalize(q.float().view(1, -1), dim=-1)
+    mask = mask64.float()
+    comp = 1.0 - mask
+
+    z_in = (feats * mask).sum(dim=(2, 3)) / mask.sum(dim=(2, 3)).clamp_min(1e-6)
+    z_out = (feats * comp).sum(dim=(2, 3)) / comp.sum(dim=(2, 3)).clamp_min(1e-6)
+
+    z_in = F.normalize(z_in, dim=-1)
+    z_out = F.normalize(z_out, dim=-1)
+    return (z_in @ q.T).squeeze(-1) - beta * (z_out @ q.T).squeeze(-1)
+
+
+def frame_iou(pred_logits, gt_masks):
+    pred = (torch.sigmoid(pred_logits.float()) > 0.4).float()
+    gt = gt_masks.float()
+    if pred.ndim == 4:
+        pred = pred.squeeze(1)
+    inter = (pred * gt).sum(dim=(1, 2))
+    union = torch.maximum(pred, gt).sum(dim=(1, 2))
+    num_pixels = pred.shape[-1] * pred.shape[-2]
+    no_obj = gt.sum(dim=(1, 2)) == 0
+    inter_no_obj = ((1.0 - pred) * (1.0 - gt)).sum(dim=(1, 2))
+    inter = torch.where(no_obj, inter_no_obj, inter)
+    union = torch.where(no_obj, torch.full_like(union, float(num_pixels)), union)
+    return inter / union.clamp_min(1e-7)
+
+
+def frame_fscore_proxy(pred_logits, gt_masks):
+    pred = (torch.sigmoid(pred_logits.float()) > 0.4).float()
+    gt = gt_masks.float()
+    if pred.ndim == 4:
+        pred = pred.squeeze(1)
+    tp = (pred * gt).sum(dim=(1, 2))
+    precision = tp / pred.sum(dim=(1, 2)).clamp_min(1e-7)
+    recall = tp / gt.sum(dim=(1, 2)).clamp_min(1e-7)
+    beta2 = 0.3
+    fscore = (1 + beta2) * precision * recall / (beta2 * precision + recall).clamp_min(1e-7)
+    no_obj = gt.sum(dim=(1, 2)) == 0
+    return torch.where(no_obj, torch.zeros_like(fscore), fscore)
+
+
+def parse_betas():
+    raw = os.environ.get("D2_BETAS", "0.5")
+    return [float(x.strip()) for x in raw.split(",") if x.strip()]
+
+
+def collect_q_pool(model, tokenizer, limit):
+    q_pool = []
+    loader = make_loader(tokenizer)
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        q = get_q(model, batch)
+        q_pool.append(
+            {
+                "sample_idx": sample_idx,
+                "vid": batch["vids"][0],
+                "ref": batch["refs"][0][0],
+                "fid": int(batch["fids"][0][0]),
+                "q": q.cpu(),
+            }
+        )
+        print(f"Collected q {sample_idx}: vid={q_pool[-1]['vid']} ref={q_pool[-1]['ref']}")
+    if not q_pool:
+        raise RuntimeError("No q vectors collected. Is the selected split empty?")
+    return q_pool
+
+
+def choose_shuffled_idx(sample_idx, q_pool):
+    if len(q_pool) <= 1:
+        return None
+    return (sample_idx + 1) % len(q_pool)
+
+
+def choose_wrong_ref_idx(sample_idx, q_pool):
+    current = q_pool[sample_idx]
+    for item in q_pool:
+        if item["sample_idx"] == sample_idx:
+            continue
+        if item["vid"] == current["vid"] and item["fid"] != current["fid"]:
+            return item["sample_idx"]
+    for item in q_pool:
+        if item["sample_idx"] == sample_idx:
+            continue
+        if item["vid"] == current["vid"] and item["ref"] != current["ref"]:
+            return item["sample_idx"]
+    return None
+
+
+def run_d2(model, tokenizer, q_pool, betas, limit):
+    rows = []
+    loader = make_loader(tokenizer)
+    q_lookup = {item["sample_idx"]: item for item in q_pool}
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(1234)
+
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        item = q_lookup[sample_idx]
+        real_q = item["q"].cuda()
+
+        low_res_masks, iou_predictions = decode_low_res(model, batch, real_q)
+        pred_mask64 = masks_to_64(torch.sigmoid(low_res_masks))
+        gt_masks = batch["masks"][0][0].float()
+        gt_mask64 = masks_to_64(gt_masks)
+        image_embeddings = batch["image_feats"][0].float()
+
+        pred_logits_hr = model.get_model().visual_model.postprocess_masks(
+            low_res_masks.to(batch["image_feats"][0].dtype),
+            input_size=batch["resizes"][0],
+            original_size=batch["orgsizes"][0],
+        ).squeeze(1)
+
+        frame_ious = frame_iou(pred_logits_hr, gt_masks)
+        frame_fscores = frame_fscore_proxy(pred_logits_hr, gt_masks)
+        pred_area = (torch.sigmoid(pred_logits_hr.float()) > 0.4).float().mean(dim=(1, 2))
+        gt_area = gt_masks.float().mean(dim=(1, 2))
+
+        shuffled_idx = choose_shuffled_idx(sample_idx, q_pool)
+        wrong_ref_idx = choose_wrong_ref_idx(sample_idx, q_pool)
+        q_controls = [
+            ("real", real_q, sample_idx),
+            ("random", torch.randn(real_q.shape, device=real_q.device, generator=generator), None),
+        ]
+        if shuffled_idx is not None:
+            q_controls.append(("shuffled", q_lookup[shuffled_idx]["q"].cuda(), shuffled_idx))
+        if wrong_ref_idx is not None:
+            q_controls.append(("wrong_ref", q_lookup[wrong_ref_idx]["q"].cuda(), wrong_ref_idx))
+
+        for beta in betas:
+            for q_type, q, q_source_idx in q_controls:
+                pred_scores = d2_scores(image_embeddings, pred_mask64, q, beta)
+                gt_scores = d2_scores(image_embeddings, gt_mask64, q, beta)
+                base_info = {
+                    "sample_idx": sample_idx,
+                    "vid": item["vid"],
+                    "ref": item["ref"],
+                    "fid": item["fid"],
+                    "split": args.eval_split,
+                    "frame_iou": math.nan,
+                    "frame_fscore_proxy": math.nan,
+                    "iou_pred": math.nan,
+                    "pred_area": math.nan,
+                    "gt_area": math.nan,
+                }
+                for frame_idx in range(pred_scores.shape[0]):
+                    base_info_frame = dict(base_info)
+                    base_info_frame.update(
+                        {
+                            "frame_iou": frame_ious[frame_idx].item(),
+                            "frame_fscore_proxy": frame_fscores[frame_idx].item(),
+                            "iou_pred": iou_predictions[frame_idx].item(),
+                            "pred_area": pred_area[frame_idx].item(),
+                            "gt_area": gt_area[frame_idx].item(),
+                        }
+                    )
+                    row = dict(base_info_frame)
+                    row.update(
+                        {
+                            "frame": frame_idx,
+                            "q_type": q_type,
+                            "beta": beta,
+                            "s_pred": pred_scores[frame_idx].item(),
+                            "s_gt": gt_scores[frame_idx].item(),
+                            "q_source_idx": q_source_idx if q_source_idx is not None else "",
+                        }
+                    )
+                    rows.append(row)
+
+        real_rows = [
+            r for r in rows if r["sample_idx"] == sample_idx and r["q_type"] == "real" and r["beta"] == betas[0]
+        ]
+        s_pred_values = [r["s_pred"] for r in real_rows]
+        print(
+            f"D2 {sample_idx}: vid={item['vid']} ref={item['ref']} "
+            f"mean_s_pred={np.mean(s_pred_values):.4f} min_s_pred={np.min(s_pred_values):.4f} "
+            f"mean_iou={frame_ious.mean().item():.4f}"
+        )
+
+    return rows
+
+
+def print_summary(rows):
+    real_rows = [r for r in rows if r["q_type"] == "real"]
+    if not real_rows:
+        return
+    by_beta = sorted(set(r["beta"] for r in real_rows))
+    print("\nSummary")
+    print(f"rows: {len(rows)}")
+    for beta in by_beta:
+        beta_rows = [r for r in rows if r["beta"] == beta]
+        print(f"\nbeta={beta}")
+        for q_type in sorted(set(r["q_type"] for r in beta_rows)):
+            qr = [r for r in beta_rows if r["q_type"] == q_type]
+            print(
+                f"{q_type:10s} "
+                f"mean_s_pred={np.mean([r['s_pred'] for r in qr]):+.4f} "
+                f"mean_s_gt={np.mean([r['s_gt'] for r in qr]):+.4f}"
+            )
+        real_beta = [r for r in beta_rows if r["q_type"] == "real"]
+        s_pred = np.array([r["s_pred"] for r in real_beta])
+        frame_iou_values = np.array([r["frame_iou"] for r in real_beta])
+        if len(s_pred) > 1 and np.std(s_pred) > 1e-8 and np.std(frame_iou_values) > 1e-8:
+            corr = np.corrcoef(s_pred, frame_iou_values)[0, 1]
+            print(f"corr(real s_pred, frame_iou)={corr:+.4f}")
+        else:
+            print("corr(real s_pred, frame_iou)=nan")
+
+
+def main():
+    set_seed(42)
+    torch.set_grad_enabled(False)
+    betas = parse_betas()
+    tokenizer, seg_token_idx = build_tokenizer()
+    limit = args.max_eval_rows if args.max_eval_rows > 0 else 30
+    print(f"Split: {args.eval_split} | samples: {limit} | betas: {betas}")
+
+    model = build_model(tokenizer, seg_token_idx)
+    q_pool = collect_q_pool(model, tokenizer, limit)
+    rows = run_d2(model, tokenizer, q_pool, betas, limit)
+    print_summary(rows)
+
+    csv_path = os.environ.get("D2_BASIC_CSV", f"/workspace/SimToken/d2_basic_{args.eval_split}_{limit}.csv")
+    os.makedirs(os.path.dirname(os.path.abspath(csv_path)), exist_ok=True)
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        writer.writeheader()
+        writer.writerows(rows)
+    print(f"\nSaved CSV: {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

d2_llm_space.py

@@ -0,0 +1,314 @@
+import csv
+import math
+import os
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import transformers
+from torch.utils.data import DataLoader
+
+from configs import args
+from datasets import REFAVS
+from decoder_invariance_check import build_model, set_seed
+from d2_basic import frame_fscore_proxy, frame_iou
+from load_model import collate_fn, dict_to_cuda
+
+
+def build_tokenizer():
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+    return tokenizer, seg_token_idx
+
+
+def make_loader(tokenizer):
+    dataset = REFAVS(args.eval_split, args, tokenizer, input_type="refer")
+    return DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
+
+
+def forward_for_hidden_and_q(model, batch):
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        output = model.forward(
+            images=batch["images"],
+            images_clip=batch["images_clip"],
+            audio_features=batch["audio_feats"],
+            image_features=batch["image_feats"],
+            input_ids=batch["input_ids"],
+            labels=batch["labels"],
+            attention_masks=batch["attention_masks"],
+            masks_list=batch["masks"],
+            resize_list=batch["resizes"],
+            orgsize_list=batch["orgsizes"],
+            conversation_list=batch["convs"],
+            refs_num=batch["refs_num"],
+            fids=batch["fids"],
+            vids=batch["vids"],
+            contrast=args.ct_weight,
+            ref_ids=batch["ref_ids"],
+            inference=True,
+        )
+    h_seg = output["seg_hidden_states"][0][0].float()
+    q = output["seg_embeddings"][0][0].float()
+    return h_seg, q
+
+
+def decode_low_res(model, batch, q):
+    visual_model = model.get_model().visual_model
+    sparse, dense = visual_model.prompt_encoder(
+        points=None,
+        boxes=None,
+        masks=None,
+        text_embeds=q.view(1, 1, -1).to(next(visual_model.parameters()).dtype),
+    )
+    sparse = sparse.to(q.dtype)
+    dense = dense.to(q.dtype)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        low_res_masks, iou_predictions = visual_model.mask_decoder(
+            image_embeddings=batch["image_feats"][0],
+            image_pe=visual_model.prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse,
+            dense_prompt_embeddings=dense,
+            multimask_output=False,
+        )
+    return low_res_masks.float(), iou_predictions.float().squeeze(-1)
+
+
+def clip_projected_tokens(model, batch):
+    images = torch.cat(batch["images_clip"], dim=0)
+    with torch.no_grad():
+        clip_tokens = model.encode_images(images)
+        projector = model.get_model().mm_projector
+        clip_tokens = clip_tokens.to(projector.weight.dtype)
+        llm_tokens = projector(clip_tokens).float()
+    return llm_tokens
+
+
+def infer_square_grid(num_tokens):
+    grid = int(math.sqrt(num_tokens))
+    if grid * grid != num_tokens:
+        raise ValueError(f"Expected square patch-token grid, got {num_tokens} tokens")
+    return grid
+
+
+def masks_to_token_grid(mask_logits_or_binary, num_tokens):
+    if mask_logits_or_binary.ndim == 3:
+        mask_logits_or_binary = mask_logits_or_binary.unsqueeze(1)
+    grid = infer_square_grid(num_tokens)
+    return F.interpolate(
+        mask_logits_or_binary.float(),
+        size=(grid, grid),
+        mode="bilinear",
+        align_corners=False,
+    ).flatten(2).transpose(1, 2).clamp(0.0, 1.0)
+
+
+def d2_scores_llm(llm_tokens, mask_tokens, h_seg, beta):
+    if llm_tokens.shape[:2] != mask_tokens.shape[:2]:
+        raise ValueError(f"Token/mask mismatch: {llm_tokens.shape} vs {mask_tokens.shape}")
+    h = F.normalize(h_seg.float().view(1, -1), dim=-1)
+    tokens = llm_tokens.float()
+    mask = mask_tokens.float()
+    comp = 1.0 - mask
+
+    z_in = (tokens * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1e-6)
+    z_out = (tokens * comp).sum(dim=1) / comp.sum(dim=1).clamp_min(1e-6)
+
+    z_in = F.normalize(z_in, dim=-1)
+    z_out = F.normalize(z_out, dim=-1)
+    return (z_in @ h.T).squeeze(-1) - beta * (z_out @ h.T).squeeze(-1)
+
+
+def parse_betas():
+    raw = os.environ.get("D2_BETAS", "0.5")
+    return [float(x.strip()) for x in raw.split(",") if x.strip()]
+
+
+def collect_hidden_pool(model, tokenizer, limit):
+    pool = []
+    loader = make_loader(tokenizer)
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        h_seg, q = forward_for_hidden_and_q(model, batch)
+        pool.append(
+            {
+                "sample_idx": sample_idx,
+                "vid": batch["vids"][0],
+                "ref": batch["refs"][0][0],
+                "fid": int(batch["fids"][0][0]),
+                "h": h_seg.cpu(),
+                "q": q.cpu(),
+            }
+        )
+        print(f"Collected h {sample_idx}: vid={pool[-1]['vid']} ref={pool[-1]['ref']}")
+    if not pool:
+        raise RuntimeError("No hidden states collected. Is the selected split empty?")
+    return pool
+
+
+def choose_shuffled_idx(sample_idx, pool):
+    if len(pool) <= 1:
+        return None
+    return (sample_idx + 1) % len(pool)
+
+
+def choose_wrong_ref_idx(sample_idx, pool):
+    current = pool[sample_idx]
+    for item in pool:
+        if item["sample_idx"] == sample_idx:
+            continue
+        if item["vid"] == current["vid"] and item["fid"] != current["fid"]:
+            return item["sample_idx"]
+    for item in pool:
+        if item["sample_idx"] == sample_idx:
+            continue
+        if item["vid"] == current["vid"] and item["ref"] != current["ref"]:
+            return item["sample_idx"]
+    return None
+
+
+def run_d2_llm(model, tokenizer, pool, betas, limit):
+    rows = []
+    lookup = {item["sample_idx"]: item for item in pool}
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(1234)
+    loader = make_loader(tokenizer)
+
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        item = lookup[sample_idx]
+        h_real = item["h"].cuda()
+        q_real = item["q"].cuda()
+
+        low_res_masks, iou_predictions = decode_low_res(model, batch, q_real)
+        llm_tokens = clip_projected_tokens(model, batch)
+        pred_mask_tokens = masks_to_token_grid(torch.sigmoid(low_res_masks), llm_tokens.shape[1])
+        gt_masks = batch["masks"][0][0].float()
+        gt_mask_tokens = masks_to_token_grid(gt_masks, llm_tokens.shape[1])
+
+        pred_logits_hr = model.get_model().visual_model.postprocess_masks(
+            low_res_masks.to(batch["image_feats"][0].dtype),
+            input_size=batch["resizes"][0],
+            original_size=batch["orgsizes"][0],
+        ).squeeze(1)
+        frame_ious = frame_iou(pred_logits_hr, gt_masks)
+        frame_fscores = frame_fscore_proxy(pred_logits_hr, gt_masks)
+        pred_area = (torch.sigmoid(pred_logits_hr.float()) > 0.4).float().mean(dim=(1, 2))
+        gt_area = gt_masks.float().mean(dim=(1, 2))
+
+        shuffled_idx = choose_shuffled_idx(sample_idx, pool)
+        wrong_ref_idx = choose_wrong_ref_idx(sample_idx, pool)
+        controls = [
+            ("real", h_real, sample_idx),
+            ("random", torch.randn(h_real.shape, device=h_real.device, generator=generator), None),
+        ]
+        if shuffled_idx is not None:
+            controls.append(("shuffled", lookup[shuffled_idx]["h"].cuda(), shuffled_idx))
+        if wrong_ref_idx is not None:
+            controls.append(("wrong_ref", lookup[wrong_ref_idx]["h"].cuda(), wrong_ref_idx))
+
+        for beta in betas:
+            for h_type, h, h_source_idx in controls:
+                pred_scores = d2_scores_llm(llm_tokens, pred_mask_tokens, h, beta)
+                gt_scores = d2_scores_llm(llm_tokens, gt_mask_tokens, h, beta)
+                for frame_idx in range(pred_scores.shape[0]):
+                    rows.append(
+                        {
+                            "sample_idx": sample_idx,
+                            "vid": item["vid"],
+                            "ref": item["ref"],
+                            "fid": item["fid"],
+                            "split": args.eval_split,
+                            "frame": frame_idx,
+                            "h_type": h_type,
+                            "beta": beta,
+                            "s_pred": pred_scores[frame_idx].item(),
+                            "s_gt": gt_scores[frame_idx].item(),
+                            "h_source_idx": h_source_idx if h_source_idx is not None else "",
+                            "frame_iou": frame_ious[frame_idx].item(),
+                            "frame_fscore_proxy": frame_fscores[frame_idx].item(),
+                            "iou_pred": iou_predictions[frame_idx].item(),
+                            "pred_area": pred_area[frame_idx].item(),
+                            "gt_area": gt_area[frame_idx].item(),
+                        }
+                    )
+
+        real_rows = [
+            r for r in rows if r["sample_idx"] == sample_idx and r["h_type"] == "real" and r["beta"] == betas[0]
+        ]
+        s_pred_values = [r["s_pred"] for r in real_rows]
+        print(
+            f"D2-LLM {sample_idx}: vid={item['vid']} ref={item['ref']} "
+            f"mean_s_pred={np.mean(s_pred_values):.4f} min_s_pred={np.min(s_pred_values):.4f} "
+            f"mean_iou={frame_ious.mean().item():.4f}"
+        )
+
+    return rows
+
+
+def print_summary(rows):
+    print("\nSummary")
+    print(f"rows: {len(rows)}")
+    for beta in sorted(set(r["beta"] for r in rows)):
+        beta_rows = [r for r in rows if r["beta"] == beta]
+        print(f"\nbeta={beta}")
+        for h_type in sorted(set(r["h_type"] for r in beta_rows)):
+            hr = [r for r in beta_rows if r["h_type"] == h_type]
+            print(
+                f"{h_type:10s} "
+                f"mean_s_pred={np.mean([r['s_pred'] for r in hr]):+.4f} "
+                f"mean_s_gt={np.mean([r['s_gt'] for r in hr]):+.4f}"
+            )
+        real_rows = [r for r in beta_rows if r["h_type"] == "real"]
+        s_pred = np.array([r["s_pred"] for r in real_rows])
+        frame_iou_values = np.array([r["frame_iou"] for r in real_rows])
+        if len(s_pred) > 1 and np.std(s_pred) > 1e-8 and np.std(frame_iou_values) > 1e-8:
+            corr = np.corrcoef(s_pred, frame_iou_values)[0, 1]
+            print(f"corr(real s_pred, frame_iou)={corr:+.4f}")
+        else:
+            print("corr(real s_pred, frame_iou)=nan")
+
+
+def main():
+    set_seed(42)
+    torch.set_grad_enabled(False)
+    betas = parse_betas()
+    tokenizer, seg_token_idx = build_tokenizer()
+    limit = args.max_eval_rows if args.max_eval_rows > 0 else 30
+    print(f"Split: {args.eval_split} | samples: {limit} | betas: {betas}")
+
+    model = build_model(tokenizer, seg_token_idx)
+    pool = collect_hidden_pool(model, tokenizer, limit)
+    rows = run_d2_llm(model, tokenizer, pool, betas, limit)
+    print_summary(rows)
+
+    csv_path = os.environ.get("D2_LLM_CSV", f"/workspace/SimToken/d2_llm_{args.eval_split}_{limit}.csv")
+    os.makedirs(os.path.dirname(os.path.abspath(csv_path)), exist_ok=True)
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        writer.writeheader()
+        writer.writerows(rows)
+    print(f"\nSaved CSV: {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

decoder_invariance_check.py

@@ -0,0 +1,256 @@
+import csv
+import os
+import random
+from functools import partial
+
+import numpy as np
+import torch
+import transformers
+from peft import LoraConfig, get_peft_model
+from torch.utils.data import DataLoader
+from transformers import AutoConfig
+
+from configs import args
+from datasets import REFAVS
+from load_model import collate_fn, dict_to_cuda
+from models.avs_model import Simtoken_ForCausalLM
+
+
+def set_seed(seed=42):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def find_lora_target_modules(model, target_modules=("q_proj", "v_proj")):
+    modules = set()
+    excluded = [
+        "visual_model",
+        "vision_tower",
+        "mm_projector",
+        "text_hidden_fcs",
+        "audio_feature_layer",
+    ]
+    for name, module in model.named_modules():
+        if not isinstance(module, torch.nn.Linear):
+            continue
+        if any(x in name for x in excluded):
+            continue
+        if any(x in name for x in target_modules):
+            modules.add(name)
+    return sorted(modules)
+
+
+def build_model(tokenizer, seg_token_idx):
+    model_args = {
+        "train_mask_decoder": True,
+        "out_dim": 256,
+        "ce_loss_weight": 1.0,
+        "dice_loss_weight": 0.5,
+        "bce_loss_weight": 2.0,
+        "seg_token_idx": seg_token_idx,
+        "vision_pretrained": args.vision_pretrained,
+        "vision_tower": args.vision_tower,
+        "use_im_start_end": False,
+        "compress": args.compress,
+        "start": args.start,
+    }
+
+    model = Simtoken_ForCausalLM.from_pretrained(
+        args.mllm,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        **model_args,
+    )
+
+    model.config.eos_token_id = tokenizer.eos_token_id
+    model.config.bos_token_id = tokenizer.bos_token_id
+    model.config.pad_token_id = tokenizer.pad_token_id
+
+    model.get_model().initialize_vision_modules(model.get_model().config)
+    vision_tower = model.get_model().get_vision_tower()
+    vision_tower.to(dtype=torch.float32, device="cuda")
+
+    model_args_from_pt = AutoConfig.from_pretrained(args.mllm)
+    model_args_from_pt.use_cluster = True
+    model_args_from_pt.freeze = False
+    model_args_from_pt.mm_tune = True
+    model_args_from_pt.spatial_cluster_rate0 = 64
+    model_args_from_pt.spatial_cluster_rate1 = 32
+    model_args_from_pt.spatial_cluster_rate2 = 16
+    model_args_from_pt.temporal_cluster_rate = 0.0625
+    model_args_from_pt.vision_tune = False
+    model.get_model().initialize_cluster_modules(model_args_from_pt)
+    model.get_model().initialize_lisa_modules(model.get_model().config)
+
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules=find_lora_target_modules(model),
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora_config)
+    model = model.to("cuda")
+    model.resize_token_embeddings(len(tokenizer))
+
+    state = torch.load(args.saved_model, map_location="cpu")
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    print(f"Loaded checkpoint: {args.saved_model}")
+    print(f"Missing keys: {len(missing)} | Unexpected keys: {len(unexpected)}")
+
+    model.eval()
+    return model
+
+
+def get_seg_embedding(model, batch):
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        output = model.forward(
+            images=batch["images"],
+            images_clip=batch["images_clip"],
+            audio_features=batch["audio_feats"],
+            image_features=batch["image_feats"],
+            input_ids=batch["input_ids"],
+            labels=batch["labels"],
+            attention_masks=batch["attention_masks"],
+            masks_list=batch["masks"],
+            resize_list=batch["resizes"],
+            orgsize_list=batch["orgsizes"],
+            conversation_list=batch["convs"],
+            refs_num=batch["refs_num"],
+            fids=batch["fids"],
+            vids=batch["vids"],
+            contrast=args.ct_weight,
+            ref_ids=batch["ref_ids"],
+            inference=True,
+        )
+    return output["seg_embeddings"][0][0:1]
+
+
+def check_one_sample(model, batch):
+    q = get_seg_embedding(model, batch)
+    image_embeddings = batch["image_feats"][0]
+
+    visual_model = model.get_model().visual_model
+    sparse, dense = visual_model.prompt_encoder(
+        points=None,
+        boxes=None,
+        masks=None,
+        text_embeds=q.unsqueeze(1),
+    )
+    sparse = sparse.to(q.dtype)
+    dense = dense.to(q.dtype)
+
+    decoder = visual_model.mask_decoder
+    image_pe = visual_model.prompt_encoder.get_dense_pe()
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        full_masks, full_iou = decoder(
+            image_embeddings=image_embeddings,
+            image_pe=image_pe,
+            sparse_prompt_embeddings=sparse,
+            dense_prompt_embeddings=dense,
+            multimask_output=False,
+        )
+
+        rows = []
+        for t in range(image_embeddings.shape[0]):
+            single_masks, single_iou = decoder(
+                image_embeddings=image_embeddings[t : t + 1],
+                image_pe=image_pe,
+                sparse_prompt_embeddings=sparse,
+                dense_prompt_embeddings=dense,
+                multimask_output=False,
+            )
+
+            diff = (full_masks[t : t + 1] - single_masks).float().abs()
+            iou_diff = (full_iou[t : t + 1] - single_iou).float().abs()
+            rows.append(
+                {
+                    "vid": batch["vids"][0],
+                    "ref": batch["refs"][0][0],
+                    "frame": t,
+                    "max_abs_diff": diff.max().item(),
+                    "mean_abs_diff": diff.mean().item(),
+                    "iou_pred_diff": iou_diff.max().item(),
+                }
+            )
+    return rows
+
+
+def main():
+    set_seed(42)
+    torch.set_grad_enabled(False)
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+
+    dataset = REFAVS(args.eval_split, args, tokenizer, input_type="refer")
+    loader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
+
+    limit = args.max_eval_rows if args.max_eval_rows > 0 else 1
+    print(f"Split: {args.eval_split} | samples to check: {limit}")
+
+    model = build_model(tokenizer, seg_token_idx)
+
+    all_rows = []
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        rows = check_one_sample(model, batch)
+        all_rows.extend(rows)
+
+        print(f"\nSample {sample_idx}: vid={rows[0]['vid']} ref={rows[0]['ref']}")
+        print("frame | max_abs_diff | mean_abs_diff | iou_pred_diff")
+        for row in rows:
+            print(
+                f"{row['frame']:02d} | "
+                f"{row['max_abs_diff']:.8e} | "
+                f"{row['mean_abs_diff']:.8e} | "
+                f"{row['iou_pred_diff']:.8e}"
+            )
+
+    if not all_rows:
+        raise RuntimeError("No rows were checked. Is the selected split empty?")
+
+    max_diff = max(row["max_abs_diff"] for row in all_rows)
+    mean_diff = sum(row["mean_abs_diff"] for row in all_rows) / len(all_rows)
+    max_iou_diff = max(row["iou_pred_diff"] for row in all_rows)
+
+    print("\nSummary")
+    print(f"checked frames: {len(all_rows)}")
+    print(f"global max_abs_diff: {max_diff:.8e}")
+    print(f"average mean_abs_diff: {mean_diff:.8e}")
+    print(f"global max_iou_pred_diff: {max_iou_diff:.8e}")
+
+    csv_path = os.environ.get("DECODER_INVARIANCE_CSV")
+    if csv_path:
+        os.makedirs(os.path.dirname(os.path.abspath(csv_path)), exist_ok=True)
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=list(all_rows[0].keys()))
+            writer.writeheader()
+            writer.writerows(all_rows)
+        print(f"Saved CSV: {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

load_model.py

@@ -1,12 +1,3 @@
-# Compatibility: transformers==4.30.2 calls hf_hub_download(use_auth_token=...),
-# removed in huggingface_hub>=0.20.  Patch before importing transformers so the
-# bound reference inside transformers.utils.hub picks up the fixed version.
-import huggingface_hub as _hfhub
-_hfhub_orig = _hfhub.hf_hub_download
-def _hfhub_compat(*args, use_auth_token=None, token=None, **kwargs):
-    return _hfhub_orig(*args, token=token or use_auth_token, **kwargs)
-_hfhub.hf_hub_download = _hfhub_compat
-
 import transformers
 
 from torch.cuda.amp import autocast, GradScaler
@@ -217,7 +208,7 @@ def collate_fn(batch, tokenizer=None):
 
 import torch.multiprocessing as mp
 if __name__ == "__main__":
-    mp.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn")
     set_seed(42)
     tokenizer = transformers.AutoTokenizer.from_pretrained(
         args.mllm,
@@ -233,9 +224,17 @@ if __name__ == "__main__":
     print("seg_token_idx: ", seg_token_idx)
 
 
-    _split = args.eval_split
-    _dataset = REFAVS(_split, args, tokenizer, input_type='refer')
-    _dataloader = DataLoader(_dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
+    if args.eval_split not in {"test_s", "test_u", "test_n"}:
+        raise ValueError(f"Unsupported eval_split: {args.eval_split}")
+
+    val_dataset = REFAVS(args.eval_split, args, tokenizer, input_type='refer')
+    val_dataloader = DataLoader(
+        val_dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=4,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
 
 
 
@@ -341,8 +340,12 @@ if __name__ == "__main__":
     model = model.to("cuda")
     model.resize_token_embeddings(len(tokenizer))
 
-    model.load_state_dict(torch.load(args.saved_model), strict=False)
-    print("saved model loaded")
+    missing, unexpected = model.load_state_dict(
+        torch.load(args.saved_model, map_location="cpu"),
+        strict=False,
+    )
+    print(f"saved model loaded: {args.saved_model}")
+    print(f"missing keys: {len(missing)} | unexpected keys: {len(unexpected)}")
 
 
     save_root = args.visualization_root
@@ -401,16 +404,15 @@ if __name__ == "__main__":
         print("visualization finished")
 
 
-    def valuate(model, dataloader, name, max_rows=-1):
+    def valuate(model, dataloader, name):
         model.eval()
 
         total_iou = 0
         total_fscore = 0
         count = 0
 
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on {name}", total=_total)):
-            if 0 < max_rows <= i:
+        for batch_idx, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on {name}")):
+            if args.max_eval_rows > 0 and batch_idx >= args.max_eval_rows:
                 break
             input_dict = dict_to_cuda(batch)
 
@@ -445,39 +447,40 @@ if __name__ == "__main__":
                 total_fscore += fscore * num_seg * T
                 count += num_seg * T
 
+        if count == 0:
+            raise RuntimeError(f"No samples were evaluated for {name}")
+
         print(f"\n  valuate on {name}:  miou: {total_iou/count}  fscore: {total_fscore/count}")
 
 
-    def valuate_Null(model, dataloader, max_rows=-1):
+    def valuate_Null(model, dataloader):
         model.eval()
 
         total_metric = 0
         count = 0
 
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on Null", total=_total)):
-            if 0 < max_rows <= i:
+        for batch_idx, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on Null")):
+            if args.max_eval_rows > 0 and batch_idx >= args.max_eval_rows:
                 break
             input_dict = dict_to_cuda(batch)
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-                with torch.no_grad():
-                    output_dict = model.forward(images=input_dict["images"],
-                                                images_clip=input_dict["images_clip"],
-                                                audio_features=input_dict["audio_feats"],
-                                                image_features=input_dict["image_feats"],
-                                                input_ids=input_dict["input_ids"],
-                                                labels=input_dict["labels"],
-                                                attention_masks=input_dict["attention_masks"],
-                                                masks_list=input_dict["masks"],
-                                                resize_list=input_dict["resizes"],
-                                                orgsize_list=input_dict["orgsizes"],
-                                                conversation_list=input_dict["convs"],
-                                                refs_num=input_dict["refs_num"],
-                                                fids=input_dict["fids"],
-                                                vids=input_dict["vids"],
-                                                contrast=args.ct_weight,
-                                                ref_ids=input_dict["ref_ids"],
-                                                inference=True)
+            with torch.no_grad():
+                output_dict = model.forward(images=input_dict["images"],
+                                            images_clip=input_dict["images_clip"],
+                                            audio_features=input_dict["audio_feats"],
+                                            image_features=input_dict["image_feats"],
+                                            input_ids=input_dict["input_ids"],
+                                            labels=input_dict["labels"],
+                                            attention_masks=input_dict["attention_masks"],
+                                            masks_list=input_dict["masks"],
+                                            resize_list=input_dict["resizes"],
+                                            orgsize_list=input_dict["orgsizes"],
+                                            conversation_list=input_dict["convs"],
+                                            refs_num=input_dict["refs_num"],
+                                            fids=input_dict["fids"],
+                                            vids=input_dict["vids"],
+                                            contrast=args.ct_weight,
+                                            ref_ids=input_dict["ref_ids"],
+                                            inference=True)
             pred_masks = output_dict["pred_masks"]  # list[B]:[num_seg, T, H, W]
             gt_masks = output_dict["gt_masks"]  # list[B]:[num_seg, T, H, W]
             for i in range(len(pred_masks)):
@@ -488,637 +491,13 @@ if __name__ == "__main__":
                 total_metric += null_metric * num_seg * T
                 count += num_seg * T
 
-        print(f"\n  valuate on test_n_refer, metric: {total_metric / count}")
-
-
+        if count == 0:
+            raise RuntimeError("No samples were evaluated for test_n")
 
+        print(f"\n  valuate on test_n_refer, metric: {total_metric / count}")
 
-    from seg_ltpo import (
-        LTPOConfig, ltpo_optimize, best_of_2_optimize, decode_full_video,
-        get_sam_model, get_anchor_indices,
-        QLTPOConfig, q_ltpo_autograd, check_grad_connectivity,
-        reset_q_ltpo_stats, get_q_ltpo_stats,
-        q_ltpo_frame_adaptive, decode_full_video_adaptive,
-        _compute_avt_proxy_reward,
-    )
-
-    def print_q_ltpo_stats(name: str) -> None:
-        stats = get_q_ltpo_stats()
-        if not stats:
-            return
-        n = len(stats)
-        acc_rate       = sum(s["accepted"]         for s in stats) / n
-        mean_gain      = sum(s["reward_gain"]       for s in stats) / n
-        mean_drift     = sum(s["drift"]             for s in stats) / n
-        clip_rate      = sum(s["hit_clip"]          for s in stats) / n
-        mean_iou_init  = sum(s["R_iou_pred_init"]   for s in stats) / n
-        mean_iou_best  = sum(s["R_iou_pred_best"]   for s in stats) / n
-        mean_area_init = sum(s["area_hard_init"]    for s in stats) / n
-        mean_area_best = sum(s["area_hard_best"]    for s in stats) / n
-        # Null safety: reward improved but predicted area grew >20 %
-        null_risk = sum(
-            1 for s in stats
-            if s["reward_gain"] > 0 and s["area_hard_best"] > s["area_hard_init"] * 1.2
-        ) / n
-        gains = sorted(s["reward_gain"] for s in stats)
-        def _pct(v, p): return v[max(0, int(len(v) * p / 100) - 1)]
-        mean_e0 = sum(s["e0"] for s in stats) / n
-        mean_mask_iou      = sum(s.get("mask_soft_iou",      0.0) for s in stats) / n
-        mean_iou_contrib   = sum(s.get("R_iou_contrib_gain", 0.0) for s in stats) / n
-        mean_soft_area_init = sum(s.get("r_area_soft_init",  0.0) for s in stats) / n
-        mean_soft_area_best = sum(s.get("r_area_soft_best",  0.0) for s in stats) / n
-        # B1 activation diagnostics
-        b1_excesses    = sorted(s.get("b1_peak_excess", 0.0) for s in stats)
-        b1_act_rate    = sum(1 for v in b1_excesses if v > 1e-8) / n
-        b1_mean_excess = sum(b1_excesses) / n
-        print(f"\n  [q-LTPO stats | {name} | n={n}]")
-        print(f"    acceptance rate      : {acc_rate:.3f}")
-        print(f"    mean e0 (exist prior): {mean_e0:.4f}  ← should differ Null vs Seen")
-        print(f"    mean reward gain     : {mean_gain:+.4f}")
-        print(f"    reward_gain p10/50/90: {_pct(gains,10):+.4f} / {_pct(gains,50):+.4f} / {_pct(gains,90):+.4f}")
-        print(f"    mean drift ‖q−q₀‖   : {mean_drift:.4f}")
-        print(f"    hit-clip ratio       : {clip_rate:.3f}")
-        print(f"    R_iou_pred init→best : {mean_iou_init:.4f} → {mean_iou_best:.4f}")
-        print(f"    R_iou_contrib_gain   : {mean_iou_contrib:+.4f}  ← λ_iou·e0·Δiou")
-        print(f"    mask soft-IoU(init,best): {mean_mask_iou:.4f}  ← 1.0=mask不变")
-        print(f"    area (hard) init→best: {mean_area_init:.4f} → {mean_area_best:.4f}")
-        print(f"    soft area init→best  : {mean_soft_area_init:.4f} → {mean_soft_area_best:.4f}")
-        print(f"    B1 activation rate   : {b1_act_rate:.3f}  ← frac(peak_area > e0)")
-        print(f"    B1 mean excess       : {b1_mean_excess:.5f}  ← mean ReLU(peak_area - e0)")
-        print(f"    B1 excess p10/50/90  : {_pct(b1_excesses,10):.5f} / {_pct(b1_excesses,50):.5f} / {_pct(b1_excesses,90):.5f}")
-        print(f"    reward↑ & area+20%↑  : {null_risk:.3f}  ← Null safety indicator")
-        # Direction II: frame-adaptive delta diagnostics
-        delta_norms = [s.get("delta_norm", 0.0) for s in stats]
-        if any(v > 0 for v in delta_norms):
-            print(f"    mean delta ‖Δ‖       : {sum(delta_norms)/n:.4f}  ← per-anchor residual norm")
-
-    def valuate_ltpo(model, dataloader, name, ltpo_cfg, optimize_fn=None,
-                     max_rows=-1, multimask=False, use_edge=False):
-        if optimize_fn is None:
-            optimize_fn = ltpo_optimize
-        """
-        Evaluate with SEG-LTPO test-time optimisation + optional boundary refinement.
-
-        decode_mode:
-          multimask=False, use_edge=False : original single-mask decode (default)
-          multimask=True,  use_edge=False : 3 candidates, SAM iou_pred selection (step 1a)
-          multimask=True,  use_edge=True  : 3 candidates, boundary-edge score (step 1b)
-        """
-        model.eval()
-        sam_model    = get_sam_model(model)
-        model_dtype  = torch.bfloat16
-        num_frames   = 10
-        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)
-
-        total_iou    = 0
-        total_fscore = 0
-        count        = 0
-
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(tqdm(dataloader, desc=f"LTPO Evaluating on {name}", total=_total)):
-            if 0 < max_rows <= i:
-                break
-            input_dict = dict_to_cuda(batch)
-
-            # ── Step 1: standard forward pass (LLM + SAM decode) ──────────
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-                with torch.no_grad():
-                    output_dict = model.forward(
-                        images=input_dict["images"],
-                        images_clip=input_dict["images_clip"],
-                        audio_features=input_dict["audio_feats"],
-                        image_features=input_dict["image_feats"],
-                        input_ids=input_dict["input_ids"],
-                        labels=input_dict["labels"],
-                        attention_masks=input_dict["attention_masks"],
-                        masks_list=input_dict["masks"],
-                        resize_list=input_dict["resizes"],
-                        orgsize_list=input_dict["orgsizes"],
-                        conversation_list=input_dict["convs"],
-                        refs_num=input_dict["refs_num"],
-                        fids=input_dict["fids"],
-                        vids=input_dict["vids"],
-                        contrast=args.ct_weight,
-                        ref_ids=input_dict["ref_ids"],
-                        inference=True,
-                    )
-
-            gt_masks          = output_dict["gt_masks"]         # list[B]:[num_seg, T, H, W]
-            seg_emb_list      = output_dict["seg_embeddings"]   # list[B]:[num_seg, 256]
-
-            for b in range(len(input_dict["images"])):
-                image_embeds_b = input_dict["image_feats"][b]   # [T, 256, 64, 64]
-                resize_b       = input_dict["resizes"][b]
-                orgsize_b      = input_dict["orgsizes"][b]
-                rgb_b = input_dict["images"][b] if use_edge else None  # [T,3,H,W]
-
-                # Convert initial Fseg to float32 for stable optimisation.
-                # seg_emb_list[b]: [num_seg, 256] in bfloat16
-                F_init_b = seg_emb_list[b].detach().float()     # [num_seg, 256]
-
-                pred_masks_ltpo = []
-                for seg_idx in range(F_init_b.shape[0]):
-                    fseg_init = F_init_b[seg_idx : seg_idx + 1]  # [1, 256]
-
-                    # ── Step 2: optimisation (float32, outside autocast) ──────
-                    best_fseg = optimize_fn(
-                        fseg_init, image_embeds_b, anchor_indices,
-                        sam_model, model_dtype, ltpo_cfg,
-                    )  # [1, 256] float32
-
-                    # ── Step 3: decode full video with best Fseg ──────────────
-                    pred_mask = decode_full_video(
-                        best_fseg, image_embeds_b, sam_model,
-                        resize_b, orgsize_b, model_dtype,
-                        rgb_frames=rgb_b, multimask=multimask,
-                    )  # [T, H, W]
-                    pred_masks_ltpo.append(pred_mask)
-
-                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)  # [num_seg, T, H, W]
-
-                num_seg = pred_masks_b.shape[0]
-                T_      = pred_masks_b.shape[1]
-                iou     = utility.mask_iou(pred_masks_b, gt_masks[b])
-                fscore  = utility.Eval_Fmeasure(pred_masks_b, gt_masks[b], None)
-
-                total_iou    += iou    * num_seg * T_
-                total_fscore += fscore * num_seg * T_
-                count        += num_seg * T_
-
-        print(f"\n  LTPO valuate on {name}:  miou: {total_iou/count:.4f}  fscore: {total_fscore/count:.4f}")
-
-
-    def valuate_ltpo_null(model, dataloader, ltpo_cfg, optimize_fn=None, max_rows=-1):
-        if optimize_fn is None:
-            optimize_fn = ltpo_optimize
-        """LTPO evaluation for Null split: measures S metric (lower = fewer false-positive masks)."""
-        model.eval()
-        sam_model      = get_sam_model(model)
-        model_dtype    = torch.bfloat16
-        num_frames     = 10
-        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)
-
-        total_metric = 0
-        count        = 0
-
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(tqdm(dataloader, desc="LTPO Evaluating on Null", total=_total)):
-            if 0 < max_rows <= i:
-                break
-            input_dict = dict_to_cuda(batch)
-
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-                with torch.no_grad():
-                    output_dict = model.forward(
-                        images=input_dict["images"],
-                        images_clip=input_dict["images_clip"],
-                        audio_features=input_dict["audio_feats"],
-                        image_features=input_dict["image_feats"],
-                        input_ids=input_dict["input_ids"],
-                        labels=input_dict["labels"],
-                        attention_masks=input_dict["attention_masks"],
-                        masks_list=input_dict["masks"],
-                        resize_list=input_dict["resizes"],
-                        orgsize_list=input_dict["orgsizes"],
-                        conversation_list=input_dict["convs"],
-                        refs_num=input_dict["refs_num"],
-                        fids=input_dict["fids"],
-                        vids=input_dict["vids"],
-                        contrast=args.ct_weight,
-                        ref_ids=input_dict["ref_ids"],
-                        inference=True,
-                    )
-
-            seg_emb_list = output_dict["seg_embeddings"]  # list[B]:[num_seg, 256]
-
-            for b in range(len(input_dict["images"])):
-                image_embeds_b = input_dict["image_feats"][b]
-                resize_b       = input_dict["resizes"][b]
-                orgsize_b      = input_dict["orgsizes"][b]
-                F_init_b       = seg_emb_list[b].detach().float()
-
-                pred_masks_ltpo = []
-                for seg_idx in range(F_init_b.shape[0]):
-                    fseg_init = F_init_b[seg_idx : seg_idx + 1]
-                    best_fseg = optimize_fn(
-                        fseg_init, image_embeds_b, anchor_indices,
-                        sam_model, model_dtype, ltpo_cfg,
-                    )
-                    pred_mask = decode_full_video(
-                        best_fseg, image_embeds_b, sam_model,
-                        resize_b, orgsize_b, model_dtype,
-                    )
-                    pred_masks_ltpo.append(pred_mask)
-
-                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)  # [num_seg, T, H, W]
-                num_seg      = pred_masks_b.shape[0]
-                T_           = pred_masks_b.shape[1]
-                null_metric  = utility.metric_s_for_null(pred_masks_b)
-
-                total_metric += null_metric * num_seg * T_
-                count        += num_seg * T_
-
-        print(f"\n  LTPO valuate on Null:  S metric: {total_metric/count:.4f}")
-
-
-    def valuate_ltpo_adaptive(model, dataloader, name, ltpo_cfg, max_rows=-1):
-        """Evaluate with Direction II frame-adaptive token optimization."""
-        model.eval()
-        sam_model      = get_sam_model(model)
-        model_dtype    = torch.bfloat16
-        num_frames     = 10
-        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)
-
-        total_iou    = 0
-        total_fscore = 0
-        count        = 0
-
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(tqdm(dataloader, desc=f"FA-LTPO Evaluating on {name}", total=_total)):
-            if 0 < max_rows <= i:
-                break
-            input_dict = dict_to_cuda(batch)
-
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-                with torch.no_grad():
-                    output_dict = model.forward(
-                        images=input_dict["images"],
-                        images_clip=input_dict["images_clip"],
-                        audio_features=input_dict["audio_feats"],
-                        image_features=input_dict["image_feats"],
-                        input_ids=input_dict["input_ids"],
-                        labels=input_dict["labels"],
-                        attention_masks=input_dict["attention_masks"],
-                        masks_list=input_dict["masks"],
-                        resize_list=input_dict["resizes"],
-                        orgsize_list=input_dict["orgsizes"],
-                        conversation_list=input_dict["convs"],
-                        refs_num=input_dict["refs_num"],
-                        fids=input_dict["fids"],
-                        vids=input_dict["vids"],
-                        contrast=args.ct_weight,
-                        ref_ids=input_dict["ref_ids"],
-                        inference=True,
-                    )
-
-            gt_masks     = output_dict["gt_masks"]       # list[B]:[num_seg, T, H, W]
-            seg_emb_list = output_dict["seg_embeddings"] # list[B]:[num_seg, 256]
-
-            for b in range(len(input_dict["images"])):
-                image_embeds_b = input_dict["image_feats"][b]
-                resize_b       = input_dict["resizes"][b]
-                orgsize_b      = input_dict["orgsizes"][b]
-                F_init_b       = seg_emb_list[b].detach().float()
-
-                pred_masks_ltpo = []
-                for seg_idx in range(F_init_b.shape[0]):
-                    fseg_init = F_init_b[seg_idx : seg_idx + 1]
-
-                    q_global, delta = q_ltpo_frame_adaptive(
-                        fseg_init, image_embeds_b, anchor_indices,
-                        sam_model, model_dtype, ltpo_cfg,
-                    )
-
-                    pred_mask = decode_full_video_adaptive(
-                        q_global, delta, anchor_indices,
-                        image_embeds_b, sam_model,
-                        resize_b, orgsize_b, model_dtype,
-                    )
-                    pred_masks_ltpo.append(pred_mask)
-
-                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)
-                num_seg = pred_masks_b.shape[0]
-                T_      = pred_masks_b.shape[1]
-                iou     = utility.mask_iou(pred_masks_b, gt_masks[b])
-                fscore  = utility.Eval_Fmeasure(pred_masks_b, gt_masks[b], None)
-
-                total_iou    += iou    * num_seg * T_
-                total_fscore += fscore * num_seg * T_
-                count        += num_seg * T_
-
-        print(f"\n  FA-LTPO valuate on {name}:  miou: {total_iou/count:.4f}  fscore: {total_fscore/count:.4f}")
-
-    # ── Step A0: reward–metric correlation study ─────────────────────────
-
-    def _print_correlation_report(per_sample: list) -> None:
-        import numpy as np
-        n = len(per_sample)
-        if n == 0:
-            return
-
-        r_iou   = np.array([s["reward_gain"]   for s in per_sample], dtype=float)
-        r_avt   = np.array([s["r_avt_gain"]    for s in per_sample], dtype=float)
-        r_avt_c = np.array([s["r_avt_c_gain"]  for s in per_sample], dtype=float)
-        dm      = np.array([s["delta_miou"]     for s in per_sample], dtype=float)
-        df      = np.array([s["delta_f"]        for s in per_sample], dtype=float)
-
-        def pearson(x, y):
-            x = x - x.mean(); y = y - y.mean()
-            denom = np.sqrt((x ** 2).sum() * (y ** 2).sum())
-            return float((x * y).sum() / (denom + 1e-12))
-
-        def wrong_frac(gains, deltas):
-            return sum(1 for g, d in zip(gains, deltas) if g > 0 and d < 0) / n
-
-        print(f"\n  [Step A0: Reward–Metric Correlation | n={n}]")
-        print(f"    mean ΔmIoU : {dm.mean():+.4f}  (std {dm.std():.4f})")
-        print(f"    mean ΔF    : {df.mean():+.4f}  (std {df.std():.4f})")
-        print(f"\n    Pearson r  with ΔmIoU :")
-        print(f"      R_iou_pred_gain   : {pearson(r_iou,   dm):+.3f}   ← current proxy")
-        print(f"      R_avt_gain        : {pearson(r_avt,   dm):+.3f}   ← cos(z_in, q_init)")
-        print(f"      R_avt_c_gain      : {pearson(r_avt_c, dm):+.3f}   ← cos(z_in,q)-β·cos(z_out,q)")
-        print(f"\n    Pearson r  with ΔF :")
-        print(f"      R_iou_pred_gain   : {pearson(r_iou,   df):+.3f}")
-        print(f"      R_avt_gain        : {pearson(r_avt,   df):+.3f}")
-        print(f"      R_avt_c_gain      : {pearson(r_avt_c, df):+.3f}")
-        print(f"\n    Wrong direction  (gain>0  but  Δ<0):")
-        print(f"      R_iou / ΔmIoU : {wrong_frac(r_iou,   dm):.3f}")
-        print(f"      R_avt / ΔmIoU : {wrong_frac(r_avt,   dm):.3f}")
-        print(f"      R_iou / ΔF    : {wrong_frac(r_iou,   df):.3f}")
-        print(f"      R_avt / ΔF    : {wrong_frac(r_avt,   df):.3f}")
-
-    def valuate_ltpo_correlation_study(model, dataloader, ltpo_cfg, max_rows=-1):
-        """Step A0: per-sample reward–metric correlation study.
-
-        For each (video, segment) sample runs:
-          1. Baseline decode  (q_init → mask → IoU/F)
-          2. q-LTPO s1        (q_best → mask → IoU/F)
-        Records reward signals and ΔmIoU / ΔF per sample, then prints
-        Pearson correlation table to identify which reward best predicts
-        actual metric improvement.
-        """
-        model.eval()
-        sam_model      = get_sam_model(model)
-        model_dtype    = torch.bfloat16
-        anchor_indices = get_anchor_indices(10, ltpo_cfg.num_anchors)
-
-        per_sample = []
-
-        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
-        for i, batch in enumerate(
-            tqdm(dataloader, desc="Correlation study (s1)", total=_total)
-        ):
-            if 0 < max_rows <= i:
-                break
-            input_dict = dict_to_cuda(batch)
-
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-                with torch.no_grad():
-                    output_dict = model.forward(
-                        images=input_dict["images"],
-                        images_clip=input_dict["images_clip"],
-                        audio_features=input_dict["audio_feats"],
-                        image_features=input_dict["image_feats"],
-                        input_ids=input_dict["input_ids"],
-                        labels=input_dict["labels"],
-                        attention_masks=input_dict["attention_masks"],
-                        masks_list=input_dict["masks"],
-                        resize_list=input_dict["resizes"],
-                        orgsize_list=input_dict["orgsizes"],
-                        conversation_list=input_dict["convs"],
-                        refs_num=input_dict["refs_num"],
-                        fids=input_dict["fids"],
-                        vids=input_dict["vids"],
-                        contrast=args.ct_weight,
-                        ref_ids=input_dict["ref_ids"],
-                        inference=True,
-                    )
-
-            gt_masks     = output_dict["gt_masks"]       # list[B]:[num_seg, T, H, W]
-            seg_emb_list = output_dict["seg_embeddings"] # list[B]:[num_seg, 256]
-
-            for b in range(len(input_dict["images"])):
-                image_embeds_b = input_dict["image_feats"][b]
-                resize_b       = input_dict["resizes"][b]
-                orgsize_b      = input_dict["orgsizes"][b]
-                F_init_b       = seg_emb_list[b].detach().float()
-
-                for seg_idx in range(F_init_b.shape[0]):
-                    q_init  = F_init_b[seg_idx : seg_idx + 1]          # [1, 256]
-                    gt_seg  = gt_masks[b][seg_idx : seg_idx + 1]        # [1, T, H, W]
-
-                    # Baseline decode (q_init, no LTPO)
-                    with torch.no_grad():
-                        pred_base = decode_full_video(
-                            q_init, image_embeds_b, sam_model,
-                            resize_b, orgsize_b, model_dtype,
-                        ).unsqueeze(0)                                  # [1, T, H, W]
-                    iou_base = utility.mask_iou(pred_base, gt_seg)
-                    f_base   = utility.Eval_Fmeasure(pred_base, gt_seg, None)
-
-                    # LTPO (s1) — also computes r_avt inside q_ltpo_autograd
-                    reset_q_ltpo_stats()
-                    q_best = q_ltpo_autograd(
-                        q_init, image_embeds_b, anchor_indices,
-                        sam_model, model_dtype, ltpo_cfg,
-                    )
-                    stat = get_q_ltpo_stats()[0]
-
-                    with torch.no_grad():
-                        pred_ltpo = decode_full_video(
-                            q_best, image_embeds_b, sam_model,
-                            resize_b, orgsize_b, model_dtype,
-                        ).unsqueeze(0)
-                    iou_ltpo = utility.mask_iou(pred_ltpo, gt_seg)
-                    f_ltpo   = utility.Eval_Fmeasure(pred_ltpo, gt_seg, None)
-
-                    per_sample.append({
-                        "reward_gain":  stat["reward_gain"],
-                        "r_avt_gain":   stat.get("r_avt_gain",   0.0),
-                        "r_avt_c_gain": stat.get("r_avt_c_gain", 0.0),
-                        "e0":           stat["e0"],
-                        "accepted":     stat["accepted"],
-                        "delta_miou":   float(iou_ltpo  - iou_base),
-                        "delta_f":      float(f_ltpo    - f_base),
-                    })
-
-        _print_correlation_report(per_sample)
-
-    # ── Stage 0: gradient connectivity check ─────────────────────────────
-    # Loads one image_embed directly from disk — no dataloader, no gt_mask,
-    # no media frames required.  F_init is a unit-scale random vector that
-    # mimics the distribution of Fseg (SAM prompt embeddings are in ℝ^256
-    # with per-dim std ≈ 0.05–0.3; we use std=0.1 as a neutral initialisation).
-    def run_stage0_check():
-        import glob
-        sam_model   = get_sam_model(model)
-        model_dtype = torch.bfloat16
-
-        embed_files = sorted(glob.glob(os.path.join(args.data_dir, "image_embed", "*.pt")))
-        if not embed_files:
-            print("[Stage 0] ERROR: no .pt files found in data/image_embed/")
-            return False
-
-        img_embs = torch.load(embed_files[0], map_location="cuda")  # [T, 256, 64, 64]
-        if img_embs.dim() == 3:                                      # [256,64,64] → [1,256,64,64]
-            img_embs = img_embs.unsqueeze(0)
-
-        torch.manual_seed(42)
-        F_init = torch.randn(1, 256, device="cuda") * 0.1           # [1, 256] float32
-
-        anchors = get_anchor_indices(img_embs.shape[0], 4)
-        diag = check_grad_connectivity(F_init, img_embs, anchors, sam_model, model_dtype)
-        print("\n[Stage 0] Gradient connectivity check:")
-        print(f"  file used          : {os.path.basename(embed_files[0])}")
-        print(f"  gradient_connected : {diag['gradient_connected']}")
-        print(f"  grad_norm (step 0) : {diag['grad_norm_step0']:.6f}")
-        print(f"  reward trajectory  : {[f'{r:.4f}' for r in diag['reward_trajectory']]}")
-        return diag["gradient_connected"]
-
-    # ── Bypass equivalence test ───────────────────────────────────────────
-    # Three controlled tests to verify that fseg.unsqueeze(1) (bypass) is
-    # numerically equivalent to prompt_encoder(text_embeds=fseg.unsqueeze(1)):
-    #   Test 1 — dense_emb dtype: dense_A.to(bfloat16) vs dense_emb_bf16 (exact 0?)
-    #   Test 2 — matched-prec anchor decode: same decoder, same inputs, both bfloat16
-    #   Test 3 — full-video (all T frames) matched-prec decode
-    # If all pass, delta_bypass_init = 0 and the +4.22% is purely from optimization.
-    def run_bypass_test():
-        from seg_ltpo import _precompute_dense_emb
-
-        sam_model   = get_sam_model(model)
-        pe          = sam_model.prompt_encoder
-        mask_dec    = sam_model.mask_decoder
-        model_dtype = torch.bfloat16
-
-        # Get one real Fseg via a standard forward pass on the first batch
-        batch = next(iter(_dataloader))
-        input_dict = dict_to_cuda(batch)
-        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
-            with torch.no_grad():
-                output_dict = model.forward(
-                    images=input_dict["images"],
-                    images_clip=input_dict["images_clip"],
-                    audio_features=input_dict["audio_feats"],
-                    image_features=input_dict["image_feats"],
-                    input_ids=input_dict["input_ids"],
-                    labels=input_dict["labels"],
-                    attention_masks=input_dict["attention_masks"],
-                    masks_list=input_dict["masks"],
-                    resize_list=input_dict["resizes"],
-                    orgsize_list=input_dict["orgsizes"],
-                    conversation_list=input_dict["convs"],
-                    refs_num=input_dict["refs_num"],
-                    fids=input_dict["fids"],
-                    vids=input_dict["vids"],
-                    contrast=args.ct_weight,
-                    ref_ids=input_dict["ref_ids"],
-                    inference=True,
-                )
 
-        fseg         = output_dict["seg_embeddings"][0][0:1].detach()  # [1,256] bfloat16
-        image_embeds = input_dict["image_feats"][0]                    # [T,256,64,64] float32
-        device       = fseg.device
-
-        anchor_indices = get_anchor_indices(image_embeds.shape[0], 4)
-        img_anc        = image_embeds[anchor_indices]                          # [A,256,64,64] float32
-        dense_emb_bf16 = _precompute_dense_emb(sam_model, model_dtype, device)  # [1,256,64,64] bfloat16
-        dense_pe       = pe.get_dense_pe().to(device)                          # float32
-
-        def _decode(img, sparse_emb, dense_emb):
-            return mask_dec(
-                image_embeddings=img,
-                image_pe=dense_pe,
-                sparse_prompt_embeddings=sparse_emb,
-                dense_prompt_embeddings=dense_emb,
-                multimask_output=False,
-            )
-
-        def _check(label, tensor_a, tensor_b, exact=False):
-            err = (tensor_a.float() - tensor_b.float()).abs().max().item()
-            tol = 0.0 if exact else 1e-4
-            status = "PASS" if err <= tol else "FAIL"
-            print(f"  [{status}]  {label:50s}  max|A-B| = {err:.2e}")
-            return err <= tol
-
-        print(f"\n[Bypass Test]  fseg dtype={fseg.dtype}  norm={fseg.float().norm().item():.4f}")
-
-        with torch.no_grad():
-            # Get prompt_encoder outputs (called outside autocast → float32)
-            sparse_A, dense_A = pe(points=None, boxes=None, masks=None,
-                                   text_embeds=fseg.unsqueeze(1))
-            sparse_B = fseg.unsqueeze(1)   # bypass sparse: identical tensor
-
-            # ── Test 1: dense_emb dtype artifact ──────────���─────────────────────
-            # Hypothesis: dense_A (float32) and dense_emb_bf16 differ only because
-            # no_mask_embed.weight is float32; casting to bfloat16 should give exact 0.
-            print("\n  [Test 1] dense_emb dtype artifact (expected: exact 0)")
-            t1 = _check("dense_A.to(bfloat16) vs dense_emb_bf16",
-                        dense_A.to(torch.bfloat16), dense_emb_bf16, exact=True)
-
-            # ── Test 2: matched-precision decode on anchors ──────────────────────
-            # Both paths use bfloat16 sparse + bfloat16 dense.
-            # If sparse_emb is identical and dense_emb is identical (per Test 1),
-            # masks and iou_preds must be identical (same decoder, same inputs).
-            print("\n  [Test 2] matched-precision anchor decode (expected: exact 0)")
-            dense_A_bf16 = dense_A.to(model_dtype)
-            masks_A, iou_A = _decode(img_anc, sparse_A,  dense_A_bf16)
-            masks_B, iou_B = _decode(img_anc, sparse_B,  dense_emb_bf16)
-            _check("sparse_emb",                    sparse_A, sparse_B,  exact=True)
-            t2m = _check("masks   (anchors, matched prec)", masks_A,  masks_B,  exact=True)
-            t2i = _check("iou_preds (anchors, matched prec)", iou_A,  iou_B,    exact=True)
-            t2  = t2m and t2i
-
-            # ── Test 3: full-video bypass-init baseline (all T frames) ──────────
-            # Extend Test 2 to all T frames; quantifies delta_bypass_init over
-            # the complete video rather than just the 4 anchor frames.
-            print(f"\n  [Test 3] full-video matched-precision decode (T={image_embeds.shape[0]} frames)")
-            masks_full_A, _ = _decode(image_embeds, sparse_A,  dense_A_bf16)
-            masks_full_B, _ = _decode(image_embeds, sparse_B,  dense_emb_bf16)
-            t3 = _check("masks (all frames, matched prec)", masks_full_A, masks_full_B, exact=True)
-
-        print("\n  ── Verdict ──────────────────────────────────────────────────────")
-        if t1 and t2 and t3:
-            print("  ALL PASS — bypass is algebraically and numerically equivalent to")
-            print("  prompt_encoder path under matched precision.  delta_bypass_init = 0.")
-            print("  The +4.22% mIoU improvement is purely from q-LTPO optimization.")
-        else:
-            failures = []
-            if not t1: failures.append("Test 1 (dense dtype)")
-            if not t2: failures.append("Test 2 (anchor decode)")
-            if not t3: failures.append("Test 3 (full-video decode)")
-            print(f"  FAIL in: {', '.join(failures)}")
-            print("  delta_bypass_init ≠ 0; need per-sample mIoU comparison to quantify.")
-
-    # ── Run evaluation ────────────────────────────────────────────────────
-
-    ltpo_cfg          = LTPOConfig()
-    q_ltpo_cfg_s1     = QLTPOConfig(stage=1)
-    q_ltpo_cfg_s2     = QLTPOConfig(stage=2)
-    q_ltpo_cfg_s21    = QLTPOConfig(stage=21)   # P1a: tether probe
-    q_ltpo_cfg_s22    = QLTPOConfig(stage=22)   # P1b: faithful ext-ref
-
-    # ── Direction B: boundary precision probes ──────────────────────────────
-    q_ltpo_cfg_b1_w03 = QLTPOConfig(stage=1, lambda_area_inc=0.3, area_inc_tau=0.0)
-    q_ltpo_cfg_b1_w10 = QLTPOConfig(stage=1, lambda_area_inc=1.0, area_inc_tau=0.0)
-
-    # ── Direction II: Frame-adaptive token optimization ─────────────────────
-    # fa_c03: delta clipped at 0.3×‖q_init‖ — moderate constraint.
-    #   First probe to answer: "does constrained frame-adaptive beat shared q?"
-    #   If yes → ablate tighter/looser constraints and smoothness in follow-up.
-    q_ltpo_cfg_fa_c03  = QLTPOConfig(stage=1, lambda_residual=0.001, lambda_smooth_temp=0.0, max_delta_drift_scale=0.3)
-
-    max_rows          = args.max_eval_rows       # -1 = all rows
-
-    # --max_eval_rows 0  → Stage 0 + bypass equivalence check, then exit
-    if max_rows == 0:
-        run_stage0_check()
-        run_bypass_test()
-    elif _split == 'test_n':
-        # Null safety check: baseline + Stage 1 + frame-adaptive
-        valuate_Null(model, _dataloader, max_rows=max_rows)
-        for cfg_name, cfg in [("s1", q_ltpo_cfg_s1)]:
-            reset_q_ltpo_stats()
-            valuate_ltpo_null(model, _dataloader, cfg,
-                              optimize_fn=q_ltpo_autograd, max_rows=max_rows)
-            print_q_ltpo_stats(f"null_q_ltpo_{cfg_name}")
-        reset_q_ltpo_stats()
-        valuate_ltpo_adaptive(model, _dataloader, "null_fa_c03",
-                              q_ltpo_cfg_fa_c03, max_rows=max_rows)
-        print_q_ltpo_stats("null_fa_c03")
+    if args.eval_split == "test_n":
+        valuate_Null(model, val_dataloader)
     else:
-        valuate(model, _dataloader, _split, max_rows=max_rows)
-        # Step A0: reward–metric correlation study (s1 + AVT proxy signals)
-        valuate_ltpo_correlation_study(
-            model, _dataloader, q_ltpo_cfg_s1, max_rows=max_rows
-        )
-
+        valuate(model, val_dataloader, args.eval_split)

save_audio_feats.py

@@ -80,4 +80,3 @@ for vid in vids:
     # print(f"{vid}: {audio_embed.shape}")
     torch.save(audio_embed, f'{save_dir}/{vid}.pt')
     print(f'{vid} embedding saved {audio_embed.shape}')
-

setup_simtoken.md

@@ -1,12 +1,22 @@
 # SimToken Setup
 
+本文档用于在新机器上重建 SimToken 环境，并准备后续 A-min 实验。
+
 ---
 
 ## 1. Create Environment
 
+先确认 GPU 和 CUDA driver 状态：
+
 ```bash
-conda create -n simtoken python=3.10 -y
-conda activate simtoken
+nvidia-smi
+```
+
+创建 conda 环境：
+
+```bash
+/opt/miniforge3/condabin/conda create -n simtoken python=3.10 -y
+/opt/miniforge3/condabin/conda activate simtoken
 
 python -m pip install --upgrade pip wheel "setuptools<81"
 
@@ -34,17 +44,63 @@ pip install \
   huggingface_hub
 ```
 
+快速验证：
+
+```bash
+python - <<'PY'
+import torch
+print("torch:", torch.__version__)
+print("cuda available:", torch.cuda.is_available())
+print("device:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu")
+PY
+```
+
 ---
 
-## 2. Download from HuggingFace（新机器初始化）
+## 2. Check Workspace After Migration
+
+使用服务器平台的迁移工具完成目录迁移后，在新机器上确认关键文件：
+
+```bash
+cd /workspace/SimToken
+
+ls -lh checkpoints/simtoken_pretrained.pth
+ls -lh models/segment_anything/sam_vit_h_4b8939.pth
+ls -d data/image_embed data/gt_mask data/audio_embed data/media
+```
+
+如果迁移后只有压缩包而没有解压目录，重新解压：
+
+```bash
+cd /workspace/SimToken/data
+
+tar -xf image_embed.tar
+tar -xzf gt_mask.tar.gz
+tar -xzf audio_embed.tar.gz
+tar -xf media.tar
+```
 
-登录 HuggingFace（token 在 https://huggingface.co/settings/tokens 生成）：
+清理迁移中不需要的缓存：
+
+```bash
+cd /workspace/SimToken
+find . -name "__pycache__" -prune -exec rm -rf {} +
+find . -name ".pytest_cache" -prune -exec rm -rf {} +
+find . -name ".cache" -prune -exec rm -rf {} +
+find . -name "*.pyc" -delete
+```
+
+---
+
+## 3. Download from HuggingFace
+
+如果新机器不使用迁移工具，而是从 HuggingFace 重新初始化，先登录：
 
 ```bash
 huggingface-cli login
 ```
 
-下载完整 repo（代码 + 权重 + 压缩数据包，共约 190G）：
+下载完整 repo：
 
 ```bash
 mkdir -p /workspace/SimToken
@@ -56,97 +112,108 @@ huggingface-cli download yfan07/SimToken \
   --local-dir-use-symlinks False
 ```
 
-下载完成后解压数据包：
+下载完成后解压数据：
 
 ```bash
 cd /workspace/SimToken/data
 
-tar -xf image_embed.tar      # ~5–10 分钟
+tar -xf image_embed.tar
 tar -xzf gt_mask.tar.gz
 tar -xzf audio_embed.tar.gz
 tar -xf media.tar
 ```
 
-
 ---
 
-## 3. Pre-download Model Weights（首次使用必做）
+## 4. Pre-download Model Weights
 
-`transformers==4.30.2` 与新版 `huggingface_hub` 存在 API 不兼容（`use_auth_token` 已移除）。
-解决方案：先用 CLI 将模型下载到本地缓存，之后运行实验时加 `TRANSFORMERS_OFFLINE=1`，跳过所有网络请求。
+`transformers==4.30.2` 与新版 `huggingface_hub` 可能存在网络/API 兼容问题。建议先用 CLI 将模型下载到本地缓存，实验时再加 `TRANSFORMERS_OFFLINE=1`。
 
 ```bash
-# Chat-UniVi-7B（~14G）
+# Chat-UniVi-7B
 huggingface-cli download Chat-UniVi/Chat-UniVi-7B-v1.5
 
-# CLIP ViT-L（~1.6G）
+# CLIP ViT-L
 huggingface-cli download openai/clip-vit-large-patch14
 ```
 
-下载完成后即永久缓存，新 session 无需重复下载。
+下载完成后做离线验证：
+
+```bash
+cd /workspace/SimToken
+
+TRANSFORMERS_OFFLINE=1 /opt/miniforge3/condabin/conda run -n simtoken \
+  python -m py_compile train.py load_model.py decoder_invariance_check.py
+```
 
 ---
 
-## 4. Example Evaluation
+## 5. Smoke Test
 
-所有评测命令统一加 `TRANSFORMERS_OFFLINE=1`：
+先跑一个轻量 sanity check，确认 checkpoint、数据和离线模型缓存都能正常读取：
 
 ```bash
 cd /workspace/SimToken
 
-# Unseen split（全量 1656 样本）
-TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u
+TRANSFORMERS_OFFLINE=1 /opt/miniforge3/condabin/conda run -n simtoken \
+  python decoder_invariance_check.py \
+  --eval_split test_s \
+  --max_eval_rows 1
+```
 
-# Seen split
-TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_s
+如果可以正常加载模型并输出 per-frame diff，就可以启动完整 A-min 训练：
 
-# Null split（S metric，越低越好）
-TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_n
+```bash
+cd /workspace/SimToken
+mkdir -p log checkpoints
+
+TRANSFORMERS_OFFLINE=1 /opt/miniforge3/condabin/conda run -n simtoken \
+  python -W ignore train.py \
+  --name amin_full_e1 \
+  --init_from_saved_model \
+  --epochs 1 \
+  --batch_size 2 \
+  --lr 1e-4 \
+  --saved_model /workspace/SimToken/checkpoints/simtoken_pretrained.pth \
+  --log_root /workspace/SimToken/log \
+  --checkpoint_root /workspace/SimToken/checkpoints
+```
 
-# 限制样本数（快速验证）
-TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u --max_eval_rows 50
+启动日志中应出现：
 
-# Stage 0 梯度连通性 + bypass 等价性检查（仅诊断）
-TRANSFORMERS_OFFLINE=1 python -W ignore load_model.py --eval_split test_u --max_eval_rows 0
+```text
+initialized training from saved model: /workspace/SimToken/checkpoints/simtoken_pretrained.pth
+missing keys: ... | unexpected keys: ...
 ```
 
-每次评估依次输出：Baseline + q-LTPO Stage 1 两组结果及诊断统计。
-
 ---
 
-## 5. Upload to HuggingFace（实验结束后）
-
-数据目录以压缩包形式存储，可大幅减少文件数量，避免 HuggingFace commit 频率限制。
+## 6. Upload to HuggingFace
 
-**第一步：将数据目录压缩为归档文件（如尚未压缩）**
+实验结束后，如需重新上传到 HuggingFace，先将数据目录压缩为归档文件，减少文件数量：
 
 ```bash
 cd /workspace/SimToken/data
 
-tar -cf image_embed.tar image_embed/     # 不压缩（.pt 已是二进制）
+tar -cf image_embed.tar image_embed/
 tar -czf gt_mask.tar.gz gt_mask/
 tar -czf audio_embed.tar.gz audio_embed/
 tar -cf media.tar media/
 
-# 确认压缩包存在后删除原始目录
 ls -lh *.tar*
 rm -rf image_embed/ gt_mask/ audio_embed/ media/
 ```
 
-**第二步：清理缓存并上传**
+清理缓存并上传：
 
 ```bash
-find /workspace/SimToken -name "__pycache__" -exec rm -rf {} + 2>/dev/null
-find /workspace/SimToken -name "*.pyc" -delete
+cd /workspace/SimToken
 
-huggingface-cli login   # token 在 https://huggingface.co/settings/tokens 生成（需 Write 权限）
+find . -name "__pycache__" -prune -exec rm -rf {} +
+find . -name ".pytest_cache" -prune -exec rm -rf {} +
+find . -name ".cache" -prune -exec rm -rf {} +
+find . -name "*.pyc" -delete
 
-cd /workspace/SimToken
+huggingface-cli login
 python upload_hf.py --repo yfan07/SimToken
 ```
-
-**注意事项：**
-- 建议在 `tmux` 里运行，防止 SSH 断开：`tmux new -s upload`，完成后 `Ctrl+B D` detach
-- 支持断点续传：中断后重新执行同一命令会自动跳过已上传文件
-- 遇到 rate limit（HTTP 429）时脚本会自动等待约 1 小时后重试
-- 监控进度：`tail -f /workspace/SimToken/upload.log`

simtoken_experiment.md

@@ -0,0 +1,369 @@
+# SimToken 实验路线文档
+
+## 0. 当前状态
+
+前置诊断已经完成，路线收敛到 **A-min dynamic referent gate training**。
+
+已确认结论：
+
+1. **SAM decoder 下游是逐帧 batch-parallel 解码**  
+   `mask_decoder(image_embeddings[0:T])[t]` 与 `mask_decoder(image_embeddings[t:t+1])[0]` 只有混合精度数值噪声差异。旧的 decoder-level joint-frame competition 假设关闭。
+
+2. **target_frame sweep 基本无效**  
+   不同 target frame 生成的 q 几乎相同，`cos_to_q5` 通常在 `0.997+`；Seen/Null 上 oracle gain 约 `+0.0009`。这条 TTO 路线关闭。
+
+3. **raw SAM-space D2 失效**  
+   256 维 `q/Fseg` 与 SAM image embedding 不在可直接 cosine 的语义空间，`real q ≈ shuffled/wrong_ref q`，甚至 random q 更高。该定义关闭。
+
+4. **LLM-space D2 有弱诊断信号，但不适合作为主 reward**  
+   用 4096 维 `[SEG]` hidden state 与 `mm_projector(CLIP patch tokens)` 后的视觉 token 计算 D2，可以得到正相关：
+   - `corr(s_pred, frame_iou) ≈ +0.316`
+   - bottom 20% `s_pred` 中 failure rate 相比随机 baseline 约 `1.60x`
+   - 控制 `iou_pred` / `pred_area` 后偏相关约 `+0.14`
+
+   结论：`s_pred(beta=1.0)` 可以作为诊断信号或 frame-aware gate 的候选输入，但不能作为核心 TTO reward。
+
+5. **margin-D2 无效**  
+   离线 `s_margin = s(real) - max(s(shuffled), s(wrong_ref))` 的 failure enrichment 约 `0.93x`，会抵消掉有用的通用可见性/质量信号。该路线关闭。
+
+当前最干净的解释是：
+
+> q 本身通常是稳定的 referent anchor；主要瓶颈不在 q 生成，也不在简单 q selection，而在 SAM decoder 如何使用已有的 `mask_token -> q` sparse self-attention path。
+
+2026-04-22 更新：
+
+完整训练每个 epoch 约 2-4 小时，瓶颈主要在 7B MLLM forward，而不在 gate 本身。因此当前实验策略已调整为：
+
+1. 先缓存固定 checkpoint 下的 `q = seg_embeddings`；
+2. 在 cached q + cached SAM image embeddings 上训练 gate-only；
+3. 用 cached eval split 快速判断 gate 是否有泛化收益；
+4. 只有 gate-only 泛化信号成立后，再跑完整 A-min 联合训练。
+
+---
+
+## 1. A-min 当前实现
+
+已在代码中加入 A-min dynamic referent gate：
+
+- 文件：`models/segment_anything/modeling/transformer.py`
+- 模块：`ReferentGate`
+- 插入位置：`TwoWayAttentionBlock` 的 sparse self-attention + `norm1` 之后，token-to-image cross-attention 之前
+- 作用对象：只作用于 `mask_tokens`
+- 不作用于：`iou_token` 和 `q/sparse_prompt` 本身
+
+SAM token index：
+
+```python
+tokens = [iou_token, mask_tokens..., sparse_prompt(q)]
+```
+
+因此：
+
+```python
+iou_token index: 0
+mask token range: 1 : 1 + num_mask_tokens
+q token index: 1 + num_mask_tokens
+```
+
+A-min gate 形式：
+
+```python
+alpha = sigmoid(Linear([mask_token, q, cos(mask_token, q)]))
+mask_token = mask_token + alpha * Linear(q)
+```
+
+为保证旧 checkpoint 初始行为不变，`proj(q)` 分支使用零初始化。当前也将 `gate` 分支零初始化，使 alpha 有干净观测基线：
+
+```python
+nn.init.zeros_(self.gate.weight)
+nn.init.zeros_(self.gate.bias)
+nn.init.zeros_(self.proj.weight)
+nn.init.zeros_(self.proj.bias)
+```
+
+初始时 gate 为 identity：
+
+```text
+max_abs_diff(gate(mask, q), mask) = 0.0
+alpha_mean = 0.5
+alpha_std = 0.0
+```
+
+当前训练 forward 保持完整链路：`prepare_inputs_labels_for_multimodal -> MLLM forward -> text_hidden_fcs -> SAM mask decoder -> loss`。`--gate_only` 只控制参数冻结范围，不再改变 forward 语义。
+
+---
+
+## 2. 当前新增工具
+
+### 2.1 训练脚本增强
+
+`train.py` 已加入：
+
+- `--max_steps`
+- `--overfit_samples`
+- `--log_gate_stats_every`
+- `--skip_eval_after_train`
+- `--eval_train_only`
+
+启动时会打印 referent gate 参数是否 trainable、是否进入 optimizer，以及初始 `proj_norm/gate_norm`。
+
+### 2.2 cached q 路线
+
+新增脚本：
+
+- `cache_q_features.py`
+  - 离线缓存 `q = seg_embeddings`
+  - cache 文件很小，因为只保存 q 和少量 metadata
+  - `image_embeddings` 仍使用已有 `data/image_embed/{vid}.pt`
+  - `gt_masks` 仍使用已有 `data/gt_mask/...`
+
+- `train_cached_gate.py`
+  - 加载 base model 和 cached q
+  - 冻结全部参数，只训练 `referent_gate`
+  - 支持 `--eval_only`、`--disable_gate`
+  - 支持 `--save_gate_only`，只保存 gate 参数，checkpoint 约 1.6MB
+  - 支持 `--gate_checkpoint`，在 base checkpoint 上 overlay gate-only checkpoint
+  - gate stats 会记录：
+
+```text
+batch_miou
+batch_fscore
+proj_norm
+gate_norm
+proj_grad_norm
+gate_grad_norm
+alpha_mean / alpha_std / alpha_min / alpha_max
+```
+
+cached 解码已优化：一个 dataloader batch 会展平成 paired frame batch 调用 `mask_decoder.forward_modified_v3`，避免逐 sample 调 decoder 的主要开销，同时不会产生 prompt/image cross product。
+
+---
+
+## 3. 已完成实验结果
+
+### 3.1 cached identity 与原始 pipeline 一致性
+
+先用 `test_s` 前 10 条验证 cached pipeline 是否与原始 `load_model.py` 对齐：
+
+```text
+cached identity:
+mIoU   = 0.9686462879
+Fscore = 0.9868578851
+
+original load_model.py:
+mIoU   = 0.9686277151
+Fscore = 0.9868472159
+
+diff:
+mIoU   = +0.0000186
+Fscore = +0.0000107
+```
+
+结论：差异远小于 0.001，cached q pipeline 与原始 eval pipeline 一致，可以用于 gate-only 快速验证。
+
+### 3.2 gate probe：梯度路径与 alpha 分化
+
+在 cached train128 上跑 50 optimizer steps：
+
+```text
+step 5:
+proj_norm=0.074015
+gate_norm=0.064479
+proj_grad_norm=0.052291
+gate_grad_norm=0.000170
+alpha_mean=0.4999
+alpha_std=0.0019
+
+step 50:
+proj_norm=0.428711
+gate_norm=0.523223
+proj_grad_norm=0.022453
+gate_grad_norm=0.000504
+alpha_mean=0.5063
+alpha_std=0.0112
+```
+
+结论：
+
+- `proj_norm` 从 0 稳定增长，注入分支有梯度；
+- `gate_norm` 也开始增长，alpha 控制分支参与学习；
+- `alpha_std` 从 0 增长，说明 gate 对不同输入有分化响应；
+- 计算图、冻结范围、optimizer param groups 均正常。
+
+### 3.3 overfit32：表达能力验证
+
+cached train32 identity baseline：
+
+```text
+mIoU   = 0.8814558
+Fscore = 0.9375512
+```
+
+cached gate overfit32，200 steps，lr=1e-4：
+
+```text
+mIoU   = 0.9085821
+Fscore = 0.9444574
+```
+
+提升：
+
+```text
+mIoU   = +0.0271263
+Fscore = +0.0069063
+```
+
+结论：在 q、SAM image embeddings、mask decoder 原始参数均固定时，仅训练 A-min gate 就能明显提高训练集 mIoU，说明 gate 机制有表达能力，梯度路径通畅。
+
+### 3.4 overfit32 泛化评估
+
+对 cached eval split 前 200 条，identity baseline：
+
+```text
+test_s mIoU   = 0.7390979
+test_s Fscore = 0.8190672
+
+test_u mIoU   = 0.6732285
+test_u Fscore = 0.7734924
+
+test_n metric = 0.0606105
+```
+
+overfit32 gate checkpoint：
+
+```text
+test_s mIoU   = 0.7199481
+test_s Fscore = 0.8045849
+
+test_u mIoU   = 0.6672303
+test_u Fscore = 0.7663978
+
+test_n metric = 0.0648588
+```
+
+delta：
+
+```text
+test_s mIoU   = -0.0191498
+test_s Fscore = -0.0144823
+
+test_u mIoU   = -0.0059983
+test_u Fscore = -0.0070946
+
+test_n metric = +0.0042483
+```
+
+结论：
+
+- overfit32 gate 没有泛化；
+- Null metric 略升，说明小样本过拟合有轻微放大前景的倾向；
+- 这不是方法失败，而是 32 个样本不足以学到泛化 referent anchoring 的预期结果；
+- 下一步应扩大 cached train 样本量，并降低 lr。
+
+---
+
+## 4. 当前下一步实验：cached train256 gate-only
+
+用户已经完成 train256 的 q 缓存。下一步用 train256 跑更保守的 gate-only 泛化实验。
+
+### Step 1：训练 cached gate-only train256
+
+```bash
+cd /workspace/SimToken
+mkdir -p log checkpoints
+
+TRANSFORMERS_OFFLINE=1 python -u -W ignore train_cached_gate.py \
+  --cache_split train \
+  --cache_root /workspace/SimToken/cache_q \
+  --name cached_gate_train256_s300_lr3e5 \
+  --epochs 20 \
+  --max_steps 300 \
+  --batch_size 8 \
+  --lr 3e-5 \
+  --saved_model /workspace/SimToken/checkpoints/simtoken_pretrained.pth \
+  --log_root /workspace/SimToken/log \
+  --checkpoint_root /workspace/SimToken/checkpoints \
+  --log_gate_stats_every 50 \
+  --skip_eval_after_train \
+  --save_gate_only \
+  2>&1 | tee /workspace/SimToken/log/cached_gate_train256_s300_lr3e5.stdout
+```
+
+训练中重点观察：
+
+```text
+batch_miou / batch_fscore 是否逐步改善
+proj_norm 是否持续增长
+alpha_std 是否温和分化
+Null 风险：alpha 是否出现极端偏移
+```
+
+如果 `proj_norm` 在前 100 steps 仍接近 0，说明 lr=3e-5 可能过小，可以改回 1e-4 或使用分层 lr。
+
+### Step 2：评估 cached train256 gate checkpoint
+
+```bash
+for split in test_s test_u test_n; do
+  TRANSFORMERS_OFFLINE=1 python -u -W ignore train_cached_gate.py \
+    --cache_split $split \
+    --cache_root /workspace/SimToken/cache_q \
+    --batch_size 8 \
+    --saved_model /workspace/SimToken/checkpoints/simtoken_pretrained.pth \
+    --gate_checkpoint /workspace/SimToken/checkpoints/cached_gate_train256_s300_lr3e5.pth \
+    --eval_only \
+    --name cached_gate_train256_s300_lr3e5_${split}_200 \
+    2>&1 | tee /workspace/SimToken/log/cached_gate_train256_s300_lr3e5_${split}_200.stdout
+done
+```
+
+对比 baseline 使用 3.4 中 identity 200 条结果。
+
+### Step 3：根据结果决策
+
+判断标准：
+
+- Seen / Unseen 都提升：进入更大 cached train 或完整 A-min；
+- Seen 提升、Unseen 不提升：gate 仍可能学 dataset pattern，需要更多 train cache 或更强正则；
+- Seen / Unseen 都下降：不要跑完整 A-min，先调 lr、正则或 gate 容量；
+- Null metric 保持 `< 0.07`：暂不加 area penalty；
+- Null metric 超过 `0.10`：强危险信号，需要 area penalty 或约束预测面积。
+
+如果 train256 有弱正收益但幅度小，先看 alpha 分布和 hard/easy frames，而不是立刻扩大。若 alpha 在所有帧上几乎一致，可能只是全局偏置；若 hard frames alpha 系统性更高，说明更像 referent anchoring。
+
+---
+
+## 5. 成功标准
+
+A-min 成功不能只看总体 mIoU，需要同时满足：
+
+1. Seen / Unseen mIoU 稳定提升；
+2. Unseen 至少不弱于 Seen 的提升趋势；
+3. Null 指标不恶化，预测面积不膨胀；
+4. hard frames 改善更明显；
+5. 如果记录 gate alpha，hard frames 的 alpha 应系统性高于 easy frames。
+
+失败解释：
+
+- 如果 Seen 提升、Unseen 不提升：可能是 gate 学到数据集模式，而不是 referent anchoring；
+- 如果 Null 恶化：gate 可能放大了通用前景显著性；
+- 如果 gate-only 无变化但完整 A-min 有收益：说明 gate 需要与 mask decoder / text projection 协同适配；
+- 如果全 split 下降：gate 插入位置、初始化或学习率需要重新检查。
+
+---
+
+## 6. 后续机制分析
+
+如果 A-min 有正收益，再做 hook 分析：
+
+1. sparse self-attention 中 `mask_token -> q`；
+2. token-to-image attention 中 mask token 对 image tokens 的关注；
+3. A-min 前后 hard/easy frames 的 gate alpha；
+4. `s_pred(beta=1.0)` 与 gate alpha 的关系。
+
+这部分用于论文解释，不作为当前阻塞项。
+
+---
+
+## 7. 当前一句话结论
+
+> A-min gate 的梯度路径、表达能力和 cached pipeline 一致性已经通过验证；overfit32 能显著提升训练集但不能泛化。当前主线是用更大 cached train set（已完成 train256 cache）验证 gate-only 泛化，再决定是否投入完整 A-min 联合训练。

target_frame_sweep.py

@@ -0,0 +1,265 @@
+import csv
+import os
+import random
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import transformers
+from torch.utils.data import DataLoader
+
+from configs import args
+from datasets import REFAVS
+from decoder_invariance_check import build_model, set_seed
+from load_model import collate_fn, dict_to_cuda
+from utils import utility
+
+
+def decode_with_q(model, batch, q):
+    visual_model = model.get_model().visual_model
+    image_embeddings = batch["image_feats"][0]
+
+    sparse, dense = visual_model.prompt_encoder(
+        points=None,
+        boxes=None,
+        masks=None,
+        text_embeds=q.unsqueeze(1),
+    )
+    sparse = sparse.to(q.dtype)
+    dense = dense.to(q.dtype)
+
+    low_res_masks, iou_predictions = visual_model.mask_decoder(
+        image_embeddings=image_embeddings,
+        image_pe=visual_model.prompt_encoder.get_dense_pe(),
+        sparse_prompt_embeddings=sparse,
+        dense_prompt_embeddings=dense,
+        multimask_output=False,
+    )
+    pred_masks = visual_model.postprocess_masks(
+        low_res_masks,
+        input_size=batch["resizes"][0],
+        original_size=batch["orgsizes"][0],
+    ).squeeze(1)
+    return pred_masks.unsqueeze(0), iou_predictions.squeeze(-1)
+
+
+def get_q_for_target_frame(model, batch, target_frame):
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        output = model.forward(
+            images=batch["images"],
+            images_clip=batch["images_clip"],
+            audio_features=batch["audio_feats"],
+            image_features=batch["image_feats"],
+            input_ids=batch["input_ids"],
+            labels=batch["labels"],
+            attention_masks=batch["attention_masks"],
+            masks_list=batch["masks"],
+            resize_list=batch["resizes"],
+            orgsize_list=batch["orgsizes"],
+            conversation_list=batch["convs"],
+            refs_num=batch["refs_num"],
+            fids=batch["fids"],
+            vids=batch["vids"],
+            contrast=args.ct_weight,
+            ref_ids=batch["ref_ids"],
+            inference=True,
+            target_frame=target_frame,
+        )
+    return output["seg_embeddings"][0][0:1]
+
+
+def mask_area(pred_masks):
+    return (torch.sigmoid(pred_masks) > 0.4).float().mean().item()
+
+
+def mean_mask_iou_to_others(mask, other_masks):
+    if not other_masks:
+        return 1.0
+    binary = (torch.sigmoid(mask) > 0.4).float()
+    other_binary = [(torch.sigmoid(m) > 0.4).float() for m in other_masks]
+    vals = []
+    for other in other_binary:
+        inter = (binary * other).sum()
+        union = torch.maximum(binary, other).sum()
+        vals.append((inter / (union + 1e-7)).item())
+    return float(np.mean(vals))
+
+
+def evaluate_one_sample(model, batch, sample_idx):
+    rows = []
+    qs = []
+    pred_masks_by_tf = []
+
+    gt_masks = batch["masks"][0]
+    vid = batch["vids"][0]
+    ref = batch["refs"][0][0]
+
+    for target_frame in range(args.frame_n):
+        q = get_q_for_target_frame(model, batch, target_frame)
+        qs.append(q.float().squeeze(0))
+
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            pred_masks, iou_predictions = decode_with_q(model, batch, q)
+        pred_masks_by_tf.append(pred_masks.detach())
+
+        miou = utility.mask_iou(pred_masks.float(), gt_masks.float())
+        fscore = utility.Eval_Fmeasure(pred_masks.float(), gt_masks.float(), None)
+        null_metric = utility.metric_s_for_null(pred_masks.float())
+        area = mask_area(pred_masks)
+        mean_iou_pred = iou_predictions.float().mean().item()
+
+        rows.append(
+            {
+                "sample_idx": sample_idx,
+                "vid": vid,
+                "ref": ref,
+                "target_frame": target_frame,
+                "mean_iou_pred": mean_iou_pred,
+                "mask_area": area,
+                "null_metric": float(null_metric),
+                "miou": miou,
+                "fscore": fscore,
+                "cos_to_q5": 0.0,
+                "mean_cos_to_other_q": 0.0,
+                "mean_mask_iou_to_other_tf": 0.0,
+            }
+        )
+
+    q_stack = F.normalize(torch.stack(qs, dim=0), dim=-1)
+    q_cos = q_stack @ q_stack.T
+    q5_idx = min(5, len(qs) - 1)
+
+    for i, row in enumerate(rows):
+        other = [j for j in range(len(rows)) if j != i]
+        row["cos_to_q5"] = q_cos[i, q5_idx].item()
+        row["mean_cos_to_other_q"] = q_cos[i, other].mean().item()
+        row["mean_mask_iou_to_other_tf"] = mean_mask_iou_to_others(
+            pred_masks_by_tf[i], [pred_masks_by_tf[j] for j in other]
+        )
+
+    return rows
+
+
+def print_sample_summary(rows):
+    print(f"\nSample {rows[0]['sample_idx']}: vid={rows[0]['vid']} ref={rows[0]['ref']}")
+    print("tf | miou | fscore | null_s | iou_pred | area | cos_to_q5 | mean_q_cos")
+    for row in rows:
+        print(
+            f"{row['target_frame']:02d} | "
+            f"{row['miou']:.4f} | "
+            f"{row['fscore']:.4f} | "
+            f"{row['null_metric']:.4f} | "
+            f"{row['mean_iou_pred']:.4f} | "
+            f"{row['mask_area']:.4f} | "
+            f"{row['cos_to_q5']:.4f} | "
+            f"{row['mean_cos_to_other_q']:.4f}"
+        )
+
+    best_miou = max(rows, key=lambda x: x["miou"])
+    best_iou_pred = max(rows, key=lambda x: x["mean_iou_pred"])
+    fixed = rows[min(5, len(rows) - 1)]
+    miou_values = [row["miou"] for row in rows]
+    q5_values = [row["cos_to_q5"] for row in rows]
+    print(
+        "Best miou tf="
+        f"{best_miou['target_frame']} ({best_miou['miou']:.4f}); "
+        "best iou_pred tf="
+        f"{best_iou_pred['target_frame']} ({best_iou_pred['mean_iou_pred']:.4f}); "
+        f"fixed tf=5 miou={fixed['miou']:.4f}"
+    )
+    print(
+        f"target-frame miou range={max(miou_values) - min(miou_values):.4f}; "
+        f"min cos_to_q5={min(q5_values):.4f}"
+    )
+
+
+def main():
+    set_seed(42)
+    torch.set_grad_enabled(False)
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+
+    dataset = REFAVS(args.eval_split, args, tokenizer, input_type="refer")
+    loader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        collate_fn=partial(collate_fn, tokenizer=tokenizer),
+    )
+
+    limit = args.max_eval_rows if args.max_eval_rows > 0 else 1
+    print(f"Split: {args.eval_split} | samples to sweep: {limit}")
+
+    model = build_model(tokenizer, seg_token_idx)
+
+    all_rows = []
+    for sample_idx, batch in enumerate(loader):
+        if sample_idx >= limit:
+            break
+        batch = dict_to_cuda(batch)
+        rows = evaluate_one_sample(model, batch, sample_idx)
+        all_rows.extend(rows)
+        print_sample_summary(rows)
+
+    if not all_rows:
+        raise RuntimeError("No rows were checked. Is the selected split empty?")
+
+    fixed_rows = [r for r in all_rows if r["target_frame"] == min(5, args.frame_n - 1)]
+    oracle_by_sample = {}
+    iou_pred_by_sample = {}
+    for row in all_rows:
+        key = row["sample_idx"]
+        if key not in oracle_by_sample or row["miou"] > oracle_by_sample[key]["miou"]:
+            oracle_by_sample[key] = row
+        if key not in iou_pred_by_sample or row["mean_iou_pred"] > iou_pred_by_sample[key]["mean_iou_pred"]:
+            iou_pred_by_sample[key] = row
+
+    fixed_miou = np.mean([r["miou"] for r in fixed_rows])
+    fixed_null_metric = np.mean([r["null_metric"] for r in fixed_rows])
+    oracle_miou = np.mean([r["miou"] for r in oracle_by_sample.values()])
+    iou_pred_selected_miou = np.mean([r["miou"] for r in iou_pred_by_sample.values()])
+    min_cos_to_q5 = np.mean(
+        [min(r["cos_to_q5"] for r in all_rows if r["sample_idx"] == sample_idx) for sample_idx in oracle_by_sample]
+    )
+    mean_miou_range = np.mean(
+        [
+            max(r["miou"] for r in all_rows if r["sample_idx"] == sample_idx)
+            - min(r["miou"] for r in all_rows if r["sample_idx"] == sample_idx)
+            for sample_idx in oracle_by_sample
+        ]
+    )
+
+    print("\nSummary")
+    print(f"samples: {len(fixed_rows)}")
+    print(f"fixed target_frame=5 mean miou: {fixed_miou:.4f}")
+    print(f"fixed target_frame=5 mean null_s: {fixed_null_metric:.4f}")
+    print(f"oracle best-target-frame mean miou: {oracle_miou:.4f}")
+    print(f"best-by-iou_pred selected mean miou: {iou_pred_selected_miou:.4f}")
+    print(f"oracle gain over fixed: {oracle_miou - fixed_miou:+.4f}")
+    print(f"iou_pred-selection gain over fixed: {iou_pred_selected_miou - fixed_miou:+.4f}")
+    print(f"mean target-frame miou range: {mean_miou_range:.4f}")
+    print(f"mean sample min cos_to_q5: {min_cos_to_q5:.4f}")
+
+    csv_path = os.environ.get("TARGET_FRAME_SWEEP_CSV")
+    if csv_path:
+        os.makedirs(os.path.dirname(os.path.abspath(csv_path)), exist_ok=True)
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=list(all_rows[0].keys()))
+            writer.writeheader()
+            writer.writerows(all_rows)
+        print(f"Saved CSV: {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -1,7 +1,7 @@
 import transformers
 from datasets import REFAVS
 from configs import args
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Subset
 from functools import partial
 from models.llava import conversation as conversation_lib
 # from  models.avs_model import VISAForCausalLM
@@ -21,6 +21,7 @@ import numpy as np
 import re
 import time
 import os
+import sys
 
 
 import warnings
@@ -235,11 +236,19 @@ if __name__ == "__main__":
     val_dataset_u_refer = REFAVS('test_u', args, tokenizer, input_type='refer')
     val_dataset_n_refer = REFAVS('test_n', args, tokenizer, input_type='refer')
 
+    if args.overfit_samples > 0:
+        overfit_n = min(args.overfit_samples, len(train_dataset))
+        train_dataset = Subset(train_dataset, list(range(overfit_n)))
+        print(f"overfit_samples enabled: using first {overfit_n} train samples")
+
+    train_eval_dataset = train_dataset
+
 
     g = torch.Generator()
     g.manual_seed(42)
 
     train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=8, worker_init_fn=seed_worker,collate_fn=partial(collate_fn, tokenizer=tokenizer), generator=g)
+    train_eval_dataloader = DataLoader(train_eval_dataset, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
 
     val_dataloader_s_refer = DataLoader(val_dataset_s_refer, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
     val_dataloader_u_refer = DataLoader(val_dataset_u_refer, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
@@ -349,6 +358,11 @@ if __name__ == "__main__":
     model = model.to("cuda")
     model.resize_token_embeddings(len(tokenizer))
 
+    if args.init_from_saved_model or args.gate_only:
+        state = torch.load(args.saved_model, map_location="cpu")
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        print(f"initialized training from saved model: {args.saved_model}")
+        print(f"missing keys: {len(missing)} | unexpected keys: {len(unexpected)}")
 
     for name, param in model.audio_feature_layer.named_parameters():
         param.requires_grad = True
@@ -366,9 +380,113 @@ if __name__ == "__main__":
         ):
             p.requires_grad = True
 
+    if args.gate_only:
+        for p in model.parameters():
+            p.requires_grad = False
+        for n, p in model.named_parameters():
+            if "referent_gate" in n:
+                p.requires_grad = True
+        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in model.parameters())
+        print(f"gate_only enabled: trainable params {trainable} / {total}")
 
     print("will save train model")
 
+    def _total_norm(values):
+        if not values:
+            return 0.0
+        return float(sum(v * v for v in values) ** 0.5)
+
+    def collect_referent_gate_stats(model):
+        gate_modules = [(n, m) for n, m in model.named_modules() if n.endswith("referent_gate")]
+        proj_norms = []
+        gate_norms = []
+        proj_grad_norms = []
+        gate_grad_norms = []
+        alpha_tensors = []
+
+        for _, module in gate_modules:
+            proj_norms.append(module.proj.weight.detach().float().norm().item())
+            gate_norms.append(module.gate.weight.detach().float().norm().item())
+            if module.proj.weight.grad is not None:
+                proj_grad_norms.append(module.proj.weight.grad.detach().float().norm().item())
+            if module.gate.weight.grad is not None:
+                gate_grad_norms.append(module.gate.weight.grad.detach().float().norm().item())
+            if module.last_alpha is not None:
+                alpha_tensors.append(module.last_alpha.detach().float().reshape(-1))
+
+        stats = {
+            "modules": len(gate_modules),
+            "proj_norm": _total_norm(proj_norms),
+            "gate_norm": _total_norm(gate_norms),
+            "proj_grad_norm": _total_norm(proj_grad_norms),
+            "gate_grad_norm": _total_norm(gate_grad_norms),
+        }
+
+        if alpha_tensors:
+            alpha = torch.cat(alpha_tensors)
+            stats.update(
+                {
+                    "alpha_mean": alpha.mean().item(),
+                    "alpha_std": alpha.std(unbiased=False).item(),
+                    "alpha_min": alpha.min().item(),
+                    "alpha_max": alpha.max().item(),
+                }
+            )
+        else:
+            stats.update(
+                {
+                    "alpha_mean": float("nan"),
+                    "alpha_std": float("nan"),
+                    "alpha_min": float("nan"),
+                    "alpha_max": float("nan"),
+                }
+            )
+
+        return stats
+
+    def print_referent_gate_optimizer_sanity(model, optimizer):
+        optimizer_param_ids = {id(p) for group in optimizer.param_groups for p in group["params"]}
+        gate_params = [(n, p) for n, p in model.named_parameters() if "referent_gate" in n]
+        trainable_gate = [(n, p) for n, p in gate_params if p.requires_grad]
+        optimizer_gate = [(n, p) for n, p in gate_params if id(p) in optimizer_param_ids]
+        optimizer_trainable_gate = [
+            (n, p) for n, p in gate_params if p.requires_grad and id(p) in optimizer_param_ids
+        ]
+        print(
+            "referent_gate sanity: "
+            f"params={sum(p.numel() for _, p in gate_params)} | "
+            f"trainable={sum(p.numel() for _, p in trainable_gate)} | "
+            f"in_optimizer={sum(p.numel() for _, p in optimizer_gate)} | "
+            f"trainable_in_optimizer={sum(p.numel() for _, p in optimizer_trainable_gate)}"
+        )
+
+        stats = collect_referent_gate_stats(model)
+        print(
+            "referent_gate init stats: "
+            f"modules={stats['modules']} | "
+            f"proj_norm={stats['proj_norm']:.6f} | "
+            f"gate_norm={stats['gate_norm']:.6f}"
+        )
+
+    def log_referent_gate_stats(global_step, loss_value):
+        stats = collect_referent_gate_stats(model)
+        message = (
+            f"gate_stats step={global_step} "
+            f"loss={loss_value:.6f} "
+            f"proj_norm={stats['proj_norm']:.6f} "
+            f"gate_norm={stats['gate_norm']:.6f} "
+            f"proj_grad_norm={stats['proj_grad_norm']:.6f} "
+            f"gate_grad_norm={stats['gate_grad_norm']:.6f} "
+            f"alpha_mean={stats['alpha_mean']:.4f} "
+            f"alpha_std={stats['alpha_std']:.4f} "
+            f"alpha_min={stats['alpha_min']:.4f} "
+            f"alpha_max={stats['alpha_max']:.4f}"
+        )
+        print(message)
+        with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
+            f.write(message + "\n")
+
     def valuate(model, dataloader, args, name):
         model.eval()
 
@@ -420,11 +538,17 @@ if __name__ == "__main__":
     epochs = args.epochs
     print("init lr:", args.lr)
     optimizer = AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
+    print_referent_gate_optimizer_sanity(model, optimizer)
 
-    gradient_accumulation_steps = int(16 // args.batch_size)
-    step_per_epoch = len(train_dataloader) // gradient_accumulation_steps
-    total_steps = epochs * step_per_epoch
+    gradient_accumulation_steps = max(1, int(16 // args.batch_size))
+    step_per_epoch = max(1, len(train_dataloader) // gradient_accumulation_steps)
+    full_total_steps = epochs * step_per_epoch
+    total_steps = min(args.max_steps, full_total_steps) if args.max_steps > 0 else full_total_steps
     warmup_steps = int(total_steps * 0.1)
+    print(
+        f"training schedule: grad_accum={gradient_accumulation_steps} | "
+        f"step_per_epoch={step_per_epoch} | total_optimizer_steps={total_steps}"
+    )
 
     scheduler = get_cosine_schedule_with_warmup(
         optimizer,
@@ -433,6 +557,9 @@ if __name__ == "__main__":
     )
 
 
+    optimizer_step_count = 0
+    stop_training = False
+
     for epoch in range(epochs):
 
         model.train()
@@ -441,6 +568,9 @@ if __name__ == "__main__":
 
         loop = tqdm(train_dataloader, desc=f"Training Epoch {epoch + 1}/{epochs}")
         for step, batch in enumerate(loop):
+            if args.max_steps > 0 and optimizer_step_count >= args.max_steps:
+                stop_training = True
+                break
             input_dict = dict_to_cuda(batch)
             output_dict = model.forward(images=input_dict["images"],
                                         images_clip=input_dict["images_clip"],
@@ -459,6 +589,7 @@ if __name__ == "__main__":
                                         contrast=args.ct_weight,
                                         ref_ids=input_dict["ref_ids"],
                                         epoch=epoch,
+                                        gate_only=args.gate_only,
                                         inference=False)
 
             loss = output_dict["loss"]
@@ -468,6 +599,15 @@ if __name__ == "__main__":
 
 
             if (step + 1) % gradient_accumulation_steps == 0:
+                optimizer_step_count += 1
+                if (
+                    args.log_gate_stats_every > 0
+                    and optimizer_step_count % args.log_gate_stats_every == 0
+                ):
+                    log_referent_gate_stats(
+                        optimizer_step_count,
+                        loss.item() * gradient_accumulation_steps,
+                    )
                 optimizer.step()
                 scheduler.step()
                 optimizer.zero_grad()
@@ -475,16 +615,33 @@ if __name__ == "__main__":
                 current_lr = scheduler.get_lr()[0]
                 loop.set_postfix(lr=current_lr, loss=running_loss / ((step + 1) / gradient_accumulation_steps))
 
-        print(f"  Epoch {epoch + 1}, Loss:{running_loss / ((step + 1) / gradient_accumulation_steps) :.4f}, Learning Rate:{scheduler.get_last_lr()[0]:.6f}")
+                if args.max_steps > 0 and optimizer_step_count >= args.max_steps:
+                    stop_training = True
+                    break
+
+        denom = max(1, optimizer_step_count)
+        print(f"  Epoch {epoch + 1}, Loss:{running_loss / denom :.4f}, Learning Rate:{scheduler.get_last_lr()[0]:.6f}")
 
 
         with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
-            f.write(f"Epoch {epoch}: running_loss {running_loss / len(train_dataloader) * gradient_accumulation_steps}  Learning Rate:{scheduler.get_last_lr()[0]:.6f}\n")
+            f.write(f"Epoch {epoch}: running_loss {running_loss / denom}  Learning Rate:{scheduler.get_last_lr()[0]:.6f}\n")
+
+        if stop_training:
+            print(f"stopped early at optimizer step {optimizer_step_count}")
+            break
 
 
     torch.save(model.state_dict(), os.path.join(args.checkpoint_root, f"{args.name}.pth"))
     print(f"trained model saved as {args.name}.pth")
 
+    if args.skip_eval_after_train:
+        print("skip_eval_after_train enabled: exiting after checkpoint save")
+        sys.exit(0)
+
+    if args.eval_train_only:
+        valuate(model, train_eval_dataloader, args, 'train_overfit')
+        sys.exit(0)
+
     # ---------------test on seen & unseen ------------------------------------------
     model.eval()
 
@@ -531,4 +688,4 @@ if __name__ == "__main__":
     print(f"\n  valuate on test_n_refer, metric: {total_metric/count}")
 
     with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
-        f.write(f"\n valuate on  test_n_refer:   metric {total_metric/count} \n")
+        f.write(f"\n valuate on  test_n_refer:   metric {total_metric/count} \n")

train_cached_gate.py

@@ -0,0 +1,439 @@
+import json
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+import transformers
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, Dataset, Subset
+from tqdm import tqdm
+
+from configs import args
+from decoder_invariance_check import build_model, set_seed
+from models.avs_model import dice_loss, sigmoid_ce_loss
+from utils import utility
+
+
+def _total_norm(values):
+    if not values:
+        return 0.0
+    return float(sum(v * v for v in values) ** 0.5)
+
+
+def collect_referent_gate_stats(model):
+    gate_modules = [(n, m) for n, m in model.named_modules() if n.endswith("referent_gate")]
+    proj_norms = []
+    gate_norms = []
+    proj_grad_norms = []
+    gate_grad_norms = []
+    alpha_tensors = []
+
+    for _, module in gate_modules:
+        proj_norms.append(module.proj.weight.detach().float().norm().item())
+        gate_norms.append(module.gate.weight.detach().float().norm().item())
+        if module.proj.weight.grad is not None:
+            proj_grad_norms.append(module.proj.weight.grad.detach().float().norm().item())
+        if module.gate.weight.grad is not None:
+            gate_grad_norms.append(module.gate.weight.grad.detach().float().norm().item())
+        if module.last_alpha is not None:
+            alpha_tensors.append(module.last_alpha.detach().float().reshape(-1))
+
+    stats = {
+        "modules": len(gate_modules),
+        "proj_norm": _total_norm(proj_norms),
+        "gate_norm": _total_norm(gate_norms),
+        "proj_grad_norm": _total_norm(proj_grad_norms),
+        "gate_grad_norm": _total_norm(gate_grad_norms),
+    }
+
+    if alpha_tensors:
+        alpha = torch.cat(alpha_tensors)
+        stats.update(
+            {
+                "alpha_mean": alpha.mean().item(),
+                "alpha_std": alpha.std(unbiased=False).item(),
+                "alpha_min": alpha.min().item(),
+                "alpha_max": alpha.max().item(),
+            }
+        )
+    else:
+        stats.update(
+            {
+                "alpha_mean": float("nan"),
+                "alpha_std": float("nan"),
+                "alpha_min": float("nan"),
+                "alpha_max": float("nan"),
+            }
+        )
+
+    return stats
+
+
+def zero_referent_gate(model):
+    with torch.no_grad():
+        for _, module in model.named_modules():
+            if not _.endswith("referent_gate"):
+                continue
+            module.gate.weight.zero_()
+            module.gate.bias.zero_()
+            module.proj.weight.zero_()
+            module.proj.bias.zero_()
+            module.last_alpha = None
+
+
+def referent_gate_state_dict(model):
+    return {
+        name: param.detach().cpu()
+        for name, param in model.state_dict().items()
+        if "referent_gate" in name
+    }
+
+
+def load_referent_gate_checkpoint(model, path):
+    checkpoint = torch.load(path, map_location="cpu")
+    if isinstance(checkpoint, dict) and checkpoint.get("type") == "referent_gate_only":
+        checkpoint = checkpoint["state_dict"]
+    gate_state = {k: v for k, v in checkpoint.items() if "referent_gate" in k}
+    if not gate_state:
+        raise RuntimeError(f"No referent_gate parameters found in {path}")
+    current = model.state_dict()
+    missing_shape = [
+        k
+        for k, v in gate_state.items()
+        if k not in current or tuple(current[k].shape) != tuple(v.shape)
+    ]
+    if missing_shape:
+        raise RuntimeError(f"Gate checkpoint has incompatible keys: {missing_shape[:5]}")
+    current.update(gate_state)
+    model.load_state_dict(current, strict=True)
+    print(f"loaded referent gate checkpoint: {path} ({len(gate_state)} tensors)")
+
+
+def log_gate_stats(model, step, loss_value, batch_metrics=None):
+    stats = collect_referent_gate_stats(model)
+    metric_text = ""
+    if batch_metrics is not None:
+        metric_text = (
+            f"batch_miou={batch_metrics['miou']:.4f} "
+            f"batch_fscore={batch_metrics['fscore']:.4f} "
+        )
+    message = (
+        f"gate_stats step={step} "
+        f"loss={loss_value:.6f} "
+        f"{metric_text}"
+        f"proj_norm={stats['proj_norm']:.6f} "
+        f"gate_norm={stats['gate_norm']:.6f} "
+        f"proj_grad_norm={stats['proj_grad_norm']:.6f} "
+        f"gate_grad_norm={stats['gate_grad_norm']:.6f} "
+        f"alpha_mean={stats['alpha_mean']:.4f} "
+        f"alpha_std={stats['alpha_std']:.4f} "
+        f"alpha_min={stats['alpha_min']:.4f} "
+        f"alpha_max={stats['alpha_max']:.4f}"
+    )
+    print(message)
+    os.makedirs(args.log_root, exist_ok=True)
+    with open(os.path.join(args.log_root, f"{args.name}.txt"), "a") as f:
+        f.write(message + "\n")
+
+
+class CachedQDataset(Dataset):
+    def __init__(self, split, cfg):
+        self.split = split
+        self.cfg = cfg
+        self.root = os.path.join(cfg.cache_root, split)
+        self.index_path = os.path.join(self.root, "index.jsonl")
+        if not os.path.exists(self.index_path):
+            raise FileNotFoundError(f"Missing cache index: {self.index_path}")
+        with open(self.index_path) as f:
+            self.rows = [json.loads(line) for line in f if line.strip()]
+
+    def __len__(self):
+        return len(self.rows)
+
+    def _load_masks(self, vid, fids):
+        masks = []
+        for fid in fids:
+            frames = []
+            for frame_idx in range(self.cfg.frame_n):
+                path = os.path.join(
+                    self.cfg.data_dir,
+                    "gt_mask",
+                    vid,
+                    f"fid_{int(fid)}",
+                    f"0000{frame_idx}.png",
+                )
+                mask = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+                if mask is None:
+                    raise FileNotFoundError(path)
+                frames.append(torch.as_tensor(mask > 0, dtype=torch.float32))
+            masks.append(torch.stack(frames, dim=0))
+        return torch.stack(masks, dim=0)
+
+    def __getitem__(self, idx):
+        row = self.rows[idx]
+        cache = torch.load(os.path.join(self.root, row["path"]), map_location="cpu")
+        vid = cache["vid"]
+        return {
+            "sample_idx": cache["sample_idx"],
+            "vid": vid,
+            "refs": cache["refs"],
+            "fids": cache["fids"],
+            "q": cache["q"].float(),
+            "image_embeddings": torch.load(
+                os.path.join(self.cfg.data_dir, "image_embed", f"{vid}.pt"),
+                map_location="cpu",
+            ).float(),
+            "gt_masks": self._load_masks(vid, cache["fids"]),
+            "resize": tuple(cache["resize"]),
+            "orgsize": tuple(cache["orgsize"]),
+        }
+
+
+def collate_cached(batch):
+    return batch
+
+
+def decode_batch(visual_model, batch, device):
+    image_pe = visual_model.prompt_encoder.get_dense_pe()
+    frame_qs = []
+    frame_image_embeddings = []
+    prompt_spans = []
+
+    for sample_idx, sample in enumerate(batch):
+        q = sample["q"].to(device=device, dtype=torch.float32)
+        image_embeddings = sample["image_embeddings"].to(device=device, dtype=torch.float32)
+        frames = image_embeddings.shape[0]
+        for prompt_idx in range(q.shape[0]):
+            start = len(frame_qs) * frames
+            frame_qs.append(q[prompt_idx].unsqueeze(0).expand(frames, -1))
+            frame_image_embeddings.append(image_embeddings)
+            prompt_spans.append((sample_idx, prompt_idx, start, start + frames))
+
+    if not frame_qs:
+        raise RuntimeError("No cached prompts were provided for decoding.")
+
+    frame_qs = torch.cat(frame_qs, dim=0)
+    frame_image_embeddings = torch.cat(frame_image_embeddings, dim=0)
+    sparse_embeddings, dense_embeddings = visual_model.prompt_encoder(
+        points=None,
+        boxes=None,
+        masks=None,
+        text_embeds=frame_qs.unsqueeze(1),
+    )
+    sparse_embeddings = sparse_embeddings.to(frame_qs.dtype)
+    dense_embeddings = dense_embeddings.to(frame_qs.dtype)
+
+    low_res_masks = visual_model.mask_decoder.forward_modified_v3(
+        image_embeddings=frame_image_embeddings,
+        image_pe=image_pe,
+        sparse_prompt_embeddings=sparse_embeddings,
+        dense_prompt_embeddings=dense_embeddings,
+    ).unsqueeze(1)
+
+    pred_by_sample = [[] for _ in batch]
+    for sample_idx, _, start, end in prompt_spans:
+        sample = batch[sample_idx]
+        pred_mask = visual_model.postprocess_masks(
+            low_res_masks[start:end],
+            input_size=sample["resize"],
+            original_size=sample["orgsize"],
+        )
+        pred_by_sample[sample_idx].append(pred_mask.squeeze(1))
+
+    return [torch.stack(pred_masks, dim=0) for pred_masks in pred_by_sample]
+
+
+def decode_sample(visual_model, sample, device):
+    return decode_batch(visual_model, [sample], device)[0]
+
+
+def compute_mask_loss(pred_masks, gt_masks):
+    mask_bce_loss = 0.0
+    mask_dice_loss = 0.0
+    num_masks = 0
+
+    for pred_mask, gt_mask in zip(pred_masks, gt_masks):
+        gt_mask = gt_mask.to(device=pred_mask.device, dtype=pred_mask.dtype)
+        num_seg, frames, height, width = gt_mask.shape
+        gt_flat = gt_mask.view(num_seg * frames, height, width)
+        pred_flat = pred_mask.view(num_seg * frames, height, width)
+
+        mask_bce_loss = mask_bce_loss + (
+            sigmoid_ce_loss(pred_flat, gt_flat, num_masks=gt_flat.shape[0])
+            * gt_flat.shape[0]
+        )
+        mask_dice_loss = mask_dice_loss + (
+            dice_loss(pred_flat, gt_flat, num_masks=gt_flat.shape[0])
+            * gt_flat.shape[0]
+        )
+        num_masks += gt_flat.shape[0]
+
+    mask_bce_loss = 2.0 * mask_bce_loss / (num_masks + 1e-8)
+    mask_dice_loss = 0.5 * mask_dice_loss / (num_masks + 1e-8)
+    return mask_bce_loss + mask_dice_loss
+
+
+def compute_batch_metrics(pred_masks, gt_masks):
+    total_iou = 0.0
+    total_fscore = 0.0
+    count = 0
+    for pred_mask, gt_mask in zip(pred_masks, gt_masks):
+        gt_mask = gt_mask.to(device=pred_mask.device, dtype=pred_mask.dtype)
+        num_seg, frames = pred_mask.shape[:2]
+        weight = num_seg * frames
+        total_iou += utility.mask_iou(pred_mask.detach().float(), gt_mask.float()) * weight
+        total_fscore += utility.Eval_Fmeasure(pred_mask.detach().float(), gt_mask.float(), None) * weight
+        count += weight
+    return {
+        "miou": total_iou / max(1, count),
+        "fscore": total_fscore / max(1, count),
+    }
+
+
+def evaluate(model, loader):
+    model.eval()
+    visual_model = model.get_model().visual_model
+    total_iou = 0.0
+    total_fscore = 0.0
+    total_null = 0.0
+    count = 0
+
+    with torch.no_grad():
+        for batch in tqdm(loader, desc=f"Cached eval {args.cache_split}"):
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                batch_pred = decode_batch(visual_model, batch, "cuda")
+            for sample, pred in zip(batch, batch_pred):
+                gt = sample["gt_masks"].to(device=pred.device, dtype=pred.dtype)
+                num_seg, frames = pred.shape[:2]
+                weight = num_seg * frames
+                if args.cache_split == "test_n":
+                    total_null += float(utility.metric_s_for_null(pred.float())) * weight
+                else:
+                    total_iou += utility.mask_iou(pred.float(), gt.float()) * weight
+                    total_fscore += utility.Eval_Fmeasure(pred.float(), gt.float(), None) * weight
+                count += weight
+
+    if count == 0:
+        raise RuntimeError("No cached samples were evaluated.")
+
+    if args.cache_split == "test_n":
+        print(f"cached valuate on test_n_refer, metric: {total_null / count}")
+    else:
+        print(
+            f"cached valuate on {args.cache_split}: "
+            f"miou: {total_iou / count} fscore: {total_fscore / count}"
+        )
+
+
+def train(model, loader):
+    if args.disable_gate:
+        raise ValueError("--disable_gate is only valid with --eval_only")
+
+    for p in model.parameters():
+        p.requires_grad = False
+    for name, p in model.named_parameters():
+        if "referent_gate" in name:
+            p.requires_grad = True
+
+    gate_params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = AdamW(gate_params, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
+    stats = collect_referent_gate_stats(model)
+    print(
+        "cached gate init: "
+        f"modules={stats['modules']} "
+        f"proj_norm={stats['proj_norm']:.6f} "
+        f"gate_norm={stats['gate_norm']:.6f} "
+        f"trainable_params={sum(p.numel() for p in gate_params)}"
+    )
+
+    visual_model = model.get_model().visual_model
+    step = 0
+    for epoch in range(args.epochs):
+        model.train()
+        order_loader = loader
+        for batch in tqdm(order_loader, desc=f"Cached gate train {epoch + 1}/{args.epochs}"):
+            if args.max_steps > 0 and step >= args.max_steps:
+                break
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                pred_masks = decode_batch(visual_model, batch, "cuda")
+            gt_masks = [sample["gt_masks"] for sample in batch]
+
+            loss = compute_mask_loss(pred_masks, gt_masks)
+            optimizer.zero_grad()
+            loss.backward()
+            step += 1
+            if args.log_gate_stats_every > 0 and step % args.log_gate_stats_every == 0:
+                batch_metrics = compute_batch_metrics(pred_masks, gt_masks)
+                log_gate_stats(model, step, loss.item(), batch_metrics)
+            optimizer.step()
+
+        if args.max_steps > 0 and step >= args.max_steps:
+            print(f"stopped early at cached optimizer step {step}")
+            break
+
+    os.makedirs(args.checkpoint_root, exist_ok=True)
+    ckpt_path = os.path.join(args.checkpoint_root, f"{args.name}.pth")
+    if args.save_gate_only:
+        torch.save(
+            {
+                "type": "referent_gate_only",
+                "base_model": args.saved_model,
+                "state_dict": referent_gate_state_dict(model),
+            },
+            ckpt_path,
+        )
+    else:
+        torch.save(model.state_dict(), ckpt_path)
+    print(f"cached gate model saved as {ckpt_path}")
+
+
+def main():
+    set_seed(42)
+    random.seed(42)
+    np.random.seed(42)
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        args.mllm,
+        cache_dir=None,
+        model_max_length=2048,
+        padding_side="right",
+        use_fast=False,
+    )
+    tokenizer.pad_token = tokenizer.unk_token
+    tokenizer.add_tokens("[SEG]")
+    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+
+    dataset = CachedQDataset(args.cache_split, args)
+    if args.overfit_samples > 0:
+        n = min(args.overfit_samples, len(dataset))
+        dataset = Subset(dataset, list(range(n)))
+        print(f"cached overfit_samples enabled: using first {n} samples")
+
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=not args.eval_only,
+        num_workers=4,
+        collate_fn=collate_cached,
+    )
+
+    model = build_model(tokenizer, seg_token_idx)
+    if args.gate_checkpoint:
+        load_referent_gate_checkpoint(model, args.gate_checkpoint)
+    if args.disable_gate:
+        zero_referent_gate(model)
+        print("disable_gate enabled: referent gate forced to identity")
+
+    if args.eval_only:
+        evaluate(model, loader)
+        return
+
+    train(model, loader)
+    if not args.skip_eval_after_train:
+        evaluate(model, loader)
+
+
+if __name__ == "__main__":
+    main()