Finding Visual Task Vectors

TL;DR

This work probes how visual in-context learning operates in vision models by identifying Visual Task Vectors—task-encoding activations that can be patched into transformer attention heads. It builds a pipeline to discover these vectors by computing mean activations per task, scoring activations for task relevance, and using REINFORCE to select patch positions that maximize task performance in MAE-VQGAN. Empirically, patched Task Vectors enable solving downstream image-to-image tasks without input-output demonstrations, often outperforming the original one-shot prompts while reducing forward computation. The study also provides ablations, qualitative analyses, and cross-domain explorations (e.g., Llama2-7B and x-ray data), highlighting the distributed encoder-decoder nature of visual task representations and showing vector arithmetic can yield composable tasks.

Abstract

Visual Prompting is a technique for teaching models to perform a visual task via in-context examples, without any additional training. In this work, we analyze the activations of MAE-VQGAN, a recent Visual Prompting model, and find task vectors, activations that encode task-specific information. Equipped with this insight, we demonstrate that it is possible to identify the task vectors and use them to guide the network towards performing different tasks without providing any input-output examples. To find task vectors, we compute the average intermediate activations per task and use the REINFORCE algorithm to search for the subset of task vectors. The resulting task vectors guide the model towards performing a task better than the original model without the need for input-output examples.

Figures (17)

• Figure 1: Visual Prompting models like MAE-VQGAN bar2022visual require input-output example(s) to describe the desired task in their forward pass. We analyze the model activations and find Task Vectors, activations that encode task information that can be reused to control the task the model performs (see Figure \ref{['fig:teaser']}a). Specifically, we tap into activations of individual attention heads and replace their outputs with Task Vectors to guide the model to the desired task (see Figure \ref{['fig:teaser']}b). Surprisingly, the resulting models perform better than the original model while removing the need for input-output examples. This confirms that Task Vectors exist in the network activation space and they can guide the model to perform the desired task.
• Figure 2: Activation Scoring Analysis. Individual scores (${\rho_{token}}(i)$) aggregated per Attention Head (left) for the encoder and decoder. Individual token scores of specific heads, along with the t-SNE JMLR:v9:vandermaaten08a clustering of head activations of different tasks (right).
• Figure 3: Qualitative Examples. We qualitatively compare the task-specific variants of our methodology's results with the original model and the CMA and GRS baselines. Our patching methodology performs better than the original MAE-VQGAN model.
• Figure 4: Qualitative Examples. We qualitatively compare the task-specific and multi-task variants of our methodology with the CMA and GRS baselines. Our patching methodology performs better than the original MAE-VQGAN model.
• Figure 5: Out of domain evaluation on x-ray images.