Upload 26 files

milk10k_new_collapse_research/README.md

@@ -0,0 +1,120 @@
+# MILK10k Collapse Research
+
+## Why this folder exists
+
+The main training pipeline (`milk10k_effb2_metadata`) does full end-to-end dual-encoder training —
+hours per run, heavy GPU, lots of hyperparameters. That's fine for final experiments, but terrible
+for iterating on new ideas.
+
+This folder is a **lightweight research sandbox**. The idea is:
+
+1. Extract frozen features from a foundation model once (DINOv2, ConvNeXt, etc.)
+2. Train cheap sklearn probes (minutes, CPU) on those features
+3. Test collapse-mitigation ideas fast without re-training the encoder
+
+Results go to `results/new_collapse_research/` and never interfere with the controlled
+dual-encoder runs.
+
+## The problem: class collapse
+
+MILK10k has a severe long-tail distribution:
+
+- Head: NV, MEL, BCC, BKL
+- Tail: **BEN_OTH**, **INF**, **MAL_OTH** (called "tail gate" classes)
+
+The tail classes get crushed — the model rarely predicts them. This manifests as low
+recall on BEN_OTH/INF/MAL_OTH even when overall metrics look fine.
+
+The `BASELINE_F1_MACRO = 0.5828` is the reference point to beat.
+
+## The pipeline
+
+```
+extract_foundation_features.py    # Step 1: freeze a backbone, dump .npz features
+          │
+          ├── linear_probe.py           # Quick baseline: logistic/MLP on frozen features
+          ├── train_hierarchical.py     # Two-stage: group → within-group expert
+          ├── train_multimodal_ssl.py   # Contrastive pretrain on pairs, then linear eval
+          │
+analyze_decision_policy.py        # Post-hoc: optimize tail-class logit biases
+```
+
+### Feature extraction (`extract_foundation_features.py`)
+
+Runs a frozen model (DINOv2, or any timm model) on all clinical/dermoscopic pairs and
+saves features as `.npz`. Features are L2-normalized by default.
+
+Output for each split: `{train,val}_features.npz` containing `clinical`, `dermoscopic`,
+`lesion_id`, and `label` arrays.
+
+### Linear probe (`linear_probe.py`)
+
+Trains a `LogisticRegression` or `MLPClassifier` on frozen features. Supports three
+feature modes:
+
+- `pair` — concatenates `[clinical, dermoscopic, |clinical-dermoscopic|, clinical×dermoscopic]`
+- `clinical` — clinical image features only
+- `dermoscopic` — dermoscopic image features only
+
+Useful for: establishing a feature-quality baseline, comparing backbones, testing
+whether the frozen DINOv2 features already contain useful signal.
+
+### Hierarchical probe (`train_hierarchical.py`)
+
+Two-stage classification:
+
+1. **Group classifier** — predicts one of three groups:
+   - `melanocytic` (MEL, NV)
+   - `keratinocyte_like` (AKIEC, BCC, BKL, SCCKA)
+   - `rare_other` (BEN_OTH, INF, MAL_OTH, DF, VASC)
+2. **Per-group expert** — predicts the fine class within the predicted group
+
+The final probability = P(group) × P(class | group), renormalized.
+
+Useful for: testing whether partitioning the problem reduces tail-class confusion.
+A rare-class like BEN_OTH only competes within `rare_other` instead of all 11 classes.
+
+### Multimodal SSL (`train_multimodal_ssl.py`)
+
+Contrastively pretrains an encoder on clinical-dermoscopic **pairs** (positive pairs
+are clinical and dermoscopic images of the same lesion) using a symmetric CLIP loss.
+After pretraining, evaluates via linear probe on frozen encoder features.
+
+Useful for: learning representations that are invariant to the clinical/dermoscopic
+modality shift, potentially improving downstream classification.
+
+### Decision policy (`analyze_decision_policy.py`)
+
+Post-hoc optimization: searches for logit biases on tail classes
+(BEN_OTH, INF, MAL_OTH) that maximize macro F1 on the validation set.
+
+Useful for: squeezing the last bit of performance from an already-trained model
+without re-training. The `--tail-only` flag restricts optimization to just the three
+tail gate classes.
+
+## How to run
+
+```bash
+# Activate your environment first
+conda activate ml2
+
+# Full pipeline (features → probes → SSL → policy)
+DATA_DIR=/path/to/milk10k bash milk10k_new_collapse_research/run_new_ideas.sh
+
+# Only feature extraction with a different backbone
+MODEL_NAME=dinov2_vits14 bash milk10k_new_collapse_research/run_new_ideas.sh
+
+# Only decision policy on an existing prediction CSV
+PREDICTIONS_CSV=/path/to/val_predictions.csv bash milk10k_new_collapse_research/run_new_ideas.sh
+```
+
+Each step skips if its outputs already exist, so you can rerun selectively.
+
+## Key metrics file
+
+Every script writes to `results/new_collapse_research/{experiment}/`:
+
+- `metrics.json` — overall metrics (f1_macro, balanced_accuracy, etc.)
+- `per_class_metrics.csv` — per-class recall, precision, f1
+- `confusion_matrix.csv` — full 11×11 confusion matrix
+- `collapse_report.md` — human-readable report focused on tail-gate classes

milk10k_new_collapse_research/__init__.py

@@ -0,0 +1,2 @@
+"""Experimental MILK10k collapse research tools."""
+

milk10k_new_collapse_research/compat.py

@@ -0,0 +1,15 @@
+"""Compatibility helpers for reusing the existing local MILK10k package."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+def ensure_legacy_package_path() -> None:
+    repo_root = Path(__file__).resolve().parents[1]
+    legacy_root = repo_root / "milk10k_effb2_metadata"
+    legacy_path = str(legacy_root)
+    if legacy_root.is_dir() and legacy_path not in sys.path:
+        sys.path.insert(0, legacy_path)
+

milk10k_new_collapse_research/config.py

@@ -0,0 +1,36 @@
+"""Shared defaults for experimental MILK10k collapse research."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+CLASS_NAMES = [
+    "AKIEC",
+    "BCC",
+    "BEN_OTH",
+    "BKL",
+    "DF",
+    "INF",
+    "MAL_OTH",
+    "MEL",
+    "NV",
+    "SCCKA",
+    "VASC",
+]
+
+DEFAULT_HIERARCHY = {
+    "melanocytic": ["MEL", "NV"],
+    "keratinocyte_like": ["AKIEC", "BCC", "BKL", "SCCKA"],
+    "rare_other": ["BEN_OTH", "INF", "MAL_OTH", "DF", "VASC"],
+}
+
+TAIL_GATE_CLASSES = ["BEN_OTH", "INF", "MAL_OTH"]
+RESULTS_ROOT = Path("results/new_collapse_research")
+BASELINE_F1_MACRO = 0.5828
+
+
+def class_to_idx(class_names: list[str] | None = None) -> dict[str, int]:
+    names = class_names or CLASS_NAMES
+    return {name: idx for idx, name in enumerate(names)}
+

milk10k_new_collapse_research/features.py

@@ -0,0 +1,129 @@
+"""Frozen feature extraction utilities."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from timm.data import create_transform, resolve_data_config
+from torch.utils.data import DataLoader, Dataset
+
+
+DINOV2_MODELS = {"dinov2_vits14", "dinov2_vitb14", "dinov2_vitl14", "dinov2_vitg14"}
+
+
+@dataclass
+class FeatureBatch:
+    lesion_ids: list[str]
+    labels: list[str]
+    clinical: np.ndarray
+    dermoscopic: np.ndarray
+
+
+class PairImageDataset(Dataset):
+    def __init__(self, df, transform) -> None:
+        self.df = df.reset_index(drop=True)
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def _load(self, path: str) -> torch.Tensor:
+        with Image.open(path) as image:
+            image = image.convert("RGB")
+        return self.transform(image)
+
+    def __getitem__(self, idx: int) -> dict[str, object]:
+        row = self.df.iloc[idx]
+        return {
+            "lesion_id": str(row["lesion_id"]),
+            "label": str(row["label"]),
+            "clinical": self._load(row["clinical_path"]),
+            "dermoscopic": self._load(row["dermoscopic_path"]),
+        }
+
+
+def build_feature_model(model_name: str, device: torch.device, pretrained: bool = True):
+    if model_name in DINOV2_MODELS:
+        if not pretrained:
+            raise ValueError("DINOv2 torch.hub models require pretrained weights; use a timm model for --no-pretrained smoke tests.")
+        model = torch.hub.load("facebookresearch/dinov2", model_name)
+        image_size = 224
+        transform = dinov2_transform(image_size)
+    else:
+        import timm
+
+        model = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool="avg")
+        cfg = resolve_data_config({}, model=model)
+        image_size = int(cfg.get("input_size", (3, 224, 224))[-1])
+        transform = create_transform(**cfg, is_training=False)
+    model.to(device).eval()
+    return model, transform, image_size
+
+
+def dinov2_transform(image_size: int):
+    from torchvision import transforms
+
+    return transforms.Compose(
+        [
+            transforms.Resize(int(image_size * 1.15), interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ]
+    )
+
+
+@torch.no_grad()
+def extract_pair_features(model, loader: DataLoader, device: torch.device, l2_normalize: bool = True) -> FeatureBatch:
+    lesion_ids: list[str] = []
+    labels: list[str] = []
+    clinical_features: list[np.ndarray] = []
+    dermoscopic_features: list[np.ndarray] = []
+    for batch in loader:
+        clinical = batch["clinical"].to(device, non_blocking=True)
+        dermoscopic = batch["dermoscopic"].to(device, non_blocking=True)
+        clinical_z = model(clinical)
+        dermoscopic_z = model(dermoscopic)
+        if l2_normalize:
+            clinical_z = F.normalize(clinical_z, dim=1)
+            dermoscopic_z = F.normalize(dermoscopic_z, dim=1)
+        clinical_features.append(clinical_z.cpu().numpy())
+        dermoscopic_features.append(dermoscopic_z.cpu().numpy())
+        lesion_ids.extend(str(item) for item in batch["lesion_id"])
+        labels.extend(str(item) for item in batch["label"])
+    return FeatureBatch(
+        lesion_ids=lesion_ids,
+        labels=labels,
+        clinical=np.concatenate(clinical_features, axis=0),
+        dermoscopic=np.concatenate(dermoscopic_features, axis=0),
+    )
+
+
+def save_feature_npz(path: Path, batch: FeatureBatch, model_name: str, split: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        path,
+        lesion_id=np.asarray(batch.lesion_ids, dtype=object),
+        label=np.asarray(batch.labels, dtype=object),
+        clinical=batch.clinical,
+        dermoscopic=batch.dermoscopic,
+        model_name=np.asarray(model_name),
+        split=np.asarray(split),
+    )
+
+
+def make_feature_matrix(data: dict[str, np.ndarray], mode: str = "pair") -> np.ndarray:
+    clinical = data["clinical"]
+    dermoscopic = data["dermoscopic"]
+    if mode == "clinical":
+        return clinical
+    if mode == "dermoscopic":
+        return dermoscopic
+    if mode == "pair":
+        return np.concatenate([clinical, dermoscopic, np.abs(clinical - dermoscopic), clinical * dermoscopic], axis=1)
+    raise ValueError(f"Unsupported feature mode: {mode}")

milk10k_new_collapse_research/hierarchy.py

@@ -0,0 +1,78 @@
+"""Hierarchical class mapping helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from milk10k_new_collapse_research.config import CLASS_NAMES, DEFAULT_HIERARCHY
+
+
+@dataclass(frozen=True)
+class HierarchySpec:
+    groups: dict[str, list[str]]
+    class_names: list[str]
+
+    @property
+    def group_names(self) -> list[str]:
+        return list(self.groups)
+
+    @property
+    def class_to_group(self) -> dict[str, str]:
+        mapping: dict[str, str] = {}
+        for group, labels in self.groups.items():
+            for label in labels:
+                if label in mapping:
+                    raise ValueError(f"Class {label!r} appears in multiple hierarchy groups.")
+                mapping[label] = group
+        missing = sorted(set(self.class_names) - set(mapping))
+        extra = sorted(set(mapping) - set(self.class_names))
+        if missing:
+            raise ValueError(f"Hierarchy missing classes: {missing}")
+        if extra:
+            raise ValueError(f"Hierarchy references unknown classes: {extra}")
+        return mapping
+
+    def group_index(self) -> dict[str, int]:
+        return {name: idx for idx, name in enumerate(self.group_names)}
+
+    def group_labels_for_class_labels(self, labels: list[str]) -> np.ndarray:
+        c2g = self.class_to_group
+        g2i = self.group_index()
+        return np.asarray([g2i[c2g[label]] for label in labels], dtype=np.int64)
+
+    def class_group_indices(self) -> list[int]:
+        c2g = self.class_to_group
+        g2i = self.group_index()
+        return [g2i[c2g[label]] for label in self.class_names]
+
+
+def default_hierarchy(class_names: list[str] | None = None) -> HierarchySpec:
+    names = class_names or CLASS_NAMES
+    name_set = set(names)
+    groups = {
+        group: [label for label in labels if label in name_set]
+        for group, labels in DEFAULT_HIERARCHY.items()
+    }
+    groups = {group: labels for group, labels in groups.items() if labels}
+    return HierarchySpec(groups=groups, class_names=names)
+
+
+def combine_group_and_expert_probabilities(
+    group_prob: np.ndarray,
+    expert_prob_by_group: dict[str, np.ndarray],
+    spec: HierarchySpec,
+) -> np.ndarray:
+    """Return flat class probabilities from group and in-group probabilities."""
+    flat = np.zeros((group_prob.shape[0], len(spec.class_names)), dtype=np.float64)
+    group_to_idx = spec.group_index()
+    class_to_idx = {label: idx for idx, label in enumerate(spec.class_names)}
+    for group_name, labels in spec.groups.items():
+        expert_prob = expert_prob_by_group[group_name]
+        if expert_prob.shape[1] != len(labels):
+            raise ValueError(f"Expert {group_name!r} has {expert_prob.shape[1]} columns, expected {len(labels)}.")
+        for local_idx, label in enumerate(labels):
+            flat[:, class_to_idx[label]] = group_prob[:, group_to_idx[group_name]] * expert_prob[:, local_idx]
+    row_sum = flat.sum(axis=1, keepdims=True)
+    return flat / np.clip(row_sum, 1e-12, None)

milk10k_new_collapse_research/metrics_ext.py

@@ -0,0 +1,126 @@
+"""Metrics and reporting helpers for collapse-focused experiments."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+from milk10k_new_collapse_research.compat import ensure_legacy_package_path
+
+ensure_legacy_package_path()
+from milk10k_effb2_metadata.metrics import compute_metrics, save_predictions
+from milk10k_new_collapse_research.config import BASELINE_F1_MACRO, TAIL_GATE_CLASSES
+
+
+def write_standard_outputs(
+    output_dir: Path,
+    val_df: pd.DataFrame,
+    y_true: np.ndarray,
+    y_prob: np.ndarray,
+    class_names: list[str],
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    metrics, per_class_df, cm = compute_metrics(y_true, y_prob, class_names)
+    if extra:
+        metrics.update(extra)
+    (output_dir / "metrics.json").write_text(json.dumps(json_safe(metrics), indent=2), encoding="utf-8")
+    per_class_df.to_csv(output_dir / "per_class_metrics.csv", index=False)
+    pd.DataFrame(cm, index=class_names, columns=class_names).to_csv(output_dir / "confusion_matrix.csv")
+    save_predictions(val_df, y_true, y_prob, class_names, output_dir)
+    (output_dir / "collapse_report.md").write_text(
+        build_collapse_report(metrics, cm, class_names),
+        encoding="utf-8",
+    )
+    return metrics
+
+
+def build_collapse_report(metrics: dict[str, Any], cm: np.ndarray, class_names: list[str]) -> str:
+    class_to_idx = {name: idx for idx, name in enumerate(class_names)}
+    lines = [
+        "# Collapse Report",
+        "",
+        f"- f1_macro: `{metrics.get('f1_macro', 0.0):.4f}`",
+        f"- baseline_f1_macro: `{BASELINE_F1_MACRO:.4f}`",
+        f"- balanced_accuracy: `{metrics.get('balanced_accuracy', 0.0):.4f}`",
+        f"- accuracy: `{metrics.get('accuracy', 0.0):.4f}`",
+        "",
+        "## Tail Gates",
+        "",
+        "| Class | Support | Correct | Recall | F1 | Gate |",
+        "|---|---:|---:|---:|---:|---|",
+    ]
+    per_class = {row["class"]: row for row in metrics.get("per_class", [])}
+    for label in TAIL_GATE_CLASSES:
+        if label not in class_to_idx:
+            continue
+        idx = class_to_idx[label]
+        row = per_class.get(label, {})
+        support = int(cm[idx].sum())
+        correct = int(cm[idx, idx])
+        recall = float(row.get("recall_sensitivity", 0.0))
+        f1 = float(row.get("f1", 0.0))
+        gate = tail_gate_status(label, correct, f1)
+        lines.append(f"| `{label}` | {support} | {correct} | {recall:.4f} | {f1:.4f} | {gate} |")
+
+    lines.extend(["", "## Key Confusions", ""])
+    for src, dst in [("MEL", "NV"), ("NV", "MEL"), ("INF", "BCC")]:
+        if src in class_to_idx and dst in class_to_idx:
+            lines.append(f"- `{src} -> {dst}`: `{int(cm[class_to_idx[src], class_to_idx[dst]])}`")
+    for src in ["AKIEC", "BCC", "BKL", "SCCKA"]:
+        if src not in class_to_idx:
+            continue
+        idx = class_to_idx[src]
+        cluster = [label for label in ["AKIEC", "BCC", "BKL", "SCCKA"] if label in class_to_idx and label != src]
+        total = sum(int(cm[idx, class_to_idx[dst]]) for dst in cluster)
+        lines.append(f"- `{src}` errors into keratinocyte-like cluster: `{total}`")
+
+    return "\n".join(lines) + "\n"
+
+
+def tail_gate_status(label: str, correct: int, f1: float) -> str:
+    if label == "BEN_OTH":
+        return "PASS" if correct > 2 or f1 > 0.3077 else "FAIL"
+    if label == "INF":
+        return "PASS" if correct >= 5 else "FAIL"
+    if label == "MAL_OTH":
+        return "PASS" if correct >= 1 else "UNRESOLVED"
+    return "CHECK"
+
+
+def load_prediction_csv(path: Path) -> tuple[pd.DataFrame, np.ndarray, np.ndarray, list[str]]:
+    df = pd.read_csv(path)
+    prob_cols = [col for col in df.columns if col.startswith("prob_")]
+    class_names = [col.removeprefix("prob_") for col in prob_cols]
+    true_col = "y_true" if "y_true" in df.columns else "true_label"
+    if true_col == "true_label" and not np.issubdtype(df[true_col].dtype, np.number):
+        label_true_col = "label_true" if "label_true" in df.columns else "true_label"
+        label_to_idx = {label: idx for idx, label in enumerate(class_names)}
+        y_true = df[label_true_col].map(label_to_idx).to_numpy(dtype=np.int64)
+    else:
+        y_true = df[true_col].to_numpy(dtype=np.int64)
+        return df, y_true, df[prob_cols].to_numpy(dtype=np.float64), class_names
+
+    return df, y_true, df[prob_cols].to_numpy(dtype=np.float64), class_names
+
+
+def json_safe(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {str(k): json_safe(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [json_safe(v) for v in value]
+    if isinstance(value, tuple):
+        return [json_safe(v) for v in value]
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, (np.integer,)):
+        return int(value)
+    if isinstance(value, (np.floating,)):
+        return float(value)
+    if isinstance(value, (np.bool_,)):
+        return bool(value)
+    return value

milk10k_new_collapse_research/run_new_ideas.sh

@@ -0,0 +1,179 @@
+#!/usr/bin/env bash
+
+set -Eeuo pipefail
+
+DATA_DIR="${DATA_DIR:-/marimo/milk10k}"
+OUT_ROOT="${OUT_ROOT:-results/new_collapse_research}"
+MODEL_NAME="${MODEL_NAME:-dinov2_vitb14}"
+PREDICTIONS_CSV="${PREDICTIONS_CSV:-}"
+
+run_py() {
+  python -m "$@"
+}
+
+###### Feature extraction ######
+
+FEATURES_DIR="$OUT_ROOT/features_$MODEL_NAME"
+TRAIN_FEATURES="$FEATURES_DIR/train_features.npz"
+VAL_FEATURES="$FEATURES_DIR/val_features.npz"
+
+COMMON_FEATURE_ARGS=(
+  --data-dir "$DATA_DIR"
+  --model-name "$MODEL_NAME"
+  --output-dir "$FEATURES_DIR"
+)
+
+run_feature_config() {
+  local run_name="$1"
+  shift
+  local log="$OUT_ROOT/${run_name}.log"
+
+  mkdir -p "$FEATURES_DIR"
+  echo "START $run_name"
+
+  if [[ -s "$TRAIN_FEATURES" && -s "$VAL_FEATURES" ]]; then
+    echo "SKIP feature extraction; both train_features.npz and val_features.npz already exist"
+  else
+    run_py milk10k_new_collapse_research.scripts.extract_foundation_features \
+      "${COMMON_FEATURE_ARGS[@]}" "$@" >"$log" 2>&1
+  fi
+
+  for artifact in "$TRAIN_FEATURES" "$VAL_FEATURES"; do
+    test -s "$artifact" || { echo "FATAL: missing $artifact" >&2; exit 3; }
+  done
+  echo "DONE $run_name"
+}
+
+run_feature_config extract_features
+
+###### Linear probe ######
+
+PROBE_DIR="$OUT_ROOT/linear_probe_$MODEL_NAME"
+PROBE_METRICS="$PROBE_DIR/metrics.json"
+
+COMMON_PROBE_ARGS=(
+  --train-features "$TRAIN_FEATURES"
+  --val-features "$VAL_FEATURES"
+  --output-dir "$PROBE_DIR"
+)
+
+run_probe_config() {
+  local run_name="$1"
+  shift
+  local log="$OUT_ROOT/${run_name}.log"
+
+  mkdir -p "$PROBE_DIR"
+  echo "START $run_name"
+
+  if [[ -s "$PROBE_METRICS" ]]; then
+    echo "SKIP probe; metrics.json already exists at $PROBE_DIR"
+  else
+    run_py milk10k_new_collapse_research.scripts.linear_probe \
+      "${COMMON_PROBE_ARGS[@]}" "$@" >"$log" 2>&1
+  fi
+
+  test -s "$PROBE_METRICS" || { echo "FATAL: missing $PROBE_METRICS" >&2; exit 3; }
+  echo "DONE $run_name"
+}
+
+run_probe_config linear_probe_baseline --feature-mode pair --probe logistic
+run_probe_config linear_probe_clinical --feature-mode clinical --probe logistic
+run_probe_config linear_probe_dermoscopic --feature-mode dermoscopic --probe logistic
+run_probe_config linear_probe_mlp --feature-mode pair --probe mlp
+
+###### Hierarchical probe ######
+
+HIER_DIR="$OUT_ROOT/hierarchical_$MODEL_NAME"
+HIER_METRICS="$HIER_DIR/metrics.json"
+
+COMMON_HIER_ARGS=(
+  --train-features "$TRAIN_FEATURES"
+  --val-features "$VAL_FEATURES"
+  --output-dir "$HIER_DIR"
+)
+
+run_hier_config() {
+  local run_name="$1"
+  shift
+  local log="$OUT_ROOT/${run_name}.log"
+
+  mkdir -p "$HIER_DIR"
+  echo "START $run_name"
+
+  if [[ -s "$HIER_METRICS" ]]; then
+    echo "SKIP hierarchical; metrics.json already exists at $HIER_DIR"
+  else
+    run_py milk10k_new_collapse_research.scripts.train_hierarchical \
+      "${COMMON_HIER_ARGS[@]}" "$@" >"$log" 2>&1
+  fi
+
+  test -s "$HIER_METRICS" || { echo "FATAL: missing $HIER_METRICS" >&2; exit 3; }
+  echo "DONE $run_name"
+}
+
+run_hier_config hierarchical_probe --feature-mode pair
+run_hier_config hierarchical_probe_clinical --feature-mode clinical
+run_hier_config hierarchical_probe_dermoscopic --feature-mode dermoscopic
+
+###### Multimodal SSL ######
+
+SSL_DIR="$OUT_ROOT/multimodal_ssl"
+SSL_METRICS="$SSL_DIR/metrics.json"
+
+run_ssl_config() {
+  local run_name="$1"
+  shift
+  local log="$OUT_ROOT/${run_name}.log"
+
+  mkdir -p "$SSL_DIR"
+  echo "START $run_name"
+
+  if [[ -s "$SSL_METRICS" ]]; then
+    echo "SKIP SSL; metrics.json already exists at $SSL_DIR"
+  else
+    run_py milk10k_new_collapse_research.scripts.train_multimodal_ssl \
+      --data-dir "$DATA_DIR" \
+      --output-dir "$SSL_DIR" \
+      "$@" >"$log" 2>&1
+  fi
+
+  test -s "$SSL_METRICS" || { echo "FATAL: missing $SSL_METRICS" >&2; exit 3; }
+  echo "DONE $run_name"
+}
+
+run_ssl_config multimodal_ssl_convnext --backbone convnext_base --epochs 20
+
+###### Decision policy (optional, requires a predictions CSV) ######
+
+if [[ -n "$PREDICTIONS_CSV" && -s "$PREDICTIONS_CSV" ]]; then
+  POLICY_DIR="$OUT_ROOT/decision_policy"
+  POLICY_METRICS="$POLICY_DIR/metrics.json"
+
+  run_policy_config() {
+    local run_name="$1"
+    shift
+    local log="$OUT_ROOT/${run_name}.log"
+
+    mkdir -p "$POLICY_DIR"
+    echo "START $run_name"
+
+    if [[ -s "$POLICY_METRICS" ]]; then
+      echo "SKIP policy; metrics.json already exists at $POLICY_DIR"
+    else
+      run_py milk10k_new_collapse_research.scripts.analyze_decision_policy \
+        --predictions "$PREDICTIONS_CSV" \
+        --output-dir "$POLICY_DIR" \
+        "$@" >"$log" 2>&1
+    fi
+
+    test -s "$POLICY_METRICS" || { echo "FATAL: missing $POLICY_METRICS" >&2; exit 3; }
+    echo "DONE $run_name"
+  }
+
+  run_policy_config decision_policy_tail --tail-only
+  run_policy_config decision_policy_all
+else
+  echo "SKIP decision policy; set PREDICTIONS_CSV to a valid prediction CSV to enable"
+fi
+
+echo "ALL RUNS COMPLETED"

milk10k_new_collapse_research/scripts/__init__.py

@@ -0,0 +1,2 @@
+"""CLI entrypoints for experimental MILK10k collapse research."""
+

milk10k_new_collapse_research/scripts/analyze_decision_policy.py

@@ -0,0 +1,106 @@
+"""Optimize rare-class logit bias on an existing validation prediction CSV."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from milk10k_new_collapse_research.compat import ensure_legacy_package_path
+
+ensure_legacy_package_path()
+from milk10k_effb2_metadata.metrics import apply_class_bias, compute_metrics, optimize_class_bias
+from milk10k_new_collapse_research.config import RESULTS_ROOT, TAIL_GATE_CLASSES
+from milk10k_new_collapse_research.metrics_ext import build_collapse_report, json_safe, load_prediction_csv
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--predictions", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=RESULTS_ROOT / "decision_policy")
+    parser.add_argument("--metric", choices=["f1_macro", "dice_macro"], default="f1_macro")
+    parser.add_argument("--max-bias", type=float, default=2.0)
+    parser.add_argument("--step", type=float, default=0.25)
+    parser.add_argument("--passes", type=int, default=3)
+    parser.add_argument("--tail-only", action="store_true", help="Optimize only BEN_OTH/INF/MAL_OTH biases.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    df, y_true, y_prob, class_names = load_prediction_csv(args.predictions)
+    base_metrics, _, base_cm = compute_metrics(y_true, y_prob, class_names)
+    if args.tail_only:
+        bias, best_score = optimize_tail_bias(y_true, y_prob, class_names, args)
+    else:
+        bias, best_score = optimize_class_bias(
+            y_true,
+            y_prob,
+            class_names,
+            metric_name=args.metric,
+            max_bias=args.max_bias,
+            step=args.step,
+            passes=args.passes,
+        )
+    adjusted_prob = apply_class_bias(y_prob, bias)
+    adjusted_metrics, per_class_df, adjusted_cm = compute_metrics(y_true, adjusted_prob, class_names)
+    out_df = df.copy()
+    pred = adjusted_prob.argmax(axis=1)
+    out_df["policy_y_pred"] = pred
+    out_df["policy_label_pred"] = [class_names[idx] for idx in pred]
+    for idx, label in enumerate(class_names):
+        out_df[f"policy_prob_{label}"] = adjusted_prob[:, idx]
+    out_df.to_csv(args.output_dir / "val_predictions.csv", index=False)
+    per_class_df.to_csv(args.output_dir / "per_class_metrics.csv", index=False)
+    pd.DataFrame(adjusted_cm, index=class_names, columns=class_names).to_csv(args.output_dir / "confusion_matrix.csv")
+    payload = {
+        "experiment": "decision_policy",
+        "source_predictions": str(args.predictions),
+        "metric": args.metric,
+        "tail_only": args.tail_only,
+        "base_metric": base_metrics[args.metric],
+        "best_search_metric": best_score,
+        "bias": {label: float(bias[idx]) for idx, label in enumerate(class_names)},
+        "adjusted": adjusted_metrics,
+        "base": base_metrics,
+    }
+    (args.output_dir / "metrics.json").write_text(json.dumps(json_safe(payload), indent=2), encoding="utf-8")
+    (args.output_dir / "collapse_report.md").write_text(
+        "# Base Collapse Report\n\n"
+        + build_collapse_report(base_metrics, base_cm, class_names)
+        + "\n# Adjusted Collapse Report\n\n"
+        + build_collapse_report(adjusted_metrics, adjusted_cm, class_names),
+        encoding="utf-8",
+    )
+
+
+def optimize_tail_bias(y_true: np.ndarray, y_prob: np.ndarray, class_names: list[str], args: argparse.Namespace):
+    bias = np.zeros(len(class_names), dtype=np.float32)
+    tail_indices = [class_names.index(label) for label in TAIL_GATE_CLASSES if label in class_names]
+    best_score = float(compute_metrics(y_true, y_prob, class_names)[0][args.metric])
+    deltas = np.arange(-args.max_bias, args.max_bias + args.step * 0.5, args.step, dtype=np.float32)
+    for _ in range(max(1, args.passes)):
+        improved = False
+        for idx in tail_indices:
+            current = float(bias[idx])
+            best_local = current
+            for delta in deltas:
+                trial = bias.copy()
+                trial[idx] = current + float(delta)
+                score = float(compute_metrics(y_true, apply_class_bias(y_prob, trial), class_names)[0][args.metric])
+                if score > best_score + 1e-12:
+                    best_score = score
+                    best_local = float(trial[idx])
+                    improved = True
+            bias[idx] = best_local
+        if not improved:
+            break
+    return bias, best_score
+
+
+if __name__ == "__main__":
+    main()

milk10k_new_collapse_research/scripts/extract_foundation_features.py

@@ -0,0 +1,77 @@
+"""Extract frozen paired clinical/dermoscopic features."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from milk10k_new_collapse_research.compat import ensure_legacy_package_path
+from milk10k_new_collapse_research.config import RESULTS_ROOT
+from milk10k_new_collapse_research.features import (
+    PairImageDataset,
+    build_feature_model,
+    extract_pair_features,
+    save_feature_npz,
+)
+
+ensure_legacy_package_path()
+from milk10k_effb2_metadata.data import lesion_split, load_paired_dataframe
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data-dir", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=RESULTS_ROOT / "features")
+    parser.add_argument("--model-name", default="dinov2_vitb14")
+    parser.add_argument("--val-size", type=float, default=0.20)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--limit", type=int, default=0, help="Optional per-split row cap for smoke tests.")
+    parser.add_argument("--no-l2-normalize", action="store_true")
+    parser.add_argument("--no-pretrained", action="store_true", help="Only for local smoke tests with timm models.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = torch.device(args.device)
+    df = load_paired_dataframe(args.data_dir)
+    train_df, val_df = lesion_split(df, args.val_size, args.seed)
+    if args.limit > 0:
+        train_df = train_df.head(args.limit).copy()
+        val_df = val_df.head(args.limit).copy()
+
+    model, transform, image_size = build_feature_model(args.model_name, device, pretrained=not args.no_pretrained)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    for split, split_df in [("train", train_df), ("val", val_df)]:
+        dataset = PairImageDataset(split_df, transform)
+        loader = DataLoader(
+            dataset,
+            batch_size=args.batch_size,
+            shuffle=False,
+            num_workers=args.num_workers,
+            pin_memory=torch.cuda.is_available(),
+        )
+        batch = extract_pair_features(model, loader, device, l2_normalize=not args.no_l2_normalize)
+        save_feature_npz(args.output_dir / f"{split}_features.npz", batch, args.model_name, split)
+
+    manifest = {
+        "model_name": args.model_name,
+        "image_size": image_size,
+        "data_dir": str(args.data_dir),
+        "train_rows": int(len(train_df)),
+        "val_rows": int(len(val_df)),
+        "val_size": args.val_size,
+        "seed": args.seed,
+    }
+    (args.output_dir / "feature_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

milk10k_new_collapse_research/scripts/linear_probe.py

@@ -0,0 +1,110 @@
+"""Train a linear or small MLP probe on frozen paired features."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import LogisticRegression
+from sklearn.neural_network import MLPClassifier
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+from milk10k_new_collapse_research.config import CLASS_NAMES, RESULTS_ROOT
+from milk10k_new_collapse_research.features import make_feature_matrix
+from milk10k_new_collapse_research.metrics_ext import write_standard_outputs
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train-features", type=Path, required=True)
+    parser.add_argument("--val-features", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=RESULTS_ROOT / "linear_probe")
+    parser.add_argument("--feature-mode", choices=["pair", "clinical", "dermoscopic"], default="pair")
+    parser.add_argument("--probe", choices=["logistic", "mlp"], default="logistic")
+    parser.add_argument("--max-iter", type=int, default=1000)
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    train = dict(np.load(args.train_features, allow_pickle=True))
+    val = dict(np.load(args.val_features, allow_pickle=True))
+    class_names = infer_class_names(train, val)
+    label_to_idx = {label: idx for idx, label in enumerate(class_names)}
+    x_train = make_feature_matrix(train, args.feature_mode)
+    x_val = make_feature_matrix(val, args.feature_mode)
+    y_train = np.asarray([label_to_idx[str(label)] for label in train["label"]], dtype=np.int64)
+    y_val = np.asarray([label_to_idx[str(label)] for label in val["label"]], dtype=np.int64)
+
+    if args.probe == "logistic":
+        clf = LogisticRegression(
+            max_iter=args.max_iter,
+            class_weight="balanced",
+            solver="lbfgs",
+            multi_class="auto",
+            random_state=args.seed,
+        )
+    else:
+        clf = MLPClassifier(
+            hidden_layer_sizes=(512,),
+            alpha=1e-4,
+            batch_size=64,
+            max_iter=args.max_iter,
+            early_stopping=True,
+            random_state=args.seed,
+        )
+    model = make_pipeline(StandardScaler(), clf)
+    model.fit(x_train, y_train)
+    y_prob = model.predict_proba(x_val)
+    y_prob = align_probabilities(y_prob, model.classes_, len(class_names))
+    val_df = prediction_frame(val)
+    metrics = write_standard_outputs(
+        args.output_dir,
+        val_df,
+        y_val,
+        y_prob,
+        class_names,
+        extra={
+            "experiment": "linear_probe",
+            "feature_mode": args.feature_mode,
+            "probe": args.probe,
+            "train_features": str(args.train_features),
+            "val_features": str(args.val_features),
+        },
+    )
+    (args.output_dir / "probe_config.json").write_text(json.dumps({"metrics_f1_macro": metrics["f1_macro"]}, indent=2))
+
+
+def infer_class_names(*arrays: dict[str, np.ndarray]) -> list[str]:
+    observed = sorted({str(label) for data in arrays for label in data["label"].tolist()})
+    preferred = [label for label in CLASS_NAMES if label in observed]
+    extras = [label for label in observed if label not in preferred]
+    return preferred + extras
+
+
+def align_probabilities(y_prob: np.ndarray, classes: np.ndarray, n_classes: int) -> np.ndarray:
+    aligned = np.zeros((y_prob.shape[0], n_classes), dtype=np.float64)
+    for src_idx, class_idx in enumerate(classes):
+        aligned[:, int(class_idx)] = y_prob[:, src_idx]
+    row_sum = aligned.sum(axis=1, keepdims=True)
+    return aligned / np.clip(row_sum, 1e-12, None)
+
+
+def prediction_frame(data: dict[str, np.ndarray]) -> pd.DataFrame:
+    return pd.DataFrame(
+        {
+            "lesion_id": [str(item) for item in data["lesion_id"]],
+            "clinical_path": [""] * len(data["label"]),
+            "dermoscopic_path": [""] * len(data["label"]),
+        }
+    )
+
+
+if __name__ == "__main__":
+    main()
+

milk10k_new_collapse_research/scripts/train_hierarchical.py

@@ -0,0 +1,114 @@
+"""Train a two-stage hierarchy on frozen paired features."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+from milk10k_new_collapse_research.config import CLASS_NAMES, RESULTS_ROOT
+from milk10k_new_collapse_research.features import make_feature_matrix
+from milk10k_new_collapse_research.hierarchy import combine_group_and_expert_probabilities, default_hierarchy
+from milk10k_new_collapse_research.metrics_ext import write_standard_outputs
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train-features", type=Path, required=True)
+    parser.add_argument("--val-features", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=RESULTS_ROOT / "hierarchical")
+    parser.add_argument("--feature-mode", choices=["pair", "clinical", "dermoscopic"], default="pair")
+    parser.add_argument("--max-iter", type=int, default=1000)
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    train = dict(np.load(args.train_features, allow_pickle=True))
+    val = dict(np.load(args.val_features, allow_pickle=True))
+    class_names = infer_class_names(train, val)
+    spec = default_hierarchy(class_names)
+    label_to_idx = {label: idx for idx, label in enumerate(class_names)}
+    x_train = make_feature_matrix(train, args.feature_mode)
+    x_val = make_feature_matrix(val, args.feature_mode)
+    train_labels = [str(label) for label in train["label"]]
+    val_labels = [str(label) for label in val["label"]]
+    y_train = np.asarray([label_to_idx[label] for label in train_labels], dtype=np.int64)
+    y_val = np.asarray([label_to_idx[label] for label in val_labels], dtype=np.int64)
+    y_train_group = spec.group_labels_for_class_labels(train_labels)
+
+    group_model = classifier(args.max_iter, args.seed)
+    group_model.fit(x_train, y_train_group)
+    group_prob = align_probabilities(group_model.predict_proba(x_val), group_model.classes_, len(spec.group_names))
+
+    expert_prob_by_group: dict[str, np.ndarray] = {}
+    for group_name, labels in spec.groups.items():
+        local_label_to_idx = {label: idx for idx, label in enumerate(labels)}
+        mask = np.asarray([label in local_label_to_idx for label in train_labels], dtype=bool)
+        if len(labels) == 1:
+            expert_prob_by_group[group_name] = np.ones((x_val.shape[0], 1), dtype=np.float64)
+            continue
+        expert = classifier(args.max_iter, args.seed)
+        expert.fit(x_train[mask], np.asarray([local_label_to_idx[label] for label in np.asarray(train_labels)[mask]]))
+        expert_prob_by_group[group_name] = align_probabilities(expert.predict_proba(x_val), expert.classes_, len(labels))
+
+    y_prob = combine_group_and_expert_probabilities(group_prob, expert_prob_by_group, spec)
+    val_df = pd.DataFrame(
+        {
+            "lesion_id": [str(item) for item in val["lesion_id"]],
+            "clinical_path": [""] * len(val_labels),
+            "dermoscopic_path": [""] * len(val_labels),
+        }
+    )
+    metrics = write_standard_outputs(
+        args.output_dir,
+        val_df,
+        y_val,
+        y_prob,
+        class_names,
+        extra={"experiment": "hierarchical_probe", "feature_mode": args.feature_mode, "groups": spec.groups},
+    )
+    (args.output_dir / "hierarchy_config.json").write_text(
+        json.dumps({"groups": spec.groups, "metrics_f1_macro": metrics["f1_macro"]}, indent=2),
+        encoding="utf-8",
+    )
+
+
+def classifier(max_iter: int, seed: int):
+    return make_pipeline(
+        StandardScaler(),
+        LogisticRegression(
+            max_iter=max_iter,
+            class_weight="balanced",
+            solver="lbfgs",
+            multi_class="auto",
+            random_state=seed,
+        ),
+    )
+
+
+def infer_class_names(*arrays: dict[str, np.ndarray]) -> list[str]:
+    observed = sorted({str(label) for data in arrays for label in data["label"].tolist()})
+    preferred = [label for label in CLASS_NAMES if label in observed]
+    extras = [label for label in observed if label not in preferred]
+    return preferred + extras
+
+
+def align_probabilities(y_prob: np.ndarray, classes: np.ndarray, n_classes: int) -> np.ndarray:
+    aligned = np.zeros((y_prob.shape[0], n_classes), dtype=np.float64)
+    for src_idx, class_idx in enumerate(classes):
+        aligned[:, int(class_idx)] = y_prob[:, src_idx]
+    row_sum = aligned.sum(axis=1, keepdims=True)
+    return aligned / np.clip(row_sum, 1e-12, None)
+
+
+if __name__ == "__main__":
+    main()
+

milk10k_new_collapse_research/scripts/train_multimodal_ssl.py

@@ -0,0 +1,222 @@
+"""Contrastively pretrain a dual-use encoder on clinical/dermoscopic MILK10k pairs."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+from timm.data import create_transform, resolve_data_config
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+
+from milk10k_new_collapse_research.compat import ensure_legacy_package_path
+from milk10k_new_collapse_research.config import CLASS_NAMES, RESULTS_ROOT
+from milk10k_new_collapse_research.metrics_ext import write_standard_outputs
+
+ensure_legacy_package_path()
+from milk10k_effb2_metadata.data import lesion_split, load_paired_dataframe
+
+
+class SslPairDataset(Dataset):
+    def __init__(self, df: pd.DataFrame, transform) -> None:
+        self.df = df.reset_index(drop=True)
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def _load(self, path: str) -> torch.Tensor:
+        with Image.open(path) as image:
+            image = image.convert("RGB")
+        return self.transform(image)
+
+    def __getitem__(self, idx: int) -> dict[str, object]:
+        row = self.df.iloc[idx]
+        return {
+            "clinical": self._load(row["clinical_path"]),
+            "dermoscopic": self._load(row["dermoscopic_path"]),
+            "label": str(row["label"]),
+            "lesion_id": str(row["lesion_id"]),
+        }
+
+
+class ContrastiveEncoder(nn.Module):
+    def __init__(self, backbone_name: str, projection_dim: int, pretrained: bool) -> None:
+        super().__init__()
+        import timm
+
+        self.encoder = timm.create_model(backbone_name, pretrained=pretrained, num_classes=0, global_pool="avg")
+        feature_dim = int(self.encoder.num_features)
+        self.projector = nn.Sequential(
+            nn.LayerNorm(feature_dim),
+            nn.Linear(feature_dim, feature_dim),
+            nn.GELU(),
+            nn.Linear(feature_dim, projection_dim),
+        )
+
+    def forward_features(self, images: torch.Tensor) -> torch.Tensor:
+        return self.encoder(images)
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        return F.normalize(self.projector(self.forward_features(images)), dim=1)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data-dir", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=RESULTS_ROOT / "multimodal_ssl")
+    parser.add_argument("--backbone", default="convnext_base")
+    parser.add_argument("--imagenet-pretrained", action="store_true")
+    parser.add_argument("--image-size", type=int, default=224)
+    parser.add_argument("--projection-dim", type=int, default=256)
+    parser.add_argument("--epochs", type=int, default=20)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument("--val-size", type=float, default=0.20)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--limit", type=int, default=0, help="Optional per-split row cap for smoke tests.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = torch.device(args.device)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    df = load_paired_dataframe(args.data_dir)
+    train_df, val_df = lesion_split(df, args.val_size, args.seed)
+    if args.limit > 0:
+        train_df = train_df.head(args.limit).copy()
+        val_df = val_df.head(args.limit).copy()
+
+    model = ContrastiveEncoder(args.backbone, args.projection_dim, args.imagenet_pretrained).to(device)
+    train_transform, eval_transform = make_transforms(model.encoder, args.image_size)
+    train_loader = DataLoader(
+        SslPairDataset(train_df, train_transform),
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+        drop_last=True,
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    history = []
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        losses = []
+        for batch in train_loader:
+            clinical = batch["clinical"].to(device, non_blocking=True)
+            dermoscopic = batch["dermoscopic"].to(device, non_blocking=True)
+            optimizer.zero_grad(set_to_none=True)
+            z1 = model(clinical)
+            z2 = model(dermoscopic)
+            loss = symmetric_clip_loss(z1, z2, args.temperature)
+            loss.backward()
+            optimizer.step()
+            losses.append(float(loss.detach().cpu().item()))
+        history.append({"epoch": epoch, "ssl_loss": float(np.mean(losses)) if losses else None})
+        print(f"epoch={epoch} ssl_loss={history[-1]['ssl_loss']}")
+
+    torch.save(
+        {
+            "backbone": args.backbone,
+            "model_state": model.state_dict(),
+            "projection_dim": args.projection_dim,
+            "image_size": args.image_size,
+            "history": history,
+        },
+        args.output_dir / "ssl_encoder.pt",
+    )
+
+    x_train, y_train, class_names = extract_features_for_probe(model, train_df, eval_transform, args, device)
+    x_val, y_val, _ = extract_features_for_probe(model, val_df, eval_transform, args, device, class_names=class_names)
+    probe = make_pipeline(
+        StandardScaler(),
+        LogisticRegression(max_iter=1000, class_weight="balanced", solver="lbfgs", multi_class="auto", random_state=args.seed),
+    )
+    probe.fit(x_train, y_train)
+    y_prob = align_probabilities(probe.predict_proba(x_val), probe.classes_, len(class_names))
+    metrics = write_standard_outputs(
+        args.output_dir,
+        val_df,
+        y_val,
+        y_prob,
+        class_names,
+        extra={"experiment": "multimodal_ssl_linear_eval", "ssl_history": history},
+    )
+    (args.output_dir / "ssl_config.json").write_text(
+        json.dumps({"backbone": args.backbone, "metrics_f1_macro": metrics["f1_macro"], "history": history}, indent=2),
+        encoding="utf-8",
+    )
+
+
+def make_transforms(model: nn.Module, image_size: int):
+    cfg = resolve_data_config({}, model=model)
+    cfg["input_size"] = (3, image_size, image_size)
+    train_transform = create_transform(**cfg, is_training=True)
+    eval_transform = create_transform(**cfg, is_training=False)
+    return train_transform, eval_transform
+
+
+def symmetric_clip_loss(z1: torch.Tensor, z2: torch.Tensor, temperature: float) -> torch.Tensor:
+    logits = z1 @ z2.T / temperature
+    labels = torch.arange(z1.size(0), device=z1.device)
+    return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels))
+
+
+@torch.no_grad()
+def extract_features_for_probe(
+    model: ContrastiveEncoder,
+    df: pd.DataFrame,
+    transform,
+    args: argparse.Namespace,
+    device: torch.device,
+    class_names: list[str] | None = None,
+):
+    dataset = SslPairDataset(df, transform)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
+    model.eval()
+    features = []
+    labels = []
+    for batch in loader:
+        clinical = batch["clinical"].to(device, non_blocking=True)
+        dermoscopic = batch["dermoscopic"].to(device, non_blocking=True)
+        c = F.normalize(model.forward_features(clinical), dim=1)
+        d = F.normalize(model.forward_features(dermoscopic), dim=1)
+        pair = torch.cat([c, d, torch.abs(c - d), c * d], dim=1)
+        features.append(pair.cpu().numpy())
+        labels.extend(str(label) for label in batch["label"])
+    if class_names is None:
+        observed = sorted(set(labels))
+        class_names = [label for label in CLASS_NAMES if label in observed] + [label for label in observed if label not in CLASS_NAMES]
+    label_to_idx = {label: idx for idx, label in enumerate(class_names)}
+    y = np.asarray([label_to_idx[label] for label in labels], dtype=np.int64)
+    return np.concatenate(features, axis=0), y, class_names
+
+
+def align_probabilities(y_prob: np.ndarray, classes: np.ndarray, n_classes: int) -> np.ndarray:
+    aligned = np.zeros((y_prob.shape[0], n_classes), dtype=np.float64)
+    for src_idx, class_idx in enumerate(classes):
+        aligned[:, int(class_idx)] = y_prob[:, src_idx]
+    row_sum = aligned.sum(axis=1, keepdims=True)
+    return aligned / np.clip(row_sum, 1e-12, None)
+
+
+if __name__ == "__main__":
+    main()
+