|
| 1 | +"""Joint intra-modal + cross-modal attention attribution. |
| 2 | +
|
| 3 | +Combines two complementary attention signals that a ViT-based VLM exposes: |
| 4 | +
|
| 5 | +- **Cross-modal attention** (``VLMOutput.attention_weights``): language tokens |
| 6 | + attending to visual patches — the signal already used by the single-step |
| 7 | + attention explainer. |
| 8 | +- **Intra-visual attention** (``VLMOutput.intra_visual_weights``): visual |
| 9 | + patches attending to each other — captures local texture / spatial grouping |
| 10 | + information that cross-modal attention under-weights. |
| 11 | +
|
| 12 | +The joint saliency map is the convex combination |
| 13 | +
|
| 14 | + joint = α · intra_visual + (1 − α) · cross_modal |
| 15 | +
|
| 16 | +then min-max normalised to ``[0, 1]``. When ``intra_visual_weights`` is |
| 17 | +absent the method degrades gracefully to pure cross-modal attention. |
| 18 | +
|
| 19 | +References |
| 20 | +---------- |
| 21 | +arXiv:2509.22415 — cross-attention captures 52–75 % of VLM saliency; the |
| 22 | +remaining variance lives in the intra-visual pathway. |
| 23 | +""" |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +import logging |
| 27 | +from dataclasses import dataclass |
| 28 | + |
| 29 | +import numpy as np |
| 30 | + |
| 31 | +from miru.attention.extractor import AttentionExtractor |
| 32 | +from miru.models.base import VLMBackend |
| 33 | + |
| 34 | +logger = logging.getLogger(__name__) |
| 35 | + |
| 36 | +DEFAULT_INTRA_WEIGHT: float = 0.4 |
| 37 | +MIN_INTRA_WEIGHT: float = 0.0 |
| 38 | +MAX_INTRA_WEIGHT: float = 1.0 |
| 39 | + |
| 40 | + |
| 41 | +@dataclass(frozen=True) |
| 42 | +class JointAttributionResult: |
| 43 | + """Output of the joint attribution explainer. |
| 44 | +
|
| 45 | + Attributes: |
| 46 | + joint_grid: Blended saliency map, float32 ``(grid_h, grid_w)`` in |
| 47 | + ``[0, 1]``. |
| 48 | + intra_weight: Actual weight applied to the intra-visual signal. |
| 49 | + cross_weight: Actual weight applied to the cross-modal signal |
| 50 | + (``1 − intra_weight``). |
| 51 | + used_intra: Whether intra-visual weights were available and used. |
| 52 | + grid_h: Height of the output grid. |
| 53 | + grid_w: Width of the output grid. |
| 54 | + """ |
| 55 | + |
| 56 | + joint_grid: np.ndarray |
| 57 | + intra_weight: float |
| 58 | + cross_weight: float |
| 59 | + used_intra: bool |
| 60 | + grid_h: int |
| 61 | + grid_w: int |
| 62 | + |
| 63 | + |
| 64 | +class JointAttribution: |
| 65 | + """Joint intra-modal + cross-modal attribution explainer. |
| 66 | +
|
| 67 | + Args: |
| 68 | + intra_weight: Weight for the intra-visual signal ``α ∈ [0, 1]``. |
| 69 | + The cross-modal signal receives weight ``1 − α``. |
| 70 | + Defaults to ``0.4`` (60 % cross-modal, 40 % intra-visual). |
| 71 | + extractor: :class:`~miru.attention.extractor.AttentionExtractor` |
| 72 | + for normalisation and grid resizing. |
| 73 | + """ |
| 74 | + |
| 75 | + def __init__( |
| 76 | + self, |
| 77 | + intra_weight: float = DEFAULT_INTRA_WEIGHT, |
| 78 | + extractor: AttentionExtractor | None = None, |
| 79 | + ) -> None: |
| 80 | + if not MIN_INTRA_WEIGHT <= intra_weight <= MAX_INTRA_WEIGHT: |
| 81 | + raise ValueError( |
| 82 | + f"intra_weight must be in [{MIN_INTRA_WEIGHT}, {MAX_INTRA_WEIGHT}]," |
| 83 | + f" got {intra_weight}" |
| 84 | + ) |
| 85 | + self._intra_weight = intra_weight |
| 86 | + self._extractor = extractor or AttentionExtractor() |
| 87 | + |
| 88 | + def explain( |
| 89 | + self, |
| 90 | + backend: VLMBackend, |
| 91 | + image_array: np.ndarray, |
| 92 | + question: str, |
| 93 | + ) -> JointAttributionResult: |
| 94 | + """Compute joint attribution for *image_array*. |
| 95 | +
|
| 96 | + Args: |
| 97 | + backend: Any registered :class:`~miru.models.base.VLMBackend`. |
| 98 | + image_array: float32 ``(H, W, 3)`` image in ``[0, 1]``. |
| 99 | + question: Natural-language question. |
| 100 | +
|
| 101 | + Returns: |
| 102 | + :class:`JointAttributionResult` with the blended saliency grid. |
| 103 | + """ |
| 104 | + out = backend.infer(image_array, question) |
| 105 | + cross_grid = self._extractor.extract(out.attention_weights) |
| 106 | + |
| 107 | + used_intra = out.intra_visual_weights is not None |
| 108 | + if not used_intra: |
| 109 | + logger.warning( |
| 110 | + "joint_attribution: backend '%s' did not supply " |
| 111 | + "intra_visual_weights; falling back to cross-modal only.", |
| 112 | + backend.name, |
| 113 | + ) |
| 114 | + joint = cross_grid |
| 115 | + actual_intra_w = 0.0 |
| 116 | + else: |
| 117 | + intra_grid = self._extractor.extract(out.intra_visual_weights) # type: ignore[arg-type] |
| 118 | + joint = ( |
| 119 | + self._intra_weight * intra_grid |
| 120 | + + (1.0 - self._intra_weight) * cross_grid |
| 121 | + ).astype(np.float32) |
| 122 | + actual_intra_w = self._intra_weight |
| 123 | + |
| 124 | + lo, hi = float(joint.min()), float(joint.max()) |
| 125 | + if hi - lo > 1e-8: |
| 126 | + joint = ((joint - lo) / (hi - lo)).astype(np.float32) |
| 127 | + |
| 128 | + h, w = joint.shape |
| 129 | + return JointAttributionResult( |
| 130 | + joint_grid=joint, |
| 131 | + intra_weight=actual_intra_w, |
| 132 | + cross_weight=1.0 - actual_intra_w, |
| 133 | + used_intra=used_intra, |
| 134 | + grid_h=h, |
| 135 | + grid_w=w, |
| 136 | + ) |
0 commit comments