|
| 1 | +"""Attention rollout for multi-layer Transformer VLMs. |
| 2 | +
|
| 3 | +Computes a multi-layer saliency map by propagating attention through all |
| 4 | +transformer layers using the geometric-mean ("attention flow") aggregation |
| 5 | +from Abnar & Zuidema (2020). |
| 6 | +
|
| 7 | +Algorithm |
| 8 | +--------- |
| 9 | +Given *L* per-layer attention maps ``A_1, β¦, A_L`` (each ``(H, W)`` float32, |
| 10 | +non-negative, from the first to the last encoder layer): |
| 11 | +
|
| 12 | +1. Min-max normalise each map: ``a_l = normalise(A_l)``. |
| 13 | +2. Add a residual identity term (models the skip-connection): |
| 14 | + ``Γ’_l = (1 β residual_weight) * a_l + residual_weight * uniform``, |
| 15 | + where ``uniform = ones(H,W)/(H*W)`` and ``residual_weight`` defaults |
| 16 | + to ``0.5``. |
| 17 | +3. Compute the geometric mean across layers in log-space: |
| 18 | + ``rollout = exp(mean_l [log(Γ’_l + Ξ΅)])`` |
| 19 | +4. Min-max normalise the result to ``[0, 1]``. |
| 20 | +
|
| 21 | +The geometric mean is preferred over the arithmetic mean because it |
| 22 | +preserves sparsity: a cell must receive *consistent* attention across |
| 23 | +all layers to score high, matching the rollout intuition that information |
| 24 | +must flow through every layer. |
| 25 | +
|
| 26 | +When ``layer_attention_weights`` is ``None`` (backend does not expose |
| 27 | +per-layer weights), the explainer falls back to the single final-layer |
| 28 | +attention map with a warning. |
| 29 | +
|
| 30 | +References |
| 31 | +---------- |
| 32 | +Abnar, S., & Zuidema, W. (2020). *Quantifying Attention Flow in |
| 33 | +Transformers*. ACL 2020. arXiv:2005.00928. |
| 34 | +""" |
| 35 | +from __future__ import annotations |
| 36 | + |
| 37 | +import logging |
| 38 | +from dataclasses import dataclass |
| 39 | + |
| 40 | +import numpy as np |
| 41 | + |
| 42 | +from miru.attention.extractor import AttentionExtractor |
| 43 | +from miru.models.base import VLMBackend |
| 44 | + |
| 45 | +logger = logging.getLogger(__name__) |
| 46 | + |
| 47 | +DEFAULT_RESIDUAL_WEIGHT: float = 0.5 |
| 48 | +_LOG_EPS: float = 1e-8 |
| 49 | + |
| 50 | + |
| 51 | +@dataclass(frozen=True) |
| 52 | +class RolloutResult: |
| 53 | + """Output of the attention rollout explainer. |
| 54 | +
|
| 55 | + Attributes: |
| 56 | + rollout_grid: Multi-layer saliency map, float32 |
| 57 | + ``(grid_h, grid_w)`` in ``[0, 1]``. |
| 58 | + n_layers: Number of layers actually used. |
| 59 | + used_layer_weights: Whether per-layer weights were available. |
| 60 | + residual_weight: Residual identity weight applied. |
| 61 | + grid_h: Height of the output grid. |
| 62 | + grid_w: Width of the output grid. |
| 63 | + """ |
| 64 | + |
| 65 | + rollout_grid: np.ndarray |
| 66 | + n_layers: int |
| 67 | + used_layer_weights: bool |
| 68 | + residual_weight: float |
| 69 | + grid_h: int |
| 70 | + grid_w: int |
| 71 | + |
| 72 | + |
| 73 | +class AttentionRollout: |
| 74 | + """Multi-layer attention rollout explainer. |
| 75 | +
|
| 76 | + Args: |
| 77 | + residual_weight: Weight for the uniform identity term ``β [0, 1]``. |
| 78 | + Represents the Transformer skip-connection contribution. |
| 79 | + Defaults to ``0.5``. |
| 80 | + extractor: :class:`~miru.attention.extractor.AttentionExtractor` |
| 81 | + for normalisation and grid resizing. |
| 82 | + """ |
| 83 | + |
| 84 | + def __init__( |
| 85 | + self, |
| 86 | + residual_weight: float = DEFAULT_RESIDUAL_WEIGHT, |
| 87 | + extractor: AttentionExtractor | None = None, |
| 88 | + ) -> None: |
| 89 | + if not 0.0 <= residual_weight <= 1.0: |
| 90 | + raise ValueError( |
| 91 | + f"residual_weight must be in [0, 1], got {residual_weight}" |
| 92 | + ) |
| 93 | + self._residual_weight = residual_weight |
| 94 | + self._extractor = extractor or AttentionExtractor() |
| 95 | + |
| 96 | + def explain( |
| 97 | + self, |
| 98 | + backend: VLMBackend, |
| 99 | + image_array: np.ndarray, |
| 100 | + question: str, |
| 101 | + ) -> RolloutResult: |
| 102 | + """Compute attention rollout for *image_array*. |
| 103 | +
|
| 104 | + Args: |
| 105 | + backend: Any registered :class:`~miru.models.base.VLMBackend`. |
| 106 | + image_array: float32 ``(H, W, 3)`` image in ``[0, 1]``. |
| 107 | + question: Natural-language question. |
| 108 | +
|
| 109 | + Returns: |
| 110 | + :class:`RolloutResult` with the rolled-out saliency grid. |
| 111 | + """ |
| 112 | + out = backend.infer(image_array, question) |
| 113 | + |
| 114 | + if out.layer_attention_weights is not None and len(out.layer_attention_weights) > 0: |
| 115 | + raw_layers = out.layer_attention_weights |
| 116 | + used_layer_weights = True |
| 117 | + else: |
| 118 | + logger.warning( |
| 119 | + "rollout: backend '%s' did not supply layer_attention_weights;" |
| 120 | + " falling back to single-layer attention.", |
| 121 | + backend.name, |
| 122 | + ) |
| 123 | + raw_layers = [out.attention_weights] |
| 124 | + used_layer_weights = False |
| 125 | + |
| 126 | + resolution = self._extractor.resolution |
| 127 | + n_cells = resolution * resolution |
| 128 | + uniform = np.full((resolution, resolution), 1.0 / n_cells, dtype=np.float32) |
| 129 | + |
| 130 | + log_sum = np.zeros((resolution, resolution), dtype=np.float64) |
| 131 | + for raw in raw_layers: |
| 132 | + norm = self._extractor.extract(raw) |
| 133 | + blended = ( |
| 134 | + (1.0 - self._residual_weight) * norm |
| 135 | + + self._residual_weight * uniform |
| 136 | + ).astype(np.float64) |
| 137 | + log_sum += np.log(blended + _LOG_EPS) |
| 138 | + |
| 139 | + geom_mean = np.exp(log_sum / len(raw_layers)).astype(np.float32) |
| 140 | + |
| 141 | + lo, hi = float(geom_mean.min()), float(geom_mean.max()) |
| 142 | + if hi - lo > 1e-8: |
| 143 | + geom_mean = ((geom_mean - lo) / (hi - lo)).astype(np.float32) |
| 144 | + |
| 145 | + h, w = geom_mean.shape |
| 146 | + return RolloutResult( |
| 147 | + rollout_grid=geom_mean, |
| 148 | + n_layers=len(raw_layers), |
| 149 | + used_layer_weights=used_layer_weights, |
| 150 | + residual_weight=self._residual_weight, |
| 151 | + grid_h=h, |
| 152 | + grid_w=w, |
| 153 | + ) |
0 commit comments