Skip to content

MaskInversion: Localized Embeddings via Optimization of Explainability Maps

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Ferrari/Rupprecht (Google/Oxford); data-free localized CLIP embeddings via mask optimization

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MaskInversion generates a localized region embedding from a frozen CLIP model by optimizing a single token—called a Localized Embedding Token (LET)—so that its explainability map matches a binary query mask. The optimization runs at test time without modifying model weights, making it a drop-in replacement for the [CLS] token in any CLIP-based pipeline. It outperforms prior region-representation methods on referring expression retrieval, zero-shot class retrieval, localized captioning, and localized image generation.

Problem

Contrastive VLMs like CLIP are trained with global image-text alignment, producing a single [CLS] token that pools over the full image. Naive region-embedding strategies—cropping or averaging patch tokens over a mask—lose context or produce misaligned representations. Existing fixes (AlphaCLIP, FGVP, ReCLIP) either require retraining with mask supervision or modify the input image heuristically, limiting flexibility and generalization.

Method

Given image \(x\) and binary query mask \(m \in \{0,1\}^{W \times H}\), MaskInversion learns a Localized Embedding Token \(\text{LET}_m \in \mathbb{R}^d\) via test-time optimization over \(K\) gradient descent steps, with the frozen CLIP encoder.

  1. Initialization: \(\text{LET}_m^{(0)} = z_0\) (the global [CLS] token).
  2. Explainability-guided optimization: At each step \(k\), the activation score is \(s^{(k)} = \cos(\text{LET}_m^{(k)}, \bar{z})\), where \(\bar{z}\) averages [CLS] and patch tokens. The explainability map \(E^{(k)} \in [0,1]^{W \times H}\) is derived via LeGrad (last-layer attention gradients). The token is updated to minimize the soft Dice loss against \(m\): $\(\mathcal{L}_{\text{Dice}} = 1 - \frac{2 \cdot \text{intersection}(E^{(k)}, m)}{\text{union}(E^{(k)}, m) + \epsilon}\)$
  3. Optional regularization: \(\mathcal{L}_{\text{reg}} = 1 - \cos(\text{LET}_m^{(k)}, z_0^L)\) balances region vs. global context; weighted by \(\alpha\).
  4. Gradient decomposition: Rather than computing second-order derivatives each iteration, the gradient \(\nabla_A = \frac{\partial z̄}{\partial A} \cdot \text{LET}_m^{(k)}\) is decomposed so that \(\frac{\partial \bar{z}}{\partial A}\) is computed once and reused as a dot product, eliminating repeated backpropagation through the explainability map computation.

AdamW with \(K=10\) iterations is used. All model weights remain frozen throughout.

Key Contributions

  • Test-time localized embedding via explainability-map optimization—no backbone retraining or modification required.
  • Gradient decomposition strategy that converts costly second-order gradient computation into a single precomputed dot product, enabling efficient multi-mask processing on the same image.
  • A regularization loss (\(\alpha \cdot \mathcal{L}_{\text{reg}}\)) to trade off between purely local and globally contextualized representations.
  • Drop-in compatibility: the learned LET replaces [CLS] tokens in classification, captioning (CLIPCap+GPT-2), and image generation (diffusion) without task-specific finetuning.

Results

  • Referring Expression Retrieval (RefCOCO, RefCOCO+, PhraseCut): MaskInversion ViT-H/14 achieves 64.0% Acc@1 on PhraseCut vs. AlphaCLIP's 34.0%, FGVP's 35.9%, RedCircle's 21.5%; 61.2% Acc@1 on RefCOCO vs. 43.4% (AlphaCLIP), 42.6% (FGVP); 52.6% on RefCOCO+ vs. 39.7% (AlphaCLIP), 38.0% (FGVP).
  • Class Retrieval (PascalVOC, PascalContext, COCO, OpenImagesV7): MaskInversion ViT-H/14 reaches 93.5% Acc@1 on PascalVOC vs. RIS's 78.0%; 61.8% on PascalContext vs. 38.1% (RIS); 63.7% on COCO vs. 43.6% (RIS); 51.2% on OpenImagesV7 vs. 43.0% (AlphaCLIP).
  • Consistent gains hold across all four CLIP backbone sizes (ViT-B/32 through ViT-H/14).

Limitations

  • Optimization at test time (10 gradient steps per mask) adds per-region inference latency compared to single-forward-pass methods; not suitable for real-time dense-mask scenarios despite gradient decomposition.
  • Gradient decomposition is derived specifically for cosine-similarity-based activation scores with LeGrad; applicability to other explainability methods may require re-derivation.
  • Evaluated only with CLIP/OpenCLIP backbones; generality to non-contrastive VLMs (e.g., LLaVA, Flamingo) is unverified.
  • Requires a ground-truth or pre-computed binary mask as input; mask quality directly affects embedding quality.

Relevance to Vision-Language Models

MaskInversion directly addresses one of CLIP's principal weaknesses—lack of spatially grounded representations—without touching model weights, which is critical given the cost of retraining large VLMs. By turning explainability maps into a differentiable training signal, it establishes a novel paradigm where model introspection (gradient-based attribution) is repurposed for representation learning rather than post-hoc analysis. The drop-in [CLS] replacement interface means it can augment any CLIP-based downstream pipeline (captioners, classifiers, diffusion models) without architectural changes, making it highly relevant for VLM practitioners building region-aware grounding, referring expression, or compositional reasoning systems. It also complements ongoing work on open-vocabulary segmentation and grounded VLMs by providing richer per-region embeddings without additional mask-annotated pretraining.