提交其余项目实验变更
This commit is contained in:
@@ -1 +1 @@
|
||||
"""Q2 robustness and Q3 explanation-selection experiments."""
|
||||
"""Q2 multimodal emotion-recognition experiments."""
|
||||
|
||||
@@ -0,0 +1,614 @@
|
||||
"""Score the frozen EarlyConcat and MoFE checkpoints using the math-Q2 protocol.
|
||||
|
||||
This script performs no training and selects no models. It evaluates the saved
|
||||
three-seed checkpoints on the official labeled test split once, and reuses the
|
||||
fixed 42-scenario validation-mask audit as the controlled-missingness protocol.
|
||||
All generated files stay under deep_learning/Q2/outputs/followups/.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from sklearn.metrics import accuracy_score, f1_score, mean_absolute_error, mean_squared_error
|
||||
from torch import nn
|
||||
|
||||
from .data import (
|
||||
ATTACHMENT2,
|
||||
MODALITIES,
|
||||
RobustStats,
|
||||
Split,
|
||||
_ids_and_targets,
|
||||
_text_mask,
|
||||
_unpickle,
|
||||
apply_robust_stats,
|
||||
fit_robust_stats,
|
||||
load_aligned,
|
||||
)
|
||||
from .models import AlignedFusionModel
|
||||
from .mofe import MixtureOfFusionExperts
|
||||
from .train_mofe import MODEL_CONFIG, _predict, _device_for, EARLYCONCAT, MOFE7_MLP
|
||||
|
||||
|
||||
Q2_ROOT = Path(__file__).resolve().parents[1]
|
||||
REPO_ROOT = Q2_ROOT.parents[1]
|
||||
REFERENCE_DIR = Q2_ROOT / "outputs" / "followups" / "R01_selected_model_reevaluation"
|
||||
OUTPUT_DIR = Q2_ROOT / "outputs" / "followups" / "R02_math_protocol_evaluation"
|
||||
SEEDS = (42, 3407, 2026)
|
||||
BOOTSTRAP_REPS = 1000
|
||||
TEST_BOOTSTRAP_SEED = 20260925
|
||||
AURC_BOOTSTRAP_SEED = 20260926
|
||||
SCENARIO_SEED = 20261833
|
||||
METHODS = (EARLYCONCAT, MOFE7_MLP)
|
||||
CURVE_MODES = ("single", "sync", "partial", "async")
|
||||
CURVE_RATES = (0.0, 0.1, 0.3, 0.5, 0.7)
|
||||
|
||||
|
||||
def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||
if not rows:
|
||||
return
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fields = list(dict.fromkeys(key for row in rows for key in row))
|
||||
with path.open("w", newline="", encoding="utf-8-sig") as stream:
|
||||
writer = csv.DictWriter(stream, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def split_from_part(part: dict[str, Any]) -> Split:
|
||||
xs = tuple(np.asarray(part[name], dtype=np.float32) for name in MODALITIES)
|
||||
masks = [
|
||||
_text_mask(part),
|
||||
np.any(np.isfinite(xs[1]) & (xs[1] != 0), axis=-1),
|
||||
np.any(np.isfinite(xs[2]) & (xs[2] != 0), axis=-1),
|
||||
]
|
||||
ids, y_cls, y_reg = _ids_and_targets(part)
|
||||
return Split(xs, np.stack(masks, axis=-1), y_cls, y_reg, ids)
|
||||
|
||||
|
||||
def load_splits(feature_path: Path) -> dict[str, Split]:
|
||||
raw = _unpickle(feature_path)
|
||||
usual = load_aligned(feature_path)
|
||||
splits = {"train": usual["train"], "valid": usual["valid"], "test": split_from_part(raw["test"])}
|
||||
groups = {
|
||||
name: {sample_id.split("$_$", 1)[0] for sample_id in split.ids}
|
||||
for name, split in splits.items()
|
||||
}
|
||||
for first, second in (("train", "valid"), ("train", "test"), ("valid", "test")):
|
||||
overlap = groups[first] & groups[second]
|
||||
if overlap:
|
||||
raise ValueError(f"official {first}/{second} source-video groups overlap: {len(overlap)}")
|
||||
for name, split in splits.items():
|
||||
expected = np.where(split.y_reg < 0, 0, np.where(split.y_reg == 0, 1, 2))
|
||||
if not np.array_equal(expected, split.y_cls):
|
||||
raise ValueError(f"{name}: classification labels disagree with strict sign of regression labels")
|
||||
return splits
|
||||
|
||||
|
||||
def metrics(split: Split, logits: np.ndarray, intensity: np.ndarray, indices: np.ndarray | None = None) -> dict[str, float]:
|
||||
if indices is None:
|
||||
indices = np.arange(split.n)
|
||||
y_cls = split.y_cls[indices]
|
||||
y_reg = split.y_reg[indices]
|
||||
pred_cls = np.asarray(logits)[indices].argmax(axis=-1)
|
||||
pred_reg = np.clip(np.asarray(intensity).reshape(-1)[indices], -3.0, 3.0)
|
||||
return {
|
||||
"accuracy": float(accuracy_score(y_cls, pred_cls)),
|
||||
"macro_f1": float(f1_score(y_cls, pred_cls, labels=[0, 1, 2], average="macro", zero_division=0)),
|
||||
"mae": float(mean_absolute_error(y_reg, pred_reg)),
|
||||
"rmse": float(math.sqrt(mean_squared_error(y_reg, pred_reg))),
|
||||
"pearson": float(np.corrcoef(y_reg, pred_reg)[0, 1]) if np.std(y_reg) > 0 and np.std(pred_reg) > 0 else float("nan"),
|
||||
}
|
||||
|
||||
|
||||
def load_model(method: str, seed: int, dims: tuple[int, int, int], device: torch.device) -> nn.Module:
|
||||
if method == EARLYCONCAT:
|
||||
checkpoint = REFERENCE_DIR / "models" / "baselines" / "concat" / f"seed_{seed}" / "model_best.pt"
|
||||
model: nn.Module = AlignedFusionModel("concat", dims=dims).to(device)
|
||||
state = torch.load(checkpoint, map_location=device, weights_only=False)
|
||||
if state.get("kind") != "concat" or int(state.get("seed", -1)) != seed:
|
||||
raise ValueError(f"unexpected EarlyConcat checkpoint: {checkpoint}")
|
||||
elif method == MOFE7_MLP:
|
||||
checkpoint = REFERENCE_DIR / "models" / MOFE7_MLP / f"seed_{seed}" / "model_best.pt"
|
||||
model = MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device)
|
||||
state = torch.load(checkpoint, map_location=device, weights_only=False)
|
||||
if state.get("config") != MODEL_CONFIG or int(state.get("seed", -1)) != seed:
|
||||
raise ValueError(f"unexpected MoFE checkpoint: {checkpoint}")
|
||||
else:
|
||||
raise ValueError(f"unknown model {method}")
|
||||
if tuple(state.get("dims", ())) != dims:
|
||||
raise ValueError(f"feature dimensions do not match checkpoint: {checkpoint}")
|
||||
model.load_state_dict(state["state_dict"])
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def best_interval(visible: np.ndarray, wanted: int, cap: int, location: str, rng: np.random.Generator) -> tuple[int, int] | None:
|
||||
steps = len(visible)
|
||||
candidates: list[tuple[int, int, int, int]] = []
|
||||
for left in range(steps):
|
||||
hits = 0
|
||||
for right in range(left, steps):
|
||||
hits += int(visible[right])
|
||||
count = min(hits, cap)
|
||||
if count:
|
||||
candidates.append((abs(count - wanted), right - left + 1, left, right))
|
||||
if not candidates:
|
||||
return None
|
||||
best = min((error, span) for error, span, _, _ in candidates)
|
||||
tied = [(left, right) for error, span, left, right in candidates if (error, span) == best]
|
||||
if location == "start":
|
||||
return min(tied, key=lambda pair: (pair[0], pair[1]))
|
||||
if location == "end":
|
||||
return max(tied, key=lambda pair: (pair[1], pair[0]))
|
||||
if location == "middle":
|
||||
center = (steps - 1) / 2
|
||||
return min(tied, key=lambda pair: (abs((pair[0] + pair[1]) / 2 - center), pair[0]))
|
||||
if location != "random":
|
||||
raise ValueError(f"unknown interval location: {location}")
|
||||
return tied[int(rng.integers(0, len(tied)))]
|
||||
|
||||
|
||||
def spread_short_spans(visible: np.ndarray, wanted: int, cap: int) -> np.ndarray:
|
||||
positions = np.flatnonzero(visible)
|
||||
count = min(int(wanted), int(cap), len(positions))
|
||||
chosen = np.zeros(len(visible), dtype=bool)
|
||||
if count <= 0:
|
||||
return chosen
|
||||
n_spans = min(3, count)
|
||||
chunks = np.array_split(positions, n_spans)
|
||||
allocations = [count // n_spans + int(i < count % n_spans) for i in range(n_spans)]
|
||||
for chunk, amount in zip(chunks, allocations):
|
||||
if amount <= 0 or len(chunk) == 0:
|
||||
continue
|
||||
amount = min(amount, len(chunk))
|
||||
start = max(0, (len(chunk) - amount) // 2)
|
||||
chosen[chunk[start:start + amount]] = True
|
||||
return chosen
|
||||
|
||||
|
||||
def continuous_mask(
|
||||
original: np.ndarray,
|
||||
rate: float,
|
||||
mode: str,
|
||||
rng: np.random.Generator,
|
||||
*,
|
||||
modalities: tuple[int, ...] | None = None,
|
||||
location: str = "random",
|
||||
span_structure: str = "long",
|
||||
) -> np.ndarray:
|
||||
"""Reproduce math/Q2 continuous masking on this model's observed positions."""
|
||||
observed = np.asarray(original, dtype=bool)
|
||||
result = observed.copy()
|
||||
if rate <= 0 or mode == "none":
|
||||
return result
|
||||
steps, modality_count = observed.shape
|
||||
present = [m for m in range(modality_count) if observed[:, m].any()]
|
||||
if not present:
|
||||
return result
|
||||
if modalities is not None:
|
||||
selected = [int(m) for m in modalities if int(m) in present]
|
||||
if not selected:
|
||||
return result
|
||||
elif mode == "single":
|
||||
selected = [int(rng.choice(present))]
|
||||
elif mode in {"sync", "partial", "async"}:
|
||||
if len(present) == 1:
|
||||
selected = present
|
||||
else:
|
||||
count = int(rng.integers(2, min(3, len(present)) + 1))
|
||||
selected = sorted(int(v) for v in rng.choice(present, size=count, replace=False))
|
||||
else:
|
||||
raise ValueError(f"unknown mask mode: {mode}")
|
||||
|
||||
def max_hide(modality: int) -> int:
|
||||
count = int(observed[:, modality].sum())
|
||||
keep = max(1, int(math.ceil(0.2 * count)))
|
||||
return max(0, count - keep)
|
||||
|
||||
target = {m: min(max_hide(m), int(round(rate * int(observed[:, m].sum())))) for m in selected}
|
||||
if mode == "sync":
|
||||
span = max(1, int(round(rate * steps)))
|
||||
if location == "start":
|
||||
left = 0
|
||||
elif location == "end":
|
||||
left = steps - span
|
||||
elif location == "middle":
|
||||
left = (steps - span) // 2
|
||||
elif location == "random":
|
||||
left = int(rng.integers(0, max(1, steps - span + 1)))
|
||||
else:
|
||||
raise ValueError(f"unknown interval location: {location}")
|
||||
right = min(steps - 1, left + span - 1)
|
||||
for m in selected:
|
||||
candidates = np.flatnonzero(observed[left:right + 1, m]) + left
|
||||
amount = min(len(candidates), max_hide(m), target[m])
|
||||
if amount:
|
||||
offset = 0 if location != "end" else len(candidates) - amount
|
||||
result[candidates[max(0, offset):max(0, offset) + amount], m] = False
|
||||
else:
|
||||
common_span = max(1, int(round(rate * steps)))
|
||||
for rank, m in enumerate(selected):
|
||||
wanted = target[m]
|
||||
if wanted <= 0:
|
||||
continue
|
||||
cap = max_hide(m)
|
||||
if span_structure == "multi_short":
|
||||
hide = spread_short_spans(observed[:, m], wanted, cap)
|
||||
elif span_structure != "long":
|
||||
raise ValueError(f"unknown span structure: {span_structure}")
|
||||
elif mode == "single" and location != "random":
|
||||
# Place a contiguous block at the requested relative location
|
||||
# among observed positions, while keeping the selected-source
|
||||
# missing amount fixed. This avoids treating padding as time.
|
||||
interval = best_interval(observed[:, m], wanted, cap, location, rng)
|
||||
hide = np.zeros(steps, dtype=bool)
|
||||
if interval is not None:
|
||||
left, right = interval
|
||||
candidates = np.flatnonzero(observed[left:right + 1, m]) + left
|
||||
amount = min(len(candidates), wanted, cap)
|
||||
if amount:
|
||||
offset = 0 if location != "end" else len(candidates) - amount
|
||||
hide[candidates[max(0, offset):max(0, offset) + amount]] = True
|
||||
elif mode in {"partial", "async"}:
|
||||
if mode == "partial":
|
||||
base_left = int(rng.integers(0, max(1, steps - common_span + 1))) if location == "random" else (
|
||||
0 if location == "start" else steps - common_span if location == "end" else (steps - common_span) // 2
|
||||
)
|
||||
offset = int(round(rank * common_span * 0.5))
|
||||
else:
|
||||
base_left = 0 if location == "random" else (
|
||||
0 if location == "start" else steps - common_span if location == "end" else (steps - common_span) // 2
|
||||
)
|
||||
available = max(1, steps - common_span + 1)
|
||||
offsets = np.rint(np.linspace(0, max(0, available - 1), len(selected))).astype(int)
|
||||
if location == "random":
|
||||
rng.shuffle(offsets)
|
||||
offset = int(offsets[rank])
|
||||
left = min(max(0, base_left + offset), max(0, steps - common_span))
|
||||
right = min(steps - 1, left + common_span - 1)
|
||||
hide = np.zeros(steps, dtype=bool)
|
||||
candidates = np.flatnonzero(observed[left:right + 1, m]) + left
|
||||
amount = min(len(candidates), wanted, cap)
|
||||
if amount:
|
||||
hide[candidates[:amount]] = True
|
||||
else:
|
||||
interval = best_interval(observed[:, m], wanted, cap, location, rng)
|
||||
hide = np.zeros(steps, dtype=bool)
|
||||
if interval is not None:
|
||||
left, right = interval
|
||||
candidates = np.flatnonzero(observed[left:right + 1, m]) + left
|
||||
amount = min(len(candidates), wanted, cap)
|
||||
if amount:
|
||||
offset = 0 if location != "end" else len(candidates) - amount
|
||||
hide[candidates[max(0, offset):max(0, offset) + amount]] = True
|
||||
result[hide, m] = False
|
||||
return result
|
||||
|
||||
|
||||
def scenario_seed(seed: int, sample_id: str, key: str) -> int:
|
||||
return int.from_bytes(hashlib.sha256(f"{seed}:{sample_id}:{key}".encode()).digest()[:8], "little")
|
||||
|
||||
|
||||
def make_scenarios(valid: Split, seed: int = SCENARIO_SEED) -> dict[str, np.ndarray]:
|
||||
scenarios = {"0.0/none": valid.mask.copy()}
|
||||
for rate in CURVE_RATES[1:]:
|
||||
for mode in CURVE_MODES:
|
||||
key = f"{rate:.1f}/{mode}"
|
||||
scenarios[key] = np.stack([
|
||||
continuous_mask(mask, rate, mode, np.random.default_rng(scenario_seed(seed, sample_id, key)))
|
||||
for sample_id, mask in zip(valid.ids, valid.mask)
|
||||
])
|
||||
modality_sets = (((0,), "T"), ((1,), "A"), ((2,), "V"), ((0, 1), "TA"), ((0, 2), "TV"), ((1, 2), "AV"), ((0, 1, 2), "TAV"))
|
||||
for selected, label in modality_sets:
|
||||
key = f"0.3/modality_{label}"
|
||||
scenarios[key] = np.stack([
|
||||
continuous_mask(mask, 0.3, "sync", np.random.default_rng(scenario_seed(seed, sample_id, key)), modalities=selected)
|
||||
for sample_id, mask in zip(valid.ids, valid.mask)
|
||||
])
|
||||
for modality_index, label in enumerate(("T", "A", "V")):
|
||||
for location in ("start", "middle", "end"):
|
||||
key = f"0.3/location_{location}_{label}"
|
||||
scenarios[key] = np.stack([
|
||||
continuous_mask(mask, 0.3, "single", np.random.default_rng(scenario_seed(seed, sample_id, key)), modalities=(modality_index,), location=location)
|
||||
for sample_id, mask in zip(valid.ids, valid.mask)
|
||||
])
|
||||
for structure in ("long", "multi_short"):
|
||||
key = f"0.3/span_{structure}_{label}"
|
||||
scenarios[key] = np.stack([
|
||||
continuous_mask(mask, 0.3, "single", np.random.default_rng(scenario_seed(seed, sample_id, key)), modalities=(modality_index,), span_structure=structure)
|
||||
for sample_id, mask in zip(valid.ids, valid.mask)
|
||||
])
|
||||
for mode in ("sync", "partial", "async"):
|
||||
key = f"0.3/synchrony_{mode}"
|
||||
scenarios[key] = np.stack([
|
||||
continuous_mask(mask, 0.3, mode, np.random.default_rng(scenario_seed(seed, sample_id, key)), modalities=(0, 1, 2))
|
||||
for sample_id, mask in zip(valid.ids, valid.mask)
|
||||
])
|
||||
return scenarios
|
||||
|
||||
|
||||
def actual_additional_rates(base: np.ndarray, scenarios: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
|
||||
result = {}
|
||||
observed = base.sum(axis=1)
|
||||
for scenario, current in scenarios.items():
|
||||
newly_hidden = base & ~current
|
||||
hidden_count = newly_hidden.sum(axis=1)
|
||||
by_modality = np.divide(
|
||||
hidden_count,
|
||||
observed,
|
||||
out=np.full(hidden_count.shape, np.nan, dtype=np.float64),
|
||||
where=observed > 0,
|
||||
)
|
||||
result[scenario] = np.nanmean(by_modality, axis=1)
|
||||
return result
|
||||
|
||||
|
||||
def aurc_from_curve(rates: list[float], maes: list[float]) -> float:
|
||||
order = np.argsort(np.asarray(rates), kind="stable")
|
||||
x = np.asarray(rates, dtype=np.float64)[order]
|
||||
y = np.asarray(maes, dtype=np.float64)[order]
|
||||
unique_x, inverse = np.unique(x, return_inverse=True)
|
||||
unique_y = np.asarray([y[inverse == i].mean() for i in range(len(unique_x))])
|
||||
if len(unique_x) <= 1 or unique_x[-1] <= 0:
|
||||
return float(maes[0])
|
||||
return float(np.trapezoid(unique_y, unique_x) / unique_x[-1])
|
||||
|
||||
|
||||
def curve_scenarios(mode: str) -> list[str]:
|
||||
return ["0.0/none"] + [f"{rate:.1f}/{mode}" for rate in CURVE_RATES[1:]]
|
||||
|
||||
|
||||
def group_indices(ids: list[str]) -> tuple[list[str], dict[str, np.ndarray]]:
|
||||
groups = sorted({sample_id.split("$_$", 1)[0] for sample_id in ids})
|
||||
mapping = {group: np.flatnonzero(np.asarray([x.split("$_$", 1)[0] == group for x in ids])) for group in groups}
|
||||
return groups, mapping
|
||||
|
||||
|
||||
def bootstrap_clean_test(
|
||||
split: Split,
|
||||
preds: dict[tuple[str, int], dict[str, np.ndarray]],
|
||||
) -> list[dict[str, Any]]:
|
||||
groups, mapping = group_indices(split.ids)
|
||||
rng = np.random.default_rng(TEST_BOOTSTRAP_SEED)
|
||||
draws: dict[str, list[float]] = defaultdict(list)
|
||||
for _ in range(BOOTSTRAP_REPS):
|
||||
chosen = rng.choice(groups, size=len(groups), replace=True)
|
||||
indices = np.concatenate([mapping[group] for group in chosen])
|
||||
per_method = {}
|
||||
for method in METHODS:
|
||||
per_seed = [metrics(split, preds[(method, seed)]["logits"], preds[(method, seed)]["intensity"], indices) for seed in SEEDS]
|
||||
per_method[method] = {key: float(np.mean([row[key] for row in per_seed])) for key in per_seed[0]}
|
||||
for metric in per_method[EARLYCONCAT]:
|
||||
draws[metric].append(per_method[MOFE7_MLP][metric] - per_method[EARLYCONCAT][metric])
|
||||
rows = []
|
||||
for metric, values in draws.items():
|
||||
rows.append({
|
||||
"comparison": "MoFE-7 + MLP Router minus EarlyConcat + BiGRU",
|
||||
"metric": metric,
|
||||
"delta_mean_over_seeds": float(np.mean([r[metric] for r in [
|
||||
metrics(split, preds[(MOFE7_MLP, seed)]["logits"], preds[(MOFE7_MLP, seed)]["intensity"])
|
||||
for seed in SEEDS
|
||||
]]) - np.mean([r[metric] for r in [
|
||||
metrics(split, preds[(EARLYCONCAT, seed)]["logits"], preds[(EARLYCONCAT, seed)]["intensity"])
|
||||
for seed in SEEDS
|
||||
]])),
|
||||
"bootstrap_ci_2p5": float(np.quantile(values, 0.025)),
|
||||
"bootstrap_ci_97p5": float(np.quantile(values, 0.975)),
|
||||
"bootstrap_probability_delta_gt_0": float(np.mean(np.asarray(values) > 0)),
|
||||
"replicates": BOOTSTRAP_REPS,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
"seed": TEST_BOOTSTRAP_SEED,
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def bootstrap_aurc(
|
||||
valid: Split,
|
||||
predictions: dict[tuple[str, int, str], dict[str, np.ndarray]],
|
||||
scenarios: dict[str, np.ndarray],
|
||||
rates_by_sample: dict[str, np.ndarray],
|
||||
) -> list[dict[str, Any]]:
|
||||
groups, mapping = group_indices(valid.ids)
|
||||
rng = np.random.default_rng(AURC_BOOTSTRAP_SEED)
|
||||
delta_by_mode: dict[str, list[float]] = {mode: [] for mode in CURVE_MODES}
|
||||
for _ in range(BOOTSTRAP_REPS):
|
||||
chosen = rng.choice(groups, size=len(groups), replace=True)
|
||||
indices = np.concatenate([mapping[group] for group in chosen])
|
||||
for mode in CURVE_MODES:
|
||||
keys = curve_scenarios(mode)
|
||||
model_aucs: dict[str, list[float]] = {method: [] for method in METHODS}
|
||||
for method in METHODS:
|
||||
for seed in SEEDS:
|
||||
xs = [float(np.nanmean(rates_by_sample[key][indices])) for key in keys]
|
||||
ys = [float(np.abs(valid.y_reg[indices] - predictions[(method, seed, key)]["intensity"][indices]).mean()) for key in keys]
|
||||
model_aucs[method].append(aurc_from_curve(xs, ys))
|
||||
delta_by_mode[mode].append(float(np.mean(model_aucs[MOFE7_MLP]) - np.mean(model_aucs[EARLYCONCAT])))
|
||||
point = {}
|
||||
for mode in CURVE_MODES:
|
||||
model_aucs = {}
|
||||
for method in METHODS:
|
||||
model_aucs[method] = []
|
||||
for seed in SEEDS:
|
||||
keys = curve_scenarios(mode)
|
||||
xs = [float(np.nanmean(rates_by_sample[key])) for key in keys]
|
||||
ys = [float(np.abs(valid.y_reg - predictions[(method, seed, key)]["intensity"]).mean()) for key in keys]
|
||||
model_aucs[method].append(aurc_from_curve(xs, ys))
|
||||
point[mode] = float(np.mean(model_aucs[MOFE7_MLP]) - np.mean(model_aucs[EARLYCONCAT]))
|
||||
rows = []
|
||||
for mode, values in delta_by_mode.items():
|
||||
rows.append({
|
||||
"mode": mode,
|
||||
"delta_aurc_mae_mofe_minus_earlyconcat": point[mode],
|
||||
"bootstrap_ci_2p5": float(np.quantile(values, 0.025)),
|
||||
"bootstrap_ci_97p5": float(np.quantile(values, 0.975)),
|
||||
"bootstrap_probability_delta_lt_0": float(np.mean(np.asarray(values) < 0)),
|
||||
"replicates": BOOTSTRAP_REPS,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
"seed": AURC_BOOTSTRAP_SEED,
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def run(device_name: str = "auto", batch_size: int = 64, masks_only: bool = False) -> None:
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
device = _device_for(device_name)
|
||||
torch.set_num_threads(4)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
feature_path = ATTACHMENT2 / "aligned_50.pkl"
|
||||
with (REFERENCE_DIR / "run_manifest.json").open("r", encoding="utf-8") as stream:
|
||||
reference_manifest = json.load(stream)
|
||||
if sha256(feature_path) != reference_manifest["feature_sha256"]:
|
||||
raise ValueError("current official feature file hash differs from the checkpoint evaluation manifest")
|
||||
|
||||
raw_splits = load_splits(feature_path)
|
||||
train = raw_splits["train"]
|
||||
valid = raw_splits["valid"]
|
||||
test = raw_splits["test"]
|
||||
scaler_path = REFERENCE_DIR / "aligned_robust_stats.npz"
|
||||
stats = RobustStats.load(scaler_path)
|
||||
computed = fit_robust_stats(train)
|
||||
scaler_diff = max(
|
||||
max(float(np.max(np.abs(a - b))) for a, b in zip(computed.center, stats.center)),
|
||||
max(float(np.max(np.abs(a - b))) for a, b in zip(computed.scale, stats.scale)),
|
||||
)
|
||||
if scaler_diff > 1e-6:
|
||||
raise ValueError(f"checkpoint scaler is not the train-only scaler (max difference {scaler_diff})")
|
||||
valid = apply_robust_stats(valid, stats)
|
||||
test = apply_robust_stats(test, stats)
|
||||
dims = tuple(x.shape[-1] for x in train.x)
|
||||
|
||||
if not masks_only:
|
||||
# Final, clean official-test evaluation; no retraining or selection occurs here.
|
||||
test_predictions: dict[tuple[str, int], dict[str, np.ndarray]] = {}
|
||||
test_rows: list[dict[str, Any]] = []
|
||||
for method in METHODS:
|
||||
for seed in SEEDS:
|
||||
model = load_model(method, seed, dims, device)
|
||||
prediction = _predict(model, test, test.mask, device, batch_size)
|
||||
test_predictions[(method, seed)] = prediction
|
||||
test_rows.append({"method": method, "seed": seed, "n_test": test.n, **metrics(test, prediction["logits"], prediction["intensity"])})
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
summary_rows = []
|
||||
for method in METHODS:
|
||||
subset = [row for row in test_rows if row["method"] == method]
|
||||
for metric in ("accuracy", "macro_f1", "mae", "rmse", "pearson"):
|
||||
values = [float(row[metric]) for row in subset]
|
||||
summary_rows.append({"method": method, "metric": metric, "mean": float(np.mean(values)), "sd_across_seeds": float(np.std(values, ddof=1))})
|
||||
write_csv(OUTPUT_DIR / "official_test_metrics_by_seed.csv", test_rows)
|
||||
write_csv(OUTPUT_DIR / "official_test_summary.csv", summary_rows)
|
||||
write_csv(OUTPUT_DIR / "official_test_paired_bootstrap.csv", bootstrap_clean_test(test, test_predictions))
|
||||
|
||||
# Reproduce the math-Q2 42-scenario design with a per-sample stable seed,
|
||||
# while applying it to the observation masks used to train these models.
|
||||
scenario_masks = make_scenarios(valid)
|
||||
rates_by_sample = actual_additional_rates(valid.mask, scenario_masks)
|
||||
condition_predictions: dict[tuple[str, int, str], dict[str, np.ndarray]] = {}
|
||||
condition_rows: list[dict[str, Any]] = []
|
||||
for method in METHODS:
|
||||
for seed in SEEDS:
|
||||
model = load_model(method, seed, dims, device)
|
||||
for scenario, masks in scenario_masks.items():
|
||||
prediction = _predict(model, valid, masks, device, batch_size)
|
||||
condition_predictions[(method, seed, scenario)] = prediction
|
||||
values = metrics(valid, prediction["logits"], prediction["intensity"])
|
||||
condition_rows.append({
|
||||
"method": method,
|
||||
"seed": seed,
|
||||
"scenario": scenario,
|
||||
"realized_additional_global_rate": float(np.nanmean(rates_by_sample[scenario])),
|
||||
"n_valid": valid.n,
|
||||
**values,
|
||||
})
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
write_csv(OUTPUT_DIR / "controlled_metrics_by_scenario.csv", condition_rows)
|
||||
|
||||
auc_rows: list[dict[str, Any]] = []
|
||||
for method in METHODS:
|
||||
for seed in SEEDS:
|
||||
for mode in CURVE_MODES:
|
||||
keys = curve_scenarios(mode)
|
||||
xs = [float(np.nanmean(rates_by_sample[key])) for key in keys]
|
||||
ys = [float(np.abs(valid.y_reg - condition_predictions[(method, seed, key)]["intensity"]).mean()) for key in keys]
|
||||
auc_rows.append({"method": method, "seed": seed, "mask_mode": mode, "aurc_mae": aurc_from_curve(xs, ys), "rates_realized": json.dumps(xs)})
|
||||
write_csv(OUTPUT_DIR / "aurc_mae_by_mode_seed.csv", auc_rows)
|
||||
auc_summary = []
|
||||
for method in METHODS:
|
||||
for mode in CURVE_MODES:
|
||||
values = [row["aurc_mae"] for row in auc_rows if row["method"] == method and row["mask_mode"] == mode]
|
||||
auc_summary.append({"method": method, "mask_mode": mode, "mean": float(np.mean(values)), "sd_across_seeds": float(np.std(values, ddof=1))})
|
||||
write_csv(OUTPUT_DIR / "aurc_mae_summary.csv", auc_summary)
|
||||
write_csv(OUTPUT_DIR / "aurc_mae_paired_bootstrap.csv", bootstrap_aurc(valid, condition_predictions, scenario_masks, rates_by_sample))
|
||||
|
||||
manifest = {
|
||||
"experiment": "Frozen EarlyConcat vs MoFE-7 evaluation under math/Q2 test protocol",
|
||||
"created_unix": time.time(),
|
||||
"device": str(device),
|
||||
"cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
|
||||
"feature_file": str(feature_path),
|
||||
"feature_sha256": sha256(feature_path),
|
||||
"representation": "official aligned_50 ordered positions; not physical-time bins",
|
||||
"train_valid_test_counts": {name: split.n for name, split in raw_splits.items()},
|
||||
"source_video_groups": {name: len({sample_id.split("$_$", 1)[0] for sample_id in split.ids}) for name, split in raw_splits.items()},
|
||||
"official_group_splits_disjoint": True,
|
||||
"test_evaluation": (
|
||||
"one final clean evaluation on official labeled test split; no training/model selection/calibration"
|
||||
if not masks_only else "test outputs preserved from the earlier single evaluation; no test prediction was rerun"
|
||||
),
|
||||
"test_prediction_performed_this_invocation": not masks_only,
|
||||
"seeds": list(SEEDS),
|
||||
"checkpoint_source": str(REFERENCE_DIR / "models"),
|
||||
"train_only_scaler": str(scaler_path),
|
||||
"scaler_max_abs_difference_from_train_refit": scaler_diff,
|
||||
"test_labels_used_for_training_or_selection": False,
|
||||
"controlled_missingness": {
|
||||
"scenario_seed": SCENARIO_SEED,
|
||||
"scenario_design": "math/Q2 42-scenario design regenerated on the Q2 models' BERT attention-mask base",
|
||||
"scenarios": len(scenario_masks),
|
||||
"AURC": "normalized trapezoidal area of MAE over realized equal-modality-weighted added missing rate, at 0/.1/.3/.5/.7 for single/sync/partial/async",
|
||||
},
|
||||
"bootstrap": {
|
||||
"replicates": BOOTSTRAP_REPS,
|
||||
"test_seed": TEST_BOOTSTRAP_SEED,
|
||||
"aurc_seed": AURC_BOOTSTRAP_SEED,
|
||||
"unit": "source video id",
|
||||
"paired": True,
|
||||
},
|
||||
}
|
||||
(OUTPUT_DIR / "run_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
||||
print(f"wrote math-protocol comparison to {OUTPUT_DIR}")
|
||||
print(f"n_test={test.n}; n_valid={valid.n}; device={device}; scenarios={len(scenario_masks)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--masks-only", action="store_true", help="Recompute validation mask scenarios without rerunning official-test inference")
|
||||
args = parser.parse_args()
|
||||
run(masks_only=args.masks_only)
|
||||
@@ -9,7 +9,7 @@ from .train_compare import _plot, _summary, _write_csv
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Rebuild Q2 summary tables from saved validation predictions")
|
||||
parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parents[1] / "outputs" / "algorithm_selection"))
|
||||
parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parents[1] / "outputs" / "followups" / "earlyconcat_standalone"))
|
||||
args = parser.parse_args()
|
||||
output = Path(args.output_dir)
|
||||
with (output / "validation_metrics_by_condition.csv").open(encoding="utf-8-sig", newline="") as stream:
|
||||
|
||||
@@ -5,6 +5,8 @@ from torch import nn
|
||||
|
||||
|
||||
class AlignedFusionModel(nn.Module):
|
||||
"""Early concatenation + BiGRU model for the supplied aligned sequence."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kind: str,
|
||||
@@ -14,8 +16,8 @@ class AlignedFusionModel(nn.Module):
|
||||
dropout: float = 0.15,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if kind not in {"concat", "gate", "crossattn"}:
|
||||
raise ValueError(f"unknown model kind: {kind}")
|
||||
if kind != "concat":
|
||||
raise ValueError(f"only the selected EarlyConcat model is maintained; got: {kind}")
|
||||
self.kind = kind
|
||||
self.hidden = hidden
|
||||
self.projections = nn.ModuleList(
|
||||
@@ -25,31 +27,9 @@ class AlignedFusionModel(nn.Module):
|
||||
self.position = nn.Parameter(torch.randn(1, steps, hidden) * 0.02)
|
||||
self.modality = nn.Parameter(torch.randn(1, 1, 3, hidden) * 0.02)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
if kind == "concat":
|
||||
self.fusion = nn.Sequential(
|
||||
nn.Linear(hidden * 3 + 3, hidden), nn.GELU(), nn.LayerNorm(hidden), nn.Dropout(dropout)
|
||||
)
|
||||
elif kind == "gate":
|
||||
self.gate_score = nn.Sequential(nn.Linear(hidden, hidden // 2), nn.Tanh(), nn.Linear(hidden // 2, 1))
|
||||
self.fusion = nn.Sequential(
|
||||
nn.Linear(hidden + 3, hidden), nn.GELU(), nn.LayerNorm(hidden), nn.Dropout(dropout)
|
||||
)
|
||||
else:
|
||||
layer = nn.TransformerEncoderLayer(
|
||||
d_model=hidden,
|
||||
nhead=4,
|
||||
dim_feedforward=hidden * 2,
|
||||
dropout=dropout,
|
||||
activation="gelu",
|
||||
batch_first=True,
|
||||
norm_first=True,
|
||||
)
|
||||
self.cross_encoder = nn.TransformerEncoder(layer, num_layers=2, enable_nested_tensor=False)
|
||||
self.fusion = nn.Sequential(
|
||||
nn.Linear(hidden + 3, hidden), nn.GELU(), nn.LayerNorm(hidden), nn.Dropout(dropout)
|
||||
)
|
||||
|
||||
self.fusion = nn.Sequential(
|
||||
nn.Linear(hidden * 3 + 3, hidden), nn.GELU(), nn.LayerNorm(hidden), nn.Dropout(dropout)
|
||||
)
|
||||
self.temporal = nn.GRU(
|
||||
input_size=hidden,
|
||||
hidden_size=hidden // 2,
|
||||
@@ -72,38 +52,16 @@ class AlignedFusionModel(nn.Module):
|
||||
encoded.append(token)
|
||||
stack = torch.stack(encoded, dim=2) # B x T x M x D
|
||||
availability = masks.to(stack.dtype)
|
||||
gate_weights = None
|
||||
|
||||
if self.kind == "concat":
|
||||
fused = self.fusion(torch.cat((stack.flatten(2), availability), dim=-1))
|
||||
elif self.kind == "gate":
|
||||
scores = self.gate_score(stack).squeeze(-1)
|
||||
scores = scores.masked_fill(~masks, -1e4)
|
||||
gate_weights = torch.softmax(scores, dim=-1) * availability
|
||||
gate_weights = gate_weights / gate_weights.sum(dim=-1, keepdim=True).clamp_min(1e-8)
|
||||
weighted = (stack * gate_weights[..., None]).sum(dim=2)
|
||||
fused = self.fusion(torch.cat((weighted, availability), dim=-1))
|
||||
else:
|
||||
batch, steps, modalities, hidden = stack.shape
|
||||
flat = stack.reshape(batch, steps * modalities, hidden)
|
||||
valid = masks.reshape(batch, steps * modalities).clone()
|
||||
empty = ~valid.any(dim=1)
|
||||
if empty.any():
|
||||
valid[empty, 0] = True
|
||||
flat[empty, 0] = 0.0
|
||||
attended = self.cross_encoder(flat, src_key_padding_mask=~valid)
|
||||
attended = attended.reshape(batch, steps, modalities, hidden)
|
||||
observed_count = availability.sum(dim=2, keepdim=True)
|
||||
pooled = (attended * availability[..., None]).sum(dim=2) / observed_count.clamp_min(1.0)
|
||||
fused = self.fusion(torch.cat((pooled, availability), dim=-1))
|
||||
fused = self.fusion(torch.cat((stack.flatten(2), availability), dim=-1))
|
||||
|
||||
temporal, _ = self.temporal(self.dropout(fused))
|
||||
time_weight = masks.any(dim=-1).to(temporal.dtype)
|
||||
empty_time = time_weight.sum(dim=1, keepdim=True) <= 0
|
||||
if empty_time.any():
|
||||
time_weight[empty_time.squeeze(1), 0] = 1.0
|
||||
pooled = (temporal * time_weight[..., None]).sum(dim=1) / time_weight.sum(dim=1, keepdim=True).clamp_min(1.0)
|
||||
pooled = (temporal * time_weight[..., None]).sum(dim=1)
|
||||
pooled = pooled / time_weight.sum(dim=1, keepdim=True).clamp_min(1.0)
|
||||
hidden = self.head(pooled)
|
||||
logits = self.classifier(hidden)
|
||||
intensity = 3.0 * torch.tanh(self.regressor(hidden).squeeze(-1))
|
||||
return {"logits": logits, "intensity": intensity, "gate": gate_weights}
|
||||
return {"logits": logits, "intensity": intensity}
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
|
||||
SUBSETS: dict[str, tuple[int, ...]] = {
|
||||
"T": (0,),
|
||||
"A": (1,),
|
||||
"V": (2,),
|
||||
"TA": (0, 1),
|
||||
"TV": (0, 2),
|
||||
"AV": (1, 2),
|
||||
"TAV": (0, 1, 2),
|
||||
}
|
||||
EXPERT_NAMES = tuple(SUBSETS)
|
||||
EXPERT_BITS = {
|
||||
name: tuple(int(i in indices) for i in range(3))
|
||||
for name, indices in SUBSETS.items()
|
||||
}
|
||||
|
||||
|
||||
class MixtureOfFusionExperts(nn.Module):
|
||||
"""Seven-subset, hard-availability MoFE with the selected MLP router.
|
||||
|
||||
Each modality has a private projection. Experts only receive the private
|
||||
projections belonging to their subset. The weighted result is passed
|
||||
through one shared temporal backbone and one shared prediction head.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: tuple[int, int, int],
|
||||
router: str = "mlp",
|
||||
expert_names: tuple[str, ...] = EXPERT_NAMES,
|
||||
availability_mode: str = "hard",
|
||||
steps: int = 50,
|
||||
latent_dim: int = 64,
|
||||
hidden: int = 128,
|
||||
dropout: float = 0.15,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if router != "mlp":
|
||||
raise ValueError(f"only the selected MLP router is maintained; got: {router}")
|
||||
if availability_mode != "hard":
|
||||
raise ValueError(f"only hard availability masking is maintained; got: {availability_mode}")
|
||||
if tuple(expert_names) != EXPERT_NAMES:
|
||||
raise ValueError("the selected MoFE uses all seven modality-subset experts")
|
||||
|
||||
self.dims = dims
|
||||
self.router_kind = router
|
||||
self.expert_names = tuple(expert_names)
|
||||
self.availability_mode = availability_mode
|
||||
self.steps = steps
|
||||
self.latent_dim = latent_dim
|
||||
self.hidden = hidden
|
||||
|
||||
# These projections are private to each modality and are not tied.
|
||||
self.private_projections = nn.ModuleList(
|
||||
nn.Sequential(nn.Linear(size, latent_dim), nn.GELU()) for size in dims
|
||||
)
|
||||
self.experts = nn.ModuleDict()
|
||||
for name in self.expert_names:
|
||||
n_modalities = len(SUBSETS[name])
|
||||
self.experts[name] = nn.Sequential(
|
||||
nn.Linear(n_modalities * latent_dim, hidden),
|
||||
nn.GELU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden, latent_dim),
|
||||
nn.LayerNorm(latent_dim),
|
||||
)
|
||||
|
||||
router_input_dim = 9
|
||||
self.router = nn.Sequential(
|
||||
nn.Linear(router_input_dim, 16),
|
||||
nn.GELU(),
|
||||
nn.Linear(16, len(self.expert_names)),
|
||||
)
|
||||
|
||||
# Shared early-fusion projection, BiGRU, and task heads.
|
||||
self.all_missing_token = nn.Parameter(torch.zeros(1, 1, latent_dim))
|
||||
self.input_projection = nn.Sequential(
|
||||
nn.Linear(latent_dim + 3, hidden),
|
||||
nn.GELU(),
|
||||
nn.LayerNorm(hidden),
|
||||
nn.Dropout(dropout),
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.temporal = nn.GRU(
|
||||
input_size=hidden,
|
||||
hidden_size=hidden // 2,
|
||||
num_layers=1,
|
||||
batch_first=True,
|
||||
bidirectional=True,
|
||||
)
|
||||
self.head = nn.Sequential(nn.Linear(hidden, hidden // 2), nn.GELU(), nn.Dropout(dropout))
|
||||
self.classifier = nn.Linear(hidden // 2, 3)
|
||||
self.regressor = nn.Linear(hidden // 2, 1)
|
||||
|
||||
@staticmethod
|
||||
def _availability(masks: torch.Tensor, names: tuple[str, ...]) -> torch.Tensor:
|
||||
masks = masks.bool()
|
||||
columns = [masks[..., list(SUBSETS[name])].all(dim=-1) for name in names]
|
||||
return torch.stack(columns, dim=-1)
|
||||
|
||||
def _router_features(
|
||||
self,
|
||||
private: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
masks: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
observed = masks.to(dtype=private[0].dtype)
|
||||
magnitude = torch.stack(
|
||||
[torch.sqrt(x.square().mean(dim=-1) + 1e-8) for x in private], dim=-1
|
||||
)
|
||||
local_ratio = F.avg_pool1d(
|
||||
observed.transpose(1, 2), kernel_size=5, stride=1, padding=2, count_include_pad=False
|
||||
).transpose(1, 2)
|
||||
return torch.cat((observed, torch.log1p(magnitude), local_ratio), dim=-1)
|
||||
|
||||
def _route(
|
||||
self,
|
||||
router_features: torch.Tensor,
|
||||
availability: torch.Tensor,
|
||||
force_expert: str | None,
|
||||
) -> torch.Tensor:
|
||||
scores = self.router(router_features)
|
||||
scores = scores.masked_fill(~availability, -1e4)
|
||||
weights = torch.softmax(scores, dim=-1) * availability.to(scores.dtype)
|
||||
# In the full seven-expert model this is exactly the all-modalities-
|
||||
# missing case. It also safely handles ablations with no eligible set.
|
||||
has_expert = availability.any(dim=-1, keepdim=True)
|
||||
weights = weights * has_expert.to(weights.dtype)
|
||||
weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-8)
|
||||
|
||||
if force_expert is not None:
|
||||
if force_expert not in self.expert_names:
|
||||
raise ValueError(f"expert {force_expert} is not enabled in this model")
|
||||
expert_idx = self.expert_names.index(force_expert)
|
||||
forced = torch.zeros_like(weights)
|
||||
forced[..., expert_idx] = 1.0
|
||||
# Force the requested expert where its modality subset is present;
|
||||
# where it is unavailable, use the learned router over eligible
|
||||
# experts instead of replacing observed information with zeros.
|
||||
return torch.where(availability[..., expert_idx, None], forced, weights)
|
||||
|
||||
return weights
|
||||
|
||||
def forward(
|
||||
self,
|
||||
xs: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
masks: torch.Tensor,
|
||||
force_expert: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
masks = masks.bool()
|
||||
if masks.ndim != 3 or masks.shape[-1] != 3:
|
||||
raise ValueError(f"masks must have shape B x T x 3, got {tuple(masks.shape)}")
|
||||
if masks.shape[1] > self.steps:
|
||||
raise ValueError(f"sequence has {masks.shape[1]} steps, model supports {self.steps}")
|
||||
|
||||
private_values = []
|
||||
for modality, (projector, x) in enumerate(zip(self.private_projections, xs)):
|
||||
projected = projector(x)
|
||||
projected = projected * masks[..., modality, None].to(projected.dtype)
|
||||
private_values.append(projected)
|
||||
private = tuple(private_values)
|
||||
router_features = self._router_features(private, masks)
|
||||
availability = self._availability(masks, self.expert_names)
|
||||
|
||||
local_expert_outputs = []
|
||||
for name in self.expert_names:
|
||||
indices = SUBSETS[name]
|
||||
expert_input = torch.cat([private[i] for i in indices], dim=-1)
|
||||
local_expert_outputs.append(self.experts[name](expert_input))
|
||||
expert_stack = torch.stack(local_expert_outputs, dim=-2)
|
||||
|
||||
alpha_local = self._route(router_features, availability, force_expert)
|
||||
fused = (expert_stack * alpha_local[..., None]).sum(dim=-2)
|
||||
has_expert = availability.any(dim=-1)
|
||||
fused = torch.where(
|
||||
has_expert[..., None], fused, self.all_missing_token.expand_as(fused)
|
||||
)
|
||||
|
||||
# Restore a stable seven-column interface for saved diagnostics,
|
||||
# including expert-set ablations.
|
||||
alpha = masks.new_zeros((*masks.shape[:2], len(EXPERT_NAMES)), dtype=private[0].dtype)
|
||||
expert_outputs = private[0].new_zeros((*masks.shape[:2], len(EXPERT_NAMES), self.latent_dim))
|
||||
for local_idx, name in enumerate(self.expert_names):
|
||||
global_idx = EXPERT_NAMES.index(name)
|
||||
alpha[..., global_idx] = alpha_local[..., local_idx]
|
||||
expert_outputs[..., global_idx, :] = expert_stack[..., local_idx, :]
|
||||
|
||||
fused_with_masks = torch.cat((fused, masks.to(fused.dtype)), dim=-1)
|
||||
encoded = self.input_projection(fused_with_masks)
|
||||
temporal, _ = self.temporal(self.dropout(encoded))
|
||||
time_weight = masks.any(dim=-1).to(temporal.dtype)
|
||||
empty_time = time_weight.sum(dim=1, keepdim=True) <= 0
|
||||
if empty_time.any():
|
||||
time_weight[empty_time.squeeze(1), 0] = 1.0
|
||||
pooled = (temporal * time_weight[..., None]).sum(dim=1)
|
||||
pooled = pooled / time_weight.sum(dim=1, keepdim=True).clamp_min(1.0)
|
||||
hidden = self.head(pooled)
|
||||
logits = self.classifier(hidden)
|
||||
intensity = 3.0 * torch.tanh(self.regressor(hidden).squeeze(-1))
|
||||
|
||||
bits = torch.tensor(
|
||||
[EXPERT_BITS[name] for name in EXPERT_NAMES],
|
||||
dtype=alpha.dtype,
|
||||
device=alpha.device,
|
||||
)
|
||||
utility = torch.einsum("bte,em->btm", alpha, bits)
|
||||
return {
|
||||
"logits": logits,
|
||||
"intensity": intensity,
|
||||
"fused": fused,
|
||||
"alpha": alpha,
|
||||
"utility": utility,
|
||||
"availability": availability,
|
||||
"expert_outputs": expert_outputs,
|
||||
"fallback": ~has_expert,
|
||||
"router_features": router_features,
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data import RobustStats, apply_robust_stats, load_aligned
|
||||
from .mofe import EXPERT_NAMES
|
||||
from .mofe import MixtureOfFusionExperts
|
||||
from .train_mofe import (
|
||||
MOFE7_MLP,
|
||||
SEEDS,
|
||||
_conditions,
|
||||
_metric_dict,
|
||||
_predict,
|
||||
_sha256,
|
||||
_write_csv,
|
||||
)
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
REFERENCE_DIR = ROOT / "outputs" / "mofe_7experts"
|
||||
DEFAULT_DIAGNOSTIC_OUTPUT = ROOT / "outputs" / "followups" / "D0_task_preference"
|
||||
MODEL_CONFIG: dict[str, Any] = {
|
||||
"router": "mlp",
|
||||
"expert_names": EXPERT_NAMES,
|
||||
"availability_mode": "hard",
|
||||
}
|
||||
|
||||
|
||||
def _load_model(seed: int, dims: tuple[int, int, int], device: torch.device) -> MixtureOfFusionExperts:
|
||||
checkpoint = REFERENCE_DIR / "models" / MOFE7_MLP / f"seed_{seed}" / "model_best.pt"
|
||||
saved = torch.load(checkpoint, map_location=device, weights_only=False)
|
||||
if saved.get("config") != MODEL_CONFIG or tuple(saved.get("dims", ())) != dims:
|
||||
raise ValueError(f"checkpoint does not match the selected single-router model: {checkpoint}")
|
||||
if int(saved.get("seed", -1)) != seed:
|
||||
raise ValueError(f"checkpoint seed mismatch: expected {seed}, found {saved.get('seed')}")
|
||||
model = MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
return model.eval()
|
||||
|
||||
|
||||
def _rank_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
groups: dict[tuple[int, str], list[dict[str, Any]]] = {}
|
||||
for row in rows:
|
||||
if row["expert"] == "learned_router":
|
||||
continue
|
||||
groups.setdefault((int(row["seed"]), str(row["condition"])), []).append(row)
|
||||
|
||||
result: list[dict[str, Any]] = []
|
||||
for (seed, condition), values in sorted(groups.items()):
|
||||
by_name = {str(row["expert"]): row for row in values}
|
||||
ordered = [by_name[name] for name in EXPERT_NAMES]
|
||||
f1 = [float(row["macro_f1"]) for row in ordered]
|
||||
mae = [float(row["mae"]) for row in ordered]
|
||||
pearson = [float(row["pearson"]) for row in ordered]
|
||||
result.append({
|
||||
"seed": seed,
|
||||
"condition": condition,
|
||||
"spearman_macro_f1_vs_mae": _spearman(f1, mae),
|
||||
"spearman_macro_f1_vs_pearson": _spearman(f1, pearson),
|
||||
"best_macro_f1_expert": EXPERT_NAMES[int(np.argmax(f1))],
|
||||
"best_mae_expert": EXPERT_NAMES[int(np.argmin(mae))],
|
||||
"best_pearson_expert": EXPERT_NAMES[int(np.argmax(pearson))],
|
||||
"macro_f1_order_best_to_worst": ">".join(EXPERT_NAMES[i] for i in np.argsort(-np.asarray(f1), kind="stable")),
|
||||
"mae_order_best_to_worst": ">".join(EXPERT_NAMES[i] for i in np.argsort(np.asarray(mae), kind="stable")),
|
||||
"pearson_order_best_to_worst": ">".join(EXPERT_NAMES[i] for i in np.argsort(-np.asarray(pearson), kind="stable")),
|
||||
})
|
||||
return result
|
||||
|
||||
|
||||
def _spearman(left: list[float], right: list[float]) -> float:
|
||||
def average_ranks(values: list[float]) -> np.ndarray:
|
||||
array = np.asarray(values, dtype=np.float64)
|
||||
order = np.argsort(array, kind="stable")
|
||||
ranks = np.empty(len(array), dtype=np.float64)
|
||||
start = 0
|
||||
while start < len(array):
|
||||
end = start + 1
|
||||
while end < len(array) and array[order[end]] == array[order[start]]:
|
||||
end += 1
|
||||
ranks[order[start:end]] = (start + 1 + end) / 2
|
||||
start = end
|
||||
return ranks
|
||||
|
||||
left_ranks = average_ranks(left)
|
||||
right_ranks = average_ranks(right)
|
||||
if np.std(left_ranks) == 0 or np.std(right_ranks) == 0:
|
||||
return 0.0
|
||||
return float(np.corrcoef(left_ranks, right_ranks)[0, 1])
|
||||
|
||||
|
||||
def _summary_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
result = []
|
||||
for expert in (*EXPERT_NAMES, "learned_router"):
|
||||
matching = [row for row in rows if row["expert"] == expert and row["condition"] != "clean"]
|
||||
per_seed: dict[int, list[dict[str, Any]]] = {}
|
||||
for row in matching:
|
||||
per_seed.setdefault(int(row["seed"]), []).append(row)
|
||||
seed_means = []
|
||||
for seed, seed_rows in sorted(per_seed.items()):
|
||||
seed_means.append({
|
||||
metric: float(np.mean([float(row[metric]) for row in seed_rows]))
|
||||
for metric in ("macro_f1", "mae", "pearson", "available_position_fraction")
|
||||
})
|
||||
if not seed_means:
|
||||
continue
|
||||
out: dict[str, Any] = {"expert": expert, "n_seeds": len(seed_means), "conditions_averaged": len(matching) // len(seed_means)}
|
||||
for metric in ("macro_f1", "mae", "pearson", "available_position_fraction"):
|
||||
values = [item[metric] for item in seed_means]
|
||||
out[f"corrupt_{metric}_mean"] = float(np.mean(values))
|
||||
out[f"corrupt_{metric}_seed_sd"] = float(np.std(values, ddof=1)) if len(values) > 1 else 0.0
|
||||
result.append(out)
|
||||
return result
|
||||
|
||||
|
||||
def _write_readout(output: Path, summary: list[dict[str, Any]], ranks: list[dict[str, Any]]) -> None:
|
||||
rho_f1_mae = np.asarray([float(row["spearman_macro_f1_vs_mae"]) for row in ranks])
|
||||
rho_f1_pearson = np.asarray([float(row["spearman_macro_f1_vs_pearson"]) for row in ranks])
|
||||
best_f1 = {name: sum(row["best_macro_f1_expert"] == name for row in ranks) for name in EXPERT_NAMES}
|
||||
best_mae = {name: sum(row["best_mae_expert"] == name for row in ranks) for name in EXPERT_NAMES}
|
||||
best_pearson = {name: sum(row["best_pearson_expert"] == name for row in ranks) for name in EXPERT_NAMES}
|
||||
rank_by_condition: dict[str, list[dict[str, Any]]] = {}
|
||||
for row in ranks:
|
||||
rank_by_condition.setdefault(str(row["condition"]), []).append(row)
|
||||
|
||||
def winners(condition: str, column: str) -> str:
|
||||
matching = rank_by_condition[condition]
|
||||
counts = {name: sum(row[column] == name for row in matching) for name in EXPERT_NAMES}
|
||||
max_count = max(counts.values())
|
||||
names = [name for name, count in counts.items() if count == max_count]
|
||||
return ", ".join(f"{name} ({max_count}/{len(matching)})" for name in names)
|
||||
|
||||
by_name = {str(row["expert"]): row for row in summary}
|
||||
lines = [
|
||||
"# Single-router MoFE 任务偏好诊断",
|
||||
"",
|
||||
"本诊断使用保留的 single-router 检查点和验证集,用于判断是否值得增加第二个 router;它不是测试集估计。",
|
||||
"",
|
||||
"## Forced-expert 规则",
|
||||
"",
|
||||
"所选模态子集可用的位置强制使用对应 expert;该子集不可用时,由已训练 router 在其他可用 expert 中选择;全模态缺失时沿用 learned missing token。可用率表示所选 expert 能被强制使用的位置比例。",
|
||||
"",
|
||||
"## 缺失条件平均指标",
|
||||
"",
|
||||
"下表先在每个 seed 内对 15 种连续块缺失条件求平均,再汇总三个 seed;seed 标准差见 CSV。",
|
||||
"",
|
||||
"| Expert | Macro-F1 ↑ | MAE ↓ | Pearson ↑ | 可强制使用比例 |",
|
||||
"| --- | ---: | ---: | ---: | ---: |",
|
||||
]
|
||||
for name in (*EXPERT_NAMES, "learned_router"):
|
||||
row = by_name[name]
|
||||
lines.append(
|
||||
f"| {name} | {float(row['corrupt_macro_f1_mean']):.3f} | {float(row['corrupt_mae_mean']):.3f} | "
|
||||
f"{float(row['corrupt_pearson_mean']):.3f} | {float(row['corrupt_available_position_fraction_mean']):.3f} |"
|
||||
)
|
||||
lines.extend([
|
||||
"",
|
||||
"## 两个任务的 expert 偏好",
|
||||
"",
|
||||
f"在 {len(ranks)} 个 seed—条件组合中,分类 Macro-F1 与回归 MAE 的平均 Spearman ρ 为 **{rho_f1_mae.mean():.3f}**。MAE 越低越好,因此负相关表示两个指标倾向于选中相似的 expert。Macro-F1 与 Pearson 的平均 ρ 为 **{rho_f1_pearson.mean():.3f}**。",
|
||||
"",
|
||||
"六个重点条件下的相关性先按三个 seed 求平均;最优 expert 一栏显示三个 seed 中的多数结果:",
|
||||
"",
|
||||
"| 条件 | ρ(Macro-F1, MAE) | ρ(Macro-F1, Pearson) | Macro-F1 最优 | MAE 最优 | Pearson 最优 |",
|
||||
"| --- | ---: | ---: | --- | --- | --- |",
|
||||
])
|
||||
key_conditions = (
|
||||
("clean", "Clean"),
|
||||
("text_30", "Text 30%"),
|
||||
("audio_30", "Audio 30%"),
|
||||
("vision_30", "Vision 30%"),
|
||||
("audio_vision_30", "Audio+Vision 30%"),
|
||||
("all_modalities_30", "All-modal 30%"),
|
||||
)
|
||||
for condition, label in key_conditions:
|
||||
condition_rows = rank_by_condition[condition]
|
||||
rho_mae = float(np.mean([float(row["spearman_macro_f1_vs_mae"]) for row in condition_rows]))
|
||||
rho_pearson = float(np.mean([float(row["spearman_macro_f1_vs_pearson"]) for row in condition_rows]))
|
||||
lines.append(
|
||||
f"| {label} | {rho_mae:.3f} | {rho_pearson:.3f} | "
|
||||
f"{winners(condition, 'best_macro_f1_expert')} | {winners(condition, 'best_mae_expert')} | "
|
||||
f"{winners(condition, 'best_pearson_expert')} |"
|
||||
)
|
||||
lines.extend([
|
||||
"",
|
||||
f"各指标的最优 expert 次数:Macro-F1({_format_counts(best_f1)});MAE({_format_counts(best_mae)});Pearson({_format_counts(best_pearson)})。",
|
||||
"",
|
||||
"当前排名没有显示稳定的分类—回归 expert 分工:Macro-F1 较高通常同时对应较低 MAE 和较高 Pearson;文本 expert 在分类与回归指标上都是最常见的赢家。因此,这项诊断**没有提供增加第二个 router 所需的任务特异模态偏好证据**。目前保留 single-router 作为活动参照;这不代表两个任务在任何数据或设置下都不可能需要不同路由。",
|
||||
"",
|
||||
"## 结论范围",
|
||||
"",
|
||||
"输入是官方提供的 50 个有序 wordpiece 位置。结果只反映这些位置及本次缺失掩码下的任务与 expert 关系,不表示物理时间可靠性。",
|
||||
"",
|
||||
"逐条件结果见 `forced_expert_metrics.csv` 和 `rank_concordance.csv`;跨 seed 汇总见 `expert_task_preference_summary.csv`。",
|
||||
])
|
||||
(output / "task_preference_diagnostic.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def _format_counts(counts: dict[str, int]) -> str:
|
||||
return ", ".join(f"{name}: {count}" for name, count in counts.items())
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> None:
|
||||
output = args.output_dir.resolve()
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else "cpu") if args.device == "auto" else torch.device(args.device)
|
||||
torch.set_num_threads(args.threads)
|
||||
|
||||
raw = load_aligned()
|
||||
scaler_path = REFERENCE_DIR / "aligned_robust_stats.npz"
|
||||
stats = RobustStats.load(scaler_path)
|
||||
valid = apply_robust_stats(raw["valid"], stats)
|
||||
dims = tuple(int(x.shape[-1]) for x in valid.x)
|
||||
|
||||
metric_rows: list[dict[str, Any]] = []
|
||||
for seed in args.seeds:
|
||||
model = _load_model(seed, dims, device)
|
||||
conditions = _conditions(valid, seed)
|
||||
for condition, rate, masks in conditions:
|
||||
predictions = {
|
||||
"learned_router": _predict(model, valid, masks, device, args.batch_size),
|
||||
**{
|
||||
expert: _predict(model, valid, masks, device, args.batch_size, force_expert=expert)
|
||||
for expert in EXPERT_NAMES
|
||||
},
|
||||
}
|
||||
for expert, prediction in predictions.items():
|
||||
coverage = 1.0
|
||||
if expert != "learned_router":
|
||||
expert_index = EXPERT_NAMES.index(expert)
|
||||
coverage = float(model._availability(torch.as_tensor(masks, dtype=torch.bool, device=device), EXPERT_NAMES)[..., expert_index].float().mean().item())
|
||||
else:
|
||||
coverage = float(prediction["availability"].astype(bool).any(axis=-1).mean())
|
||||
metric_rows.append({
|
||||
"method": MOFE7_MLP,
|
||||
"seed": seed,
|
||||
"condition": condition,
|
||||
"missing_rate": rate,
|
||||
"expert": expert,
|
||||
"n_valid": valid.n,
|
||||
"available_position_fraction": coverage,
|
||||
**_metric_dict(valid.y_cls, valid.y_reg, prediction["logits"], prediction["intensity"]),
|
||||
})
|
||||
print(f"forced-expert diagnostic complete for seed={seed} on {device}", flush=True)
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
rank_rows = _rank_rows(metric_rows)
|
||||
summary = _summary_rows(metric_rows)
|
||||
_write_csv(output / "forced_expert_metrics.csv", metric_rows)
|
||||
_write_csv(output / "expert_task_preference_summary.csv", summary)
|
||||
_write_csv(output / "rank_concordance.csv", rank_rows)
|
||||
_write_readout(output, summary, rank_rows)
|
||||
feature_path = ROOT.parents[1] / "E题数据" / "附件2-数据集特征文件" / "aligned_50.pkl"
|
||||
metadata = {
|
||||
"diagnostic": "forced-expert task preference for the retained single-router MoFE-7",
|
||||
"checkpoint_dir": str(REFERENCE_DIR / "models" / MOFE7_MLP),
|
||||
"checkpoint_seeds": list(args.seeds),
|
||||
"feature_file": str(feature_path),
|
||||
"feature_sha256": _sha256(feature_path),
|
||||
"scaler_file": str(scaler_path),
|
||||
"scaler_sha256": _sha256(scaler_path),
|
||||
"device": str(device),
|
||||
"cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
|
||||
"python_version": sys.version,
|
||||
"torch_version": torch.__version__,
|
||||
"numpy_version": np.__version__,
|
||||
"valid_examples": valid.n,
|
||||
"conditions": [condition for condition, _, _ in _conditions(valid, args.seeds[0])],
|
||||
"corruption_seed_protocol": "seed + 13 + pattern_index*101 + int(rate*1000)",
|
||||
"forced_expert_policy": "use the requested expert where its modality subset is available; fall back to the trained single router at positions where it is unavailable; all-missing positions use the learned missing token",
|
||||
"scope_note": "50 official ordered wordpiece positions; no claim about physical-time reliability",
|
||||
"interpretation_note": "Ranking agreement is descriptive on the supplied validation split; it is a motivation diagnostic, not an unbiased test-set estimate.",
|
||||
}
|
||||
(output / "run_manifest.json").write_text(json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
print(f"saved forced-expert diagnostic to {output}", flush=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Measure classification/regression preferences across existing MoFE experts.")
|
||||
parser.add_argument("--seeds", type=int, nargs="+", default=list(SEEDS))
|
||||
parser.add_argument("--batch-size", type=int, default=128)
|
||||
parser.add_argument("--threads", type=int, default=4)
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--output-dir", type=Path, default=DEFAULT_DIAGNOSTIC_OUTPUT)
|
||||
run(parser.parse_args())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -43,7 +43,7 @@ PATTERNS = {
|
||||
"audio_vision": (1, 2),
|
||||
"all_modalities": (0, 1, 2),
|
||||
}
|
||||
KINDS = ("concat", "gate", "crossattn")
|
||||
KINDS = ("concat",)
|
||||
|
||||
|
||||
def seed_everything(seed: int) -> None:
|
||||
@@ -292,7 +292,7 @@ def _summary(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
|
||||
def _plot(summary: list[dict[str, Any]], rows: list[dict[str, Any]], path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
colors = {"concat": "#4e79a7", "gate": "#f28e2b", "crossattn": "#59a14f"}
|
||||
colors = {"concat": "#4e79a7"}
|
||||
fig, axes = plt.subplots(1, 2, figsize=(11, 4.4), constrained_layout=True)
|
||||
for row in summary:
|
||||
kind = row["method"]
|
||||
@@ -474,14 +474,14 @@ def _run(args: argparse.Namespace) -> None:
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Q2 local-missingness model and alignment transfer selection")
|
||||
parser = argparse.ArgumentParser(description="Train the EarlyConcat baseline and its alignment-transfer control")
|
||||
parser.add_argument("--seeds", type=int, nargs="+", default=[42, 3407, 2026])
|
||||
parser.add_argument("--epochs", type=int, default=32)
|
||||
parser.add_argument("--patience", type=int, default=6)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--threads", type=int, default=4)
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parents[1] / "outputs" / "algorithm_selection"))
|
||||
parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parents[1] / "outputs" / "followups" / "earlyconcat_standalone"))
|
||||
args = parser.parse_args()
|
||||
_run(args)
|
||||
|
||||
|
||||
@@ -0,0 +1,513 @@
|
||||
"""Retrain the two maintained Q2 models under the math/Q2 V2 protocol.
|
||||
|
||||
The model architectures and joint CE + SmoothL1 objective stay unchanged.
|
||||
Training masks, official splits, validation scenarios, and final-test handling
|
||||
follow the corresponding math/Q2 protocol where those choices apply.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import time
|
||||
from collections import Counter, defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sklearn.metrics import accuracy_score, f1_score, mean_absolute_error, mean_squared_error
|
||||
from torch import nn
|
||||
|
||||
from .data import ATTACHMENT2, RobustStats, Split, apply_robust_stats, fit_robust_stats
|
||||
from .evaluate_math_protocol import (
|
||||
AURC_BOOTSTRAP_SEED,
|
||||
BOOTSTRAP_REPS,
|
||||
CURVE_MODES,
|
||||
METHODS,
|
||||
SCENARIO_SEED,
|
||||
TEST_BOOTSTRAP_SEED,
|
||||
actual_additional_rates,
|
||||
aurc_from_curve,
|
||||
continuous_mask,
|
||||
curve_scenarios,
|
||||
load_splits,
|
||||
make_scenarios,
|
||||
metrics,
|
||||
scenario_seed,
|
||||
sha256,
|
||||
write_csv,
|
||||
)
|
||||
from .models import AlignedFusionModel
|
||||
from .mofe import MixtureOfFusionExperts
|
||||
from .train_mofe import EARLYCONCAT, MODEL_CONFIG, MOFE7_MLP, _predict
|
||||
from .train_compare import _loss, seed_everything
|
||||
|
||||
|
||||
Q2_ROOT = Path(__file__).resolve().parents[1]
|
||||
OUTPUT_DIR = Q2_ROOT / "outputs" / "followups" / "R03_math_protocol_retraining"
|
||||
SEED = 20260924
|
||||
TRAIN_MASK_SEED = 20261227
|
||||
BATCH_SIZE = 64
|
||||
EPOCH_LIMIT = 12
|
||||
PATIENCE = 3
|
||||
LEARNING_RATE = 3e-4
|
||||
WEIGHT_DECAY = 1e-3
|
||||
SELECTION_SCENARIOS = ("0.0/none", "0.3/single", "0.3/sync", "0.5/async")
|
||||
TRAIN_RATES = (0.0, 0.1, 0.3, 0.5, 0.7)
|
||||
TRAIN_MODES = ("single", "sync", "partial", "async")
|
||||
|
||||
|
||||
def device_for(name: str) -> torch.device:
|
||||
if name == "auto":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(name)
|
||||
|
||||
|
||||
def set_deterministic(seed: int) -> None:
|
||||
seed_everything(seed)
|
||||
torch.set_num_threads(4)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def build_model(method: str, dims: tuple[int, int, int], device: torch.device) -> nn.Module:
|
||||
if method == EARLYCONCAT:
|
||||
return AlignedFusionModel("concat", dims=dims).to(device)
|
||||
if method == MOFE7_MLP:
|
||||
return MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device)
|
||||
raise ValueError(f"unknown method: {method}")
|
||||
|
||||
|
||||
def model_state(model: nn.Module, method: str) -> dict[str, Any]:
|
||||
state: dict[str, Any] = {
|
||||
"method": method,
|
||||
"dims": tuple(int(x) for x in model_dims(model)),
|
||||
"state_dict": model.state_dict(),
|
||||
"seed": SEED,
|
||||
"protocol": "math/Q2 V2 adapted deterministic-model training",
|
||||
}
|
||||
if method == EARLYCONCAT:
|
||||
state["kind"] = "concat"
|
||||
else:
|
||||
state["config"] = MODEL_CONFIG
|
||||
return state
|
||||
|
||||
|
||||
def model_dims(model: nn.Module) -> tuple[int, int, int]:
|
||||
if isinstance(model, AlignedFusionModel):
|
||||
return tuple(layer[0].in_features for layer in model.projections) # type: ignore[return-value]
|
||||
if isinstance(model, MixtureOfFusionExperts):
|
||||
return tuple(layer[0].in_features for layer in model.private_projections) # type: ignore[return-value]
|
||||
raise TypeError(type(model))
|
||||
|
||||
|
||||
def train_masks_for_epoch(split: Split, epoch: int) -> tuple[np.ndarray, Counter[str]]:
|
||||
"""Sample reproducible math-protocol rates/patterns per training example."""
|
||||
rows: list[np.ndarray] = []
|
||||
counts: Counter[str] = Counter()
|
||||
for sample_id, observed in zip(split.ids, split.mask):
|
||||
rng = np.random.default_rng(scenario_seed(TRAIN_MASK_SEED + SEED, sample_id, f"train/{epoch}"))
|
||||
rate = float(rng.choice(TRAIN_RATES))
|
||||
mode = str(rng.choice(TRAIN_MODES))
|
||||
key = f"{rate:.1f}/{mode}"
|
||||
counts[key] += 1
|
||||
row = continuous_mask(observed, rate, mode, rng)
|
||||
rows.append(row)
|
||||
return np.stack(rows), counts
|
||||
|
||||
|
||||
def _batched_loss(
|
||||
model: nn.Module,
|
||||
split: Split,
|
||||
masks: np.ndarray,
|
||||
device: torch.device,
|
||||
batch_size: int,
|
||||
) -> float:
|
||||
model.eval()
|
||||
losses: list[float] = []
|
||||
weights: list[int] = []
|
||||
with torch.inference_mode():
|
||||
for start in range(0, split.n, batch_size):
|
||||
end = min(start + batch_size, split.n)
|
||||
xs = tuple(torch.as_tensor(x[start:end], dtype=torch.float32, device=device) for x in split.x)
|
||||
mb = torch.as_tensor(masks[start:end], dtype=torch.bool, device=device)
|
||||
y_cls = torch.as_tensor(split.y_cls[start:end], dtype=torch.long, device=device)
|
||||
y_reg = torch.as_tensor(split.y_reg[start:end], dtype=torch.float32, device=device)
|
||||
losses.append(float(_loss(model(xs, mb), y_cls, y_reg).item()))
|
||||
weights.append(end - start)
|
||||
return float(np.average(losses, weights=weights))
|
||||
|
||||
|
||||
def selection_loss(model: nn.Module, valid: Split, scenarios: dict[str, np.ndarray], device: torch.device) -> float:
|
||||
return float(np.mean([
|
||||
_batched_loss(model, valid, scenarios[key], device, BATCH_SIZE)
|
||||
for key in SELECTION_SCENARIOS
|
||||
]))
|
||||
|
||||
|
||||
def train_one(
|
||||
method: str,
|
||||
train: Split,
|
||||
valid: Split,
|
||||
valid_scenarios: dict[str, np.ndarray],
|
||||
orders: list[np.ndarray],
|
||||
output_dir: Path,
|
||||
device: torch.device,
|
||||
) -> tuple[nn.Module, int, list[dict[str, Any]], Counter[str]]:
|
||||
set_deterministic(SEED)
|
||||
model = build_model(method, tuple(x.shape[-1] for x in train.x), device)
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
|
||||
xs = tuple(torch.as_tensor(x, dtype=torch.float32, device=device) for x in train.x)
|
||||
y_cls = torch.as_tensor(train.y_cls, dtype=torch.long, device=device)
|
||||
y_reg = torch.as_tensor(train.y_reg, dtype=torch.float32, device=device)
|
||||
checkpoint_path = output_dir / "model_best.pt"
|
||||
history: list[dict[str, Any]] = []
|
||||
train_mask_counts: Counter[str] = Counter()
|
||||
best_loss = math.inf
|
||||
best_epoch = 0
|
||||
stale = 0
|
||||
|
||||
for epoch in range(1, EPOCH_LIMIT + 1):
|
||||
model.train()
|
||||
epoch_masks, epoch_counts = train_masks_for_epoch(train, epoch)
|
||||
train_mask_counts.update(epoch_counts)
|
||||
batch_losses: list[float] = []
|
||||
order = orders[epoch - 1]
|
||||
for start in range(0, train.n, BATCH_SIZE):
|
||||
indices_np = order[start:start + BATCH_SIZE]
|
||||
indices = torch.as_tensor(indices_np, dtype=torch.long, device=device)
|
||||
mb = torch.as_tensor(epoch_masks[indices_np], dtype=torch.bool, device=device)
|
||||
output = model(tuple(x.index_select(0, indices) for x in xs), mb)
|
||||
loss = _loss(output, y_cls.index_select(0, indices), y_reg.index_select(0, indices))
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||||
optimizer.step()
|
||||
batch_losses.append(float(loss.detach().item()))
|
||||
|
||||
valid_selection_loss = selection_loss(model, valid, valid_scenarios, device)
|
||||
row = {
|
||||
"method": method,
|
||||
"seed": SEED,
|
||||
"epoch": epoch,
|
||||
"train_loss": float(np.mean(batch_losses)),
|
||||
"valid_selection_loss": valid_selection_loss,
|
||||
"valid_clean_loss": _batched_loss(model, valid, valid.mask, device, BATCH_SIZE),
|
||||
}
|
||||
history.append(row)
|
||||
print(
|
||||
f"[{method}] epoch={epoch:02d} train={row['train_loss']:.4f} "
|
||||
f"valid_selection={valid_selection_loss:.4f} clean={row['valid_clean_loss']:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
if valid_selection_loss < best_loss - 1e-4:
|
||||
best_loss = valid_selection_loss
|
||||
best_epoch = epoch
|
||||
stale = 0
|
||||
torch.save(model_state(model, method) | {"best_epoch": best_epoch}, checkpoint_path)
|
||||
else:
|
||||
stale += 1
|
||||
if stale >= PATIENCE:
|
||||
break
|
||||
|
||||
saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
model.eval()
|
||||
write_csv(output_dir / "training_history.csv", history)
|
||||
return model, best_epoch, history, train_mask_counts
|
||||
|
||||
|
||||
def _group_map(ids: list[str]) -> tuple[list[str], dict[str, np.ndarray]]:
|
||||
source_ids = [sample_id.split("$_$", 1)[0] for sample_id in ids]
|
||||
groups = sorted(set(source_ids))
|
||||
mapping = {
|
||||
group: np.flatnonzero(np.asarray([source == group for source in source_ids]))
|
||||
for group in groups
|
||||
}
|
||||
return groups, mapping
|
||||
|
||||
|
||||
def test_group_bootstrap(test: Split, predictions: dict[str, dict[str, np.ndarray]]) -> list[dict[str, Any]]:
|
||||
groups, mapping = _group_map(test.ids)
|
||||
rng = np.random.default_rng(TEST_BOOTSTRAP_SEED)
|
||||
draws: dict[str, list[float]] = defaultdict(list)
|
||||
for _ in range(BOOTSTRAP_REPS):
|
||||
selected = rng.choice(groups, size=len(groups), replace=True)
|
||||
indices = np.concatenate([mapping[group] for group in selected])
|
||||
values = {
|
||||
method: metrics(test, predictions[method]["logits"], predictions[method]["intensity"], indices)
|
||||
for method in METHODS
|
||||
}
|
||||
for name in values[EARLYCONCAT]:
|
||||
draws[name].append(values[MOFE7_MLP][name] - values[EARLYCONCAT][name])
|
||||
point = {
|
||||
name: metrics(test, predictions[MOFE7_MLP]["logits"], predictions[MOFE7_MLP]["intensity"])[name]
|
||||
- metrics(test, predictions[EARLYCONCAT]["logits"], predictions[EARLYCONCAT]["intensity"])[name]
|
||||
for name in draws
|
||||
}
|
||||
return [{
|
||||
"comparison": f"{MOFE7_MLP} minus {EARLYCONCAT}",
|
||||
"metric": name,
|
||||
"delta": point[name],
|
||||
"bootstrap_ci_2p5": float(np.quantile(values, 0.025)),
|
||||
"bootstrap_ci_97p5": float(np.quantile(values, 0.975)),
|
||||
"bootstrap_probability_delta_gt_0": float(np.mean(np.asarray(values) > 0)),
|
||||
"replicates": BOOTSTRAP_REPS,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
"seed": TEST_BOOTSTRAP_SEED,
|
||||
} for name, values in draws.items()]
|
||||
|
||||
|
||||
def validation_aurc_bootstrap(
|
||||
valid: Split,
|
||||
predictions: dict[tuple[str, str], dict[str, np.ndarray]],
|
||||
rates_by_sample: dict[str, np.ndarray],
|
||||
) -> list[dict[str, Any]]:
|
||||
groups, mapping = _group_map(valid.ids)
|
||||
rng = np.random.default_rng(AURC_BOOTSTRAP_SEED)
|
||||
deltas: dict[str, list[float]] = {mode: [] for mode in CURVE_MODES}
|
||||
|
||||
def score(method: str, mode: str, indices: np.ndarray) -> float:
|
||||
keys = curve_scenarios(mode)
|
||||
xs = [float(np.nanmean(rates_by_sample[key][indices])) for key in keys]
|
||||
ys = [
|
||||
float(np.abs(valid.y_reg[indices] - predictions[(method, key)]["intensity"][indices]).mean())
|
||||
for key in keys
|
||||
]
|
||||
return aurc_from_curve(xs, ys)
|
||||
|
||||
for _ in range(BOOTSTRAP_REPS):
|
||||
selected = rng.choice(groups, size=len(groups), replace=True)
|
||||
indices = np.concatenate([mapping[group] for group in selected])
|
||||
for mode in CURVE_MODES:
|
||||
deltas[mode].append(score(MOFE7_MLP, mode, indices) - score(EARLYCONCAT, mode, indices))
|
||||
rows = []
|
||||
for mode in CURVE_MODES:
|
||||
all_indices = np.arange(valid.n)
|
||||
values = deltas[mode]
|
||||
rows.append({
|
||||
"mask_mode": mode,
|
||||
"delta_aurc_mae_mofe_minus_earlyconcat": score(MOFE7_MLP, mode, all_indices) - score(EARLYCONCAT, mode, all_indices),
|
||||
"bootstrap_ci_2p5": float(np.quantile(values, 0.025)),
|
||||
"bootstrap_ci_97p5": float(np.quantile(values, 0.975)),
|
||||
"bootstrap_probability_delta_lt_0": float(np.mean(np.asarray(values) < 0)),
|
||||
"replicates": BOOTSTRAP_REPS,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
"seed": AURC_BOOTSTRAP_SEED,
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def run(device_name: str = "auto", output_dir: Path = OUTPUT_DIR) -> None:
|
||||
if output_dir.exists() and any(output_dir.iterdir()):
|
||||
raise FileExistsError(f"refusing to overwrite non-empty result directory: {output_dir}")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
device = device_for(device_name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA was requested but is unavailable")
|
||||
|
||||
feature_path = ATTACHMENT2 / "aligned_50.pkl"
|
||||
raw_splits = load_splits(feature_path)
|
||||
train_raw, valid_raw, test_raw = raw_splits["train"], raw_splits["valid"], raw_splits["test"]
|
||||
stats = fit_robust_stats(train_raw)
|
||||
train, valid, test = (apply_robust_stats(s, stats) for s in (train_raw, valid_raw, test_raw))
|
||||
stats_path = output_dir / "aligned_robust_stats.npz"
|
||||
stats.save(stats_path)
|
||||
dims = tuple(int(x.shape[-1]) for x in train.x)
|
||||
valid_scenarios = make_scenarios(valid, SCENARIO_SEED)
|
||||
if len(valid_scenarios) != 42:
|
||||
raise ValueError(f"expected 42 controlled scenarios, got {len(valid_scenarios)}")
|
||||
rates_by_sample = actual_additional_rates(valid.mask, valid_scenarios)
|
||||
|
||||
set_deterministic(SEED)
|
||||
order_rng = np.random.default_rng(SEED + 809)
|
||||
orders = [order_rng.permutation(train.n) for _ in range(EPOCH_LIMIT)]
|
||||
best_epochs: dict[str, int] = {}
|
||||
training_rows: list[dict[str, Any]] = []
|
||||
mask_count_rows: list[dict[str, Any]] = []
|
||||
parameter_rows: list[dict[str, Any]] = []
|
||||
|
||||
for method in METHODS:
|
||||
model_dir = output_dir / "models" / method / f"seed_{SEED}"
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
model, best_epoch, history, mask_counts = train_one(
|
||||
method, train, valid, valid_scenarios, orders, model_dir, device
|
||||
)
|
||||
best_epochs[method] = best_epoch
|
||||
training_rows.extend(history)
|
||||
parameter_rows.append({
|
||||
"method": method,
|
||||
"parameters_total": sum(p.numel() for p in model.parameters()),
|
||||
"parameters_trainable": sum(p.numel() for p in model.parameters() if p.requires_grad),
|
||||
"best_epoch": best_epoch,
|
||||
})
|
||||
for key, count in sorted(mask_counts.items()):
|
||||
mask_count_rows.append({"method": method, "seed": SEED, "rate_mode": key, "sample_epoch_assignments": count})
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
write_csv(output_dir / "training_history.csv", training_rows)
|
||||
write_csv(output_dir / "training_mask_distribution.csv", mask_count_rows)
|
||||
write_csv(output_dir / "parameter_count.csv", parameter_rows)
|
||||
|
||||
# Reload the selected checkpoints, then conduct one final official-test pass.
|
||||
test_predictions: dict[str, dict[str, np.ndarray]] = {}
|
||||
test_rows: list[dict[str, Any]] = []
|
||||
condition_predictions: dict[tuple[str, str], dict[str, np.ndarray]] = {}
|
||||
condition_rows: list[dict[str, Any]] = []
|
||||
for method in METHODS:
|
||||
checkpoint_path = output_dir / "models" / method / f"seed_{SEED}" / "model_best.pt"
|
||||
saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
model = build_model(method, dims, device)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
model.eval()
|
||||
|
||||
test_prediction = _predict(model, test, test.mask, device, BATCH_SIZE)
|
||||
test_predictions[method] = test_prediction
|
||||
test_rows.append({
|
||||
"method": method,
|
||||
"seed": SEED,
|
||||
"best_epoch": best_epochs[method],
|
||||
"n_test": test.n,
|
||||
**metrics(test, test_prediction["logits"], test_prediction["intensity"]),
|
||||
})
|
||||
|
||||
for scenario, masks in valid_scenarios.items():
|
||||
prediction = _predict(model, valid, masks, device, BATCH_SIZE)
|
||||
condition_predictions[(method, scenario)] = prediction
|
||||
condition_rows.append({
|
||||
"method": method,
|
||||
"seed": SEED,
|
||||
"scenario": scenario,
|
||||
"realized_additional_global_rate": float(np.nanmean(rates_by_sample[scenario])),
|
||||
"n_valid": valid.n,
|
||||
**metrics(valid, prediction["logits"], prediction["intensity"]),
|
||||
})
|
||||
print(f"[valid/{method}] {scenario} done", flush=True)
|
||||
del model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
write_csv(output_dir / "official_test_metrics_by_seed.csv", test_rows)
|
||||
write_csv(output_dir / "official_test_paired_bootstrap.csv", test_group_bootstrap(test, test_predictions))
|
||||
write_csv(output_dir / "controlled_metrics_by_scenario.csv", condition_rows)
|
||||
|
||||
test_summary = []
|
||||
for method in METHODS:
|
||||
row = next(r for r in test_rows if r["method"] == method)
|
||||
for metric in ("accuracy", "macro_f1", "mae", "rmse", "pearson"):
|
||||
test_summary.append({"method": method, "metric": metric, "mean": row[metric], "sd_across_seeds": 0.0, "n_seeds": 1})
|
||||
write_csv(output_dir / "official_test_summary.csv", test_summary)
|
||||
|
||||
aurc_rows: list[dict[str, Any]] = []
|
||||
for method in METHODS:
|
||||
for mode in CURVE_MODES:
|
||||
keys = curve_scenarios(mode)
|
||||
xs = [float(np.nanmean(rates_by_sample[key])) for key in keys]
|
||||
ys = [
|
||||
float(np.abs(valid.y_reg - condition_predictions[(method, key)]["intensity"]).mean())
|
||||
for key in keys
|
||||
]
|
||||
aurc_rows.append({
|
||||
"method": method,
|
||||
"seed": SEED,
|
||||
"mask_mode": mode,
|
||||
"aurc_mae": aurc_from_curve(xs, ys),
|
||||
"rates_realized": json.dumps(xs),
|
||||
})
|
||||
write_csv(output_dir / "aurc_mae_by_mode_seed.csv", aurc_rows)
|
||||
write_csv(output_dir / "aurc_mae_paired_bootstrap.csv", validation_aurc_bootstrap(valid, condition_predictions, rates_by_sample))
|
||||
|
||||
manifest = {
|
||||
"experiment": "Retrained EarlyConcat and MoFE-7 + MLP Router using math/Q2 V2-compatible protocol",
|
||||
"created_unix": time.time(),
|
||||
"device": str(device),
|
||||
"cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
|
||||
"feature_file": str(feature_path),
|
||||
"feature_sha256": sha256(feature_path),
|
||||
"representation": "official aligned_50 ordered positions; not Q1 physical-time bins",
|
||||
"train_valid_test_counts": {name: split.n for name, split in raw_splits.items()},
|
||||
"source_video_groups": {name: len({sid.split("$_$", 1)[0] for sid in split.ids}) for name, split in raw_splits.items()},
|
||||
"official_group_splits_disjoint": True,
|
||||
"train_only_scaler": str(stats_path),
|
||||
"scaler_fit": "median and 1.4826*MAD on observed training rows only; zero-MAD fallback to std then 1",
|
||||
"seed": SEED,
|
||||
"model_seeds": [SEED],
|
||||
"training_configuration": {
|
||||
"epoch_limit": EPOCH_LIMIT,
|
||||
"early_stopping_patience": PATIENCE,
|
||||
"batch_size": BATCH_SIZE,
|
||||
"optimizer": "AdamW",
|
||||
"learning_rate": LEARNING_RATE,
|
||||
"weight_decay": WEIGHT_DECAY,
|
||||
"gradient_clip_norm": 1.0,
|
||||
"early_stopping_metric": "mean validation joint CE + 0.5*SmoothL1 over 0.0/none, 0.3/single, 0.3/sync, 0.5/async",
|
||||
"architecture_preserved": {
|
||||
EARLYCONCAT: "EarlyConcat + BiGRU",
|
||||
MOFE7_MLP: "MoFE-7 + MLP Router",
|
||||
},
|
||||
"objective": "cross entropy + 0.5 * SmoothL1(intensity/3, label/3); same objective for both methods",
|
||||
"training_corruption": {
|
||||
"rates": list(TRAIN_RATES),
|
||||
"patterns": list(TRAIN_MODES),
|
||||
"preserve_at_least_fraction_per_selected_modality": 0.2,
|
||||
"generator_seed": TRAIN_MASK_SEED,
|
||||
"same_sample_masks_and_batch_orders_across_models": True,
|
||||
},
|
||||
},
|
||||
"validation_protocol": {
|
||||
"scenario_seed": SCENARIO_SEED,
|
||||
"scenario_count": len(valid_scenarios),
|
||||
"same_fixed_masks_for_both_models": True,
|
||||
"scenario_design": "math/Q2 42 controlled continuous-mask scenarios regenerated on each sample's original observation mask",
|
||||
"selection_scenarios": list(SELECTION_SCENARIOS),
|
||||
"selection_note": "Deterministic-model adaptation; uses joint supervised loss instead of C5's probabilistic selection NLL.",
|
||||
"aurc": "normalized trapezoidal MAE area over realized equal-modality-weighted additional missing rate for single/sync/partial/async at 0/.1/.3/.5/.7",
|
||||
},
|
||||
"test_protocol": {
|
||||
"official_test_final_clean_passes": 1,
|
||||
"test_used_for_training_or_checkpoint_selection": False,
|
||||
"metrics": ["accuracy", "macro_f1", "mae", "rmse", "pearson"],
|
||||
"paired_group_bootstrap_replicates": BOOTSTRAP_REPS,
|
||||
"bootstrap_unit": "source video id",
|
||||
"bootstrap_seed": TEST_BOOTSTRAP_SEED,
|
||||
},
|
||||
}
|
||||
(output_dir / "run_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
||||
(output_dir / "hypothesis.md").write_text(
|
||||
"# R03: 按 math/Q2 V2 口径重训两种保留模型\n\n"
|
||||
"## 假设\n\n"
|
||||
"在保持 EarlyConcat + BiGRU 与 MoFE-7 + MLP Router 结构及共同监督目标不变的情况下,"
|
||||
"使用数学方案中的官方划分、连续块缺失训练和 42 个固定验证情景,可以公平比较两种模型的干净测试表现与缺失鲁棒性。\n\n"
|
||||
"## 唯一实验改动\n\n"
|
||||
"相对现有检查点,本轮重新训练时将缺失训练改为 0/10/30/50/70% 与 single/sync/partial/async,"
|
||||
"每个被选模态至少保留 20% 观测;训练和批次顺序在两个模型间配对。数学方案中的 C5 概率损失不适用于现有确定性分类/回归头,"
|
||||
"因此保留项目既有的 CE + 0.5 SmoothL1 联合目标。\n\n"
|
||||
"## 数据使用\n\n"
|
||||
"标准化器只在官方训练集观测行上拟合;官方验证集只用于早停与缺失评估;官方测试集在全部检查点确定后做一次干净评估。\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
print(f"wrote retraining results to {output_dir}", flush=True)
|
||||
print(f"train/valid/test={train.n}/{valid.n}/{test.n}; device={device}; best_epochs={best_epochs}", flush=True)
|
||||
for row in test_rows:
|
||||
print(
|
||||
f"{row['method']}: Acc={row['accuracy']:.4f} Macro-F1={row['macro_f1']:.4f} "
|
||||
f"MAE={row['mae']:.4f} RMSE={row['rmse']:.4f} Pearson={row['pearson']:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--device", default="auto", choices=("auto", "cuda", "cpu"))
|
||||
parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR)
|
||||
arguments = parser.parse_args()
|
||||
run(device_name=arguments.device, output_dir=arguments.output_dir)
|
||||
@@ -0,0 +1,787 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from sklearn.metrics import accuracy_score, f1_score, mean_absolute_error
|
||||
from torch import nn
|
||||
|
||||
from .data import (
|
||||
ATTACHMENT2,
|
||||
MODALITIES,
|
||||
RobustStats,
|
||||
Split,
|
||||
apply_robust_stats,
|
||||
augment_masks,
|
||||
corrupt_masks,
|
||||
fit_robust_stats,
|
||||
load_aligned,
|
||||
)
|
||||
from .models import AlignedFusionModel
|
||||
from .mofe import EXPERT_NAMES, SUBSETS, MixtureOfFusionExperts
|
||||
from .train_compare import PATTERNS, _loss, _pearson, _train_one, seed_everything
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
REFERENCE_OUTPUT = ROOT / "outputs" / "mofe_7experts"
|
||||
DEFAULT_OUTPUT = ROOT / "outputs" / "mofe_7experts"
|
||||
EARLYCONCAT = "B0_early_concat"
|
||||
MOFE7_MLP = "B5_mofe_mlp"
|
||||
SEEDS = (42, 3407, 2026)
|
||||
RATES = (0.10, 0.20, 0.30)
|
||||
HIDDEN = 128
|
||||
LATENT_DIM = 64
|
||||
MODEL_CONFIG: dict[str, Any] = {
|
||||
"router": "mlp",
|
||||
"expert_names": EXPERT_NAMES,
|
||||
"availability_mode": "hard",
|
||||
}
|
||||
SUMMARY_METRICS = (
|
||||
"corrupt_macro_f1",
|
||||
"worst_condition_macro_f1",
|
||||
"text_30_macro_f1",
|
||||
"corrupt_mae",
|
||||
"corrupt_pearson",
|
||||
)
|
||||
|
||||
|
||||
def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not rows:
|
||||
return
|
||||
fields = list(dict.fromkeys(key for row in rows for key in row))
|
||||
with path.open("w", newline="", encoding="utf-8-sig") as stream:
|
||||
writer = csv.DictWriter(stream, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def _read_csv(path: Path) -> list[dict[str, str]]:
|
||||
if not path.exists():
|
||||
return []
|
||||
with path.open("r", newline="", encoding="utf-8-sig") as stream:
|
||||
return list(csv.DictReader(stream))
|
||||
|
||||
|
||||
def _sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _device_for(name: str) -> torch.device:
|
||||
if name == "auto":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(name)
|
||||
|
||||
|
||||
def _conditions(valid: Split, seed: int) -> list[tuple[str, float, np.ndarray]]:
|
||||
rows = [("clean", 0.0, valid.mask.copy())]
|
||||
for rate in RATES:
|
||||
for pattern_idx, (pattern, modalities) in enumerate(PATTERNS.items()):
|
||||
masks = corrupt_masks(
|
||||
valid.mask,
|
||||
rate,
|
||||
modalities,
|
||||
seed + 13 + pattern_idx * 101 + int(rate * 1000),
|
||||
)
|
||||
rows.append((f"{pattern}_{int(rate * 100)}", rate, masks))
|
||||
return rows
|
||||
|
||||
|
||||
def _metric_dict(
|
||||
y_cls: np.ndarray,
|
||||
y_reg: np.ndarray,
|
||||
logits: np.ndarray,
|
||||
intensity: np.ndarray,
|
||||
) -> dict[str, float]:
|
||||
predicted_class = np.asarray(logits).argmax(axis=-1)
|
||||
predicted_intensity = np.clip(np.asarray(intensity).reshape(-1), -3.0, 3.0)
|
||||
return {
|
||||
"accuracy": float(accuracy_score(y_cls, predicted_class)),
|
||||
"macro_f1": float(f1_score(y_cls, predicted_class, labels=[0, 1, 2], average="macro", zero_division=0)),
|
||||
"mae": float(mean_absolute_error(y_reg, predicted_intensity)),
|
||||
"pearson": _pearson(y_reg, predicted_intensity),
|
||||
}
|
||||
|
||||
|
||||
def _validation_loss(model: nn.Module, valid: Split, device: torch.device, batch_size: int) -> float:
|
||||
model.eval()
|
||||
values: list[float] = []
|
||||
weights: list[int] = []
|
||||
with torch.inference_mode():
|
||||
for start in range(0, valid.n, batch_size):
|
||||
end = min(start + batch_size, valid.n)
|
||||
xs = tuple(torch.as_tensor(x[start:end], dtype=torch.float32, device=device) for x in valid.x)
|
||||
masks = torch.as_tensor(valid.mask[start:end], dtype=torch.bool, device=device)
|
||||
y_cls = torch.as_tensor(valid.y_cls[start:end], dtype=torch.long, device=device)
|
||||
y_reg = torch.as_tensor(valid.y_reg[start:end], dtype=torch.float32, device=device)
|
||||
values.append(float(_loss(model(xs, masks), y_cls, y_reg).item()))
|
||||
weights.append(end - start)
|
||||
return float(np.average(values, weights=weights))
|
||||
|
||||
|
||||
def _train_mofe(
|
||||
train: Split,
|
||||
valid: Split,
|
||||
output_dir: Path,
|
||||
device: torch.device,
|
||||
seed: int,
|
||||
epochs: int,
|
||||
patience: int,
|
||||
batch_size: int,
|
||||
reuse_checkpoint: bool,
|
||||
) -> tuple[MixtureOfFusionExperts, int, list[dict[str, Any]]]:
|
||||
dims = tuple(int(x.shape[-1]) for x in train.x)
|
||||
checkpoint_path = output_dir / "model_best.pt"
|
||||
history_path = output_dir / "training_history.csv"
|
||||
if reuse_checkpoint and checkpoint_path.exists():
|
||||
saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
if saved.get("config") != MODEL_CONFIG or tuple(saved.get("dims", ())) != dims or int(saved.get("seed", -1)) != seed:
|
||||
raise ValueError(f"cached MoFE checkpoint does not match the selected configuration: {checkpoint_path}")
|
||||
model = MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
history = [
|
||||
{"method": MOFE7_MLP, "seed": seed, **{key: float(value) for key, value in row.items() if key in {"epoch", "train_loss", "valid_clean_loss"}}}
|
||||
for row in _read_csv(history_path)
|
||||
]
|
||||
return model.eval(), int(saved.get("best_epoch", 0)), history
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
seed_everything(seed)
|
||||
model = MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device)
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-4, weight_decay=1e-4)
|
||||
xs = tuple(torch.as_tensor(x, dtype=torch.float32, device=device) for x in train.x)
|
||||
base_masks = train.mask
|
||||
y_cls = torch.as_tensor(train.y_cls, dtype=torch.long, device=device)
|
||||
y_reg = torch.as_tensor(train.y_reg, dtype=torch.float32, device=device)
|
||||
rng = np.random.default_rng(seed + 809)
|
||||
best_loss = math.inf
|
||||
best_epoch = 0
|
||||
stale_epochs = 0
|
||||
history: list[dict[str, Any]] = []
|
||||
|
||||
for epoch in range(1, epochs + 1):
|
||||
model.train()
|
||||
order = rng.permutation(train.n)
|
||||
batch_losses: list[float] = []
|
||||
for start in range(0, train.n, batch_size):
|
||||
ids_np = order[start:start + batch_size]
|
||||
ids = torch.as_tensor(ids_np, dtype=torch.long, device=device)
|
||||
masks_np = augment_masks(base_masks[ids_np], rng)
|
||||
masks = torch.as_tensor(masks_np, dtype=torch.bool, device=device)
|
||||
output = model(tuple(x.index_select(0, ids) for x in xs), masks)
|
||||
loss = _loss(output, y_cls.index_select(0, ids), y_reg.index_select(0, ids))
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||||
optimizer.step()
|
||||
batch_losses.append(float(loss.detach().item()))
|
||||
|
||||
valid_loss = _validation_loss(model, valid, device, batch_size)
|
||||
row = {
|
||||
"method": MOFE7_MLP,
|
||||
"seed": seed,
|
||||
"epoch": epoch,
|
||||
"train_loss": float(np.mean(batch_losses)),
|
||||
"valid_clean_loss": valid_loss,
|
||||
}
|
||||
history.append(row)
|
||||
print(f"[MoFE-7 MLP] seed={seed} epoch={epoch:02d} train={row['train_loss']:.4f} valid={valid_loss:.4f}", flush=True)
|
||||
if valid_loss < best_loss - 1e-4:
|
||||
best_loss = valid_loss
|
||||
best_epoch = epoch
|
||||
stale_epochs = 0
|
||||
torch.save({
|
||||
"method": MOFE7_MLP,
|
||||
"config": MODEL_CONFIG,
|
||||
"dims": dims,
|
||||
"state_dict": model.state_dict(),
|
||||
"seed": seed,
|
||||
"best_epoch": epoch,
|
||||
}, checkpoint_path)
|
||||
else:
|
||||
stale_epochs += 1
|
||||
if stale_epochs >= patience:
|
||||
break
|
||||
|
||||
saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
model.eval()
|
||||
_write_csv(history_path, history)
|
||||
return model, best_epoch, history
|
||||
|
||||
|
||||
def _load_or_train_concat(
|
||||
train: Split,
|
||||
valid: Split,
|
||||
output_dir: Path,
|
||||
device: torch.device,
|
||||
seed: int,
|
||||
epochs: int,
|
||||
patience: int,
|
||||
batch_size: int,
|
||||
reuse_checkpoint: bool,
|
||||
) -> tuple[AlignedFusionModel, int, list[dict[str, Any]]]:
|
||||
checkpoint_path = output_dir / "model_best.pt"
|
||||
dims = tuple(int(x.shape[-1]) for x in train.x)
|
||||
if reuse_checkpoint and checkpoint_path.exists():
|
||||
saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
if saved.get("kind") != "concat" or tuple(saved.get("dims", ())) != dims or int(saved.get("seed", -1)) != seed:
|
||||
raise ValueError(f"cached EarlyConcat checkpoint does not match: {checkpoint_path}")
|
||||
model = AlignedFusionModel("concat", dims=dims).to(device)
|
||||
model.load_state_dict(saved["state_dict"])
|
||||
history = [
|
||||
{"method": EARLYCONCAT, "seed": seed, **{key: float(value) for key, value in row.items() if key in {"epoch", "train_loss", "valid_clean_loss"}}}
|
||||
for row in _read_csv(output_dir / "training_history.csv")
|
||||
]
|
||||
return model.eval(), int(saved.get("best_epoch", 0)), history
|
||||
|
||||
model, best_epoch, history = _train_one(
|
||||
"concat", train, valid, output_dir, device, seed, epochs, patience, batch_size
|
||||
)
|
||||
rows = [{"method": EARLYCONCAT, "seed": seed, **row} for row in history]
|
||||
return model.eval(), best_epoch, rows
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _predict(
|
||||
model: nn.Module,
|
||||
split: Split,
|
||||
masks: np.ndarray,
|
||||
device: torch.device,
|
||||
batch_size: int,
|
||||
force_expert: str | None = None,
|
||||
) -> dict[str, np.ndarray]:
|
||||
fields = ["logits", "intensity"]
|
||||
if isinstance(model, MixtureOfFusionExperts):
|
||||
fields.extend(("alpha", "utility", "availability", "fallback"))
|
||||
chunks: dict[str, list[np.ndarray]] = {name: [] for name in fields}
|
||||
for start in range(0, split.n, batch_size):
|
||||
end = min(start + batch_size, split.n)
|
||||
xs = tuple(torch.as_tensor(x[start:end], dtype=torch.float32, device=device) for x in split.x)
|
||||
mask_batch = torch.as_tensor(masks[start:end], dtype=torch.bool, device=device)
|
||||
output = model(xs, mask_batch, force_expert=force_expert) if isinstance(model, MixtureOfFusionExperts) else model(xs, mask_batch)
|
||||
for name in fields:
|
||||
value = output[name]
|
||||
chunks[name].append(value.float().cpu().numpy())
|
||||
result = {name: np.concatenate(values, axis=0) for name, values in chunks.items()}
|
||||
result["intensity"] = np.clip(result["intensity"].reshape(-1), -3.0, 3.0)
|
||||
return result
|
||||
|
||||
|
||||
def _condition_row(
|
||||
method: str,
|
||||
seed: int,
|
||||
condition: str,
|
||||
rate: float,
|
||||
split: Split,
|
||||
prediction: dict[str, np.ndarray],
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"method": method,
|
||||
"seed": seed,
|
||||
"condition": condition,
|
||||
"missing_rate": rate,
|
||||
"n_valid": split.n,
|
||||
**_metric_dict(split.y_cls, split.y_reg, prediction["logits"], prediction["intensity"]),
|
||||
}
|
||||
|
||||
|
||||
def _diagnostics(
|
||||
seed: int,
|
||||
condition: str,
|
||||
masks: np.ndarray,
|
||||
prediction: dict[str, np.ndarray],
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
alpha = prediction["alpha"]
|
||||
availability = prediction["availability"].astype(bool)
|
||||
active = availability.any(axis=-1)
|
||||
active_alpha = alpha[active]
|
||||
if active_alpha.size:
|
||||
means = active_alpha.mean(axis=0)
|
||||
entropy = -(active_alpha * np.log(np.maximum(active_alpha, 1e-12))).sum(axis=-1) / np.log(len(EXPERT_NAMES))
|
||||
high_weight = (active_alpha.max(axis=-1) > 0.8).mean()
|
||||
else:
|
||||
means = np.zeros(len(EXPERT_NAMES), dtype=np.float64)
|
||||
entropy = np.zeros(0, dtype=np.float64)
|
||||
high_weight = 0.0
|
||||
route_row: dict[str, Any] = {
|
||||
"method": MOFE7_MLP,
|
||||
"seed": seed,
|
||||
"condition": condition,
|
||||
"active_position_fraction": float(active.mean()),
|
||||
"fallback_position_fraction": float((~active).mean()),
|
||||
"normalized_router_entropy": float(entropy.mean()) if entropy.size else 0.0,
|
||||
"fraction_active_positions_max_weight_over_0p8": float(high_weight),
|
||||
}
|
||||
for index, name in enumerate(EXPERT_NAMES):
|
||||
route_row[f"alpha_{name}_mean"] = float(means[index])
|
||||
utility = prediction["utility"]
|
||||
utility_row: dict[str, Any] = {"method": MOFE7_MLP, "seed": seed, "condition": condition}
|
||||
for modality, name in enumerate(MODALITIES):
|
||||
observed = masks[..., modality]
|
||||
utility_row[f"utility_{name}_mean"] = float(utility[..., modality][observed].mean()) if observed.any() else 0.0
|
||||
return route_row, utility_row
|
||||
|
||||
|
||||
def _summary_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
summaries: list[dict[str, Any]] = []
|
||||
for method in (EARLYCONCAT, MOFE7_MLP):
|
||||
matching = [row for row in rows if row["method"] == method]
|
||||
seeds = sorted({int(row["seed"]) for row in matching})
|
||||
conditions = list(dict.fromkeys(row["condition"] for row in matching))
|
||||
by_seed_condition = {(int(row["seed"]), row["condition"]): row for row in matching}
|
||||
clean = [by_seed_condition[(seed, "clean")] for seed in seeds]
|
||||
corrupt_conditions = [condition for condition in conditions if condition != "clean"]
|
||||
corrupt_by_seed = {
|
||||
seed: [by_seed_condition[(seed, condition)] for condition in corrupt_conditions]
|
||||
for seed in seeds
|
||||
}
|
||||
condition_f1 = {
|
||||
condition: float(np.mean([by_seed_condition[(seed, condition)]["macro_f1"] for seed in seeds]))
|
||||
for condition in corrupt_conditions
|
||||
}
|
||||
worst_condition = min(condition_f1, key=condition_f1.get)
|
||||
row: dict[str, Any] = {"method": method, "n_seeds": len(seeds), "worst_condition": worst_condition}
|
||||
for metric in ("accuracy", "macro_f1", "mae", "pearson"):
|
||||
clean_values = [float(item[metric]) for item in clean]
|
||||
corrupt_values = [float(np.mean([item[metric] for item in corrupt_by_seed[seed]])) for seed in seeds]
|
||||
row[f"clean_{metric}"] = float(np.mean(clean_values))
|
||||
row[f"clean_{metric}_sd"] = float(np.std(clean_values, ddof=1)) if len(clean_values) > 1 else 0.0
|
||||
row[f"corrupt_{metric}_mean"] = float(np.mean(corrupt_values))
|
||||
row[f"corrupt_{metric}_sd"] = float(np.std(corrupt_values, ddof=1)) if len(corrupt_values) > 1 else 0.0
|
||||
row["worst_condition_macro_f1"] = condition_f1[worst_condition]
|
||||
row["worst_single_run_macro_f1"] = min(
|
||||
item["macro_f1"] for seed in seeds for item in corrupt_by_seed[seed]
|
||||
)
|
||||
text_30 = [by_seed_condition[(seed, "text_30")] for seed in seeds]
|
||||
row["text_30_macro_f1"] = float(np.mean([item["macro_f1"] for item in text_30]))
|
||||
row["text_30_macro_f1_sd"] = float(np.std([item["macro_f1"] for item in text_30], ddof=1)) if len(text_30) > 1 else 0.0
|
||||
for condition in ("audio_30", "vision_30", "audio_vision_30", "all_modalities_30"):
|
||||
values = [by_seed_condition[(seed, condition)]["macro_f1"] for seed in seeds]
|
||||
row[f"{condition}_macro_f1"] = float(np.mean(values))
|
||||
row[f"{condition}_macro_f1_sd"] = float(np.std(values, ddof=1)) if len(values) > 1 else 0.0
|
||||
row["corrupt_macro_f1"] = row["corrupt_macro_f1_mean"]
|
||||
row["corrupt_mae"] = row["corrupt_mae_mean"]
|
||||
row["corrupt_pearson"] = row["corrupt_pearson_mean"]
|
||||
summaries.append(row)
|
||||
return summaries
|
||||
|
||||
|
||||
def _bootstrap_distributions(
|
||||
method: str,
|
||||
predictions: dict[tuple[str, int, str], dict[str, np.ndarray]],
|
||||
valid: Split,
|
||||
seeds: list[int],
|
||||
conditions: list[str],
|
||||
group_counts: np.ndarray,
|
||||
) -> dict[str, np.ndarray]:
|
||||
group_names = sorted({sample_id.split("$_$", 1)[0] for sample_id in valid.ids})
|
||||
group_index = {name: index for index, name in enumerate(group_names)}
|
||||
row_group = np.asarray([group_index[sample_id.split("$_$", 1)[0]] for sample_id in valid.ids], dtype=np.int64)
|
||||
n_groups = len(group_names)
|
||||
n_slots = len(seeds) * len(conditions)
|
||||
confusion_by_group = np.zeros((n_groups, n_slots, 9), dtype=np.float64)
|
||||
regression_by_group = np.zeros((n_groups, n_slots, 7), dtype=np.float64)
|
||||
for seed_index, seed in enumerate(seeds):
|
||||
for condition_index, condition in enumerate(conditions):
|
||||
slot = seed_index * len(conditions) + condition_index
|
||||
pred = predictions[(method, seed, condition)]
|
||||
predicted_class = pred["logits"].argmax(axis=-1)
|
||||
code = valid.y_cls * 3 + predicted_class
|
||||
np.add.at(confusion_by_group[:, slot, :], (row_group, code), 1.0)
|
||||
intensity = np.clip(pred["intensity"].reshape(-1), -3.0, 3.0)
|
||||
values = np.stack((
|
||||
np.ones(valid.n),
|
||||
np.abs(valid.y_reg - intensity),
|
||||
valid.y_reg,
|
||||
valid.y_reg ** 2,
|
||||
intensity,
|
||||
intensity ** 2,
|
||||
valid.y_reg * intensity,
|
||||
), axis=-1)
|
||||
for statistic in range(values.shape[-1]):
|
||||
np.add.at(regression_by_group[:, slot, statistic], row_group, values[:, statistic])
|
||||
|
||||
weighted_confusion = np.einsum("rg,gsk->rsk", group_counts, confusion_by_group, optimize=True)
|
||||
cm = weighted_confusion.reshape(len(group_counts), len(seeds), len(conditions), 3, 3)
|
||||
true_count = cm.sum(axis=-1)
|
||||
predicted_count = cm.sum(axis=-2)
|
||||
true_positive = np.diagonal(cm, axis1=-2, axis2=-1)
|
||||
denominator = true_count + predicted_count
|
||||
class_f1 = np.divide(2.0 * true_positive, denominator, out=np.zeros_like(true_positive), where=denominator > 0)
|
||||
macro_f1 = class_f1.mean(axis=-1)
|
||||
|
||||
weighted_regression = np.einsum("rg,gsk->rsk", group_counts, regression_by_group, optimize=True)
|
||||
regression = weighted_regression.reshape(len(group_counts), len(seeds), len(conditions), 7)
|
||||
count = np.maximum(regression[..., 0], 1.0)
|
||||
mae = regression[..., 1] / count
|
||||
sum_y, sum_y2, sum_pred, sum_pred2, sum_yp = (regression[..., index] for index in range(2, 7))
|
||||
covariance = sum_yp - sum_y * sum_pred / count
|
||||
variance_y = np.maximum(sum_y2 - sum_y ** 2 / count, 0.0)
|
||||
variance_pred = np.maximum(sum_pred2 - sum_pred ** 2 / count, 0.0)
|
||||
denominator_corr = np.sqrt(variance_y * variance_pred)
|
||||
pearson = np.divide(covariance, denominator_corr, out=np.zeros_like(covariance), where=denominator_corr > 1e-12)
|
||||
text_30_index = conditions.index("text_30")
|
||||
return {
|
||||
"corrupt_macro_f1": macro_f1[:, :, 1:].mean(axis=(1, 2)),
|
||||
"worst_condition_macro_f1": macro_f1[:, :, 1:].mean(axis=1).min(axis=1),
|
||||
"text_30_macro_f1": macro_f1[:, :, text_30_index].mean(axis=1),
|
||||
"corrupt_mae": mae[:, :, 1:].mean(axis=(1, 2)),
|
||||
"corrupt_pearson": pearson[:, :, 1:].mean(axis=(1, 2)),
|
||||
}
|
||||
|
||||
|
||||
def _paired_bootstrap(
|
||||
predictions: dict[tuple[str, int, str], dict[str, np.ndarray]],
|
||||
valid: Split,
|
||||
seeds: list[int],
|
||||
conditions: list[str],
|
||||
reps: int,
|
||||
bootstrap_seed: int,
|
||||
summaries: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
groups = sorted({sample_id.split("$_$", 1)[0] for sample_id in valid.ids})
|
||||
rng = np.random.default_rng(bootstrap_seed)
|
||||
draws = rng.integers(0, len(groups), size=(reps, len(groups)))
|
||||
group_counts = np.zeros((reps, len(groups)), dtype=np.float64)
|
||||
for rep in range(reps):
|
||||
group_counts[rep] = np.bincount(draws[rep], minlength=len(groups))
|
||||
candidate = _bootstrap_distributions(MOFE7_MLP, predictions, valid, seeds, conditions, group_counts)
|
||||
reference = _bootstrap_distributions(EARLYCONCAT, predictions, valid, seeds, conditions, group_counts)
|
||||
summary_map = {row["method"]: row for row in summaries}
|
||||
point_keys = {
|
||||
"corrupt_macro_f1": "corrupt_macro_f1",
|
||||
"worst_condition_macro_f1": "worst_condition_macro_f1",
|
||||
"text_30_macro_f1": "text_30_macro_f1",
|
||||
"corrupt_mae": "corrupt_mae",
|
||||
"corrupt_pearson": "corrupt_pearson",
|
||||
}
|
||||
rows = []
|
||||
for metric in SUMMARY_METRICS:
|
||||
delta = candidate[metric] - reference[metric]
|
||||
key = point_keys[metric]
|
||||
rows.append({
|
||||
"comparison": "MoFE-7 MLP vs EarlyConcat",
|
||||
"candidate": MOFE7_MLP,
|
||||
"reference": EARLYCONCAT,
|
||||
"metric": metric,
|
||||
"delta_candidate_minus_reference": float(summary_map[MOFE7_MLP][key] - summary_map[EARLYCONCAT][key]),
|
||||
"bootstrap_ci_2p5": float(np.quantile(delta, 0.025)),
|
||||
"bootstrap_ci_97p5": float(np.quantile(delta, 0.975)),
|
||||
"bootstrap_probability_delta_gt_0": float(np.mean(delta > 0.0)),
|
||||
"bootstrap_replicates": reps,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
"seed": bootstrap_seed,
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def _plot_summary(output: Path, summaries: list[dict[str, Any]]) -> None:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
labels = ["EarlyConcat + BiGRU", "MoFE-7 + MLP Router"]
|
||||
by_method = {row["method"]: row for row in summaries}
|
||||
methods = (EARLYCONCAT, MOFE7_MLP)
|
||||
metrics = ("clean_macro_f1", "corrupt_macro_f1", "worst_condition_macro_f1")
|
||||
names = ("Clean", "Mean corrupted", "Worst condition")
|
||||
x = np.arange(len(names))
|
||||
width = 0.34
|
||||
fig, ax = plt.subplots(figsize=(8.6, 4.8), constrained_layout=True)
|
||||
for offset, method, label, color in (
|
||||
(-width / 2, methods[0], labels[0], "#4e79a7"),
|
||||
(width / 2, methods[1], labels[1], "#f28e2b"),
|
||||
):
|
||||
values = [by_method[method][metric] for metric in metrics]
|
||||
ax.bar(x + offset, values, width, label=label, color=color)
|
||||
ax.set_xticks(x, names)
|
||||
ax.set_ylabel("Macro-F1")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_title("Q2 selected-model validation comparison")
|
||||
ax.legend(frameon=False)
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(output / "comparison_earlyconcat_mofe7.png", dpi=180)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def _parameter_rows(dims: tuple[int, int, int], device: torch.device) -> list[dict[str, Any]]:
|
||||
models: dict[str, nn.Module] = {
|
||||
EARLYCONCAT: AlignedFusionModel("concat", dims=dims),
|
||||
MOFE7_MLP: MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG),
|
||||
}
|
||||
baseline_count = sum(parameter.numel() for parameter in models[EARLYCONCAT].parameters() if parameter.requires_grad)
|
||||
rows = []
|
||||
for name, model in models.items():
|
||||
count = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
|
||||
rows.append({
|
||||
"method": name,
|
||||
"trainable_parameters": count,
|
||||
"ratio_to_earlyconcat": count / baseline_count,
|
||||
"within_2x_earlyconcat": bool(count <= 2 * baseline_count),
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def _smoke_test(train: Split, output: Path, device: torch.device, seed: int) -> dict[str, Any]:
|
||||
seed_everything(seed)
|
||||
dims = tuple(int(x.shape[-1]) for x in train.x)
|
||||
count = min(4, train.n)
|
||||
xs = tuple(torch.as_tensor(x[:count], dtype=torch.float32, device=device) for x in train.x)
|
||||
masks = torch.as_tensor(train.mask[:count].copy(), dtype=torch.bool, device=device)
|
||||
masks[0] = True
|
||||
if count > 1:
|
||||
masks[1, 5:12, 0] = False
|
||||
if count > 2:
|
||||
masks[2, 18:23, :] = False
|
||||
target_class = torch.as_tensor(train.y_cls[:count], dtype=torch.long, device=device)
|
||||
target_intensity = torch.as_tensor(train.y_reg[:count], dtype=torch.float32, device=device)
|
||||
reports: dict[str, Any] = {}
|
||||
models: dict[str, nn.Module] = {
|
||||
EARLYCONCAT: AlignedFusionModel("concat", dims=dims).to(device),
|
||||
MOFE7_MLP: MixtureOfFusionExperts(dims=dims, **MODEL_CONFIG).to(device),
|
||||
}
|
||||
for name, model in models.items():
|
||||
model.train()
|
||||
result = model(xs, masks)
|
||||
loss = _loss(result, target_class, target_intensity)
|
||||
loss.backward()
|
||||
gradient = sum(float(p.grad.detach().abs().sum().cpu()) for p in model.parameters() if p.grad is not None)
|
||||
reports[name] = {
|
||||
"logits_shape": list(result["logits"].shape),
|
||||
"intensity_shape": list(result["intensity"].shape),
|
||||
"finite_loss": bool(torch.isfinite(loss).item()),
|
||||
"gradient_l1": gradient,
|
||||
}
|
||||
mofe_model = models[MOFE7_MLP]
|
||||
mo = mofe_model(xs, masks)
|
||||
active = mo["availability"].any(dim=-1)
|
||||
alpha_sums = mo["alpha"].sum(dim=-1)
|
||||
alpha_error = float((alpha_sums[active] - 1).abs().max().cpu()) if active.any() else 0.0
|
||||
unavailable_weights = float(mo["alpha"].masked_select(~mo["availability"]).abs().max().cpu()) if (~mo["availability"]).any() else 0.0
|
||||
expert_gradients = {
|
||||
name: sum(float(parameter.grad.detach().abs().sum().cpu()) for parameter in expert.parameters() if parameter.grad is not None)
|
||||
for name, expert in mofe_model.experts.items()
|
||||
}
|
||||
router_gradient = sum(float(parameter.grad.detach().abs().sum().cpu()) for parameter in mofe_model.router.parameters() if parameter.grad is not None)
|
||||
if alpha_error > 1e-6 or unavailable_weights > 1e-8:
|
||||
raise RuntimeError(f"MoFE routing mask invariant failed: sum_error={alpha_error}, unavailable={unavailable_weights}")
|
||||
if not all(value > 0 for value in expert_gradients.values()) or router_gradient <= 0:
|
||||
raise RuntimeError(f"MoFE expert/router gradients are incomplete: {expert_gradients}; router={router_gradient}")
|
||||
report = {
|
||||
"passed": all(item["finite_loss"] and item["gradient_l1"] > 0 for item in reports.values()),
|
||||
"seed": seed,
|
||||
"device": str(device),
|
||||
"cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
|
||||
"batch_size_checked": count,
|
||||
"steps": train.steps,
|
||||
"models": reports,
|
||||
"mofe_experts": list(EXPERT_NAMES),
|
||||
"mofe_alpha_shape": list(mo["alpha"].shape),
|
||||
"mofe_max_weight_sum_error": alpha_error,
|
||||
"mofe_max_weight_on_unavailable_experts": unavailable_weights,
|
||||
"mofe_expert_gradient_l1": expert_gradients,
|
||||
"mofe_router_gradient_l1": router_gradient,
|
||||
"parameter_count": {row["method"]: row["trainable_parameters"] for row in _parameter_rows(dims, device)},
|
||||
}
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
(output / "smoke_test.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
|
||||
return report
|
||||
|
||||
|
||||
def _run(args: argparse.Namespace) -> None:
|
||||
output = args.output_dir.resolve()
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
device = _device_for(args.device)
|
||||
torch.set_num_threads(args.threads)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
raw = load_aligned()
|
||||
computed_stats = fit_robust_stats(raw["train"])
|
||||
reference_stats_path = REFERENCE_OUTPUT / "aligned_robust_stats.npz"
|
||||
if reference_stats_path.exists():
|
||||
stats = RobustStats.load(reference_stats_path)
|
||||
scaler_diff = max(
|
||||
max(float(np.max(np.abs(a - b))) for a, b in zip(computed_stats.center, stats.center)),
|
||||
max(float(np.max(np.abs(a - b))) for a, b in zip(computed_stats.scale, stats.scale)),
|
||||
)
|
||||
else:
|
||||
stats = computed_stats
|
||||
scaler_diff = 0.0
|
||||
train = apply_robust_stats(raw["train"], stats)
|
||||
valid = apply_robust_stats(raw["valid"], stats)
|
||||
stats.save(output / "aligned_robust_stats.npz")
|
||||
dims = tuple(int(x.shape[-1]) for x in train.x)
|
||||
feature_path = ATTACHMENT2 / "aligned_50.pkl"
|
||||
if not feature_path.exists():
|
||||
raise FileNotFoundError(f"official aligned feature file not found: {feature_path}")
|
||||
|
||||
if args.phase == "smoke":
|
||||
report = _smoke_test(train, output, device, args.seeds[0])
|
||||
report["scaler_max_abs_difference_from_reference"] = scaler_diff
|
||||
(output / "smoke_test.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
|
||||
print(f"selected-model smoke: passed={report['passed']} device={device}", flush=True)
|
||||
return
|
||||
|
||||
seeds = list(args.seeds)
|
||||
metrics_rows: list[dict[str, Any]] = []
|
||||
predictions: dict[tuple[str, int, str], dict[str, np.ndarray]] = {}
|
||||
router_rows: list[dict[str, Any]] = []
|
||||
utility_rows: list[dict[str, Any]] = []
|
||||
expert_rows: list[dict[str, Any]] = []
|
||||
history_rows: list[dict[str, Any]] = []
|
||||
best_epochs: dict[str, int] = {}
|
||||
condition_names: list[str] = []
|
||||
|
||||
for seed in seeds:
|
||||
baseline_dir = output / "models" / "baselines" / "concat" / f"seed_{seed}"
|
||||
baseline, baseline_epoch, baseline_history = _load_or_train_concat(
|
||||
train, valid, baseline_dir, device, seed, args.epochs, args.patience,
|
||||
args.batch_size, args.reuse_checkpoints and not args.force_retrain,
|
||||
)
|
||||
mofe_dir = output / "models" / MOFE7_MLP / f"seed_{seed}"
|
||||
mofe, mofe_epoch, mofe_history = _train_mofe(
|
||||
train, valid, mofe_dir, device, seed, args.epochs, args.patience,
|
||||
args.batch_size, args.reuse_checkpoints and not args.force_retrain,
|
||||
)
|
||||
best_epochs[f"{EARLYCONCAT}_seed_{seed}"] = baseline_epoch
|
||||
best_epochs[f"{MOFE7_MLP}_seed_{seed}"] = mofe_epoch
|
||||
history_rows.extend(baseline_history)
|
||||
history_rows.extend(mofe_history)
|
||||
|
||||
conditions = _conditions(valid, seed)
|
||||
names = [condition for condition, _, _ in conditions]
|
||||
if condition_names and names != condition_names:
|
||||
raise RuntimeError("validation condition ordering changed between seeds")
|
||||
condition_names = names
|
||||
for method, model in ((EARLYCONCAT, baseline), (MOFE7_MLP, mofe)):
|
||||
for condition, rate, masks in conditions:
|
||||
prediction = _predict(model, valid, masks, device, args.batch_size)
|
||||
predictions[(method, seed, condition)] = prediction
|
||||
metrics_rows.append(_condition_row(method, seed, condition, rate, valid, prediction))
|
||||
if method == MOFE7_MLP:
|
||||
route_row, utility_row = _diagnostics(seed, condition, masks, prediction)
|
||||
router_rows.append(route_row)
|
||||
utility_rows.append(utility_row)
|
||||
for expert in EXPERT_NAMES:
|
||||
forced = _predict(model, valid, masks, device, args.batch_size, force_expert=expert)
|
||||
observed = masks[..., list(SUBSETS[expert])].all(axis=-1)
|
||||
expert_rows.append({
|
||||
"method": MOFE7_MLP,
|
||||
"seed": seed,
|
||||
"condition": condition,
|
||||
"expert": expert,
|
||||
"available_position_fraction": float(observed.mean()),
|
||||
**_metric_dict(valid.y_cls, valid.y_reg, forced["logits"], forced["intensity"]),
|
||||
})
|
||||
print(f"evaluated {method}/seed{seed}", flush=True)
|
||||
del baseline, mofe
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
summaries = _summary_rows(metrics_rows)
|
||||
paired = _paired_bootstrap(
|
||||
predictions, valid, seeds, condition_names, args.bootstrap_reps,
|
||||
args.bootstrap_seed, summaries,
|
||||
) if args.bootstrap_reps > 0 else []
|
||||
parameter_rows = _parameter_rows(dims, device)
|
||||
output_rows = {
|
||||
"metrics_by_condition.csv": metrics_rows,
|
||||
"summary.csv": summaries,
|
||||
"paired_bootstrap.csv": paired,
|
||||
"parameter_count.csv": parameter_rows,
|
||||
"router_weights_by_condition.csv": router_rows,
|
||||
"routing_entropy.csv": router_rows,
|
||||
"modality_utility_by_condition.csv": utility_rows,
|
||||
"expert_condition_matrix.csv": expert_rows,
|
||||
"training_history.csv": history_rows,
|
||||
}
|
||||
for filename, rows in output_rows.items():
|
||||
_write_csv(output / filename, rows)
|
||||
_plot_summary(output / "figures", summaries)
|
||||
|
||||
manifest = {
|
||||
"experiment": "Q2 selected models: EarlyConcat + BiGRU and MoFE-7 + MLP Router",
|
||||
"created_unix": time.time(),
|
||||
"python_version": sys.version,
|
||||
"torch_version": torch.__version__,
|
||||
"numpy_version": np.__version__,
|
||||
"device": str(device),
|
||||
"cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
|
||||
"feature_file": str(feature_path),
|
||||
"feature_sha256": _sha256(feature_path),
|
||||
"feature_dimensions": dict(zip(MODALITIES, dims)),
|
||||
"sequence_length": train.steps,
|
||||
"representation_note": "official ordered 50-wordpiece positions; not 50 physical-time bins",
|
||||
"train_examples": train.n,
|
||||
"valid_examples": valid.n,
|
||||
"train_source_video_groups": len({sample_id.split("$_$", 1)[0] for sample_id in train.ids}),
|
||||
"valid_source_video_groups": len({sample_id.split("$_$", 1)[0] for sample_id in valid.ids}),
|
||||
"train_only_scaler": str(output / "aligned_robust_stats.npz"),
|
||||
"scaler_max_abs_difference_from_reference": scaler_diff,
|
||||
"test_labels_used": False,
|
||||
"seeds": seeds,
|
||||
"epochs_max": args.epochs,
|
||||
"patience": args.patience,
|
||||
"batch_size": args.batch_size,
|
||||
"optimizer": "AdamW(lr=1.5e-4, weight_decay=1e-4), gradient clip 1.0",
|
||||
"training_mask_augmentation": "same contiguous-block augment_masks protocol for both models",
|
||||
"validation_conditions": condition_names,
|
||||
"validation_corruption_seed": "seed + 13 + pattern_index*101 + int(rate*1000)",
|
||||
"loss": "cross_entropy + 0.5*SmoothL1(intensity/3, regression_label/3)",
|
||||
"models": {
|
||||
EARLYCONCAT: "project modalities independently, concatenate features and masks, then BiGRU",
|
||||
MOFE7_MLP: {
|
||||
"experts": list(EXPERT_NAMES),
|
||||
"router": "MLP over per-position observed values and local observation statistics",
|
||||
"availability": "hard mask; unavailable expert weights are zero",
|
||||
"shared_temporal_backbone": "one BiGRU after position-wise expert mixture",
|
||||
},
|
||||
},
|
||||
"best_epochs": best_epochs,
|
||||
"paired_bootstrap": {
|
||||
"replicates": args.bootstrap_reps,
|
||||
"seed": args.bootstrap_seed,
|
||||
"resampling_unit": "source video id",
|
||||
"paired": True,
|
||||
},
|
||||
}
|
||||
(output / "run_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
||||
print(f"selected-model results saved to {output}", flush=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Train and compare the two retained Q2 models.")
|
||||
parser.add_argument("--phase", choices=("smoke", "full"), default="full")
|
||||
parser.add_argument("--epochs", type=int, default=32)
|
||||
parser.add_argument("--patience", type=int, default=6)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--threads", type=int, default=4)
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--seeds", type=int, nargs="+", default=list(SEEDS))
|
||||
parser.add_argument("--bootstrap-reps", type=int, default=1000)
|
||||
parser.add_argument("--bootstrap-seed", type=int, default=20260924)
|
||||
parser.add_argument("--reuse-checkpoints", action="store_true")
|
||||
parser.add_argument("--force-retrain", action="store_true")
|
||||
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
|
||||
_run(parser.parse_args())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user