Empowering Segmentation Ability to Multi-modal Large Language Models

Reasoning Segmentation is first proposed in  [17], which aims to parse the question-oriented or query-focused visual region in an image and output a segmentation mask to locate it. Although this task shares similarities with referring segmentation task [12], it requires more knowledge about the real world and common sense to reason about the target region from intricate expression.

The Pipeline of MLLMs The MLLMs take multi-modal information as input that includes an image $\mathbf{I}\in\mathcal{R}^{H\times W\times 3}$ and text $\mathcal{T}$ of random length. The image $\mathbf{I}$ is encoded with the multi-modal encoder $v$ to share the same embedding space with the text embedding [4]. The text can be further divided into two parts. The first part is the user query $\mathcal{T}_{q}$, which is related to the requirements of the user. The second part is the pre-designed task prompt $\mathcal{T}_{p}$, which is hand-crafted or generated automatically. Although the prompt is not given by the user, a proper prompt can instruct the MLLMs to generate a response with higher quality [16]. Formally, the MLLMs $f_{l}$ generate the answer $\mathcal{T}_{a}$ by

where $\{\cdot\}$ indicates the concatenation of text input. The MLLMs also support multi-turn conversation by concatenating the former input and response with another query and prompt, which can be formulated as:

where $\mathcal{T}_{p_{i}}$ and $\mathcal{T}_{q_{i}}$ indicate the task prompt and user query in the $i$-th turn.

Reason Prompting Step: The first step is to prompt MLLMs to reason about the target region in the image from the user query. The reasoning step can be formulated as a VQA task, so we rephrase the user query into a question for a specific object or area in this image. Concretely, we design the following prompt for this step:

Prompt: What is the object or part that is [USER QUERY] in this image?

Output: It is fire in the fire pit. The fire is hot and …

[USER QUERY] is filled with the user query. When the user query is a question, we directly use it as the input and do not add other prompts. Notice that we do not instruct the MLLMs to output the rationale explicitly. We find that the MLLMs tend to answer the question followed by the explanation, although no explicit prompt is given. Additionally, forcing the model to output the rationale while not providing any knowledge for it may result in severe hallucination [43]. Given these concerns above, we leave the reason prompt as a question and do not prompt the MLLMs to give a detailed explanation explicitly.

Target Prompting Step: The MLLMs output both the name of the target region and the explanation including many unrelated descriptions (e.g. hot, light…) in the first step. These descriptions may disturb the MLLMs to generate exact visual attributions in the attribute prompting step, and it is not trivial to find the exact target from the response. In the second step, we input the user query, response, and image to the MLLMs to provide sufficient information so that the MLLMs can analyze the target. Specifically, the target prompt is formulated as follows:

Prompt: Here is the conversation:

The question is: [QUESTION]

The answer is: [ANSWER]

Please analyze the conversation and identify the distinct physical objects or areas that the user wants from the image.

Output: The user wants the fire from the image.

[QUESTION] and [ANSWER] is the question and answer in the first step, respectively.

Attribute Prompting Step: In this step, we prompt the MLLMs to provide visual information towards the target. Although it seems straightforward to prompt the MLLMs to output the location of the target object in the image, previous work [3] shows that the MLLMs cannot output accurate spatial positions in zero-shot settings. To address this issue, we choose to extract the visual attributes rather than the specific location to aid the downstream segmentation model. The visual attributes include color, shape, and the relative position to the other objects in the image. We prompt the MLLMs to extract the visual attributes as follows:

Prompt: Here is the target:

The target is: [TARGET]

Briefly describe the target entity or part’s visual attributes that can discriminate them from the image. Each visual attribute can be color, shape, and relative position to other objects in the image.

Output: The fire can be discriminated from the image by its bright orange color and the fact that it is emitting heat and light. The fire is surrounded by …

[TARGET] is the target analyzed in the second step. For simplicity, we present the most critical parts of the prompts. Note that we prompt the MLLMs to generate the discriminative visual attributes which are more helpful in segmenting the target region from the image.

Segmentation Pipeline: We input the target name with the visual attributes $\mathbf{E}$ and the image $\mathbf{I}$ into the segmentation model $f_{s}$ to generate the final mask $\mathbf{M}$, which can be formulated as:

In our method, we directly use MLLMs’ token embeddings corresponding to the target name and visual attributes and don’t utilize extra text encoders like BERT [13]. We add an MLP to project these textual embeddings to the visual feature space. Given the projected textual embeddings $\mathbf{E}$, to perform multi-scale interaction between $\mathbf{E}$ and the visual features in the segmentation model, we introduce several language-aware adapters in different layers. Motivated by [42], we design our language-aware module based on the attention mechanism. Specifically, we leverage the cross-attention between the textual embedding and the visual features from the segmentation model to inject the visual attributes into the segmentation model.

Training Objectives: During training, we freeze all the parameters of the segmentation model except for the adapter and decoder. Different from LISA [17], our model is only trained with the segmentation loss $\mathcal{L}_{mask}$. Specifically, our loss $\mathcal{L}_{mask}$ is a combination of per-pixel binary cross-entropy (BCE) loss and DICE loss, which can be formulated as

where $\hat{\mathbf{M}}$ is the ground-truth and $\mathbf{M}$ is the prediction. $\lambda_{bce}$ and $\lambda_{dice}$ denote the loss weights.

Datasets: Following LISA [17], our method adopts a composed dataset containing multiple datasets from different tasks. These datasets are collected from three categories as mentioned in Sec. 3.4. The first category is semantic segmentation datasets, including ADE20k [45], COCO-Stuff [2], and part semantic segmentation dataset PACO-LVIS [33], and PASCAL-Part [5]. The second is referring segmentation datasets, including RefCOCO [12], RefCOCO+ [12], RefCOCOg [27], Refclef [12]. The datasets mentioned above provide a large number of visual samples with different granularity to promote our LLaVASeg with strong segmentation ability and generalization abilities. The third category is the reasoning segmentation dataset ReasonSeg [17]. We evaluate the performance of the conversational ability based on the reasoning segmentation dataset by evaluating the MLLMs’ response in the reason prompting step. Furthermore, we exclude the COCO samples in the refCOCO(+/g) validation set to avoid data leakage.

Implementation Details: For all the experiments, we use LLaVA [23] as the architecture of MLLMs if not specified. The backbone of LLaVA is set as Llama2 [36]. For the segmentation model, we choose the backbone of ViT-H SAM [15] as the vision backbone. For model training, AdamW [25] is used as our optimizer. The overall learning rate is set as 0.0001, and the weight decay is set to 0.0001. $\lambda_{bce}$ and $\lambda_{dice}$ are set to $1$ and $0.5$, respectively. The overall batch size is set to 160, where each sample contains a pair of text and a referred mask. We train our model for 12000 iterations in total.

Evaluation Metrics: Similar to LISA [17], we adopt gIoU (the average of all per-image intersection-over-unions) and cIoU (the cumulative intersection over the cumulative union) as evaluation metrics on reasoning segmentation. We evaluate our framework’s dialogue quality based on the CIDEr metric [37] and ROUGE-L metric [20]. ROUGE-L is a metric commonly used in natural language processing for evaluating the quality of summaries by calculating the length of the longest common subsequences between the generated summary and the reference summary. CIDEr [37] is a metric specifically designed for evaluating the quality of image captions. It takes into account consensus and diversity among reference captions, providing a more comprehensive evaluation of the generated image captions. We adopt these two evaluation metrics for a more complete and comprehensive evaluation of the generated response from our framework and LISA.

Segmentation: We evaluate the segmentation performance of our LLaVASeg and compare it with the previous methods on the ReasonSeg dataset. For a fair comparison with LISA, we utilize the ‘LISA-13B-explain‘ model which can output both the explanation and the segmentation mask. The result is shown in Tab. 1. It can be seen that our LLaVASeg achieves superior performance on segmentation, which is 1.8% higher than LISA in terms of the gIoU metric. Furthermore, we should note that our LLaVASeg does not fine-tune LLaVA on the Reasonseg dataset, and inherently makes mistakes in reasoning and results in lower cIoU than LISA. We present a failure case study as shown in Fig. 3. In this case, when the LLaVA model fails to reason about the correct answer, both LLaVASeg and LISA predict the wrong area. However, due to LLaVASeg’s superior segmentation ability, it predicts the area referred by the LLaVA response more completely than LISA. Consequently, such samples result in higher cumulative union and lower cIoU. Besides, we also observe that gIoU is much more stable than cIoU since cIoU is largely biased towards large-area target [17]. This experiment shows that our model can achieve competitive segmentation performance without fine-tuning MLLMs.

Explanation Generation: To investigate the impact of fine-tuning MLLMs for segmentation tasks on their original reasoning ability and dialogue ability, we evaluate the explanation performance. Specifically, we prompt both the LISA and the off-the-shelf LLaVA used in our LLaVASeg, to produce the answer and the explanation to the user query. The experiment’s results, presented in Tab. 1, reveal a notable decline in the dialogue performance when comparing LISA with LLaVA. These findings support our motivation and prove the significance of our work that empowers the MLLMs with segmentation ability while preserving their reasoning ability. We also present visualization results in Fig. 4 for an intuitive understanding of our method.

Vanilla Referring Segmentation: We evaluate our method on the referring segmentation datasets to show that the proposed LLaVASeg framework can also handle the vanilla referring segmentation task. As shown in Tab. 2, our method achieves competitive performance compared with previous works, which outperforms the previous works on most of the dataset splits. This demonstrates the superior segmentation ability of our LLaVASeg.

Different Prompts: To show the effectiveness of the proposed chain-of-thought prompting, we ablate on the performance of our LLaVASeg when using different prompts. Specifically, we use only one or two steps of our chain-of-thought prompting in this ablation for comparison. The results are shown in Tab. LABEL:table:different_prompts. It can be seen from the results that when only the first step is leveraged, the performance drops significantly. When the targeting prompting step is added, the performance increases. This indicates that the unrelated rationale and explanation hinder the segmentation severely. When we further add the attribute prompting step, the performance also increases by a large margin, which demonstrates the effectiveness of the visual attributes. This underscores the effectiveness of our chain-of-thought prompting in maximizing the reasoning segmentation accuracy.

Effectiveness of Multi-Scale Prompting: We ablate on the impact of the multi-scale features on the reasoning segmentation performance. The results are shown in Tab. LABEL:table:multi_scale. It can be seen that the setting using the multi-scale features achieves superior performance than using single-scale features. the performance of both gIoU and cIoU increases accordingly. These results demonstrate that the extracted vision attributes including both low-level and high-level information can benefit from the multi-scale features.

Replacing Chain-of-Thought with LISA: In this experiment, we show the effectiveness of the explicitly generated visual attributes. We employed LISA, which extracts the visual attributes implicitly in the <SEG> token, to undertake the second and third steps in the chain-of-thought prompting as the comparison. Specifically, we give the MLLMs response from the first step to the LISA as the instruction for segmentation. The results are shown in Tab. LABEL:table:agent_structure. We can see that LISA achieves a gIoU of 54.3% and a cIoU of 50.9%, and LLaVASeg outperforms it with a gIoU of 59.1% and a cIoU of 52.8%. We argue that the significant improvements mainly come from the explicit visual attributes extraction in our chain-of-thought prompting when comparing with LISA. With the same reasoning response as input, our LLaVASeg leverages an explicitly generated visual attribute to aid the segmentation model, while LISA does this implicitly in the <SEG> token.

Ablation on Generalization Ability: The experiments are evaluated on the ReasonSeg validation set. To test the generalization ability of LLaVASeg, we train LLaVASeg on the dataset excluding the ReasonSeg training set. The results in Tab. LABEL:table:vision_ft show that there is a performance drop when the ReasonSeg training set is excluded. Although the performance degrades, our LLaVASeg without training on the ReasonSeg dataset still achieves competitive performance, which further supports its generalization ability.