Adding two sequence-level and consistency-based uncertainty measures: Predictive Kernel Entropy and Spectral uncertainty

.gitignore

@@ -179,3 +179,6 @@ cython_debug/
 
 # Agent instructions
 **/AGENTS.md
+
+#vscode
+.vscode/

README.md

@@ -186,6 +186,8 @@ UE methods such as `EigValLaplacian()` support fully blackbox LLMs that do not p
 | Eccentricity (Ecc) [(Lin et al., 2023)](https://arxiv.org/abs/2305.19187)                                                                                                      | Black-box   | Meaning Diversity   | High    | Low    |         No          | sequence      |
 | Lexical similarity (LexSim) [(Fomicheva et al., 2020a)](https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00330/96475/Unsupervised-Quality-Estimation-for-Neural-Machine) | Black-box   | Meaning Diversity   | High    | Low    |         No          | sequence      |
 | Kernel Language Entropy [(Nikitin et al., 2024)](https://arxiv.org/pdf/2405.20003)                                                                                             | Black-box | Meaning Diversity | High | Low | No | sequence      |
+| Predictive Kernel Entropy (PKE) [(Gruber et al., 2024)](https://arxiv.org/abs/2310.05833) | Black-box | Meaning Diversity | High | Low | No | sequence |
+| Spectral Uncertainty [(Walha et al., 2025)](https://arxiv.org/abs/2509.22272) | Black-box | Meaning Diversity | High | Low | No | sequence |
 | LUQ [(Zhang et al., 2024)](https://aclanthology.org/2024.emnlp-main.299/)                                                                                                      | Black-box   | Meaning diversity   | High    | Low    |         No          | sequence      |
 | Verbalized Uncertainty 1S [(Tian et al., 2023)](https://arxiv.org/abs/2305.14975)                                                                                              | Black-box   | Reflexive           | Low     | Low    |         No          | sequence      |
 | Verbalized Uncertainty 2S [(Tian et al., 2023)](https://arxiv.org/abs/2305.14975)                                                                                              | Black-box   | Reflexive           | Medium  | Low    |         No          | sequence      |

examples/basic_example_visual.ipynb

@@ -98,7 +98,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "6958a441",
    "metadata": {},
    "outputs": [
@@ -115,7 +115,12 @@
     "%load_ext autoreload\n",
     "%autoreload 2\n",
     "\n",
-    "from transformers import AutoModelForVision2Seq, AutoProcessor\n",
+    "from transformers import AutoProcessor\n",
+    "try:\n",
+    "    from transformers import AutoModelForVision2Seq\n",
+    "except ImportError:\n",
+    "    # transformers >= 5.0 renamed AutoModelForVision2Seq → AutoModelForImageTextToText\n",
+    "    from transformers import AutoModelForImageTextToText as AutoModelForVision2Seq\n",
     "from lm_polygraph.model_adapters.visual_whitebox_model import VisualWhiteboxModel\n",
     "from lm_polygraph import estimate_uncertainty\n",
     "from lm_polygraph.estimators import MaximumTokenProbability, MaximumSequenceProbability, SemanticEntropy, EigValLaplacian"

examples/configs/estimators/default_estimators.yaml

@@ -9,6 +9,8 @@
 - name: PTrueSampling
 - name: MonteCarloSequenceEntropy
 - name: MonteCarloNormalizedSequenceEntropy
+- name: PredictiveKernelEntropy
+- name: SpectralUncertainty
 - name: LexicalSimilarity
   cfg:
     metric: "rouge1"

examples/configs/model/default_visual.py

@@ -1,4 +1,10 @@
-from transformers import AutoModelForVision2Seq, AutoProcessor
+from transformers import AutoProcessor
+
+try:
+    from transformers import AutoModelForVision2Seq
+except ImportError:
+    # transformers >= 5.0 renamed AutoModelForVision2Seq → AutoModelForImageTextToText
+    from transformers import AutoModelForImageTextToText as AutoModelForVision2Seq
 
 
 def load_model(model_path: str, device_map: str):

src/lm_polygraph/defaults/register_default_stat_calculators.py

@@ -56,6 +56,19 @@ def _register(
         "device": None,
     }
 
+    # Shared sentence embedding model defaultconfig
+    sentence_embedding_model_cfg = {
+        "model_name": "all-mpnet-base-v2",
+        "device": None,
+        "cache_folder": None,
+        "encoding_arguments": {
+            "batch_size": 256,
+            "convert_to_numpy": True,
+            "convert_to_tensor": False,
+            "show_progress_bar": False,
+        },
+    }
+
     _register(InitialStateCalculator)
     _register(RawInputCalculator)
     _register(
@@ -74,7 +87,11 @@ def _register(
         {"nli_model": nli_model_cfg},
     )
     _register(SemanticClassesCalculator)
-
+    _register(
+        SampleSentenceEmbeddingsCalculator,
+        "lm_polygraph.defaults.stat_calculator_builders.default_SampleSentenceEmbeddingsCalculator",
+        {"sentence_embedding_model": sentence_embedding_model_cfg},
+    )
     if model_type == "Blackbox":
         _register(BlackboxGreedyTextsCalculator)
         _register(BlackboxSamplingGenerationCalculator)

src/lm_polygraph/defaults/stat_calculator_builders/default_SampleSentenceEmbeddingsCalculator.py

@@ -0,0 +1,12 @@
+from lm_polygraph.stat_calculators.sample_sentence_embeddings import (
+    SampleSentenceEmbeddingsCalculator,
+)
+from .utils import load_sentence_embedding_model
+
+
+def load_stat_calculator(config, builder):
+    if not hasattr(builder, "sentence_embedding_model"):
+        builder.sentence_embedding_model = load_sentence_embedding_model(
+            **config.sentence_embedding_model
+        )
+    return SampleSentenceEmbeddingsCalculator(builder.sentence_embedding_model)

src/lm_polygraph/defaults/stat_calculator_builders/utils.py

@@ -1,6 +1,8 @@
 import logging
+from typing import Dict, Optional
 
 from lm_polygraph.utils.deberta import Deberta, MultilingualDeberta
+from lm_polygraph.utils.sentence_embedder import SentenceEmbedder
 
 log = logging.getLogger("lm_polygraph")
 
@@ -21,3 +23,19 @@ def load_nli_model(
         f"Initialized {nli_model.deberta_path} on {nli_model.device} with batch_size={nli_model.batch_size}"
     )
     return nli_model
+
+
+def load_sentence_embedding_model(
+    model_name: str = "all-mpnet-base-v2",
+    device: Optional[str] = None,
+    cache_folder: Optional[str] = None,
+    encoding_arguments: Optional[Dict] = None,
+) -> SentenceEmbedder:
+    model = SentenceEmbedder(
+        model_name,
+        device=device,
+        cache_folder=cache_folder,
+        encoding_arguments=encoding_arguments,
+    )
+    log.info(f"Initialized SentenceEmbedder({model_name}) on {model.device}")
+    return model

src/lm_polygraph/estimators/__init__.py

@@ -90,3 +90,5 @@
 from .rauq import RAUQ
 from .csl import CSL
 from .semantic_density import SemanticDensity
+from .predictive_kernel_entropy import PredictiveKernelEntropy
+from .spectral_uncertainty import SpectralUncertainty

src/lm_polygraph/estimators/predictive_kernel_entropy.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from typing import Dict
+
+from .estimator import Estimator
+import sklearn.metrics
+
+
+class PredictiveKernelEntropy(Estimator):
+    """
+    Predictive Kernel Entropy (PKE) uncertainty estimator.
+
+    For each input, n answers are sampled and encoded with an external sentence
+    embedding model. Uncertainty is the negative mean pairwise kernel similarity:
+
+        PKE = -1 / (n*(n-1)) * sum_{i != j} k(X_i, X_j)
+
+    Higher values indicate higher uncertainty.
+
+    Reference: https://arxiv.org/abs/2310.05833
+
+    Parameters:
+        kernel (str): Kernel function — one of "rbf", "laplacian", "cosine". Default "rbf".
+        gamma (float): Gamma parameter for the RBF and the laplacian kernels. Ignored for the cosine kernel.
+    """
+
+    def __init__(self, kernel: str = "rbf", gamma: float = 1.0):
+        super().__init__(["sample_sentence_embeddings"], "sequence")
+        if kernel not in ("rbf", "laplacian", "cosine"):
+            raise ValueError(
+                f"Unknown kernel '{kernel}'. Choose from: rbf, laplacian, cosine."
+            )
+        self.kernel = kernel
+        self.gamma = gamma
+        if kernel == "rbf":
+            self.kernel_function = lambda x: sklearn.metrics.pairwise.rbf_kernel(
+                x, gamma=gamma
+            )
+        elif kernel == "laplacian":
+            self.kernel_function = lambda x: sklearn.metrics.pairwise.laplacian_kernel(
+                x, gamma=gamma
+            )
+        elif kernel == "cosine":
+            self.kernel_function = sklearn.metrics.pairwise.linear_kernel
+
+    def __str__(self) -> str:
+        return f"PredictiveKernelEntropy(kernel={self.kernel}, gamma={self.gamma})"
+
+    def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray:
+        scores = []
+        for embeddings in stats["sample_sentence_embeddings"]:
+            E = np.array(embeddings)  # (n, d)
+            n = len(E)
+            if n < 2:
+                scores.append(0.0)
+                continue
+            K = self.kernel_function(E)  # (n, n)
+            scores.append((K.diagonal().sum() - K.sum()) / (n * (n - 1)))
+        return np.array(scores)

src/lm_polygraph/estimators/spectral_uncertainty.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+from typing import Dict
+
+from .estimator import Estimator
+import sklearn.metrics
+import scipy
+
+
+class SpectralUncertainty(Estimator):
+    """
+    Spectral Uncertainty estimator.
+
+    For each input, n answers are sampled and encoded with an external sentence
+    embedding model. The empirical kernel matrix K is built from pairwise kernel
+    values between embeddings, the kernel matrix is normalized to unit trace, and
+    as a result, eigenvalues of K can be interpreted as a probability distribution
+    over latent semantic directions. Finally, uncertainty is computed as the Shannon
+    entropy of that distribution:
+
+        H = -Σ λᵢ log(λᵢ)
+
+    Higher values indicate higher uncertainty (more spread in the spectrum).
+
+    Reference: https://arxiv.org/abs/2509.22272
+
+    Parameters:
+        kernel (str): Kernel function — one of "rbf", "laplacian", "cosine". Default "rbf".
+        gamma (float): Gamma parameter for rbf/laplacian kernels. Ignored for cosine.
+        eps (float): Small value to be added to the diagonal to avoid numerical issues.
+    """
+
+    def __init__(self, kernel: str = "rbf", gamma: float = 1.0, eps: float = 1e-6):
+        super().__init__(["sample_sentence_embeddings"], "sequence")
+        if kernel not in ("rbf", "laplacian", "cosine"):
+            raise ValueError(
+                f"Unknown kernel '{kernel}'. Choose from: rbf, laplacian, cosine."
+            )
+        self.kernel = kernel
+        self.gamma = gamma
+        self.eps = eps
+        if kernel == "rbf":
+            self.kernel_function = lambda x: sklearn.metrics.pairwise.rbf_kernel(
+                x, gamma=gamma
+            )
+        elif kernel == "laplacian":
+            self.kernel_function = lambda x: sklearn.metrics.pairwise.laplacian_kernel(
+                x, gamma=gamma
+            )
+        elif kernel == "cosine":
+            self.kernel_function = sklearn.metrics.pairwise.linear_kernel
+
+    def __str__(self) -> str:
+        return f"SpectralUncertainty(kernel={self.kernel}, gamma={self.gamma})"
+
+    def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray:
+        scores = []
+        for embeddings in stats["sample_sentence_embeddings"]:
+            E = np.array(embeddings)  # (n, d)
+            n = len(E)
+            if n < 2:
+                scores.append(0.0)
+                continue
+            K = self.kernel_function(E)  # (n, n) symmetric
+            K /= n  # Normalizing to trace =1. Here we assume that the embeddings are normalized, so the trace is n.
+            try:
+                eigenvalues, _ = scipy.linalg.eigh(K)
+            except np.linalg.LinAlgError:
+                print(
+                    f"Linalg error happened because of condition {np.linalg.cond(K)}. Trying numpy function."
+                )
+                try:
+                    eigenvalues, _ = np.linalg.eigh(K)
+                except np.linalg.LinAlgError:
+                    epsilon = 1e-6  # You can tune this
+                    K += epsilon * np.eye(K.shape[0])
+                    print(
+                        f"Numy function failed. Adding {epsilon} to the diagonal. New condition is {np.linalg.cond(K)}"
+                    )
+                    eigenvalues, _ = scipy.linalg.eigh(K)
+            # Due to numerical issues, sometimes negative eigenvalues appear (although theoretically impossible because PSD). These are very marginal and are clipped to 0.
+            eigenvalues = np.array([(ev if ev >= 0 else 0.0) for ev in eigenvalues])
+            H = scipy.stats.entropy(eigenvalues)
+            scores.append(H)
+        return np.array(scores)

src/lm_polygraph/stat_calculators/__init__.py

@@ -62,3 +62,4 @@
 from .semantic_classes import SemanticClassesCalculator
 from .attention_forward_pass_visual import AttentionForwardPassCalculatorVisual
 from .vllm_logprobs_extraction import VLLMLogprobsExtractionCalculator
+from .sample_sentence_embeddings import SampleSentenceEmbeddingsCalculator

src/lm_polygraph/stat_calculators/sample_sentence_embeddings.py

@@ -0,0 +1,39 @@
+import numpy as np
+
+from typing import Dict, List, Tuple
+
+from .stat_calculator import StatCalculator
+from lm_polygraph.utils.model import Model
+from lm_polygraph.utils.sentence_embedder import SentenceEmbedder
+
+
+class SampleSentenceEmbeddingsCalculator(StatCalculator):
+    """
+    Encodes each sampled generation with an external sentence embedding model.
+    Produces `sample_sentence_embeddings`: List[np.ndarray] of shape (n_samples, embedding_dim).
+    Used by PredictiveKernelEntropy and any other estimator that needs dense
+    embeddings of sampled texts rather than the LLM's own hidden states.
+    """
+
+    @staticmethod
+    def meta_info() -> Tuple[List[str], List[str]]:
+        return ["sample_sentence_embeddings"], ["sample_texts"]
+
+    def __init__(self, model: SentenceEmbedder):
+        super().__init__()
+        self._embedder = model
+
+    def __call__(
+        self,
+        dependencies: Dict[str, np.ndarray],
+        texts: List[str],
+        model: Model,
+        max_new_tokens: int = 100,
+        **kwargs,
+    ) -> Dict[str, np.ndarray]:
+        return {
+            "sample_sentence_embeddings": [
+                self._embedder.encode(samples)
+                for samples in dependencies["sample_texts"]
+            ]
+        }

src/lm_polygraph/utils/factory_estimator.py

@@ -60,6 +60,8 @@ def load_simple_estimators(name: str, config):
         CSL,
         SemanticDensity,
         BoostedProbSequence,
+        PredictiveKernelEntropy,
+        SpectralUncertainty,
     ]
 
     try:

src/lm_polygraph/utils/sentence_embedder.py

@@ -0,0 +1,53 @@
+import numpy as np
+import torch
+
+from typing import Dict, List, Optional
+from sentence_transformers import SentenceTransformer
+
+
+class SentenceEmbedder:
+    """
+    Wrapper around SentenceTransformer that manages model loading and encoding
+    arguments. Analogous to the Deberta wrapper — intended to be constructed
+    once and shared across stat calculators via BuilderEnvironmentStatCalculator.
+    """
+
+    _DEFAULT_ENCODING_ARGUMENTS: Dict = {
+        "batch_size": 256,
+        "convert_to_numpy": True,
+        "convert_to_tensor": False,
+        "show_progress_bar": False,
+    }
+
+    def __init__(
+        self,
+        model_name: str = "all-mpnet-base-v2",
+        device: Optional[str] = None,
+        cache_folder: Optional[str] = None,
+        encoding_arguments: Optional[Dict] = None,
+    ):
+        self.model_name = model_name
+        if device is None:
+            self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+        self.cache_folder = cache_folder
+        self.encoding_arguments = {
+            **self._DEFAULT_ENCODING_ARGUMENTS,
+            **(encoding_arguments or {}),
+        }
+        self._model = None
+        self.setup()
+
+    def setup(self):
+        if self._model is not None:
+            return
+
+        self._model = SentenceTransformer(
+            self.model_name,
+            device=self.device,
+            cache_folder=self.cache_folder,
+        )
+
+    def encode(self, samples: List[str]) -> np.ndarray:
+        return self._model.encode(samples, **self.encoding_arguments)