Batch extra-info disease attribution stats

evaluate_extra_info_attribution.py

@@ -212,38 +212,6 @@ def death_distribution_parameters(
     return "weibull", torch.stack([nan, scale, shape], dim=1)
 
 
-def disease_distribution_parameters(
-    model,
-    hidden: torch.Tensor,
-    *,
-    token_ids: Sequence[int],
-    dist_mode: str,
-    logits: torch.Tensor | None = None,
-    rho: torch.Tensor | None = None,
-    eps: float = 1e-8,
-) -> tuple[str, torch.Tensor]:
-    ids = [int(x) for x in token_ids]
-    if not ids:
-        empty = hidden.new_empty((hidden.shape[0], 0, 3))
-        return "none", empty
-
-    all_logits = model.calc_risk(hidden) if logits is None else logits
-    disease_lambda = F.softplus(all_logits[:, ids]) + float(eps)
-
-    if dist_mode in {"exponential", "mixed"}:
-        nan = torch.full_like(disease_lambda, float("nan"))
-        return "exponential", torch.stack([disease_lambda, nan, nan], dim=2)
-
-    if dist_mode == "weibull":
-        all_rho = model.calc_weibull_rho(hidden) if rho is None else rho
-        shape = all_rho[:, ids].to(dtype=disease_lambda.dtype).clamp_min(float(eps))
-        scale = torch.pow(disease_lambda.clamp_min(float(eps)), -1.0 / shape)
-        nan = torch.full_like(disease_lambda, float("nan"))
-        return "weibull", torch.stack([nan, scale, shape], dim=2)
-
-    raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
-
-
 def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
     return torch.stack(
         [
@@ -258,19 +226,72 @@ def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch
     )
 
 
-def disease_parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
+def all_disease_parameter_pair_block(
+    *,
+    original_logits: torch.Tensor,
+    ablated_logits: torch.Tensor,
+    dist_mode: str,
+    original_rho: torch.Tensor | None = None,
+    ablated_rho: torch.Tensor | None = None,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    original_lambda = F.softplus(original_logits) + float(eps)
+    ablated_lambda = F.softplus(ablated_logits) + float(eps)
+
+    if dist_mode in {"exponential", "mixed"}:
+        nan = torch.full_like(original_lambda, float("nan"))
         return torch.stack(
             [
-            original[:, :, 0],
+                original_lambda,
-            ablated[:, :, 0],
+                ablated_lambda,
-            original[:, :, 1],
+                nan,
-            ablated[:, :, 1],
+                nan,
-            original[:, :, 2],
+                nan,
-            ablated[:, :, 2],
+                nan,
             ],
             dim=2,
         )
 
+    if dist_mode == "weibull":
+        if original_rho is None or ablated_rho is None:
+            raise ValueError("rho tensors are required for weibull disease parameters")
+        original_shape = original_rho.to(dtype=original_lambda.dtype).clamp_min(float(eps))
+        ablated_shape = ablated_rho.to(dtype=ablated_lambda.dtype).clamp_min(float(eps))
+        original_scale = torch.pow(original_lambda.clamp_min(float(eps)), -1.0 / original_shape)
+        ablated_scale = torch.pow(ablated_lambda.clamp_min(float(eps)), -1.0 / ablated_shape)
+        nan = torch.full_like(original_lambda, float("nan"))
+        return torch.stack(
+            [
+                nan,
+                nan,
+                original_scale,
+                ablated_scale,
+                original_shape,
+                ablated_shape,
+            ],
+            dim=2,
+        )
+
+    raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
+
+
+def grouped_parameter_stats(
+    values: torch.Tensor,
+    group_token_mask: torch.Tensor,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    finite = torch.isfinite(values)
+    values64 = values.to(dtype=torch.float64)
+    safe_values = torch.where(finite, values64, torch.zeros_like(values64))
+    mask = group_token_mask.to(device=values.device, dtype=torch.float64)
+    sums = torch.einsum("nvc,gv->ngc", safe_values, mask)
+    sumsq = torch.einsum("nvc,gv->ngc", safe_values * safe_values, mask)
+    counts = torch.einsum("nvc,gv->ngc", finite.to(dtype=torch.float64), mask)
+    return (
+        sums.detach().cpu().numpy().astype(np.float64, copy=False),
+        sumsq.detach().cpu().numpy().astype(np.float64, copy=False),
+        counts.detach().cpu().numpy().astype(np.float64, copy=False),
+    )
+
 
 def build_extra_info_ablated_slice(
     batch: dict[str, torch.Tensor],
@@ -401,24 +422,27 @@ def update_death_summary(
             acc[f"sumsq__{column}"] += float(np.square(vals).sum())
 
 
-def update_disease_parameter_summary(
+def update_disease_parameter_summary_from_group_stats(
     summary: dict[tuple[Any, ...], dict[str, float]],
     *,
     key_rows: pd.DataFrame,
-    target_group: str,
+    group_names: Sequence[str],
-    target_group_label: str,
+    group_labels: Sequence[str],
-    values: np.ndarray,
+    sums: np.ndarray,
+    sumsq: np.ndarray,
+    counts: np.ndarray,
 ) -> None:
-    if key_rows.empty or values.size == 0:
+    if key_rows.empty or sums.size == 0:
         return
-    table = key_rows.reset_index(drop=True).copy()
+    rows = key_rows.reset_index(drop=True)
-    grouped = table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False)
+    for row_idx, row in rows.iterrows():
-    for key, group in grouped:
+        base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS)
-        if not isinstance(key, tuple):
+        for group_idx, (group, label) in enumerate(zip(group_names, group_labels)):
-            key = (key,)
+            count_row = counts[int(row_idx), int(group_idx)]
-        full_key = (*key, str(target_group), str(target_group_label))
+            n_add = float(np.nanmax(count_row)) if count_row.size else 0.0
-        idx = group.index.to_numpy(dtype=np.int64)
+            if n_add <= 0:
-        vals_2d = values[idx].reshape(-1, values.shape[-1])
+                continue
+            full_key = (*base_key, str(group), str(label))
             acc = summary.setdefault(
                 full_key,
                 {
@@ -428,12 +452,14 @@ def update_disease_parameter_summary(
                     **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
                 },
             )
-        acc["n"] += float(vals_2d.shape[0])
+            acc["n"] += n_add
             for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
-            vals = finite_float64(vals_2d[:, col_idx])
+                count = float(counts[int(row_idx), int(group_idx), int(col_idx)])
-            acc[f"count__{column}"] += float(vals.size)
+                if count <= 0:
-            acc[f"sum__{column}"] += float(vals.sum())
+                    continue
-            acc[f"sumsq__{column}"] += float(np.square(vals).sum())
+                acc[f"count__{column}"] += count
+                acc[f"sum__{column}"] += float(sums[int(row_idx), int(group_idx), int(col_idx)])
+                acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
 
 
 def write_death_summary_csv(
@@ -624,6 +650,8 @@ def main() -> None:
         "all_modeled_diseases": "All modeled diseases",
         **organ_labels,
     }
+    group_names = list(risk_groups.keys())
+    group_labels = [str(risk_group_labels[group]) for group in group_names]
 
     state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
     dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
@@ -635,6 +663,20 @@ def main() -> None:
     load_model_state(model, state_dict)
     model.eval()
 
+    group_token_mask = torch.zeros(
+        (len(group_names), int(dataset.vocab_size)),
+        dtype=torch.float32,
+        device=device,
+    )
+    for group_idx, group in enumerate(group_names):
+        valid_tokens = [
+            int(token)
+            for token in risk_groups[group]
+            if 0 <= int(token) < int(dataset.vocab_size) and int(token) != death_idx
+        ]
+        if valid_tokens:
+            group_token_mask[group_idx, torch.as_tensor(valid_tokens, dtype=torch.long, device=device)] = 1.0
+
     batch_size = int(cfg_get(args, cfg, "batch_size", 128))
     attribution_batch_size = int(
         cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
@@ -675,6 +717,9 @@ def main() -> None:
 
     death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
     disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {}
+    death_key_chunks: list[pd.DataFrame] = []
+    death_value_chunks: list[np.ndarray] = []
+    disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
 
     for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True):
         batch_dev = {
@@ -697,19 +742,6 @@ def main() -> None:
             )
             original_logits = model.calc_risk(hidden)
             original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
-            original_disease_params_by_group = {}
-            disease_distribution_name_by_group = {}
-            for group, tokens in risk_groups.items():
-                disease_distribution, original_disease_params = disease_distribution_parameters(
-                    model,
-                    hidden,
-                    token_ids=tokens,
-                    dist_mode=dist_mode,
-                    logits=original_logits,
-                    rho=original_rho,
-                )
-                disease_distribution_name_by_group[group] = disease_distribution
-                original_disease_params_by_group[group] = original_disease_params
 
         for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches(
             batch_dev,
@@ -751,37 +783,39 @@ def main() -> None:
                 original_death_params[row_tensor],
                 ablated_death_params,
             ).detach().cpu().numpy()
+            death_key_chunks.append(key_table)
+            death_value_chunks.append(value_block)
+
+            disease_values = all_disease_parameter_pair_block(
+                original_logits=original_logits[row_tensor],
+                ablated_logits=ablated_logits,
+                dist_mode=dist_mode,
+                original_rho=None if original_rho is None else original_rho[row_tensor],
+                ablated_rho=ablated_rho,
+            )
+            sums, sumsq, counts = grouped_parameter_stats(
+                disease_values,
+                group_token_mask,
+            )
+            disease_stat_chunks.append((key_table, sums, sumsq, counts))
+
+    death_summary.clear()
+    for key_rows, values in zip(death_key_chunks, death_value_chunks):
         update_death_summary(
             death_summary,
-                key_rows=key_table,
+            key_rows=key_rows,
-                values=value_block,
+            values=values,
         )
 
-            for group, tokens in risk_groups.items():
+    for key_rows, sums, sumsq, counts in disease_stat_chunks:
-                disease_distribution, ablated_disease_params = disease_distribution_parameters(
+        update_disease_parameter_summary_from_group_stats(
-                    model,
-                    ablated_hidden,
-                    token_ids=tokens,
-                    dist_mode=dist_mode,
-                    logits=ablated_logits,
-                    rho=ablated_rho,
-                )
-                if disease_distribution_name_by_group[group] != disease_distribution:
-                    raise RuntimeError(
-                        "Disease distribution changed between original and ablated passes "
-                        f"for group {group!r}: {disease_distribution_name_by_group[group]!r} "
-                        f"vs {disease_distribution!r}"
-                    )
-                value_block = disease_parameter_pair_block(
-                    original_disease_params_by_group[group][row_tensor],
-                    ablated_disease_params,
-                ).detach().cpu().numpy()
-                update_disease_parameter_summary(
             disease_parameter_summary,
-                    key_rows=key_table,
+            key_rows=key_rows,
-                    target_group=group,
+            group_names=group_names,
-                    target_group_label=risk_group_labels[group],
+            group_labels=group_labels,
-                    values=value_block,
+            sums=sums,
+            sumsq=sumsq,
+            counts=counts,
         )
 
     death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"