STABLE DIFFUSION MODELS ARE SECRETLY GOOD AT VISUAL IN-CONTEXT LEARNING

Generalization of the model

We appreciate your suggestion to explore additional tasks to further demonstrate the model's generalization.

Before addressing this, it is important to distinguish between:

• Generalist/multi-task models: These models (e.g., [8-10]) are trained on multiple tasks and encode each task using unique embeddings or tokens. This often limits their generalization to unseen tasks.
• V-ICL models: These models infer task and context during inference through the relationships between the query image, prompt image, and prompt groundtruth.

We think that the Painter model is inclined towards multi-task models than V-ICL due to a few factors: (1) it is trained using annotated datasets (2) encodes the tasks using a color scheme in the prompt groundtruth, instead of inferring it from the relationship between the prompt image and its groundtruth during inference, (3) constrained to the tasks seen during training. In all fairness, the authors acknowledge that they generalize only to unseen categories of seen tasks, indicating that it has limited abilities to generalize to unseen tasks. The fact that the tasks are encoded in the coloring scheme, and that Painter collapses to a multi-task model has been highlighted by Wang et al. in [11] - Sec 3.1. This was observed in our experiments as well, where we saw that Painter struggles with novel tasks or different coloring schemes. For instance, different coloring in keypoint detection results in segmentation outputs, and edge detection using NYUv2 images results in depth maps, which implies poor task inference, overfitting to the training tasks/datasets, and dependence on the predefined coloring scheme (see Section E of revised supplementary for details and visualizations).

Hence, while we agree that generalization to more tasks enhances the practicality of V-ICL, we emphasize that the proposed approach has better generalization ability due to the following reasons:

• The proposed approach is able to perform better than other methods like Visual Prompting, IMProv, Painter, and LVM on six novel tasks.

• We do not use a color scheme to encode the task, but rather model the relationship between prompt image and its groundtruth to infer the task, as intended in visual in-context learning.

• Based on the reviewer's suggestion, we included an additional generative task of low-light enhancement (see Table 6) on the LOL dataset [12] where we show that the proposed method yields comparable results to the Painter network, despite the fact that Painter network uses the LOL dataset in training. With this expanded evaluation, we show the superiority of SD-VICL for a total of seven tasks, out of which two (colorization and low-light enhancement) are generative tasks.

Table 6: Evaluation on low-light enhancement on LOL

Usage of text seems to be overlooked

While we agree that text descriptions could supplement in-context learning, in this work we exclusively focus on visual in-context learning, where we demonstrate that V-ICL properties could be uncovered from an off-the-shelf Stable Diffusion model, contrary to existing work which needs to train foundation models specifically targeting V-ICL. Hence, we intentionally did not use text conditioning to remove any language dependence for task inference, as ideally, V-ICL should infer the task and context from the prompt images. Additionally, we show in Tables 1 and 2 of the main paper that the proposed method without any text guidance is able to perform better than models like IMProv [4] that use text prompting in addition to visual prompting.

Nonetheless, we acknowledge the importance of extending this paradigm to multi-modal contexts in future research. Based on the insights of our work, we hope future research will extend the proposed training-free paradigm to multi-modal prompting (e.g., images + text) and other multi-modal foundation models such as EMU [13], as highlighted by the reviewer.

Thank you for the valuable feedback. In what follows, we address the reviewer's concerns in detail.

Limited comparisons

Since this concern is similar to the one raised by reviewer HFYR, we refer the reader to our comments in response to reviewer HFYR's comments (see Table 4 for quantitative comparisons). To summarize, we have added new experiments which show that the proposed approach outperforms additional recent methods such as Painter [1] and LVM [2]. Additionally, we also observed that these recent models show a tendency to overfit to the tasks/datasets that they are trained on and hence are limited in terms of generalization performance. This is in contrast to our approach, which is training free and has better generalization capabilities. (See Section E of revised supplementary for details and visualizations.) We further compare our method against task-specific unsupervised models that build on the Stable Diffusion priors as described in Tables 1 and 2 in the response to reviewer bhAU’s comments.

Paper does not target best-demonstration selection

We agree with the reviewer's observation about the scope of this work. V-ICL research generally takes two broader directions:

1. Enabling in-context learning in vision foundation models [3,4].
2. Prompt engineering — finding the best prompt demonstrating the task [5-7].

The scope of this work explores along the first direction.

For prompt selection, we employed UnSupPR [5], which retrieves nearest neighbors based on CLIP features without task-specific training. This preserves our training-free generalization, ensuring alignment with the V-ICL paradigm. In Table 5, we observe that the proposed method with unsupervised prompt retrieval performs better than other V-ICL approaches such as VPR-SupPR [5] and InMeMO [7] that use supervised prompt retrieval. Considering that these approaches show that the performance is enhanced by improving the prompt selection, we believe it applies to our approach as well, where better prompts can only improve the results further.

Table 5: Comparison against prompt retrieval (PR) methods based on MAE-VQGAN, evaluated using foreground segmentation on Pascal-5i

[1] Wang et al., "Images speak in images: A generalist painter for in-context visual learning" [CVPR'23]

[2] Bai et al., "LVM: Sequential modeling enables scalable learning for large vision models" [CVPR'24]

[3] Bar et al., "Visual prompting via image inpainting" [NeurIPS'22]

[4] Xu et al., "IMProv: Inpainting-based Multimodal Prompting for Computer Vision Tasks" [arXiv'23]

[5] Zhang et al., "What Makes Good Examples for Visual In-Context Learning?" [NeurIPS'23]

[6] Sun et al., "Exploring effective factors for improving visual in-context learning" [arXiv'23]

[7] Zhang et al., "Instruct me more! random prompting for visual in-context learning" [WACV'24]

[8] Lu et al., "Unified-IO: A unified model for vision, language, and multi-modal tasks" [ICLR'22]

[9] Wang et al., "OFA: Unifying architectures, tasks, and modalities through a simple sequence-to-sequence learning framework" [ICML'22]

[10] Kolesnikov et al., "UViM: A unified modeling approach for vision with learned guiding codes" [NeurIPS'22]

[11] Wang et al., "SegGPT: Segmenting everything in context" [ICCV'23]

[12] Wei et al., "Deep Retinex Decomposition for Low-Light Enhancement" [BMVC'18]

[13] Sun et al., "EMU: Generative pretraining in multimodality" [ICLR'24]

Thanks for author’s feedback. From your feedback, I still have concerns that 1) the 1-shot performance is not competitive to other V-ICL models. 2) The ensemble mechanism you proposed is actually designed for addressing the uniform weight during the feature ensemble. In other words, your implementation is could be used for SegGPT since you can directly use your proposed weights for its mean calculation (you can split those composite images in the feature space, as have done in SegGPT. I understand it would be difficult to implement such a method in MAE-VQGAN since images are combined and resized into 1 image. But for SegGPT and Painter, they just concatenate two images at the H dimension. So technically speaking, they are not composited as same as MAE-VQGAN. Making the use of internal representation to construct qkv weights should not be restricted as you mentioned. Even as you claim, such an ensembling mechanism would face a limited application. 3) ICL is also a special form of few-shot learning. So the rebuttal from the aspect of other field like Few-shot Learning is weird. 4) No experimental results of intuitive VAE-oriented task vector baseline is provided, so the advantage of using iterative denoising is unknown, let alone such an iterative process would bring extra costs. 5) In my opinion, no use of text is a weaken exploration on SD.

Overall. Based on the all responses from authors, I feel this work is still lacking some important exploration in using SD for ICL, which could be refined for next-round submission, thus I vote for rejection.

We appreciate the reviewer's thoughtful follow-up comments and the opportunity to clarify further. Below, we address each of the raised concerns in detail.

Clarification on the number of prompts in the updated tables

We apologize for not explicitly clarifying the number of prompts used in the updated table. As mentioned in our response to Reviewer bhAU (see https://openreview.net/forum?id=fKrFTGnoXY&noteId=x4d5tHpbN9), all results in the updated table were reported using five prompts, as this configuration effectively demonstrates the combined potential of our training-free V-ICL framework and the IWPE ensembling mechanism.
Note that we have revised tables that now include results with a single prompt which still outperforms the related methods on both tasks.

Table 1: Evaluation of semantic segmentation on Cityscapes

Table 2: Evaluation of keypoint detection on DeepFashion

Effectiveness of the proposed IWPE ensembling mechanism

Existing V-ICL models such as Painter and LVM are trained on curated datasets, relying on the model to implicitly learn the context and task relationships from the pre-defined input patterns provided. These methods, by design, use fixed input formats (e.g., composite images or sequences) that do not allow modifications to latent-based prompt ensembling or weighting methods, which we argue is suboptimal (Section 2.3 of the main paper). To clarify, adding IWPE to existing approaches would require re-training their models which is in contrast to the goal of our work to make visual in-context learning methods training-free.

In contrast, our framework allows for flexibility in prompt weighting strategies. Hence, in Section 3.2 of the main paper, we evaluate our Implicitly Weighted Prompt Ensembling (IWPE) mechanism against the feature ensembling (FE) approach proposed by SegGPT, the only other latent-based ensembling method. IWPE shows substantial improvements over FE, enhancing mIoU by 4.0% and accuracy by 2.9%. These results highlight the effectiveness of IWPE in dynamically weighing prompts to improve task performance.

Clarification on not using text conditioning

We would like to clarify that the decision to focus exclusively on visual inputs in our work is deliberate and aligns with the goals of this paper. While Stable Diffusion is widely recognized for its text-to-image generative capabilities, our goal is to uncover its potential for visual in-context learning (V-ICL) using purely visual prompts, where the model infers task and context relationships solely from input-output pairs.

This focus allows us to explore the feasibility of V-ICL without relying on text conditioning. Previous works, such as [14-16], have similarly omitted text conditioning of Stable Diffusion in their pipelines to demonstrate task-specific performance without dependence on textual descriptions. Additionally, the complexity of vision tasks often makes it challenging to design unambiguous text descriptions, further supporting our decision to rely solely on visual prompts.

Potential of cross-attention-based task vector in VAE-based models

We acknowledge the reviewer’s suggestion that simpler generative models like VAEs could potentially achieve similar ICL properties using cross-attention-based task vectors. However, our work demonstrates that Stable Diffusion’s iterative denoising mechanism is particularly advantageous for uncovering task and context relationships.

Unlike in VAEs, the iterative nature of the denoising process gradually aligns the task and context relationships with the target domain, avoiding destabilization of the pre-trained generation process. Modifying an off-the-shelf foundation model to generate outputs in a new target domain requires mechanisms to smoothly integrate these changes into the generative pipeline, which the denoising process facilitates.

While our insights could inspire future work to adapt similar techniques to VAE-based models, the iterative denoising framework provides a distinct advantage for the proposed training-free V-ICL pipeline.

Clarification on "such cross attention operation to implement in-context learning is quite a common''

We would like to clarify that the comment from Reviewer bhAU on the use of cross-attention being common was with regards to works like metric learning [17], few-shot segmentation [18,19], and multi-modal representation learning [20,21]. As it can be noted from Reviewer bhAU's latest comments, they acknowledge that the specific form of cross-attention that we use has not been explored earlier for visual in-context learning. As detailed in our response to Reviewer bhAU (see https://openreview.net/forum?id=fKrFTGnoXY&noteId=FoxGQ4LQZf), while the concept of cross-attention has been used in diverse tasks, our specific formulation differs significantly:

• To the best of our knowledge, this specific form of cross-attention which models task and context relationships, is being utilized for the first time, particularly to facilitate in-context learning

• Our method uses three distinct sources for the query, key, and value vectors (from the query image, prompt image, and prompt groundtruth, respectively). Unlike works such as TriBERT [21], which pairs key and value vectors within the same source, our formulation explicitly pairs keys and values from distinct images ($A$ and $B$) (further clarified in https://openreview.net/forum?id=fKrFTGnoXY&noteId=oGE8jLTFcA)

• Existing works cited by Reviewer bhAU require further training after introducing cross-attention to capture correspondences between inputs. In contrast, our method operates in a completely training-free manner, as acknowledged by Reviewer bhAU.

These distinctions are critical and demonstrate the novelty of our approach in leveraging cross-attention for V-ICL in generative diffusion models.

Clarification on the contributions of our work as it stands in the research of Visual In-Context Learning

We wish to contextualize our contributions within the field of V-ICL research. Unlike ICL research in NLP, visual ICL is still in its early stages. While ICL properties emerge naturally in LLMs, to date, no prior work has demonstrated that current vision foundation models are inherently capable of V-ICL. Existing attempts, such as Visual Prompting (NeurIPS’22), IMProv, Painter (CVPR’23), and LVM (CVPR’24), enable V-ICL through training on curated or uncurated datasets, undermining the core benefits of ICL: the ability to adapt to novel tasks without further training or model updates, and limited generalization performance (as we highlight in Figure 4 of the main paper and Section E of the supplementary).

In contrast, our work addresses this gap by uncovering V-ICL properties in Stable Diffusion using an entirely training-free approach for the first time, aligning more closely with the original definition of ICL in LLM research, while yielding significantly superior generalization performance (a core benefit in ICL). This represents a significant step forward in V-ICL research and demonstrates the feasibility of leveraging pre-trained generative priors for task adaptation without additional supervision.

[14] Chen et al., "Anifacediff: High-fidelity face reenactment via facial parametric conditioned diffusion models." [arXiv'24]

[15] Karras et al., "Dreampose: Fashion image-to-video synthesis via stable diffusion." [ICCV'23]

[16] Tian et al., "Diffuse, attend and segment: Unsupervised zero-shot segmentation using stable diffusion." [CVPR'24]

[17] Kotovenko et al., "Cross-Image-Attention for Conditional Embeddings in Deep Metric Learning", [CVPR'23]

[18] Lin et al., "Few Shot Medical Image Segmentation with Cross Attention Transformer", [MICCAI'23]

[19] Hossain et al. "Visual Prompting for Generalized Few-shot Segmentation: A Multi-scale Approach", [CVPR'24]

[20] Lu et al., "ViLBERT: Pretraining task-agnostic visiolinguistic representations for vision-and-language tasks" [NeurIPS'19]

[21] Rahman et al., "TriBERT: Human-centric audio-visual representation learning" [NeurIPS'21]

[1] Xu et al. "IMProv: Inpainting-based Multimodal Prompting for Computer Vision Tasks" [arXiv'23]

[2] Bar et al. "Visual prompting via image inpainting" [NeurIPS'22]

[3] Wang et al. "Images speak in images: A generalist painter for in-context visual learning" [CVPR'23]

[4] Bai et al. "LVM: Sequential modeling enables scalable learning for large vision models" [CVPR'24]

[5] Tang et al. "Emergent correspondence from image diffusion" [NeurIPS 2023]

[6] Zhang et al. "A tale of two features: Stable diffusion complements dino for zero-shot semantic correspondence" [NeurIPS'24]

[7] Barsellotti et al. "Training-free open-vocabulary segmentation with offline diffusion-augmented prototype generation" [CVPR'24]

[8] Yin et al. "Improved distribution matching distillation for fast image synthesis" [NeurIPS'24]

[9] Bolya et al. "Token merging for fast stable diffusion" [CVPR'23]

We appreciate your valuable feedback and suggestions, which have enhanced the quality of the paper. Please find below additional comparisons to existing approaches and our response to the concerns raised regarding inference speed.

Limited comparisons

Our intention of restricting the comparisons to IMProv [1] and Visual Prompting [2] models was solely to ensure fair assessments, as they were the only methods that did not rely on task specific datasets such as ours. Nonetheless, we acknowledge the reviewer's concerns and have extended the evaluations to other models like Painter [3] and LVM [4] as shown in Table 4.

• Painter: A ViT-Large model trained on annotated datasets such as COCO, ADE20K, and NYUv2.

• LVM: Built on OpenLLaMA’s 7B-parameter model and trained on the UVD-V1 dataset, which combines 50 datasets (e.g., LAION5B) with both annotated and unannotated images.

As shown in Table 4, our method outperforms both Painter and LVM across multiple tasks.

Additionally, Painter and LVM often suffer from overfitting to training tasks, leading to poor generalization when exposed to novel tasks. For example, when the Painter network was provided with images from the NYUv2 dataset with edge images as prompts, it ends up predicting depth maps instead of edge detection output, indicating the model being overfit to the dataset on which it was trained. See Section E in the revised supplementary for details and visualizations.

Additionally, LVM exhibits inconsistencies in its outputs, indicating that it is unable to infer the task appropriately based on the given input prompts. Specifically, for a given task, despite the example prompts coming from the same task, the model predicts outputs for other tasks in some cases. For example, when given prompts from foreground segmentation, the network sometimes predicts outputs that belong to other tasks such as keypoints, segmentation maps, or RGB images. See Section E in the revised supplementary for details and visualizations.

These models attempt to apply learned task-specific priors instead of inferring the task information from the relationship between the prompt image and its ground truth, as expected by in-context learning. In contrast, SD-VICL shows superior ability to infer the task and context through the query and prompt inputs, instead of any learned fixed task-specific priors. These results validate our method’s ability to handle novel tasks without overfitting or requiring extensive training on annotated datasets.

We further compare our method against task-specific unsupervised models that build on the Stable Diffusion priors, as described in Tables 1 and 2 in the response to reviewer bhAU's comments.

Table 4: Extended comparisons against Painter and LVM

Question: Possibility of a one-step adaptation approach

We appreciate this suggestion and agree that inference time is a challenge for diffusion-based methods. The possibility of one-step adaptation like [5-7] is something we can explore as future work. Separately, exploring faster diffusion strategies such as [8,9] is also a potential avenue for future work to reduce inference time.

In addition, we would like to bring the reviewer's attention to Section 3.2 [L497:527] in the main paper, where we show that using multiple prompts via IWPE is one approach to reduce inference time. For example, with five prompts, performance surpasses the single-prompt while reducing inference time by approximately 1.5×.

Thank you for your valuable feedback. In accordance with your suggestions, we have clarified the contributions, outlined the distinctions with other methodologies, and incorporated additional comparisons to related research. Please find below detailed responses to your valuable feedback.

Diffusion models, as trained on vast datasets, have shown excellent performance across various unsupervised tasks

We agree with the reviewer that Stable Diffusion models inherently capture richer priors since they have been trained on vast datasets. As recommended, we compare our approach with task-specific Stable Diffusion based unsupervised methods Tian et al., 2024 and Hedlin et al., 2024 [1,2] and list the following benefits:

• Unlike [1,2] which are task specific, the proposed approach is capable of generalizing to multiple tasks without any task-specific modifications.

• The proposed approach has the ability to benefit from example prompts enabling better overall performance.

• Stable Diffusion-based unsupervised methods primarily utilize the learned feature/attention maps in a post-processing step to predict the task output, disregarding the generative process itself to generate the task outputs. In contrast, our pipeline integrates the generative process itself to generate the task outputs, thereby enhancing the quality of the outputs.

• As presented in Table 1 and Table 2, our proposed method demonstrates superior performance compared to [1] and [2] for Semantic Segmentation and Keypoint Detection, respectively. In the semantic segmentation task, we achieve absolute improvements of 10.3% in mIoU and 6.7% in accuracy, surpassing the results of [1]. Note that we use the same evaluation protocol as [1] (i.e., Hungarian matching) to ensure a fair comparison. Similarly, in the keypoint detection task, we demonstrate improvements of 32.4% in MSE and 12.2% in PCK, outperforming the results of [2].

Table 1: Evaluation of semantic segmentation on Cityscapes

Table 2: Evaluation of keypoint detection on DeepFashion

Furthermore, we compare our approach with training-based V-ICL approaches such as Painter [3] and LVM [4], as depicted in Table 3. It is noteworthy that Painter is ViT-Large model trained on multiple datasets and LVM is based on the 7B-parameter model of OpenLLama [5] trained on a vast dataset – UVD-V1 – which is a unified dataset combining 50 datasets (including 1.5B subset of LAION5B). We demonstrate that our method yields superior generalization performance compared to these models across multiple tasks, despite the fact that Painter/LVM require some form of training to enable in-context learning.

Table 3: Extended comparisons against Painter and LVM

[1] Tian et al., "Diffuse attend and segment: Unsupervised zero-shot segmentation using stable diffusion" [CVPR'24]

[2] Hedlin et al., "Unsupervised keypoints from pretrained diffusion models" [CVPR'24]

[3] Wang et al., "Images speak in images: A generalist painter for in-context visual learning" [CVPR'23]

[4] Bai et al., "LVM: Sequential modeling enables scalable learning for large vision models" [CVPR'24]

[5] Geng & Liu, "Openllama: An open reproduction of llama, May 2023. URL https://github.com/openlm-research/open_llama"

[6] Lu et al., "ViLBERT: Pretraining task-agnostic visiolinguistic representations for vision-and-language tasks" [NeurIPS'19]

[7] Rahman et al., "TriBERT: Human-centric audio-visual representation learning" [NeurIPS'21]

[8] Xu et al., "IMProv: Inpainting-based Multimodal Prompting for Computer Vision Tasks" [arXiv'23]

[9] Wang et al., "SegGPT: Segmenting everything in context" [ICCV'23]

Use of distinct query and key-value vectors as a form of cross-attention is quite common in models like ViLBERT and TriBERT

We acknowledge that employing distinct query, key-value vectors in cross-attention is not a novel concept, as demonstrated in ViLBERT [6] and TriBERT [7]. Nevertheless, our method differs in several key aspects:

• Use for in-context learning: To the best of our knowledge, this specific form of cross-attention which models task and context relationships, is being utilized for the first time, particularly to facilitate in-context learning.

• Distinct Q, K, V sources: While prior cross-attention methods derive the key-value pair from the same source, our formulation uses Q from the query image, K from the prompt image, and V from the prompt groundtruth, which amounts to three distinct images [L244-249], enabling explicit task-context infusion. This design is discussed in L454-476, where alternative attention formulations using the key-value pair from the same source such as {$Q_C, K_B, V_B$} and {$Q_D, K_B, V_B$} yield inferior performance compared to our formulation.

• Training-free: Unlike ViLBERT/TriBERT, which fine-tune their model once the cross-attention formulation is in effect, our approach does not require any training.

To address any potential ambiguity, we will include explicit references to these prior works in the revised submission and expand our discussion in the corresponding sections to clarify how our approach uniquely applies this concept for in-context learning.

We may observe similar in-context emergent properties with multi-modal LLMs

We agree with the reviewer that in-context emergent properties are likely associated with large models being trained on large datasets. In this work, we set out to enable the use of these properties of foundational models for training-free visual in-context learning, which is a challenging task. Please note that this possibility has not been previously investigated in prior research.

Lack of technical contributions and novelty

Thank you for the feedback. We believe the contributions can be better re-phrased as follows:

In contrast to LLMs, due to the diverse nature of vision tasks and their outputs, it is more challenging to directly utilize existing vision foundation models for visual in-context learning. To address this challenge, all previous V-ICL methods require some form of training that deviates from the in-context learning paradigm prevalent in LLM works. In contrast, our method introduces the first training-free V-ICL framework using the Stable Diffusion model.

While we agree that our approach harnesses the learned generative priors of Stable Diffusion, it is important to note that uncovering V-ICL properties remains a challenging endeavor, and the Stable Diffusion model cannot be directly employed for in-context learning of diverse downstream tasks. To overcome this while enabling a training-free approach, we introduce modeling of task and context relationships through in-place attention recomputation. To the best of our knowledge, we are the first to explicitly enforce context and task inference. Unlike multi-modal prompting methods such as IMProv [8], we achieve this without language guidance, solely through self-attention modifications.

Additionally, as discussed in Section 2.3, we propose implicitly weighted prompt ensembling (IWPE) which dynamically weighs prompts based on their query correspondence. This is in contrast to SegGPT's feature ensembling (FE) [9] that uniformly weighs all prompts implying equal informativeness – which is suboptimal. The proposed IWPE yields absolute improvements of 4.0% in mIoU and 2.9% in accuracy over FE.

Question: Is the emergent property unique to diffusion models?

As demonstrated previously in the NLP domain, in-context learning properties are inherent to foundation models trained on extensive datasets. The primary objective of this work was to demonstrate that visual in-context learning using foundation models is feasible under the same stringent definitions and requirements employed in the NLP domain. While achieving this goal, we presented algorithms that incorporate pre-trained diffusion models. However, we anticipate that similar approaches can be extended to other foundation models.

I appreciate the effort of the authors for their detailed rebuttal. While the results presented are intriguing, particularly in demonstrating the generalization capabilities of the proposed training-free approach, I would like some clarification: how many example prompts were provided for each task? This is important for better understanding on my part and ensuring that the comparison is fair. Expecting further clarification from the authors.

We appreciate the reviewer's insights regarding prior works on the use of cross-attention, which we will discuss in the revised related work section. We agree with the reviewer on the concept of cross-attention being used in past literature for diverse tasks such as metric learning [10], few-shot segmentation [11,12], and multi-modal representation learning [6,7], where they formulate relationships between different inputs. As the reviewer has duly noted, these methods often require to be fine-tuned to effectively capture the correspondences.

While it is true that TriBERT utilizes three modalities in computing its cross-attention, we would like to further clarify our point regarding the use of distinct $Q$, $K$, and $V$ values in our cross-attention formulation.

TriBERT's cross-attention mechanism, as described in its paper and code, is formulated as:

$

\qquad \qquad f = \text{softmax} \left ( \frac{Q_1 (K_2 \oplus K_3)^T}{\sqrt{d}} \right ) (V_2 \oplus V_3 ) \tag{1}

$

where the subscripts of $Q$, $K$, and $V$ indicate the three modalities ($1,2,3$), and $\oplus$ denotes the concatenation operation in the feature dimension.

In contrast, our cross-attention formulation, as given in equation (7) of the main paper, is as follows:

$

\qquad \qquad f = \text{softmax} \left ( \frac{Q_C K_A ^T}{\tau \cdot \sqrt{d}} \right ) V_B \tag{2}

$

where $A$, $B$, and $C$ represent three distinct images: the prompt image, the prompt groundtruth, and the query image, respectively.

Expanding TriBERT's attention mechanism (equation (1)) reveals the following key properties:

• Attention weight computation: The concatenated keys $K_2 \oplus K_3$ are used to compute the attention scores:

$

   \qquad \qquad  [ \alpha_2 ; \\ \alpha_3 ]  = 
    \text{softmax} \left ( \frac{Q_1 (K_2 \oplus K_3)^T}{\sqrt{d}} \right ) 

$

Here, $\alpha_2$ and $\alpha_3$ represent the attention scores between the query $Q_1$ and the keys $K_2$ and $K_3$, respectively. These scores are normalized jointly via softmax across the feature dimension of all keys in $K_2 \oplus K_3$.

• Value aggregation: The computed attention weights are applied to the concatenated values of $V_2 \oplus V_3$:

$

\qquad \qquad f = \Sigma_i \alpha_{2,i} V_{2,i} + \Sigma_j \alpha_{3,j} V_{3,j} \tag{4}

$

where $\alpha_{2,i}$ and $\alpha_{3,j}$ are normalized attention weights for $K_2$ and $K_3$, respectively.

• Key-Value pairing: While the softmax enables some interaction between the keys $K_2$ and $K_3$ when normalizing the attention weights, the key-value pairing in aggregation remains modality-specific. There is no direct interaction between $K_2$ and $V_3$, or between $K_3$ and $V_2$.

This design implies that while TriBERT's attention spans multiple modalities, its direct key-value interactions (value aggregations - eq (4) ) are restricted to within the same modality. In contrast, our cross-attention mechanism explicitly links distinct inputs (i.e., $Q_C$ attend to $K_A$ and aggregates $V_B$) enabling direct interactions across three distinct images. This explicit interaction across different sources is a defining characteristic of our formulation and facilitates task-context modeling within the visual domain.

In summary, the main differences between our method and TriBERT are: (1) TriBERT’s key and value vectors directly interact only within the same modality, while our formulation explicitly uses key and value vectors from distinct images to model task and context; (2) our approach is entirely training-free, leveraging pre-trained priors, whereas TriBERT requires fine-tuning; and (3) our cross-attention formulation is specifically designed for visual in-context learning, enabling task and context inference in a novel, training-free paradigm.

We hope this clarification highlights the differences between our method and prior works like TriBERT. We will ensure these distinctions are explicitly discussed in the revised manuscript to avoid ambiguity.

[10] Kotovenko et al., "Cross-Image-Attention for Conditional Embeddings in Deep Metric Learning", [CVPR'23]

[11] Lin et al., "Few Shot Medical Image Segmentation with Cross Attention Transformer", [MICCAI'23]

[12] Hossain et al. "Visual Prompting for Generalized Few-shot Segmentation: A Multi-scale Approach", [CVPR'24]