Topological divergence

examples/configs/estimators/default_estimators.yaml

@@ -100,3 +100,7 @@
 - name: CocoaPPL
 - name: CocoaMTE
 - name: SemanticDensity
+- name: TopologicalDivergence
+  cfg:
+    heads: null
+    n_jobs: -1

examples/configs/stat_calculators/default_calculators.yaml

@@ -27,4 +27,4 @@
     - "train_embeddings"
     - "background_train_embeddings"
     - "train_greedy_log_likelihoods"
-  dependencies:
+  dependencies:

requirements.txt

@@ -35,3 +35,5 @@ evaluate>=0.4.2
 spacy>=3.4.0,<3.8.0
 fastchat
 diskcache>=5.6.3
+joblib
+ripser

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,

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,
+    )

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

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}")

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)