Source-linked AI summary

PSALM: Pixelwise SegmentAtion with Large Multi-Modal Model

Zheng Zhang, Yeyao Ma, Enming Zhang, Xiang Bai

arXiv:2403.14598v1cs.CV

TL;DR

LMMs’ text-only outputs make pixel-level segmentation difficult, while segmentation tasks also require varied inputs, outputs, and capabilities. PSALM adds a mask decoder and flexible input schema to unify tasks, with reported gains from joint training and zero-shot performance on unseen tasks.

  • Problem

    LMMs produce text rather than masks, and segmentation tasks have diverse input and output requirements that are challenging to unify in one model.

    Source: §1 · pp. 1–2
  • Method

    PSALM adds a mask decoder to an LMM and uses images, task instructions, condition prompts, and learnable mask tokens in a flexible schema for segmentation.

    Source: §1 · p. 2
  • Results

    Joint training improves performance over training tasks individually, PSALM outperforms other LLM-based pixel-reasoning methods on referring benchmarks, and it shows promising zero-shot performance on three unseen segmentation tasks.

    Source: §1 · p. 3
  • Takeaways & Limitations

    The experiments indicate potential for addressing general image segmentation tasks and some task generalization across in-domain and out-of-domain settings.

    Source: §1 · p. 3, §5 · p. 16

Abstract

from arXiv · show

PSALM is a powerful extension of the Large Multi-modal Model (LMM) to address the segmentation task challenges. To overcome the limitation of the LMM being limited to textual output, PSALM incorporates a mask decoder and a well-designed input schema to handle a variety of segmentation tasks. This schema includes images, task instructions, conditional prompts, and mask tokens, which enable the model to generate and classify segmentation masks effectively. The flexible design of PSALM supports joint training across multiple datasets and tasks, leading to improved performance and task generalization. PSALM achieves superior results on several benchmarks, such as RefCOCO/RefCOCO+/RefCOCOg, COCO Panoptic Segmentation, and COCO-Interactive, and further exhibits zero-shot capabilities on unseen tasks, such as open-vocabulary segmentation, generalized referring expression segmentation and video object segmentation, making a significant step towards a GPT moment in computer vision. Through extensive experiments, PSALM demonstrates its potential to transform the domain of image segmentation, leveraging the robust visual understanding capabilities of LMMs as seen in natural language processing. Code and models are available at https://github.com/zamling/PSALM.

1 Introduction

PSALM extends LMMs to segmentation by adding a mask decoder and a flexible input schema for unifying diverse tasks. The paper reports improved performance from joint training and zero-shot results on unseen segmentation tasks.

1 Introduction

  • LMMs’ text-only output and discrete tokens make direct mask generation difficult, while segmentation tasks demand varied inputs, outputs, and capabilities that are challenging to unify in one model.The passage gives category inputs, object IDs and class-confidence scores, referring sentences, and interactive points, scribbles, boxes, or masks as examples.Source: §1 · pp. 1–2
  • PSALM extends LMMs from text-output tasks to general segmentation by placing a mask decoder on the LMM and using a flexible input schema to unify tasks in one model.The paper presents Fig. 1 as showing representative tasks supported by the model.Source: §1 · p. 2, Figure 1 · p. 2
  • The schema combines images, task-instruction prompts, condition prompts, and learnable mask tokens; the LMM output mask tokens feed a mask generator that produces proposals.Condition prompts supply information such as category names, sentences, or visual features; condition embeddings can also classify proposals.Source: §1 · pp. 2–3
  • Joint training on COCO Panoptic, RefCOCO/RefCOCO+/RefCOCOg, and COCO-Interactive significantly improves performance over training tasks individually; PSALM also outperforms LISA [21], PixelLM [43], and GSVA [53] on referring benchmarks.The authors use Phi-1.5 1.3B [26], while the compared methods use Vicuna-7B or LLama2-13B [48].Source: §1 · p. 3
  • PSALM shows promising zero-shot performance on open-vocabulary segmentation, generalized referring expression segmentation, and video object segmentation.These are unseen tasks addressed without additional training.Source: §1 · p. 3

2 Related works

Related work covers LMMs, LMM-based mask generation, and unified segmentation frameworks. PSALM is positioned as combining LMMs with broader segmentation-task support than prior LMM-centric methods.

2.1 Large Multimodal Models

  • LLaVA and BLIP-2 [23] map visual features into LLM input space, while Flamingo uses intermediate-layer feature fusion; these approaches target text or image generation rather than pixel-level segmentation.Table 1 compares method capabilities; its caption says PSALM handles more segmentation tasks than other LLM-centric methods, which can also perform text generation.Source: §2.1 · pp. 3–4, Table 1 · p. 4

2.2 Pixel Reasoning with LMM

  • LISA [21] uses a special seg token to aggregate sentence information for mask decoding, and later methods extend LMM-based segmentation to multiple objects and richer inputs.The cited examples include u-LLaVA [55], NExT-Chat [63], GLaMM [42], PerceptionGPT [37], PixelLM [43], GSVA [53], and LISA++ [59]; these methods primarily target referring segmentation, unlike PSALM’s generalized task design.Source: §2.2 · p. 4

2.3 Unified Segmentation Framework

  • Unified segmentation frameworks such as OneFormer, UNINEXT, X-Decoder, and SEEM support varied tasks, but the cited works are vision-centric rather than LMM-based.Mask2Former uses different models for generic segmentation tasks, whereas OneFormer integrates three tasks within one model.Source: §2.3 · p. 5

3 Methods

PSALM pairs an LMM with a mask generator and flexible inputs for general segmentation. Its design uses task instructions, condition prompts, and mask tokens to support different segmentation requirements.

3.1 Large Multimodal Model and Input Schema

  • PSALM updates learnable mask tokens through the LMM before decoding them, while condition prompts provide classifier weights that separate mask generation from classification.This design is described as more flexible and adaptable to more tasks, with more efficient decoupling than LISA and PixelLM.Source: §3.1 · pp. 6–7

3.2 Mask Generator

  • The mask generator combines multilevel visual features, mask tokens, and condition embeddings to predict segmentation masks and their category probabilities.Its design follows Mask2Former, using multi-scale deformable attention in a pixel decoder and a transformer-based mask decoder.Source: §3.2 · p. 7

3.3 Design of Condition Prompts

  • Category prompts pool category-token embeddings, sentence prompts use a [REF] anchor, and visual-prior prompts pool interaction regions before their embeddings classify mask proposals.Visual-prior inputs include points, boxes, masks, or scribbles.Source: §3.3 · pp. 7–8

3.4 Training Objectives

  • Training first aligns vision and language while freezing the vision encoder and LMM, then freezes only the vision encoder and matches proposals to ground-truth masks with mask and classification losses.The second stage uses bipartite matching; the mask loss combines pixel-level BCE and Dice losses.Source: §3.4 · pp. 8–9

4 Experiments

The experiments evaluate PSALM across in-domain and unseen segmentation tasks, using joint training and ablations to assess its design. Results show strong in-domain performance and zero-shot capability, with open-vocabulary performance still below the best task-specific method.

4.1 Datasets and Tasks

  • Training uses CC3M for vision-language alignment, then jointly trains on generic, referring, interactive, and vision-language instruction tasks; evaluation covers in-domain and held-out tasks.COCO-Interactive is an in-house dataset because no established public dataset covers all four interaction types.Source: §4.1 · p. 9

4.2 Implementation Details

  • The default system uses a Swin-B visual encoder, a Phi-1.5 1.3B LMM, 100 mask tokens, 56k training iterations, and no task-specific fine-tuning.Ablations use 9k iterations; experiments run on 16 V100 GPUs.Source: §4.2 · p. 9

4.3 Ablations

  • Using mask tokens as LMM inputs performs better than feeding them directly to the mask generator, while decoupling mask prediction and classification improves semantic segmentation.Multiple mask proposals also improve RefCOCO over using a single mask.Source: §4.3 · p. 10, Table 2 · p. 10, Table 3 · p. 10
  • Condition-prompt designs work best when matched to the task: avg_pooling performs better on COCO-Pan and A150-OV, while [REF] works better on RefCOCO.Combining the designs preserves their advantages and achieves the best overall performance.Source: §4.3 · p. 11, Table 5 · p. 11
  • Without vision-language alignment, all four tasks perform worse, including drops of -1.1 mAP on A150-OV and -4.4 cIoU on RefCOCO; joint training greatly improves performance.The authors describe generic segmentation as refining referring segmentation masks and referring expressions as helping recognize unseen categories for open-vocabulary segmentation.Source: §4.3 · pp. 11–12, Table 6 · p. 11, Table 7 · p. 11

4.4 System-Level Comparison on In-Domain Tasks

  • PSALM achieves state-of-the-art performance on RefCOCO and RefCOCO+, competitive performance on RefCOCOg and COCO panoptic segmentation, and leading results on interactive point, scribble, and mask inputs.On box inputs, SAM performs better on mIoU but worse on cIoU.Source: Table 8 · p. 12, §4.4 · pp. 12–13, Table 9 · p. 13

4.5 Generalizability on Out-of-Domain Tasks

  • Without special design, PSALM achieves reasonably good open-vocabulary segmentation performance but remains below the best task-specific method, SAN.The authors report significant improvement after adding LVIS training data.Source: Table 11 · p. 14, §4.5 · p. 14
  • PSALM shows promising zero-shot performance on gRefCOCO and DAVIS-2017, addressing multi-target referring expressions and video segmentation without training on video data.For gRefCOCO, it outperforms a LISA version pretrained on gRefCOCO but without task-specific fine-tuning.Source: §4.5 · p. 15, Table 12 · p. 15, Table 13 · p. 15

5 Conclusion

PSALM extends LMMs from text output to image segmentation and unifies multiple segmentation tasks with a novel input model. It performs well on in-domain tasks and shows out-of-domain generalization.

5 Conclusion

  • PSALM extends LMM capabilities from text output to image segmentation and unifies varied segmentation tasks with a novel input model.The conclusion identifies this extension as addressing LMM output limitations.Source: §5 · p. 16
  • PSALM exhibits excellent performance across multiple in-domain segmentation tasks.Source: §5 · p. 16
  • PSALM also demonstrates generalization ability on out-of-domain tasks, underscoring its potential.Source: §5 · p. 16

6 Prompts for Different Tasks

PSALM varies task instructions and condition prompts by task, with tasks using the same condition type generally sharing an instruction prompt. Some category-conditioned tasks also include a background label.

6 Prompts for Different Tasks

  • Tasks use different instruction and condition prompts, while tasks sharing a condition type generally use the same instruction prompt.Source: §6 · p. 21
  • The schema uses task instruction prompts, condition prompts, and learnable mask tokens, with prompt details listed across tasks.Source: §6 · p. 21, Table 14 · p. 21
  • For category-conditioned tasks that need background handling, such as instance segmentation, the joint sentence appends “background” at the end.Source: §6 · p. 21

7 Training Details

PSALM is trained in two stages: vision-language alignment followed by joint task training. The final stage samples among four tasks and uses Hungarian matching to assign mask-proposal ground truth.

7 Training Details

  • Training has two stages, beginning with vision-language alignment that follows LLaVA settings except for using Phi-1.5 and Swin.The first-stage hyperparameters are listed in Table 15.Source: §7 · p. 21, Table 15 · p. 22
  • The final model trains for 56k iterations, sampling equally among generic, referring, and interactive segmentation and visual-language instruction tasks.Images are resized to 1024^2 by padding the shorter side while preserving aspect ratio; the visual encoder is frozen and other parts remain trainable.Source: §7 · p. 21
  • Hungarian matching assigns mask-proposal ground truth, using classification and mask losses for generic segmentation and mask loss alone for other tasks.Second-stage hyperparameters are listed in Table 16.Source: §7 · p. 22, Table 16 · p. 22
  • Ablation experiments use the final model’s training setting but shorten training to 9k iterations to reduce experiment cost.Source: §7 · p. 22

8 Details on building COCO-Interactivate Dataset

COCO-Interactive is built from COCO2017 instance masks and boxes, which are used to generate point, scribble, box, and mask visual prompts.

8 Details on building COCO-Interactivate Dataset

  • COCO-Interactive uses COCO2017 instance masks and bounding boxes to generate four visual-prompt types: points, scribbles, boxes, and masks.Source: §8 · p. 23
  • A point prompt is randomly sampled within a circular region centered on the instance box, with radius half the box’s short side.Source: §8 · p. 23
  • Scribbles use a jittered box with IoU above 0.5, then a randomly parameterized sine curve along one diagonal.The box scale factor ranges from 0.5 to 1.2; curve amplitude ranges from 10 to 20.Source: §8 · p. 23
  • Box prompts jitter ground-truth boxes with scale factors from 0.9 to 1.1; mask prompts blur ground-truth masks with a Gaussian filter before binarization.The Gaussian kernel standard deviation is 2.Source: §8 · p. 23

9 Datasets

The section describes datasets spanning panoptic, referring, instance, open-vocabulary, semantic, video, and cross-view segmentation. It also includes a visualization of visual prompt types.

COCO-Panoptic

  • COCO-Panoptic covers over 200,000 images, with annotations for 80 instance categories and additional semantic categories.Source: §9 · p. 23

RefCOCO

  • RefCOCO, RefCOCO+, and RefCOCOg use COCO images annotated with expressions that uniquely identify objects, with differing splits and annotation styles.Source: §9 · p. 23

LVIS and ADE20K

  • LVIS provides instance annotations for over 1,000 categories and is known for its long-tail category distribution; ADE20K provides dense pixel labels for 150 categories across things and stuff.Source: §9 · pp. 23–24

Visual prompts

Other segmentation datasets

  • Cityscapes focuses on urban scenes, while Pascal VOC and Pascal Context provide semantic labels; DAVIS-2017 evaluates video segmentation with per-pixel annotations of multiple objects.Source: §9 · p. 24

gRefCOCO, Ego-Exo4D, and CC3M

  • gRefCOCO includes multi-target, no-target, and single-target expressions; Ego-Exo4D supports cross-view object correspondence, and CC3M contains around three million image-caption pairs.Source: §9 · p. 25

10 Implement Details of Decouple Ablation

The non-decoupled design uses category condition embeddings to generate and classify masks, removing the need for a matching loss. Table 17 is captioned as reporting performance on multimodal benchmarks.

Non-decoupled design

  • The non-decoupled design omits the mask token and uses category condition embeddings as mask queries.Source: §10 · p. 25
  • In this design, the category condition embedding predicts the class and mask at the same time, so matching loss is unnecessary.Source: §10 · p. 25

Table scope

  • Table 17 is captioned as reporting performance on multi-modal benchmarks.Source: Table 17 · p. 25

11 Additional Experiment

Additional experiments evaluate PSALM on multimodal benchmarks and the Ego-Exo4D correspondence benchmark. The latter is evaluated zero-shot on the validation set because the official test set and model were unavailable.

Ego-Exo4D Correspondence Benchmark

Multi-modal benchmark evaluation

Zero-shot evaluation

  • In zero-shot VQA without training data, LISA achieves a score of 0.12; the paper also describes LISA and PixelLM as yielding suboptimal results on these tasks.Source: Multi-modal Benchmark Evaluation · p. 26

Ego-Exo4D Correspondence Benchmark

12 More Qualitative Results

The section provides additional qualitative examples for in-domain and out-of-domain segmentation tasks, as well as Ego-Exo correspondence.

In-domain tasks

Out-of-domain tasks

Ego-Exo correspondence

  • Figure 11 provides further Ego-Exo correspondence examples; the accompanying labels identify ground truth for frame 0 and predictions.Source: §12 · p. 32, Figure 11 · p. 34
Loading 2403.14598v1…