evalrx 0.1.2__py3-none-any.whl

evalrx/models/paper_methods/ifcd.py

@@ -0,0 +1,204 @@
+"""TruthX-backed IFCD components.
+
+IFCD's paper contrasts an anti-hallucination and a hallucination-inducing
+internal-representation edit.  This module ports the released TruthX MLP
+editor and exposes it through reversible modern-HF hooks.  The hooks edit the
+post-projection attention output (the current Transformers boundary), whereas
+the original LLaMA patch edits just before ``o_proj``; callers must therefore
+report this route as an architecture adaptation.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+
+def _dims(value: Any) -> list[int]:
+    text = str(value or "")
+    return [int(item) for item in text.split(",") if item]
+
+
+def _checkpoint_arg(args: Any, name: str) -> Any:
+    return args.get(name) if isinstance(args, dict) else getattr(args, name)
+
+
+class TruthXEditor:
+    """Released TruthX autoencoder with a reversible edit-strength switch."""
+
+    def __init__(self, checkpoint_path: str | Path, *, hidden_size: int, top_layers: int = 15):
+        import torch
+        from torch import nn
+
+        path = Path(checkpoint_path)
+        if not path.is_file():
+            raise FileNotFoundError(f"TruthX checkpoint does not exist: {path}")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        checkpoint = torch.load(path, map_location=device, weights_only=False)
+        args = checkpoint["args"]
+        semantic_dim = int(_checkpoint_arg(args, "semantic_latent_dim"))
+        truthful_dim = int(_checkpoint_arg(args, "truthful_latent_dim"))
+
+        def stack(in_dim: int, hidden: list[int], out_dim: int, *, normalize: bool) -> Any:
+            modules: list[Any] = []
+            current = in_dim
+            for dim in hidden:
+                modules.append(nn.Sequential(nn.Linear(current, dim), nn.LayerNorm(dim), nn.LeakyReLU()))
+                current = dim
+            modules.append(nn.Sequential(nn.Linear(current, out_dim), nn.LayerNorm(out_dim), nn.LeakyReLU()))
+            return nn.Sequential(*modules)
+
+        class Autoencoder(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.semantic_encoder = stack(
+                    hidden_size, _dims(_checkpoint_arg(args, "semantic_hidden_dims")), semantic_dim, normalize=False
+                )
+                self.truthful_encoder = stack(
+                    hidden_size, _dims(_checkpoint_arg(args, "truthful_hidden_dims")), truthful_dim, normalize=True
+                )
+                self.proj = nn.Linear(truthful_dim, semantic_dim, bias=False) if semantic_dim != truthful_dim else None
+                decoder_dims = _dims(_checkpoint_arg(args, "decoder_hidden_dims"))
+                decoder: list[Any] = []
+                current = semantic_dim
+                for dim in decoder_dims:
+                    decoder.append(nn.Sequential(nn.Linear(current, dim), nn.LayerNorm(dim), nn.LeakyReLU()))
+                    current = dim
+                self.decoder = nn.Sequential(*decoder) if decoder else nn.Identity()
+                # TruthX stores this as ``nn.Sequential(Linear(...))``;
+                # preserve the state-dict key ``final_layer.0`` exactly.
+                self.final_layer = nn.Sequential(nn.Linear(current, hidden_size))
+                self.cross_attention = nn.MultiheadAttention(embed_dim=semantic_dim, num_heads=1)
+
+            def truth(self, x: Any) -> Any:
+                return torch.nn.functional.normalize(self.truthful_encoder(x), p=2, dim=-1)
+
+            def reconstruct(self, x: Any, truth: Any) -> Any:
+                semantic = self.semantic_encoder(x)
+                if self.proj is not None:
+                    truth = self.proj(truth)
+                attended, _ = self.cross_attention(
+                    semantic.unsqueeze(0), truth.unsqueeze(0), truth.unsqueeze(0)
+                )
+                return self.final_layer(self.decoder(semantic + attended[0]))
+
+        self.model = Autoencoder().to(device)
+        self.model.load_state_dict(checkpoint["state_dict"])
+        self.model.eval()
+        self.pos_center = checkpoint["pos_center"].to(device)
+        self.neg_center = checkpoint["neg_center"].to(device)
+        # The official release serializes this as a Python list, unlike the
+        # center tensors. Keep it host-side because it only gates hook setup.
+        self.rank = [int(value) for value in checkpoint["rank"]]
+        self.top_layers = int(top_layers)
+        self.strength = 0.5
+
+    def edit(self, values: Any, *, layer_index: int, component: str) -> Any:
+        """Apply TruthX's positive or negative edit only to the final token."""
+        import torch
+        import torch.nn.functional as F
+
+        rank_index = 2 * int(layer_index) + (0 if component == "attn" else 1)
+        if rank_index >= len(self.rank) or int(self.rank[rank_index]) > self.top_layers:
+            return values
+        batch, sequence, width = values.shape
+        flat = values.contiguous().view(-1, width)
+        dtype = self.model.semantic_encoder[0][0].weight.dtype
+        x = flat.to(dtype)
+        # TruthX encodes the unflattened sequence, then flattens it inside its
+        # autoencoder forward. Keeping that order is material: its
+        # cross-attention is over tokens in the current generated sequence.
+        truth = self.model.truth(values.to(dtype))
+        pos = self.pos_center[rank_index].unsqueeze(0).to(dtype)
+        neg = self.neg_center[rank_index].unsqueeze(0).to(dtype)
+        direction = (pos - neg).unsqueeze(0)
+        positive = self.model.reconstruct(
+            x, F.normalize(truth + direction, p=2, dim=-1).reshape(-1, truth.shape[-1])
+        )
+        negative = self.model.reconstruct(
+            x, F.normalize(truth - direction, p=2, dim=-1).reshape(-1, truth.shape[-1])
+        )
+        delta = (positive - negative).view(batch, sequence, width).to(values.dtype)
+        delta = F.normalize(delta, p=2, dim=-1) * torch.linalg.vector_norm(
+            values, ord=2, dim=-1, keepdim=True
+        )
+        mask = torch.zeros((batch, sequence, 1), device=values.device, dtype=values.dtype)
+        mask[:, -1:, :] = 1
+        return values + delta.to(values.dtype) * float(self.strength) * mask
+
+
+@contextmanager
+def truthx_editing(language_model: Any, editor: TruthXEditor) -> Iterator[None]:
+    """Install temporary TruthX hooks on the LLaMA attention and MLP outputs."""
+    decoder = getattr(language_model, "model", language_model)
+    handles: list[Any] = []
+
+    def edit_output(layer_index: int, component: str):
+        def hook(_module: Any, _inputs: Any, output: Any) -> Any:
+            if isinstance(output, tuple):
+                return (editor.edit(output[0], layer_index=layer_index, component=component), *output[1:])
+            return editor.edit(output, layer_index=layer_index, component=component)
+
+        return hook
+
+    try:
+        for index, layer in enumerate(decoder.layers):
+            handles.append(layer.self_attn.register_forward_hook(edit_output(index, "attn")))
+            handles.append(layer.mlp.register_forward_hook(edit_output(index, "ffn")))
+        yield
+    finally:
+        for handle in handles:
+            handle.remove()
+
+
+class IFCDLogitsProcessor:
+    """Compute ``(1 + alpha) log p+ - alpha log p-`` with IFCD's cutoff."""
+
+    def __init__(
+        self,
+        model: Any,
+        negative_inputs: dict[str, Any],
+        editor: TruthXEditor,
+        *,
+        alpha: float = 0.1,
+        beta: float = 0.1,
+        edit_strength: float = 0.5,
+    ) -> None:
+        self.model = model
+        self.negative_inputs = negative_inputs
+        self.editor = editor
+        self.alpha = float(alpha)
+        self.beta = float(beta)
+        self.edit_strength = float(edit_strength)
+        self._negative_output: Any = None
+
+    def __call__(self, input_ids: Any, scores: Any) -> Any:
+        import math
+
+        import torch
+        import torch.nn.functional as F
+
+        positive = F.log_softmax(scores, dim=-1)
+        old_strength = self.editor.strength
+        self.editor.strength = -self.edit_strength
+        try:
+            with torch.no_grad():
+                if self._negative_output is None:
+                    self._negative_output = self.model(
+                        **self.negative_inputs, use_cache=True, return_dict=True
+                    )
+                else:
+                    self._negative_output = self.model(
+                        input_ids=input_ids[:, -1:],
+                        past_key_values=self._negative_output.past_key_values,
+                        use_cache=True,
+                        return_dict=True,
+                    )
+        finally:
+            self.editor.strength = old_strength
+        negative = F.log_softmax(self._negative_output.logits[:, -1, :], dim=-1)
+        cutoff = positive.max(dim=-1, keepdim=True).values + math.log(self.beta)
+        contrastive = (1 + self.alpha) * positive - self.alpha * negative
+        return contrastive.masked_fill(positive < cutoff, -float("inf"))

evalrx/models/paper_methods/pai.py

@@ -0,0 +1,164 @@
+"""PAI decoding executor from *Paying More Attention to Image*.
+
+PAI's released LLaVA route increases the pre-softmax attention score from the
+currently generated token to image tokens in a contiguous layer interval. The
+paired classifier-free-guidance branch uses a separate image-free KV cache.
+"""
+
+from __future__ import annotations
+
+import types
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+
+@contextmanager
+def image_attention_boost(
+    model: Any,
+    *,
+    image_start: int,
+    image_end: int,
+    alpha: float = 0.2,
+    start_layer: int = 2,
+    end_layer: int = 32,
+) -> Iterator[None]:
+    """Temporarily apply PAI's released pre-softmax image-attention formula.
+
+    ``image_end`` is exclusive.  The patch targets only LLaMA eager attention;
+    callers must restore it after generation so a candidate cannot leak into a
+    later baseline or a different repair attempt.
+    """
+    import torch
+    import torch.nn.functional as F
+    from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
+
+    decoder = getattr(model, "model", model)
+    all_layers = decoder.layers
+    layers = all_layers[start_layer : min(end_layer, len(all_layers))]
+    original_forwards = [layer.self_attn.forward for layer in layers]
+
+    def pai_forward(
+        attention: Any,
+        hidden_states: Any,
+        position_embeddings: tuple[Any, Any],
+        attention_mask: Any = None,
+        past_key_values: Any = None,
+        cache_position: Any = None,
+        **kwargs: Any,
+    ) -> tuple[Any, Any]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, attention.head_dim)
+        query_states = attention.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = attention.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = attention.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, attention.layer_idx, cache_kwargs
+            )
+        keys = repeat_kv(key_states, attention.num_key_value_groups)
+        values = repeat_kv(value_states, attention.num_key_value_groups)
+        weights = torch.matmul(query_states, keys.transpose(2, 3)) * attention.scaling
+        if attention_mask is not None:
+            weights = weights + attention_mask[:, :, :, : keys.shape[-2]]
+        # PAI (attention.py): a <- a + alpha * abs(a), at the final query
+        # position and the image token interval, before softmax.  Its CFG
+        # processor makes an image-free language-model call using this same
+        # temporary patch; source PAI disables the boost for that call.
+        stop = min(int(image_end), weights.shape[-1])
+        if not getattr(attention, "_evalrx_pai_suspend", False) and int(image_start) < stop:
+            target = weights[:, :, -1, int(image_start) : stop]
+            weights[:, :, -1, int(image_start) : stop] = target + target.abs() * float(alpha)
+        weights = F.softmax(weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        weights = F.dropout(
+            weights,
+            p=0.0 if not attention.training else attention.attention_dropout,
+            training=attention.training,
+        )
+        output = torch.matmul(weights, values).transpose(1, 2).contiguous()
+        output = output.reshape(*input_shape, -1).contiguous()
+        output = attention.o_proj(output)
+        return output, weights if kwargs.get("output_attentions", False) else None
+
+    try:
+        for layer in layers:
+            layer.self_attn.forward = types.MethodType(pai_forward, layer.self_attn)
+        yield
+    finally:
+        for layer, original in zip(layers, original_forwards):
+            layer.self_attn.forward = original
+
+
+class PAICFGLogitsProcessor:
+    """Released PAI classifier-free-guidance decoding for a LLaMA LM.
+
+    PAI derives an unconditional cache from the same prompt with the image
+    token span removed.  It combines conditional and unconditional log
+    probabilities at every generation step, while temporarily suppressing the
+    image-attention boost in the unconditional pass.  This mirrors
+    ``CFG.py`` in the authors' release, without persisting a second model.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        unconditional_input_ids: Any,
+        *,
+        guidance_scale: float = 2.0,
+        start_layer: int = 2,
+        end_layer: int = 32,
+    ) -> None:
+        # HF LLaVA stores a bare ``LlamaModel`` in ``language_model`` and the
+        # final LM head on the outer wrapper.  Use the wrapper for CFG logits,
+        # but locate the inner decoder layers to suspend PAI's attention boost.
+        self.model = model
+        self.unconditional_input_ids = unconditional_input_ids
+        self.guidance_scale = float(guidance_scale)
+        language_model = getattr(model, "language_model", model)
+        decoder = getattr(language_model, "model", language_model)
+        self.layers = decoder.layers[start_layer : min(end_layer, len(decoder.layers))]
+        self._output: Any = None
+
+    @contextmanager
+    def _without_attention_boost(self) -> Iterator[None]:
+        previous = [
+            getattr(layer.self_attn, "_evalrx_pai_suspend", False) for layer in self.layers
+        ]
+        try:
+            for layer in self.layers:
+                layer.self_attn._evalrx_pai_suspend = True
+            yield
+        finally:
+            for layer, value in zip(self.layers, previous):
+                layer.self_attn._evalrx_pai_suspend = value
+
+    def __call__(self, input_ids: Any, scores: Any) -> Any:
+        import math
+
+        import torch
+        import torch.nn.functional as F
+
+        conditional = F.log_softmax(scores, dim=-1)
+        if self.guidance_scale == 1.0:
+            return conditional
+        with self._without_attention_boost(), torch.no_grad():
+            if self._output is None:
+                self._output = self.model(
+                    input_ids=self.unconditional_input_ids,
+                    use_cache=True,
+                    return_dict=True,
+                )
+            else:
+                self._output = self.model(
+                    input_ids=input_ids[:, -1:],
+                    use_cache=True,
+                    past_key_values=self._output.past_key_values,
+                    return_dict=True,
+                )
+        unconditional = F.log_softmax(self._output.logits[:, -1, :], dim=-1)
+        cutoff = conditional.max(dim=-1, keepdim=True).values + math.log(0.1)
+        guided = self.guidance_scale * (conditional - unconditional) + unconditional
+        return guided.masked_fill(conditional < cutoff, -float("inf"))

evalrx/models/paper_methods/tcd.py

@@ -0,0 +1,202 @@
+"""Temporal Contrastive Decoding (TCD) — Li et al. 2026, arXiv:2604.15383.
+
+TCD is NOT a VCD-family port. VCD/IFCD/PAI all share one combine rule
+(``(1+alpha)*z - alpha*z_tilde`` masked below a ``max + log(beta)`` cutoff,
+see :mod:`evalrx.models.paper_methods.vcd`); TCD's is additive,
+positive-rectified, and restricted to a small candidate set, with no cutoff
+masking at all (Eq. 7-9 below). The two other structural differences from
+VCD's per-step-noisy-forward pattern:
+
+* the "corrupted" view is a **waveform** operation (Hann-window blur, Eq. 1)
+  followed by a full audio-encoder re-encode, computed ONCE before decoding
+  starts — not a per-step latent perturbation;
+* the blur window and update scale are **per-example adaptive**, derived from
+  a self-normalized stability score (Eq. 2-6) that needs the audio encoder's
+  per-layer hidden-state trajectory AND the decoder's per-layer attention to
+  audio tokens — signals VCD/IFCD/PAI never touch.
+
+This module holds pure, framework-light math only (all Eq. references are to
+the paper's Section 3 / Appendix A). The orchestration — encoding, KV-cache
+management, the two-branch decode loop — lives in
+``HFLocalModel.generate_tcd`` (``models/backends/hf_local.py``), matching
+this package's split of "paper math here, model plumbing in the backend".
+
+Hyperparameter defaults (:class:`TCDHyperparams`) are Table 6 verbatim, read
+directly from the PDF (Appendix A) — not re-derived, not taken from a
+WebFetch summary of the paper.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TCDHyperparams:
+    """Table 6 defaults. The paper uses ONE configuration across benchmarks;
+    only ``gamma_gate`` is set per backbone (2.0 for Qwen2-Audio-Instruct,
+    the paper's own hyperparameter anchor — see Appendix A.1/A.2(4))."""
+
+    l_attn: int = 4          # last decoder layers aggregated for the step-wise gate's r_t (Eq. 8)
+    tau: float = 4.0          # stability aggregation softmax temperature (Eq. 4)
+    w_min: float = 8.0        # blur window bounds, ms (Eq. 5)
+    w_max: float = 30.0
+    lam_min: float = 0.3      # update-scale bounds (Eq. 6)
+    lam_max: float = 1.5
+    k_orig: int = 16          # candidate set from the original logits (Eq. 9's Omega_t)
+    k_blur: int = 8           # candidate set from the slow-path logits
+    gamma_gate: float = 2.0   # gate multiplier; paper value for Qwen2-Audio-Instruct
+    alpha_entropy: float = 0.5   # entropy exponent in the gate (Eq. 8)
+    k_ent: int = 5            # top-K for the normalized entropy signal (Eq. 8)
+    eps: float = 1e-6         # numerical stability term (Eq. 3)
+
+
+def hann_blur_waveform(waveform: Any, sample_rate: int, window_ms: float) -> Any:
+    """Eq. 1: temporally blur the waveform with a normalized Hann window, then
+    rescale to preserve global RMS amplitude ("rescaling x_tilde to preserve
+    global amplitude", Section 3.2).
+
+    ``waveform`` is a 1-D mono float32 numpy array (matches
+    ``hf_local._resolve_audio``'s contract). Returns the blurred waveform at
+    the same length and sample rate — it is re-encoded through the SAME
+    processor path as the original audio, not treated specially downstream.
+    """
+    import numpy as np
+
+    win_len = max(1, int(round(float(window_ms) / 1000.0 * sample_rate)))
+    if win_len % 2 == 0:
+        win_len += 1  # odd length -> symmetric window, unambiguous center tap
+    if win_len <= 1:
+        return np.asarray(waveform, dtype=np.float32).copy()
+
+    window = np.hanning(win_len).astype(np.float32)
+    window_sum = window.sum()
+    if window_sum <= 0:  # pragma: no cover - hanning(n>=3) is always positive-summed
+        return np.asarray(waveform, dtype=np.float32).copy()
+    window = window / window_sum
+
+    pad = win_len // 2
+    padded = np.pad(np.asarray(waveform, dtype=np.float32), (pad, pad), mode="reflect")
+    blurred = np.convolve(padded, window, mode="valid").astype(np.float32)
+
+    orig_rms = float(np.sqrt(np.mean(np.square(waveform)))) + 1e-8
+    blur_rms = float(np.sqrt(np.mean(np.square(blurred)))) + 1e-8
+    return (blurred * (orig_rms / blur_rms)).astype(np.float32)
+
+
+def layer_stability(hidden_states: list, *, eps: float) -> Any:
+    """Eq. 2-3: per-layer magnitude ``M``, temporal flux ``F``, and stability
+    ``S_l = M_l / (M_l + F_l + eps)``.
+
+    ``hidden_states`` is a list of per-layer encoder hidden-state tensors,
+    each ``(seq_len, dim)`` with the batch dimension already stripped (one
+    entry per encoder layer, time-ordered along ``seq_len`` — the audio
+    encoder's own frame axis, NOT decoding steps). Returns a 1-D tensor of
+    per-layer stability scores, same length as ``hidden_states``.
+    """
+    import torch
+
+    magnitudes, fluxes = [], []
+    for h in hidden_states:
+        h = h.float()
+        magnitudes.append(h.norm(dim=-1).mean())
+        if h.shape[0] > 1:
+            fluxes.append((h[1:] - h[:-1]).norm(dim=-1).mean())
+        else:  # pragma: no cover - a single-frame encoder trajectory is degenerate input
+            fluxes.append(torch.zeros((), dtype=h.dtype, device=h.device))
+    M = torch.stack(magnitudes)
+    F = torch.stack(fluxes)
+    return M / (M + F + eps)
+
+
+def audio_attention_ratio(attn_layer: Any, audio_mask: Any) -> Any:
+    """Fraction of one decoder layer's attention mass landing on audio-token
+    key positions, averaged over heads and query rows.
+
+    ``attn_layer`` is ``(heads, q, k)`` (batch stripped); ``audio_mask`` is a
+    bool tensor ``(k,)``. Shared by both Eq. 4's per-layer ``r_l`` (all
+    layers, computed once from the prefill) and Eq. 8's step-wise ``r_t``
+    (last ``l_attn`` layers, recomputed at every decode step) — same
+    quantity, different layer subset and call frequency.
+    """
+    attn_layer = attn_layer.float()
+    audio_mass = attn_layer[..., audio_mask].sum(dim=-1)
+    total_mass = attn_layer.sum(dim=-1).clamp_min(1e-8)
+    return (audio_mass / total_mass).mean()
+
+
+def aggregate_stability(stability_per_layer: Any, attention_ratio_per_layer: Any, *, temperature: float) -> float:
+    """Eq. 4: softmax-weight per-layer stability by audio-attention ratio.
+
+    Both inputs are 1-D tensors of equal length. This equal-length
+    requirement is exactly the point in the paper where per-layer stability
+    (an ENCODER-side quantity) and per-layer attention ratio (a DECODER-side
+    quantity) get zipped together index-for-index — faithful only when the
+    audio encoder and text decoder have the same layer count (true for
+    Qwen2-Audio-Instruct: 32 encoder / 32 decoder layers; NOT true for e.g.
+    Qwen2.5-Omni's 32/28 split). Callers must truncate both inputs to
+    ``min(len(...), len(...))`` themselves when depths differ, and report
+    ``paper_method_fidelity("tcd")`` as ``"adapted"`` in that case.
+    """
+    import torch
+
+    if stability_per_layer.shape[0] != attention_ratio_per_layer.shape[0]:
+        raise ValueError(
+            "aggregate_stability requires equal-length per-layer inputs "
+            f"(got {stability_per_layer.shape[0]} stability vs "
+            f"{attention_ratio_per_layer.shape[0]} attention-ratio entries); "
+            "truncate to min(...) before calling for a mismatched-depth architecture"
+        )
+    weights = torch.softmax(temperature * attention_ratio_per_layer.float(), dim=0)
+    return float((weights * stability_per_layer.float()).sum())
+
+
+def adaptive_blur_params(stability: float, hp: TCDHyperparams) -> "tuple[float, float]":
+    """Eq. 5-6: map the scalar stability score to (blur window ms, update scale)."""
+    window_ms = hp.w_min + (hp.w_max - hp.w_min) * stability
+    lam = hp.lam_min + (hp.lam_max - hp.lam_min) * stability
+    return float(window_ms), float(lam)
+
+
+def topk_renormalized_entropy(logits: Any, k: int) -> float:
+    """Eq. 8's uncertainty signal: normalized entropy of softmax(logits),
+    RENORMALIZED over its own top-K probability mass first, then divided by
+    log(K) so the result is bounded in [0, 1]."""
+    import torch
+
+    probs = torch.softmax(logits.float(), dim=-1)
+    top = torch.topk(probs, min(int(k), probs.shape[-1]))
+    top_probs = top.values / top.values.sum().clamp_min(1e-12)
+    entropy = -(top_probs * top_probs.clamp_min(1e-12).log()).sum()
+    denom = math.log(top.values.shape[-1]) if top.values.shape[-1] > 1 else 1.0
+    return float(entropy / denom)
+
+
+def reliance_gate(audio_reliance: float, entropy_hat: float, hp: TCDHyperparams) -> float:
+    """Eq. 8: ``g_t = min(gamma_gate * r_t * H_hat_t ** alpha, 1.0)``."""
+    return min(hp.gamma_gate * float(audio_reliance) * (max(entropy_hat, 0.0) ** hp.alpha_entropy), 1.0)
+
+
+def fuse_logits(z: Any, z_tilde: Any, *, lam: float, gate_value: float, hp: TCDHyperparams) -> Any:
+    """Eq. 7-9: gated, candidate-restricted, positive-only logit update.
+
+    Unlike VCD/IFCD/PAI, there is NO cutoff masking here — tokens outside the
+    candidate set ``Omega_t`` are returned UNCHANGED (``z_t(j)`` verbatim),
+    they are not suppressed to ``-inf``. Only the union of each branch's
+    top-K is nudged, and only upward (``ReLU(z - z_tilde)``, Eq. 7), scaled
+    by ``lam * gate_value``.
+    """
+    import torch
+
+    positive_diff = (z - z_tilde).clamp_min(0.0)
+    orig_idx = torch.topk(z, min(hp.k_orig, z.shape[-1])).indices
+    blur_idx = torch.topk(z_tilde, min(hp.k_blur, z_tilde.shape[-1])).indices
+    candidates = torch.zeros_like(z, dtype=torch.bool)
+    candidates[orig_idx] = True
+    candidates[blur_idx] = True
+
+    fused = z.clone()
+    fused[candidates] = z[candidates] + lam * gate_value * positive_diff[candidates]
+    return fused

evalrx/models/paper_methods/vcd.py

@@ -0,0 +1,45 @@
+"""Hugging Face adapter for VCD's released per-token contrastive sampler."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+
+class VCDLogitsProcessor:
+    """Contrast clean scores with a cached noisy-image generation path.
+
+    The authors' sampling loop performs the second multimodal forward at every
+    decode step.  This processor owns the corresponding noisy KV cache, which
+    lets an unmodified Hugging Face LLaVA ``generate`` call execute the same
+    score formula.
+    """
+
+    def __init__(self, model: Any, noisy_inputs: dict[str, Any], *, alpha: float, beta: float) -> None:
+        self.model = model
+        self.noisy_inputs = noisy_inputs
+        self.alpha = float(alpha)
+        self.beta = float(beta)
+        self._noisy_output: Any = None
+
+    def __call__(self, input_ids: Any, scores: Any) -> Any:
+        import torch
+
+        with torch.no_grad():
+            if self._noisy_output is None:
+                self._noisy_output = self.model(
+                    **self.noisy_inputs,
+                    use_cache=True,
+                    return_dict=True,
+                )
+            else:
+                self._noisy_output = self.model(
+                    input_ids=input_ids[:, -1:],
+                    past_key_values=self._noisy_output.past_key_values,
+                    use_cache=True,
+                    return_dict=True,
+                )
+        noisy_scores = self._noisy_output.logits[:, -1, :]
+        cutoff = scores.max(dim=-1, keepdim=True).values + math.log(self.beta)
+        contrastive = (1.0 + self.alpha) * scores - self.alpha * noisy_scores
+        return contrastive.masked_fill(scores < cutoff, -float("inf"))