#!/usr/bin/env python3 """Run revised RAVEL staged experiments with reusable outputs. The script is resumable: a run is skipped when its metrics JSON and prediction CSV already exist, unless `--overwrite` is passed. """ from __future__ import annotations import argparse import csv import json import math import random import sys import time from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score from torch.optim import AdamW from transformers import CLIPProcessor, DebertaV2Tokenizer from transformers.utils import logging as hf_logging hf_logging.set_verbosity_error() PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from src.revised_ravel_model import token_loss @dataclass(frozen=True) class MethodSpec: key: str display: str architecture: str train_target: str enable_lora: bool = True use_coattention: bool = True use_auxiliary: bool = True use_disagreement: bool = True parameter_matched: bool = False disagreement_formulation: str = "signed_difference" additional_loss: str = "none" lambda_primary: float = 0.5 lambda_unimodal: float = 0.25 consistency_beta: float = 0.25 infonce_alpha: float = 0.10 infonce_temperature: float = 0.07 METHODS: Dict[str, MethodSpec] = { "frozen_concat": MethodSpec( key="frozen_concat", display="Frozen concat", architecture="token", train_target="no_coattn", enable_lora=False, use_coattention=False, use_auxiliary=False, use_disagreement=False, ), "lora_concat": MethodSpec( key="lora_concat", display="LoRA concat", architecture="token", train_target="no_coattn", enable_lora=True, use_coattention=False, use_auxiliary=False, use_disagreement=False, ), "frozen_token_coattn": MethodSpec( key="frozen_token_coattn", display="Frozen token co-attention", architecture="token", train_target="primary", enable_lora=False, use_coattention=True, use_auxiliary=False, use_disagreement=False, ), "legacy_global": MethodSpec( key="legacy_global", display="RAVEL-Global", architecture="legacy_global", train_target="legacy_full", enable_lora=True, use_coattention=False, use_auxiliary=False, use_disagreement=False, ), "token_coattn": MethodSpec( key="token_coattn", display="Token co-attention", architecture="token", train_target="primary", enable_lora=True, use_coattention=True, use_auxiliary=False, use_disagreement=False, ), "token_aux": MethodSpec( key="token_aux", display="Token co-attention + unimodal heads", architecture="token", train_target="primary_aux", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=False, ), "param_mlp": MethodSpec( key="param_mlp", display="Parameter-matched MLP", architecture="token", train_target="param_mlp", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=False, parameter_matched=True, ), "full_revised": MethodSpec( key="full_revised", display="Full revised RAVEL", architecture="token", train_target="full", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, ), "text_only": MethodSpec( key="text_only", display="Text-only", architecture="token", train_target="text_only", enable_lora=True, use_coattention=False, use_auxiliary=True, use_disagreement=False, ), "vision_only": MethodSpec( key="vision_only", display="Vision-only", architecture="token", train_target="vision_only", enable_lora=True, use_coattention=False, use_auxiliary=True, use_disagreement=False, ), "absolute_difference": MethodSpec( key="absolute_difference", display="Absolute posterior difference", architecture="token", train_target="disagreement_feature", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, disagreement_formulation="absolute_difference", ), "js_divergence": MethodSpec( key="js_divergence", display="Jensen-Shannon divergence", architecture="token", train_target="disagreement_feature", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, disagreement_formulation="js_divergence", ), "log_probability_ratio": MethodSpec( key="log_probability_ratio", display="Log-probability ratio", architecture="token", train_target="disagreement_feature", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, disagreement_formulation="log_probability_ratio", ), "attention_discrepancy": MethodSpec( key="attention_discrepancy", display="Attention discrepancy", architecture="token", train_target="disagreement_feature", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, disagreement_formulation="attention_discrepancy", ), "consistency_loss": MethodSpec( key="consistency_loss", display="JS consistency loss", architecture="token", train_target="consistency_loss", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=False, additional_loss="js_consistency", ), "infonce_alignment": MethodSpec( key="infonce_alignment", display="InfoNCE alignment", architecture="token", train_target="infonce_alignment", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=False, additional_loss="infonce_alignment", ), "lambda_u_0_0": MethodSpec( key="lambda_u_0_0", display="Full revised RAVEL lambda_u=0.0", architecture="token", train_target="full", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, lambda_unimodal=0.0, ), "lambda_u_0_25": MethodSpec( key="lambda_u_0_25", display="Full revised RAVEL lambda_u=0.25", architecture="token", train_target="full", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, lambda_unimodal=0.25, ), "lambda_u_0_5": MethodSpec( key="lambda_u_0_5", display="Full revised RAVEL lambda_u=0.5", architecture="token", train_target="full", enable_lora=True, use_coattention=True, use_auxiliary=True, use_disagreement=True, lambda_unimodal=0.5, ), } STAGE_B_METHODS = [ "frozen_concat", "lora_concat", "frozen_token_coattn", "legacy_global", "token_coattn", "token_aux", "param_mlp", "full_revised", ] STAGE_C_METHODS = [ "legacy_global", "text_only", "vision_only", "token_coattn", "token_aux", "param_mlp", "full_revised", ] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run revised RAVEL staged experiments.") parser.add_argument("--stage", choices=["stage_b", "stage_c", "all"], default="all") parser.add_argument("--datasets", nargs="+", default=["mvsa_multiple", "hfm_existing", "hfm_deleak"]) parser.add_argument("--seeds", nargs="+", type=int, default=[1, 3, 5, 7, 11]) parser.add_argument("--stage-b-seed", type=int, default=7) parser.add_argument("--methods", nargs="+", default=None, choices=sorted(METHODS)) parser.add_argument("--epochs", type=int, default=5) parser.add_argument("--patience", type=int, default=2) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--grad-accum-steps", type=int, default=2) parser.add_argument("--max-length", type=int, default=96) parser.add_argument("--learning-rate", type=float, default=5e-5) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--device", default="cuda") parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--output-root", default="ravel_revision_results") parser.add_argument("--overwrite", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--limit-train-samples", type=int, default=None) parser.add_argument("--limit-val-samples", type=int, default=None) parser.add_argument("--limit-test-samples", type=int, default=None) parser.add_argument( "--hfm-deleak-manifest", default="ravel_revision_results/data_audit/hfm_split_manifest_deleaked.csv", ) return parser.parse_args() def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def write_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: Sequence[str]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=list(fieldnames), extrasaction="ignore") writer.writeheader() for row in rows: writer.writerow(row) def append_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: Sequence[str]) -> None: path.parent.mkdir(parents=True, exist_ok=True) exists = path.exists() with path.open("a", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=list(fieldnames), extrasaction="ignore") if not exists: writer.writeheader() for row in rows: writer.writerow(row) def freeze_non_lora(module: nn.Module) -> None: for name, parameter in module.named_parameters(): parameter.requires_grad = "lora_" in name def set_module_trainable(module: Optional[nn.Module], trainable: bool) -> None: if module is None: return for parameter in module.parameters(): parameter.requires_grad = trainable def apply_method_trainability(model: nn.Module, method: MethodSpec) -> None: """Freeze modules that are intentionally bypassed by an ablation target.""" target = method.train_target if target == "no_coattn": for name in [ "fusion", "visual_head", "text_head", "refinement", "extra_mlp_control", ]: set_module_trainable(getattr(model, name, None), False) return if target == "primary": for name in ["visual_head", "text_head", "refinement", "extra_mlp_control"]: set_module_trainable(getattr(model, name, None), False) return if target == "primary_aux": for name in ["refinement", "extra_mlp_control"]: set_module_trainable(getattr(model, name, None), False) return if target == "param_mlp": set_module_trainable(getattr(model, "refinement", None), False) return if target == "disagreement_feature": set_module_trainable(getattr(model, "extra_mlp_control", None), False) return if target in {"consistency_loss", "infonce_alignment"}: for name in ["refinement", "extra_mlp_control"]: set_module_trainable(getattr(model, name, None), False) return def maybe_limit(samples: List[Any], limit: Optional[int]) -> List[Any]: if limit is None or limit <= 0 or len(samples) <= limit: return samples buckets: Dict[Any, List[Any]] = {} for sample in samples: label = getattr(sample, "label", getattr(sample, "combined_majority", "")) buckets.setdefault(label, []).append(sample) per_label = max(1, limit // max(1, len(buckets))) selected: List[Any] = [] for label in sorted(buckets, key=str): selected.extend(buckets[label][:per_label]) selected.extend(samples[: max(0, limit - len(selected))]) return selected[:limit] def load_dataset( dataset_key: str, seed: int, batch_size: int, max_length: int, num_workers: int, method: MethodSpec, hfm_deleak_manifest: str, limits: Tuple[Optional[int], Optional[int], Optional[int]], ) -> Tuple[Any, Dict[str, Any], Any, Any, Any, List[Any], List[Any], List[Any], int, List[str]]: if dataset_key == "mvsa_multiple": from src.mvsa_multiple_pipeline import ( CLARAModel, DEFAULT_MVSA_MULTIPLE_CONFIG, LABEL_ID_TO_NAME, MVSALoader, create_dataloaders, summarize_splits, ) cfg = dict(DEFAULT_MVSA_MULTIPLE_CONFIG) cfg.update( { "architecture": method.architecture, "enable_clip_lora": method.enable_lora, "batch_size": batch_size, "max_length": max_length, "num_workers": num_workers, "pin_memory": True, "persistent_workers": bool(num_workers > 0), "prefetch_factor": 2, "seed": seed, "learning_rate": 5e-5, "weight_decay": 0.01, "use_mixup_negative": False, "use_weighted_sampler": False, "text_unfreeze_mode": "freeze_all", "unfreeze_epoch": 0, "paper_exact_counts": True, } ) loader = MVSALoader(cfg["text_dir"], cfg["label_file"]) loader.load( preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")), require_unanimous=bool(cfg["require_unanimous"]), require_cross_agree=bool(cfg["require_cross_agree"]), paper_exact_counts=True, ) train_samples, val_samples, test_samples = loader.split( train_ratio=float(cfg["train_ratio"]), val_ratio=float(cfg["val_ratio"]), seed=seed, paper_811=True, ) train_samples = maybe_limit(train_samples, limits[0]) val_samples = maybe_limit(val_samples, limits[1]) test_samples = maybe_limit(test_samples, limits[2]) processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"]) tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"]) train_loader, val_loader, test_loader = create_dataloaders( train_samples=train_samples, val_samples=val_samples, test_samples=test_samples, clip_processor=processor, tokenizer=tokenizer, batch_size=batch_size, max_length=max_length, num_workers=num_workers, pin_memory=True, persistent_workers=bool(num_workers > 0), prefetch_factor=2, use_mixup_negative=False, mixup_alpha=0.0, negative_class_boost=1.0, min_ratio_negative=0.0, weighted_train_sampler=False, ) label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))] return ( CLARAModel, cfg, train_loader, val_loader, test_loader, train_samples, val_samples, test_samples, int(cfg["num_classes"]), label_names, ) if dataset_key == "mvsa_single": from src.mvsa_single_pipeline import ( CLARAModel, DEFAULT_MVSA_SINGLE_CONFIG, LABEL_ID_TO_NAME, MVSASingleLoader, create_dataloaders, ) cfg = dict(DEFAULT_MVSA_SINGLE_CONFIG) cfg.update( { "architecture": method.architecture, "enable_clip_lora": method.enable_lora, "enable_text_lora": method.enable_lora, "batch_size": batch_size, "max_length": max_length, "num_workers": num_workers, "pin_memory": True, "persistent_workers": bool(num_workers > 0), "prefetch_factor": 2, "seed": seed, "learning_rate": 5e-5, "weight_decay": 0.01, "use_mixup_negative": False, "use_weighted_sampler": False, "text_unfreeze_mode": "freeze_all", "unfreeze_epoch": 0, "loss_type": "ce", "ce_use_class_weights": False, "label_smoothing": 0.0, } ) loader = MVSASingleLoader(cfg["text_dir"], cfg["label_file"]) samples = loader.load( preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")), require_unanimous=bool(cfg.get("require_unanimous", False)), require_cross_agree=bool(cfg.get("require_cross_agree", False)), ) train_samples, val_samples, test_samples = loader.split( train_ratio=float(cfg["train_ratio"]), val_ratio=float(cfg["val_ratio"]), seed=seed, ) train_samples = maybe_limit(train_samples, limits[0]) val_samples = maybe_limit(val_samples, limits[1]) test_samples = maybe_limit(test_samples, limits[2]) processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"]) tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"]) train_loader, val_loader, test_loader = create_dataloaders( train_samples=train_samples, val_samples=val_samples, test_samples=test_samples, clip_processor=processor, tokenizer=tokenizer, batch_size=batch_size, max_length=max_length, num_workers=num_workers, pin_memory=True, persistent_workers=bool(num_workers > 0), prefetch_factor=2, use_mixup_negative=False, mixup_alpha=0.0, negative_class_boost=1.0, min_ratio_negative=0.0, weighted_train_sampler=False, ) label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))] _ = samples # Keep the load step explicit for provenance and error checking. return ( CLARAModel, cfg, train_loader, val_loader, test_loader, train_samples, val_samples, test_samples, int(cfg["num_classes"]), label_names, ) if dataset_key in {"hfm_existing", "hfm_deleak"}: from src.hfm_pipeline import ( CLARAModel, DEFAULT_HFM_CONFIG, LABEL_ID_TO_NAME, HFMLoader, create_dataloaders, ) cfg = dict(DEFAULT_HFM_CONFIG) cfg.update( { "architecture": method.architecture, "enable_clip_lora": method.enable_lora, "batch_size": batch_size, "max_length": max_length, "num_workers": num_workers, "pin_memory": True, "seed": seed, "learning_rate": 5e-5, "weight_decay": 0.01, "text_unfreeze_mode": "freeze_all", "num_classes": 2, } ) loader = HFMLoader(cfg["text_dir"], cfg["image_root"]) if dataset_key == "hfm_deleak": loader.load_from_manifest(hfm_deleak_manifest) else: loader.load() train_samples = maybe_limit(loader.get_split("train"), limits[0]) val_samples = maybe_limit(loader.get_split("val"), limits[1]) test_samples = maybe_limit(loader.get_split("test"), limits[2]) processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"]) tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"]) train_loader, val_loader, test_loader = create_dataloaders( train_samples=train_samples, val_samples=val_samples, test_samples=test_samples, clip_processor=processor, tokenizer=tokenizer, batch_size=batch_size, max_length=max_length, num_workers=num_workers, pin_memory=True, weighted_train_sampler=False, ) label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))] return ( CLARAModel, cfg, train_loader, val_loader, test_loader, train_samples, val_samples, test_samples, int(cfg["num_classes"]), label_names, ) raise ValueError(f"Unsupported dataset: {dataset_key}") def _js_vector_torch(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: p = p.clamp(min=eps) q = q.clamp(min=eps) m = (0.5 * (p + q)).clamp(min=eps) return 0.5 * p * (p / m).log() + 0.5 * q * (q / m).log() def _attention_entropy(attn: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: probs = attn.clamp(min=eps) entropy = -(probs * probs.log()).sum(dim=-1) support = attn.size(-1) if support > 1: entropy = entropy / math.log(float(support)) return entropy.mean(dim=tuple(range(1, entropy.ndim))) def disagreement_feature( outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor], formulation: str, num_classes: int, ) -> torch.Tensor: visual_probs = outputs["visual_probs"] text_probs = outputs["text_probs"] if formulation == "signed_difference": return visual_probs - text_probs if formulation == "absolute_difference": return (visual_probs - text_probs).abs() if formulation == "js_divergence": return _js_vector_torch(visual_probs, text_probs) if formulation == "log_probability_ratio": return (visual_probs.clamp(min=1e-8).log() - text_probs.clamp(min=1e-8).log()) if formulation == "attention_discrepancy": v2t_list = outputs.get("attention_v2t", []) t2v_list = outputs.get("attention_t2v", []) if not v2t_list or not t2v_list: raise RuntimeError("attention_discrepancy requires return_attention=True outputs.") v2t = torch.stack(v2t_list, dim=0).mean(dim=0) t2v = torch.stack(t2v_list, dim=0).mean(dim=0) scalar = (_attention_entropy(v2t) - _attention_entropy(t2v)).unsqueeze(-1) return scalar.expand(-1, num_classes) raise ValueError(f"Unsupported disagreement formulation: {formulation}") def consistency_js_loss(outputs: Dict[str, torch.Tensor]) -> torch.Tensor: return _js_vector_torch(outputs["visual_probs"], outputs["text_probs"]).sum(dim=-1).mean() def infonce_loss(outputs: Dict[str, torch.Tensor], temperature: float) -> torch.Tensor: visual = F.normalize(outputs["visual_global"].float(), dim=-1) text = F.normalize(outputs["text_global"].float(), dim=-1) logits = visual @ text.T / max(float(temperature), 1e-6) labels = torch.arange(logits.size(0), device=logits.device) return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) def logits_for_target( model: nn.Module, batch: Dict[str, torch.Tensor], method_or_target: Any, criterion: nn.Module, labels: torch.Tensor, num_classes: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: if isinstance(method_or_target, MethodSpec): method = method_or_target target = method.train_target else: method = None target = str(method_or_target) pixel_values = batch["pixel_values"] input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] if target == "legacy_full": outputs = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask) logits = outputs["logits"] return logits, criterion(logits, labels), outputs, {} if target == "no_coattn": logits = model.logits_without_coattention(pixel_values, input_ids, attention_mask) return logits, criterion(logits, labels), {}, {} if target in {"text_only", "vision_only"} and hasattr(model, "encode_modalities"): encoded = model.encode_modalities(pixel_values, input_ids, attention_mask) if target == "text_only": logits = model.text_head(encoded["text_global"]) else: logits = model.visual_head(encoded["visual_global"]) return logits, criterion(logits, labels), {}, {} outputs = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask) if target == "full": loss, parts = token_loss( outputs, labels, criterion, lambda_primary=method.lambda_primary if method is not None else 0.5, lambda_unimodal=method.lambda_unimodal if method is not None else 0.25, ) return outputs["logits"], loss, outputs, parts if target == "primary": logits = outputs["pred_logits"] return logits, criterion(logits, labels), outputs, {} if target == "primary_aux": logits = outputs["pred_logits"] loss = criterion(logits, labels) loss = loss + 0.25 * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) return logits, loss, outputs, {} if target == "param_mlp": logits = model.extra_mlp_control(outputs["fused"]) loss = criterion(logits, labels) loss = loss + 0.25 * criterion(outputs["pred_logits"], labels) loss = loss + 0.25 * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) return logits, loss, outputs, {} if target == "disagreement_feature": if method is None: raise ValueError("disagreement_feature requires a MethodSpec.") if num_classes is None: num_classes = int(outputs["visual_probs"].shape[-1]) if method.disagreement_formulation == "attention_discrepancy": outputs = model( pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask, return_attention=True, ) feature = disagreement_feature( outputs, batch, formulation=method.disagreement_formulation, num_classes=int(num_classes), ).to(dtype=outputs["fused"].dtype) logits = model.refinement(outputs["fused"], feature) loss = criterion(logits, labels) loss = loss + method.lambda_primary * criterion(outputs["pred_logits"], labels) loss = loss + method.lambda_unimodal * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) outputs = dict(outputs) outputs["logits"] = logits outputs["disagreement"] = feature return logits, loss, outputs, {} if target == "consistency_loss": if method is None: raise ValueError("consistency_loss requires a MethodSpec.") logits = outputs["pred_logits"] loss = criterion(logits, labels) loss = loss + method.lambda_unimodal * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) loss = loss + method.consistency_beta * consistency_js_loss(outputs) outputs = dict(outputs) outputs["logits"] = logits outputs["disagreement"] = torch.zeros_like(outputs["visual_probs"]) return logits, loss, outputs, {} if target == "infonce_alignment": if method is None: raise ValueError("infonce_alignment requires a MethodSpec.") logits = outputs["pred_logits"] loss = criterion(logits, labels) loss = loss + method.lambda_unimodal * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) loss = loss + method.infonce_alpha * infonce_loss(outputs, method.infonce_temperature) outputs = dict(outputs) outputs["logits"] = logits outputs["disagreement"] = torch.zeros_like(outputs["visual_probs"]) return logits, loss, outputs, {} if target == "no_disagreement": logits = model.refinement(outputs["fused"], torch.zeros_like(outputs["disagreement"])) loss = criterion(logits, labels) loss = loss + 0.5 * criterion(outputs["pred_logits"], labels) loss = loss + 0.25 * ( criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels) ) return logits, loss, outputs, {} raise ValueError(target) def expected_calibration_error(probs: np.ndarray, y_true: np.ndarray, bins: int = 15) -> float: conf = probs.max(axis=1) pred = probs.argmax(axis=1) correct = (pred == y_true).astype(float) edges = np.linspace(0.0, 1.0, bins + 1) ece = 0.0 for low, high in zip(edges[:-1], edges[1:]): mask = (conf > low) & (conf <= high) if not mask.any(): continue ece += (mask.mean()) * abs(correct[mask].mean() - conf[mask].mean()) return float(ece) def adaptive_ece(probs: np.ndarray, y_true: np.ndarray, bins: int = 15) -> float: conf = probs.max(axis=1) pred = probs.argmax(axis=1) correct = (pred == y_true).astype(float) order = np.argsort(conf) chunks = np.array_split(order, min(bins, len(order))) ece = 0.0 for chunk in chunks: if len(chunk) == 0: continue ece += (len(chunk) / len(conf)) * abs(correct[chunk].mean() - conf[chunk].mean()) return float(ece) def brier_score(probs: np.ndarray, y_true: np.ndarray, num_classes: int) -> float: one_hot = np.eye(num_classes)[y_true] return float(np.mean(np.sum((probs - one_hot) ** 2, axis=1))) def nll_score(probs: np.ndarray, y_true: np.ndarray) -> float: return float(-np.mean(np.log(np.clip(probs[np.arange(len(y_true)), y_true], 1e-12, 1.0)))) def aurc_score(probs: np.ndarray, y_true: np.ndarray) -> Tuple[float, float, float, float, float]: conf = probs.max(axis=1) pred = probs.argmax(axis=1) errors = (pred != y_true).astype(float) order = np.argsort(-conf) sorted_errors = errors[order] cum_errors = np.cumsum(sorted_errors) coverage = np.arange(1, len(errors) + 1) / max(1, len(errors)) risk = cum_errors / np.arange(1, len(errors) + 1) trapezoid = getattr(np, "trapezoid", None) if trapezoid is None: trapezoid = getattr(np, "trapz") aurc = float(trapezoid(risk, coverage)) if len(errors) > 1 else float(risk[-1]) def risk_at(cov: float) -> float: idx = max(0, min(len(errors) - 1, math.ceil(cov * len(errors)) - 1)) return float(risk[idx]) def acc_at(cov: float) -> float: return 1.0 - risk_at(cov) return aurc, risk_at(0.8), risk_at(0.9), acc_at(0.8), acc_at(0.9) def compute_metrics( logits: np.ndarray, y_true: np.ndarray, num_classes: int, ) -> Dict[str, float]: probs = torch.softmax(torch.tensor(logits, dtype=torch.float32), dim=-1).numpy() pred = probs.argmax(axis=1) metrics: Dict[str, float] = { "accuracy": float(accuracy_score(y_true, pred)), "macro_precision": float(precision_score(y_true, pred, average="macro", zero_division=0)), "macro_recall": float(recall_score(y_true, pred, average="macro", zero_division=0)), "macro_f1": float(f1_score(y_true, pred, average="macro", zero_division=0)), "weighted_f1": float(f1_score(y_true, pred, average="weighted", zero_division=0)), "raw_ece": expected_calibration_error(probs, y_true), "adaptive_ece": adaptive_ece(probs, y_true), "nll": nll_score(probs, y_true), "brier": brier_score(probs, y_true, num_classes), } aurc, r80, r90, a80, a90 = aurc_score(probs, y_true) metrics.update( { "aurc": aurc, "risk_at_80_coverage": r80, "risk_at_90_coverage": r90, "accuracy_at_80_coverage": a80, "accuracy_at_90_coverage": a90, } ) try: if num_classes == 2: metrics["auroc"] = float(roc_auc_score(y_true, probs[:, 1])) else: metrics["auroc"] = float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro")) except Exception: metrics["auroc"] = float("nan") return metrics def samples_to_ids(samples: List[Any]) -> List[str]: ids: List[str] = [] for idx, sample in enumerate(samples): ids.append(str(getattr(sample, "image_id", getattr(sample, "sample_id", idx)))) return ids def evaluate( model: nn.Module, loader: Any, samples: List[Any], method: MethodSpec, device: torch.device, criterion: nn.Module, num_classes: int, dataset_key: str, seed: int, ) -> Tuple[Dict[str, float], List[Dict[str, Any]], np.ndarray, np.ndarray]: model.eval() logits_all: List[np.ndarray] = [] labels_all: List[np.ndarray] = [] visual_probs_all: List[np.ndarray] = [] text_probs_all: List[np.ndarray] = [] primary_logits_all: List[np.ndarray] = [] refined_logits_all: List[np.ndarray] = [] with torch.no_grad(): for batch in loader: labels = batch["labels"].long().to(device) batch_gpu = { key: value.to(device, non_blocking=True) if torch.is_tensor(value) else value for key, value in batch.items() } logits, _loss, outputs, _parts = logits_for_target( model, batch_gpu, method, criterion, labels, num_classes=num_classes ) logits_all.append(logits.float().cpu().numpy()) labels_all.append(labels.cpu().numpy()) if isinstance(outputs, dict) and "visual_probs" in outputs: visual_probs_all.append(outputs["visual_probs"].float().cpu().numpy()) text_probs_all.append(outputs["text_probs"].float().cpu().numpy()) primary_logits_all.append(outputs["pred_logits"].float().cpu().numpy()) refined_logits_all.append(outputs["logits"].float().cpu().numpy()) logits_np = np.concatenate(logits_all, axis=0) y_np = np.concatenate(labels_all, axis=0) metrics = compute_metrics(logits_np, y_np, num_classes) probs = torch.softmax(torch.tensor(logits_np, dtype=torch.float32), dim=-1).numpy() pred = probs.argmax(axis=1) sample_ids = samples_to_ids(samples) visual_probs = np.concatenate(visual_probs_all, axis=0) if visual_probs_all else np.full_like(probs, np.nan) text_probs = np.concatenate(text_probs_all, axis=0) if text_probs_all else np.full_like(probs, np.nan) primary_logits = ( np.concatenate(primary_logits_all, axis=0) if primary_logits_all else np.full_like(logits_np, np.nan) ) refined_logits = ( np.concatenate(refined_logits_all, axis=0) if refined_logits_all else np.full_like(logits_np, np.nan) ) primary_probs = torch.softmax(torch.tensor(primary_logits, dtype=torch.float32), dim=-1).numpy() refined_probs = torch.softmax(torch.tensor(refined_logits, dtype=torch.float32), dim=-1).numpy() disagreement = visual_probs - text_probs rows: List[Dict[str, Any]] = [] for i in range(len(y_np)): row: Dict[str, Any] = { "sample_id": sample_ids[i] if i < len(sample_ids) else str(i), "dataset": dataset_key, "split": "test", "seed": seed, "method": method.key, "true_label": int(y_np[i]), "predicted_label": int(pred[i]), "prediction_entropy": float(-(probs[i] * np.log(np.clip(probs[i], 1e-12, 1.0))).sum()), "max_confidence": float(probs[i].max()), "primary_prediction": int(np.nanargmax(primary_probs[i])) if np.isfinite(primary_probs[i]).all() else "", "primary_confidence": float(np.nanmax(primary_probs[i])) if np.isfinite(primary_probs[i]).all() else "", "refined_prediction": int(np.nanargmax(refined_probs[i])) if np.isfinite(refined_probs[i]).all() else "", "refined_confidence": float(np.nanmax(refined_probs[i])) if np.isfinite(refined_probs[i]).all() else "", "correct_primary": int(np.nanargmax(primary_probs[i]) == y_np[i]) if np.isfinite(primary_probs[i]).all() else "", "correct_refined": int(np.nanargmax(refined_probs[i]) == y_np[i]) if np.isfinite(refined_probs[i]).all() else "", "number_of_images": 1, } for c in range(num_classes): row[f"logit_class_{c}"] = float(logits_np[i, c]) row[f"prob_class_{c}"] = float(probs[i, c]) row[f"visual_prob_class_{c}"] = float(visual_probs[i, c]) if np.isfinite(visual_probs[i, c]) else "" row[f"text_prob_class_{c}"] = float(text_probs[i, c]) if np.isfinite(text_probs[i, c]) else "" row[f"disagreement_class_{c}"] = float(disagreement[i, c]) if np.isfinite(disagreement[i, c]) else "" if np.isfinite(disagreement[i]).all(): row["disagreement_l1"] = float(np.abs(disagreement[i]).sum()) else: row["disagreement_l1"] = "" rows.append(row) return metrics, rows, logits_np, y_np def train_one( model: nn.Module, train_loader: Any, val_loader: Any, method: MethodSpec, device: torch.device, epochs: int, patience: int, grad_accum_steps: int, learning_rate: float, weight_decay: float, num_classes: int, ) -> Tuple[Dict[str, Any], Dict[str, torch.Tensor], List[Dict[str, Any]]]: criterion = nn.CrossEntropyLoss() optimizer = AdamW( [parameter for parameter in model.parameters() if parameter.requires_grad], lr=learning_rate, weight_decay=weight_decay, ) use_bf16 = bool(device.type == "cuda" and torch.cuda.is_bf16_supported()) amp_dtype = torch.bfloat16 if use_bf16 else torch.float16 def trainable_state_dict() -> Dict[str, torch.Tensor]: trainable_names = {name for name, parameter in model.named_parameters() if parameter.requires_grad} return { key: value.detach().cpu().clone() for key, value in model.state_dict().items() if key in trainable_names } best_state: Dict[str, torch.Tensor] = {} best_val = -1.0 best_epoch = 0 bad_epochs = 0 history: List[Dict[str, Any]] = [] for epoch in range(1, epochs + 1): model.train() train_losses: List[float] = [] optimizer.zero_grad(set_to_none=True) start = time.time() for step, batch in enumerate(train_loader, start=1): labels = batch["labels"].long().to(device, non_blocking=True) batch_gpu = { key: value.to(device, non_blocking=True) if torch.is_tensor(value) else value for key, value in batch.items() } with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device.type == "cuda"): _logits, loss, _outputs, _parts = logits_for_target( model, batch_gpu, method, criterion, labels, num_classes=num_classes ) loss = loss / max(1, grad_accum_steps) loss.backward() if step % max(1, grad_accum_steps) == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad(set_to_none=True) train_losses.append(float(loss.detach().item()) * max(1, grad_accum_steps)) if step % 250 == 0: print( f" epoch={epoch} step={step}/{len(train_loader)} " f"loss={float(np.mean(train_losses[-50:])):.4f}", flush=True, ) if len(train_loader) % max(1, grad_accum_steps) != 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad(set_to_none=True) val_metrics, _rows, _logits, _labels = evaluate( model=model, loader=val_loader, samples=[], method=method, device=device, criterion=criterion, num_classes=num_classes, dataset_key="val", seed=0, ) val_score = val_metrics["macro_f1"] history_row = { "epoch": epoch, "train_loss": float(np.mean(train_losses)) if train_losses else float("nan"), "val_macro_f1": val_metrics["macro_f1"], "val_weighted_f1": val_metrics["weighted_f1"], "val_accuracy": val_metrics["accuracy"], "val_raw_ece": val_metrics["raw_ece"], "elapsed_seconds": time.time() - start, } history.append(history_row) print( f"epoch={epoch} train_loss={history_row['train_loss']:.4f} " f"val_mF1={val_metrics['macro_f1']:.4f} val_ECE={val_metrics['raw_ece']:.4f}", flush=True, ) if val_score > best_val: best_val = val_score best_epoch = epoch best_state = trainable_state_dict() bad_epochs = 0 else: bad_epochs += 1 if bad_epochs >= patience: break return {"best_epoch": best_epoch, "best_val_macro_f1": best_val}, best_state, history def run_single(args: argparse.Namespace, dataset_key: str, method_key: str, seed: int) -> Optional[Dict[str, Any]]: method = METHODS[method_key] out_root = Path(args.output_root) run_dir = out_root / "runs" / dataset_key / method.key / f"seed_{seed}" metrics_path = run_dir / "metrics.json" pred_path = out_root / "predictions" / dataset_key / f"{method.key}_seed_{seed}.csv" history_path = run_dir / "history.csv" ckpt_path = run_dir / "checkpoint.pt" if metrics_path.exists() and pred_path.exists() and not args.overwrite: print(f"SKIP completed {dataset_key} {method.key} seed={seed}", flush=True) return json.loads(metrics_path.read_text(encoding="utf-8")) print(f"RUN dataset={dataset_key} method={method.key} seed={seed}", flush=True) if args.dry_run: return None set_seed(seed) device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu") ( model_cls, cfg, train_loader, val_loader, test_loader, train_samples, val_samples, test_samples, num_classes, label_names, ) = load_dataset( dataset_key=dataset_key, seed=seed, batch_size=args.batch_size, max_length=args.max_length, num_workers=args.num_workers, method=method, hfm_deleak_manifest=args.hfm_deleak_manifest, limits=(args.limit_train_samples, args.limit_val_samples, args.limit_test_samples), ) cfg.update( { "architecture": method.architecture, "enable_clip_lora": method.enable_lora, "enable_text_lora": method.enable_lora, "seed": seed, "batch_size": args.batch_size, "max_length": args.max_length, } ) model = model_cls(cfg).to(device) if hasattr(model, "vision_lora"): freeze_non_lora(model.vision_lora) if hasattr(model, "text"): freeze_non_lora(model.text) apply_method_trainability(model, method) stats = model.parameter_stats() train_info, best_state, history = train_one( model=model, train_loader=train_loader, val_loader=val_loader, method=method, device=device, epochs=args.epochs, patience=args.patience, grad_accum_steps=args.grad_accum_steps, learning_rate=args.learning_rate, weight_decay=args.weight_decay, num_classes=num_classes, ) if best_state: model.load_state_dict(best_state, strict=False) criterion = nn.CrossEntropyLoss() test_metrics, prediction_rows, _logits, _labels = evaluate( model=model, loader=test_loader, samples=test_samples, method=method, device=device, criterion=criterion, num_classes=num_classes, dataset_key=dataset_key, seed=seed, ) run_dir.mkdir(parents=True, exist_ok=True) torch.save( { "model_state": { key: value.detach().cpu() for key, value in model.state_dict().items() if key in {name for name, parameter in model.named_parameters() if parameter.requires_grad} }, "checkpoint_type": "trainable_parameters_only", "cfg": cfg, "dataset": dataset_key, "method": method.key, "seed": seed, "metrics": test_metrics, "train_info": train_info, }, ckpt_path, ) write_csv(history_path, history, ["epoch", "train_loss", "val_macro_f1", "val_weighted_f1", "val_accuracy", "val_raw_ece", "elapsed_seconds"]) pred_fields = [ "sample_id", "dataset", "split", "seed", "method", "true_label", "predicted_label", ] for c in range(num_classes): pred_fields.extend( [ f"logit_class_{c}", f"prob_class_{c}", f"visual_prob_class_{c}", f"text_prob_class_{c}", f"disagreement_class_{c}", ] ) pred_fields.extend( [ "disagreement_l1", "prediction_entropy", "max_confidence", "primary_prediction", "primary_confidence", "refined_prediction", "refined_confidence", "correct_primary", "correct_refined", "number_of_images", ] ) write_csv(pred_path, prediction_rows, pred_fields) metric_payload: Dict[str, Any] = { "dataset": dataset_key, "method": method.key, "method_display": method.display, "configuration": method.train_target, "disagreement_formulation": method.disagreement_formulation, "additional_loss": method.additional_loss, "lambda_primary": method.lambda_primary, "lambda_unimodal": method.lambda_unimodal, "consistency_beta": method.consistency_beta if method.additional_loss == "js_consistency" else None, "infonce_alpha": method.infonce_alpha if method.additional_loss == "infonce_alignment" else None, "infonce_temperature": method.infonce_temperature if method.additional_loss == "infonce_alignment" else None, "seed": seed, "num_classes": num_classes, "label_names": label_names, "train_samples": len(train_samples), "val_samples": len(val_samples), "test_samples": len(test_samples), "total_params": stats["total"], "trainable_params": stats["trainable"], "best_epoch": train_info["best_epoch"], **test_metrics, } metrics_path.write_text(json.dumps(metric_payload, indent=2), encoding="utf-8") print( f"DONE {dataset_key} {method.key} seed={seed} " f"mF1={test_metrics['macro_f1']:.4f} wF1={test_metrics['weighted_f1']:.4f} " f"ECE={test_metrics['raw_ece']:.4f}", flush=True, ) return metric_payload def aggregate(output_root: Path) -> None: metric_rows: List[Dict[str, Any]] = [] for path in sorted((output_root / "runs").glob("*/*/seed_*/metrics.json")): try: metric_rows.append(json.loads(path.read_text(encoding="utf-8"))) except Exception: continue fields = [ "dataset", "method", "method_display", "configuration", "seed", "accuracy", "macro_precision", "macro_recall", "macro_f1", "weighted_f1", "auroc", "raw_ece", "adaptive_ece", "nll", "brier", "aurc", "risk_at_80_coverage", "risk_at_90_coverage", "accuracy_at_80_coverage", "accuracy_at_90_coverage", "total_params", "trainable_params", "best_epoch", "train_samples", "val_samples", "test_samples", ] write_csv(output_root / "aggregate_results" / "main_results.csv", metric_rows, fields) grouped: Dict[Tuple[str, str], List[Dict[str, Any]]] = {} for row in metric_rows: grouped.setdefault((row["dataset"], row["method"]), []).append(row) summary_rows: List[Dict[str, Any]] = [] for (dataset, method), rows in sorted(grouped.items()): out: Dict[str, Any] = {"dataset": dataset, "method": method, "runs": len(rows)} for metric in ["accuracy", "macro_f1", "weighted_f1", "raw_ece", "nll", "brier", "aurc"]: vals = np.array([float(row[metric]) for row in rows if row.get(metric) is not None], dtype=float) if vals.size: out[f"{metric}_mean"] = float(vals.mean()) out[f"{metric}_std"] = float(vals.std(ddof=1)) if vals.size > 1 else 0.0 summary_rows.append(out) write_csv( output_root / "aggregate_results" / "summary_by_method.csv", summary_rows, [ "dataset", "method", "runs", "accuracy_mean", "accuracy_std", "macro_f1_mean", "macro_f1_std", "weighted_f1_mean", "weighted_f1_std", "raw_ece_mean", "raw_ece_std", "nll_mean", "nll_std", "brier_mean", "brier_std", "aurc_mean", "aurc_std", ], ) def planned_runs(args: argparse.Namespace) -> List[Tuple[str, str, int]]: runs: List[Tuple[str, str, int]] = [] if args.methods: method_keys = args.methods elif args.stage == "stage_b": method_keys = STAGE_B_METHODS elif args.stage == "stage_c": method_keys = STAGE_C_METHODS else: method_keys = sorted(set(STAGE_B_METHODS + STAGE_C_METHODS), key=(STAGE_B_METHODS + STAGE_C_METHODS).index) if args.stage == "stage_b": seeds = [args.stage_b_seed] else: seeds = args.seeds for dataset in args.datasets: for seed in seeds: for method in method_keys: runs.append((dataset, method, seed)) return runs def main() -> None: args = parse_args() output_root = Path(args.output_root) output_root.mkdir(parents=True, exist_ok=True) plan = planned_runs(args) print(f"Planned runs: {len(plan)}", flush=True) for dataset_key, method_key, seed in plan: run_single(args, dataset_key, method_key, seed) aggregate(output_root) aggregate(output_root) print(f"Wrote aggregate results to {output_root / 'aggregate_results'}", flush=True) if __name__ == "__main__": main()