diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index 628e1d42d..c16ccece0 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -100,3 +100,7 @@ - name: CocoaPPL - name: CocoaMTE - name: SemanticDensity +- name: TopologicalDivergence + cfg: + heads: null + n_jobs: -1 diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index 265648300..b9617b723 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -27,4 +27,4 @@ - "train_embeddings" - "background_train_embeddings" - "train_greedy_log_likelihoods" - dependencies: \ No newline at end of file + dependencies: diff --git a/requirements.txt b/requirements.txt index 5b98757a2..acb2db7ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -35,3 +35,5 @@ evaluate>=0.4.2 spacy>=3.4.0,<3.8.0 fastchat diskcache>=5.6.3 +joblib +ripser diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index b7fa94fc7..983a41622 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -148,6 +148,30 @@ def _register( }, ) _register(AttentionForwardPassCalculator) + _register( + TrainMTopDivCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator", + { + "heads_extraction_priority": "cache", + "cache_path": "./workdir/output", + "model_heads_cache": "model_heads_cache.yaml", + "max_heads": 6, + "n_jobs": -1, + "dataset": ["LM-polygraph/coqa", "continuation"], + "text_column": "input", + "label_column": "output", + "description": "", + "prompt": "", + "few_shot_split": "train", + "train_split": "train", + "load_from_disk": False, + "subsample_train_dataset": 100, + "n_shot": 0, + "batch_size": 1, + "seed": [1], + "size": None, + }, + ) elif model_type == "VisualLM": _register( GreedyProbsVisualCalculator, diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py new file mode 100644 index 000000000..6acf7bb10 --- /dev/null +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py @@ -0,0 +1,63 @@ +import os +import logging + +from lm_polygraph.utils.dataset import Dataset +from lm_polygraph.stat_calculators.train_mtopdiv import ( + TrainMTopDivCalculator, +) + + +log = logging.getLogger("lm_polygraph") + + +def load_dataset(args): + + log.info("=" * 100) + log.info("Loading train dataset...") + + train_dataset = Dataset.load( + args.dataset, + args.text_column, + getattr(args, "label_column", None), + batch_size=args.batch_size, + prompt=getattr(args, "prompt", ""), + description=getattr(args, "description", ""), + mmlu_max_subject_size=100, + n_shot=getattr(args, "n_shot", 5), + few_shot_split=args.few_shot_split, + few_shot_prompt=None, + instruct=None, + split=args.train_split, + size=args.size, + load_from_disk=args.load_from_disk, + trust_remote_code=False, + ) + + if args.subsample_train_dataset != -1: + train_dataset.subsample( + args.subsample_train_dataset, + seed=( + int(list(args.seed)[0]) if not isinstance(args.seed, int) else args.seed + ), + ) + + log.info("Done loading train data.") + return train_dataset + + +def load_stat_calculator(config, builder): + priority = config.heads_extraction_priority + cache_path = os.path.join(config.cache_path, config.model_heads_cache) + max_heads = config.max_heads + n_jobs = config.n_jobs + + def load_train_dataset_fn(): + return load_dataset(config) + + return TrainMTopDivCalculator( + priority, + load_train_dataset_fn, + cache_path, + max_heads, + n_jobs, + ) diff --git a/src/lm_polygraph/estimators/__init__.py b/src/lm_polygraph/estimators/__init__.py index e107c140b..55e6a9915 100644 --- a/src/lm_polygraph/estimators/__init__.py +++ b/src/lm_polygraph/estimators/__init__.py @@ -83,6 +83,7 @@ from .kernel_language_entropy import KernelLanguageEntropy from .luq import LUQ from .eigenscore import EigenScore +from .topological_divergence import TopologicalDivergence from .cocoa import CocoaMSP, CocoaPPL, CocoaMTE from .rauq import RAUQ from .csl import CSL diff --git a/src/lm_polygraph/estimators/mtopdiv_utils.py b/src/lm_polygraph/estimators/mtopdiv_utils.py new file mode 100644 index 000000000..edeeef8d1 --- /dev/null +++ b/src/lm_polygraph/estimators/mtopdiv_utils.py @@ -0,0 +1,187 @@ +import os +import yaml +import warnings +import numpy as np +from typing import List, Tuple, Optional +import ast + +try: + from ripser import ripser +except ImportError: + raise ImportError( + "Please install the 'ripser' package to use TopologicalDivergence estimator. " + "You can install it via 'pip install ripser'." + ) + +try: + from joblib import Parallel, delayed + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + IS_PARALLEL_AVAILABLE = True +except ImportError: + IS_PARALLEL_AVAILABLE = False + warnings.warn( + "Joblib is not installed. Parallel processing for TopologicalDivergence call will not be available. " + "Please install it via 'pip install joblib' if you want to use parallel processing." + ) + + +def transform_attention_scores_to_distances( + attention_weights: np.ndarray, +) -> np.ndarray: + """Transform attention matrix to the matrix of distances between tokens. + + Parameters + ---------- + attention_weights : np.ndarray + Attention matrices of one sample (n_heads x n_tokens x n_tokens). + + Returns + ------- + np.ndarray + Distance matrix. + + """ + attention_weights = attention_weights.astype(np.float32) + n_tokens = attention_weights.shape[-1] + distance_mx = 1 - np.clip(attention_weights, a_min=0.0, a_max=None) + zero_diag = np.ones((n_tokens, n_tokens)) - np.eye(n_tokens) + + distance_mx *= np.broadcast_to(zero_diag, distance_mx.shape) + distance_mx = np.minimum( + np.swapaxes(distance_mx, -1, -2), + distance_mx, + ) + + return distance_mx + + +def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: + """ + Compute the MTopDiv (Manifold Topology Divergence) score from a distance matrix. + + This function calculates the sum of persistence intervals in the H₀ (zero-dimensional) + persistent homology barcode, corresponding to the lifetimes of connected components + in a Vietoris–Rips filtration. + + Parameters: + distance_mx (np.ndarray): A square, symmetric distance matrix. + + Returns: + float: Sum of finite H₀ barcode lengths (birth–death), representing topological diversity. + """ + barcodes = ripser(distance_mx, distance_matrix=True, maxdim=0)["dgms"] + if len(barcodes) > 0: + return barcodes[0][:-1, 1].sum() + return 0 + + +def get_mtopdivs( + heads: List[Tuple[int, int]], + length_responses: List[int], + attention_weights_batch: np.array, + n_jobs: Optional[float] = 1, +) -> np.ndarray: + batch_size = attention_weights_batch.shape[0] + padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) + + def job(layer_head_pair): + layer, head = layer_head_pair + mtopdivs = [] + distance_matrices = transform_attention_scores_to_distances( + attention_weights_batch[:, layer, head] + ) + + for sample_id in range(batch_size): + distance_matrix = distance_matrices[sample_id] + padding_length = padding_lengths[sample_id] + response_length = length_responses[sample_id] + + if padding_length > 0: + distance_matrix = distance_matrix[:-padding_length, :-padding_length] + distance_matrix[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrix) / response_length + mtopdivs.append(mtopdiv) + + return np.array(mtopdivs, dtype=float) + + if IS_PARALLEL_AVAILABLE: + with Parallel(n_jobs=n_jobs, prefer="processes") as parallel: + mtopdivs = parallel(delayed(job)((layer, head)) for layer, head in heads) + else: + mtopdivs = [job((layer, head)) for layer, head in heads] + + mtopdivs = np.stack(mtopdivs, axis=1) + return mtopdivs + + +def load_model_heads( + cache_path: Optional[str], + model_path: str, +) -> Optional[List[Tuple[int, int]]]: + """ + Load model heads. + + Parameters + ---------- + cache_path : Optional[str] + Path to the YAML cache file. + model_path : str + Unique path to the model. + + Returns + ------- + Optional[List[Tuple[int, int]]] + List of heads for the model if found, else None. + """ + if cache_path and os.path.isfile(cache_path): + try: + with open(cache_path, "r") as f: + config = yaml.safe_load(f) + models = config.get("models", []) + for model_entry in models: + if isinstance(model_entry, dict) and model_path in model_entry: + return ast.literal_eval(model_entry[model_path]) + + print(f"Model '{model_path}' not found in cache.") + except Exception as e: + print(f"Failed to load heads from cache: {e}") + return None + + +def save_model_heads( + cache_path: str, + model_path: str, + heads: List[Tuple[int, int]], +) -> None: + """ + Save a list of heads for a model. + + Parameters + ---------- + cache_path : str + Path to the YAML file to save the model heads. + model_path : str + Unique path to the model. + heads : List[Tuple[int, int]] + List of heads to save. + """ + if os.path.isfile(cache_path): + with open(cache_path, "r") as f: + config = yaml.safe_load(f) or {} + else: + config = {} + models = config.get("models", []) + + for entry in models: + if isinstance(entry, dict) and model_path in entry: + return + + models.append({model_path: f"{heads}"}) + config["models"] = models + + os.makedirs(os.path.dirname(cache_path), exist_ok=True) + + with open(cache_path, "w") as f: + yaml.safe_dump(config, f, sort_keys=False) + print(f"Saved heads for '{model_path}' to {cache_path}") diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py new file mode 100644 index 000000000..7f5a0fdc2 --- /dev/null +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -0,0 +1,85 @@ +from typing import Dict, List, Tuple, Optional +import numpy as np + +from .estimator import Estimator +from .mtopdiv_utils import get_mtopdivs + + +class TopologicalDivergence(Estimator): + """ + Estimates the sequence-level uncertainty of a language model following the method of + "Hallucination Detection in LLMs with Topological Divergence on Attention Graphs" + as provided in the paper https://arxiv.org/abs/2504.10063. + Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Computes topological divergences between prompt and response attention + graphs to identify hallucination-indicative heads. + """ + + def __init__( + self, + heads: Optional[List[Tuple[int, int]]] = None, + n_jobs: int = 1, + ): + """ + Initializes TopologicalDivergence estimator. + + Parameters: + heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. + First integer is layer index, second is head index. + If not provided or empty, all heads will be used. + n_jobs (int): Number of jobs for parallel processing. Default: -1. + """ + if not heads: + calculators = [ + "topological_divergence_heads", + "greedy_tokens", + "forwardpass_attention_weights", + ] + else: + calculators = ["greedy_tokens", "forwardpass_attention_weights"] + + super().__init__(calculators, "sequence") + self._heads = heads + self._n_jobs = n_jobs + + def __str__(self): + return "TopologicalDivergence" + + @property + def heads(self): + return self._heads + + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: + """ + Calculates sequence-wise MTopDiv scores for selected heads of attention masks. + Returns the mean MTopDiv score for each input text. + + Parameters: + stats (Dict[str, np.ndarray]): input statistics, consisting of: + * tokenized model generations for each input text in 'greedy_tokens' + * attention scores of shape [batch_size, num_layers, num_heads, seq_len, seq_len] + in 'forwardpass_attention_weights' + Returns: + np.ndarray: float uncertainty for each sample in input statistics. + Higher values indicate more uncertain samples. + """ + length_responses = list(map(len, stats["greedy_tokens"])) + attention_weights_batch = stats["forwardpass_attention_weights"] + + if self._heads is None: + best_heads = stats["topological_divergence_heads"] + self._heads = best_heads + + # After selecting heads once, we drop train stats to prevent recomputing them in later runs. + # This assumes the manager object is reinitialized before the next estimator call. + # Alternatively, a new estimator instance can be created with the selected heads. + self.stats_dependencies = ["greedy_tokens", "forwardpass_attention_weights"] + + mtopdivs = get_mtopdivs( + self._heads, + length_responses, + attention_weights_batch, + n_jobs=self._n_jobs, + ) + + return np.mean(mtopdivs, axis=1) diff --git a/src/lm_polygraph/stat_calculators/__init__.py b/src/lm_polygraph/stat_calculators/__init__.py index 871a8c49d..1e52a54b6 100644 --- a/src/lm_polygraph/stat_calculators/__init__.py +++ b/src/lm_polygraph/stat_calculators/__init__.py @@ -60,3 +60,4 @@ from .extract_claims import ClaimsExtractor from .infer_causal_lm_calculator import InferCausalLMCalculator from .semantic_classes import SemanticClassesCalculator +from .train_mtopdiv import TrainMTopDivCalculator diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py new file mode 100644 index 000000000..9ec809216 --- /dev/null +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -0,0 +1,127 @@ +import numpy as np +from sklearn.metrics import roc_auc_score +from tqdm import tqdm +from typing import Dict, List, Tuple, Literal, Callable +from itertools import product + +from ..stat_calculators import ( + StatCalculator, + GreedyProbsCalculator, + AttentionForwardPassCalculator, +) +from ..generation_metrics import RougeMetric +from lm_polygraph.utils.model import WhiteboxModel +from lm_polygraph.estimators.mtopdiv_utils import ( + get_mtopdivs, + load_model_heads, + save_model_heads, +) + + +class TrainMTopDivCalculator(StatCalculator): + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + """ + Returns the statistics and dependencies for the calculator. + """ + + return ["topological_divergence_heads"], [] + + def __init__( + self, + priority: Literal["train", "cache"] = "cache", + load_train_dataset_fn: Callable = None, + cache_path: str = None, + max_heads: int = 6, + n_jobs: int = 1, + ): + super().__init__() + + self.priority = priority + self.cache_path = cache_path + self.load_train_dataset_fn = load_train_dataset_fn + self.max_heads = max_heads + self.n_jobs = n_jobs + + def select_heads(self, scores, labels): + grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] + deltas = hal_scores.mean(0) - grounded_scores.mean(0) + heads = sorted(range(len(deltas)), key=lambda x: deltas[x], reverse=True) + + best_auroc, n_opt = 0, 0 + for n in range(1, self.max_heads + 1): + n_best_heads = heads[:n] + predictions = scores[:, n_best_heads].mean(axis=1) + roc_auc = roc_auc_score(labels, predictions) + if roc_auc > best_auroc: + best_auroc = roc_auc + n_opt = n + return heads[:n_opt] + + def __call__( + self, + dependencies: Dict[str, np.array], + texts: List[str], + model: WhiteboxModel, + max_new_tokens: int = 100, + ) -> Dict[str, np.ndarray]: + name_or_path = model.model.config.name_or_path + heads = load_model_heads(self.cache_path, name_or_path) + if heads and self.priority == "cache": + heads = np.array(heads) + heads = heads[: self.max_heads] + return {"topological_divergence_heads": heads} + + train_dataset = self.load_train_dataset_fn() + + mtopdivs = [] + labels = [] + + greedy_calc = GreedyProbsCalculator(False, False) + attn_forward_pass_calc = AttentionForwardPassCalculator() + generation_metric = RougeMetric("rougeL") + + for input_text, target in tqdm(train_dataset): + stats = greedy_calc( + dependencies=None, + texts=input_text, + model=model, + max_new_tokens=max_new_tokens, + ) + attn_weights_batch = attn_forward_pass_calc( + dependencies=stats, + texts=input_text, + model=model, + max_new_tokens=max_new_tokens, + )["forwardpass_attention_weights"] + length_responses = list(map(len, stats["greedy_tokens"])) + + _, num_layers, num_heads, _, _ = attn_weights_batch.shape + heads = product(range(num_layers), range(num_heads)) + + mtopdivs.append( + get_mtopdivs( + heads, + length_responses, + attn_weights_batch, + n_jobs=self.n_jobs, + ) + ) + labels.append( + generation_metric( + stats=stats, + target_texts=target, + ) + ) + + mtopdivs = np.concatenate(mtopdivs, axis=0) + + labels = np.concatenate(labels) + labels = ~(labels > 0.3) + + heads = self.select_heads(mtopdivs, labels) + heads = np.unravel_index(heads, (num_layers, num_heads)) + heads = np.stack(heads, axis=1) + + save_model_heads(self.cache_path, name_or_path, heads.tolist()) + return {"topological_divergence_heads": heads} diff --git a/src/lm_polygraph/utils/dataset.py b/src/lm_polygraph/utils/dataset.py index 6d99437f0..4e4a36018 100644 --- a/src/lm_polygraph/utils/dataset.py +++ b/src/lm_polygraph/utils/dataset.py @@ -50,8 +50,8 @@ def select(self, indices: List[int]): Parameters: indices (List[int]): indices to left in the dataset.Must have the same length as input texts. """ - self.x = [self.x[i] for i in indices] - self.y = [self.y[i] for i in indices] + self.x = [self.x[int(i)] for i in indices] + self.y = [self.y[int(i)] for i in indices] return self def train_test_split(self, test_size: int, seed: int, split: str = "train"): diff --git a/src/lm_polygraph/utils/factory_estimator.py b/src/lm_polygraph/utils/factory_estimator.py index f4eb9f352..e84ad42e1 100644 --- a/src/lm_polygraph/utils/factory_estimator.py +++ b/src/lm_polygraph/utils/factory_estimator.py @@ -52,6 +52,7 @@ def load_simple_estimators(name: str, config): FocusClaim, AttentionScore, AttentionScoreClaim, + TopologicalDivergence, CocoaMSP, CocoaPPL, CocoaMTE, diff --git a/test/test_estimators.py b/test/test_estimators.py index 9bc68bdf8..b8ebed98d 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -1,5 +1,6 @@ import torch import pytest +from joblib.externals.loky import get_reusable_executor from transformers import AutoModelForCausalLM, AutoTokenizer @@ -252,6 +253,28 @@ def test_attentionscore(model): assert isinstance(ue.uncertainty, float) +def test_topological_divergence_fixed_heads(model): + estimator = TopologicalDivergence( + heads=[[0, 0], [0, 1], [1, 0], [1, 1]], + n_jobs=-1, + ) + ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) + assert isinstance(ue.uncertainty, float) + + +def test_topological_divergence_select_heads(model): + estimator = TopologicalDivergence(n_jobs=-1) + ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) + assert isinstance(ue.uncertainty, float) + # second time heads should be loaded from cache + estimator = TopologicalDivergence(n_jobs=-1) + ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) + assert isinstance(ue.uncertainty, float) + + def test_cocoamsp(model): estimator = CocoaMSP() ue = estimate_uncertainty(model, estimator, INPUT)