From e3c9ecd19f383ac984af1dafc2cdbcfaa4e947af Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 1 Jul 2026 12:47:30 +0800 Subject: [PATCH] Batch extra-info disease attribution stats --- evaluate_extra_info_attribution.py | 264 ++++++++++++++++------------- 1 file changed, 149 insertions(+), 115 deletions(-) diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py index 818c31b..85e3359 100644 --- a/evaluate_extra_info_attribution.py +++ b/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,17 +226,70 @@ def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch ) -def disease_parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: - return torch.stack( - [ - original[:, :, 0], - ablated[:, :, 0], - original[:, :, 1], - ablated[:, :, 1], - original[:, :, 2], - ablated[:, :, 2], - ], - dim=2, +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_lambda, + ablated_lambda, + nan, + nan, + nan, + 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), ) @@ -401,39 +422,44 @@ 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, - target_group_label: str, - values: np.ndarray, + group_names: Sequence[str], + group_labels: Sequence[str], + 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() - grouped = table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False) - for key, group in grouped: - if not isinstance(key, tuple): - key = (key,) - full_key = (*key, str(target_group), str(target_group_label)) - idx = group.index.to_numpy(dtype=np.int64) - vals_2d = values[idx].reshape(-1, values.shape[-1]) - acc = summary.setdefault( - full_key, - { - "n": 0.0, - **{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, - **{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, - **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, - }, - ) - acc["n"] += float(vals_2d.shape[0]) - for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS): - vals = finite_float64(vals_2d[:, col_idx]) - acc[f"count__{column}"] += float(vals.size) - acc[f"sum__{column}"] += float(vals.sum()) - acc[f"sumsq__{column}"] += float(np.square(vals).sum()) + rows = key_rows.reset_index(drop=True) + for row_idx, row in rows.iterrows(): + base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS) + for group_idx, (group, label) in enumerate(zip(group_names, group_labels)): + count_row = counts[int(row_idx), int(group_idx)] + n_add = float(np.nanmax(count_row)) if count_row.size else 0.0 + if n_add <= 0: + continue + full_key = (*base_key, str(group), str(label)) + acc = summary.setdefault( + full_key, + { + "n": 0.0, + **{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, + **{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, + **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, + }, + ) + acc["n"] += n_add + for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS): + count = float(counts[int(row_idx), int(group_idx), int(col_idx)]) + if count <= 0: + continue + 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,38 +783,40 @@ def main() -> None: original_death_params[row_tensor], ablated_death_params, ).detach().cpu().numpy() - update_death_summary( - death_summary, - key_rows=key_table, - values=value_block, - ) + death_key_chunks.append(key_table) + death_value_chunks.append(value_block) - for group, tokens in risk_groups.items(): - disease_distribution, ablated_disease_params = disease_distribution_parameters( - 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, - target_group=group, - target_group_label=risk_group_labels[group], - values=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_rows, + values=values, + ) + + for key_rows, sums, sumsq, counts in disease_stat_chunks: + update_disease_parameter_summary_from_group_stats( + disease_parameter_summary, + key_rows=key_rows, + group_names=group_names, + group_labels=group_labels, + sums=sums, + sumsq=sumsq, + counts=counts, + ) death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"