PixelLM: Pixel Reasoning with Large Multimodal Model

In plain words

A chat model that sees pictures can describe them, but it cannot outline the objects it talks about. Earlier attempts bolt on a separate, heavy segmentation model to draw the outlines. This paper keeps the chat model and adds a small decoder that turns special hidden "marker" tokens into outlines. The markers are learned, and the model writes one of them for each object it names in its answer, so a question about several things yields several outlines. The authors also built a dataset of questions that name many objects at once, since existing datasets mostly name one. On their own benchmarks the method beats the closest prior system, and the results are the authors' own.

TL;DR

PixelLM is a large multimodal model (LLaVA backbone, LoRA-tuned) that answers in text and writes a learned codebook token group for each target. A light pixel decoder reads those tokens' hidden states together with CLIP features from two layers and produces one mask per target §3.1. It is trained with a text loss, a mask BCE term and a DICE term §3.2. The authors report higher gIoU and cIoU than LISA on their MUSE benchmark and on multi-referring and refCOCO splits, at about half the TFLOPs tab. 1. The official code (Apache-2.0) does not train the target refinement loss the paper describes; see Paper vs code.

Problem

LMMs produce text for whole images or regions, but pixel-level answers such as object masks are weak §1. The closest prior method, LISA, handles reasoning segmentation with one target and depends on a pre-trained SAM model for masks, which costs compute and caps performance at SAM's level §1. Existing public datasets (RefCOCO, ReasonSeg) lack multi-target questions whose object descriptions link to masks §4.1. The task is therefore: segment an open set of targets named inside a free-form answer, with no separate segmentation network.

Method

Framework §3.1 fig. 1. Four parts: a pre-trained CLIP-ViT vision encoder; a large language model; a lightweight pixel decoder; and a segmentation codebook of learned tokens §3.1. The LLM receives the image and the query, and writes an answer that contains the codebook tokens. For each target, the hidden states of its tokens go into the decoder with the image features, and the decoder produces the mask §3.1.

Segmentation codebook §3.1 fig. 2. The codebook is a set of L groups of N tokens, one group per visual scale. The main model uses two scales, the "+2 scale" row of the ablation, which matches its headline numbers tab. 4. The ablation uses N = 3 tokens per group §5.5. The paper's Fig. 2 caption describes two scales with two tokens each fig. 2. The text explains that a single token per scale cannot cover multi-target answers, so several tokens are fused into one embedding per scale with a linear layer φ §3.1.

Vision features §3.1 app. A tab. A.3. The image passes through the frozen encoder at 448 px input (CLIP-ViT-L/14-336 with interpolated position embeddings) §5.1. The final-layer output goes to the LLM through a projection p_{V→T}. Every selected layer goes through a projection p_{V→D} into the decoder. Tab. A.3 picks layers 23 and 14 (1-based; the 23rd is the penultimate layer) tab. A.3.

Pixel decoder §3.1 eq. 1 alg. 1. The decoder has one attention block per scale. It runs from the deepest scale down. Each block's mask m^ℓ modulates the feature map of the next scale through a sigmoid gate (σ(m)+1), so the next scale attends to regions the previous scale found relevant eq. 1. The per-scale masks are then summed with learned weights γ, constrained to |γ| = 1 §3.1. Alg. 1 in the appendix gives the same steps as PyTorch-style pseudocode alg. 1 app. A.

Target refinement loss §3.2. Pixels predicted for two or more targets get weight α instead of 1 in a per-pixel binary cross-entropy §3.2.

Overall objective §3.2. A weighted sum of the text cross-entropy, the refinement loss and a DICE loss; the DICE formula is not given in the text.

# Forward pass for one query  [§3.1] [eq. 1] [alg. 1]
# Input:  image (3, 448, 448); query text
# Output: one mask per target named in the answer

def pixellm_forward(image, query):
    # 1. Vision features at the chosen layers  [§3.1] [tab. A.3]
    scale_feats = [vision_layer(image, l) for l in layers]  # L x (C, h, w)
    # 2. Last-layer features into the LLM; all scales into the decoder  [§3.1]
    llm_tokens = project_to_text(scale_feats[-1])  # (h*w, d)
    decoder_feats = [project_to_decoder(f) for f in scale_feats]  # (256, h, w)
    # 3. The LLM writes an answer with L*N codebook tokens per target  [§3.1]
    answer, hidden = llm.generate(llm_tokens, query)  # hidden: (T, d)
    # 4. Take each target's N tokens per scale and fuse them  [§3.1]
    #    The paper uses a linear layer phi; see Paper vs code
    for target in targets(answer):
        group = hidden[target]  # (L, N, d)
        seg_emb = fuse_tokens(group)  # (L, d)
        # 5. Decode scale by scale, starting from the deepest  [eq. 1]
        mask = None
        for l in range(L):
            if mask is not None:
                # gate the feature with the previous scale's mask  [eq. 1]
                decoder_feats[l] = decoder_feats[l] * (sigmoid(mask) + 1)
            mask = decode_block(seg_emb[l], decoder_feats[l])  # m^l
        # 6. Combine the scale masks with learned weights  [§3.1]
        final_mask = sum(gamma[l] * m[l] for l in range(L))
        yield final_mask

Key equations, as written in the paper:

f^{\ell\prime} = f^{L} \text{ if } \ell = L, \quad f^{\ell} \odot (\sigma(m^{\ell+1}) + 1) \text{ if } \ell < L eq. 1

m^{\ell} = \mathrm{Attn}^{\ell}(h^{\ell}, f^{\ell\prime}), \qquad \hat{M} = \sum_{\ell=1}^{L} \gamma^{\ell} m^{\ell}, \quad |\gamma| = 1 §3.1

y_{res} = \mathcal{F}(p_{V\rightarrow T}(I^{L}), x_{txt}, C_{seg}), \qquad h^{\ell} = \phi(h^{\ell}_{1}, \dots, h^{\ell}_{N}) §3.1

\mathcal{L} = \mathcal{L}_{txt} + \lambda_{ref}\mathcal{L}_{ref} + \lambda_{dice}\mathcal{L}_{dice} §3.2

\mathcal{L}_{ref} = \frac{1}{KHW} \sum_k \sum_i A_i \, \mathcal{L}_{BCE}(\hat{M}_{k_i}, M_{k_i}), \qquad A_i = \begin{cases} \alpha & \sum_k \hat{M}_{k_i} \ge 2 \\ 1 & \text{otherwise} \end{cases} §3.2

The paper says the codebook tokens are "included in the loss calculation" of L_txt §3.2.

# Target refinement loss, as the paper writes it  [§3.2]
# Input:  pred_logits (K, H, W); gt_masks (K, H, W); alpha
# Output: scalar loss

def refinement_loss(pred_logits, gt_masks, alpha):
    # Count predicted masks per pixel; the paper treats M_hat as binary  [§3.2]
    overlap_count = (pred_logits > 0).sum(dim=0)  # (H, W)
    # Weight 2 (alpha) where two or more targets overlap  [§3.2]
    weight = where(overlap_count >= 2, alpha, 1.0)  # (H, W)
    # Per-pixel binary cross-entropy for each target  [§3.2]
    bce = binary_cross_entropy(pred_logits, gt_masks)  # (K, H, W)
    return mean(weight * bce)

Hyperparameters

Values come from the paper's Tab. A.1 unless marked "code"; "not stated" means neither the paper's text nor its tables give it.

MUSE dataset §4.1 app. C.2. 246k question-answer pairs over about 910k instances from LVIS (over 1000 categories); average 3.7 targets per answer. Splits: train 239k, val 2.8k, test 4.3k §4.1.

MUSE evaluation §4.3 app. A. Predicted masks are matched to ground truth with the Hungarian algorithm, using cost BCE + DICE (App. A). Unmatched masks are paired with empty masks. The ground-truth text description replaces each mask's segment in the answer. A GPT-3.5 prompt scores each prediction (see Pitfalls for the scale). Per-prediction IoU is kept only where the score exceeds 0.5 (App. A), and the image score is the sum of per-prediction IoUs divided by P, the padded set size app. A.

Results

Author-reported. Setup: MUSE (val and test) §4.1; multi-referring and refCOCO-family benchmarks §5.2. Hardware: 8 A100 GPUs for training §5.1. TFLOPs are computed at 512 generated tokens, not measured app. A.

Multi-target reasoning segmentation, MUSE tab. 1 (gIoU / cIoU; "few" and "many" are test subsets split at three targets §4.1; w/o SAM = no separate SAM model; † = PixelLM without token fusion and refinement loss):

Source: tab. 1. The "Test few"/"Test many" column pairs are gIoU / cIoU.

Multi-referring segmentation tab. 2 (metric not stated in the caption; the paper's text compares against LISA and LISA_aug only). Values are per split:

Referring segmentation tab. 3 (the commented-out earlier caption calls this cIoU; the active caption does not state a metric):

Source: tab. 3.

Ablations, MUSE val and refCOCOg (7B model) tab. 4 (gIoU / cIoU for MUSE; val(U) / test(U) for refCOCOg):

Data generator (MUSE val only): GPT-4 with 30k pairs: 30.0 gIoU / 35.7 cIoU; GPT-4V with 30k: 35.6 / 40.0; GPT-4V with 100k: 38.9 / 45.7; GPT-4V with 200k: 42.6 / 50.7 tab. 4. Source: tab. 4 and its subtables.

Efficiency tab. 1. TFLOPs 3.57 vs 7.16 for 7B and 6.65 vs 10.24 for 13B, which the paper describes as a 50% and 35% reduction §5.3.

Paper vs code

Official code: MaverickRen/PixelLM (Apache-2.0), linked from the paper's project page, pinned at commit ebf07922745a926517a9a1d16fddd392a7544bb0. Code paths below are relative to the repository root.

• Target refinement loss is not trained. The paper's §3.2 loss and α = 2.0 tab. A.1 are in the code as overlap_loss PixelLM.py, but its call is commented out PixelLM.py and the sum is only BCE + DICE PixelLM.py. The code's bce_loss_weight, default 2.0 train_ds.py, scales plain BCE. Paper vs code on this point is a discrepancy, not an omission in the card.
• Gradient clipping. Appendix A says the authors "do not use ... gradient clip" app. A. The DeepSpeed config sets gradient_clipping to 1.0 train_ds.py.
• Token fusion. The paper fuses N tokens with a linear layer φ §3.1. The code computes a weighted sum with one scalar per token, initialised to 1/N PixelLM.py PixelLM.py. These scalars are stored in a plain Python list, which PyTorch does not register as module parameters, so they may never train. Check a checkpoint before relying on either reading.
• Scale combination. The paper learns γ with |γ| = 1 §3.1. The code combines scale masks with a list of scalars initialised to 1/L PixelLM.py PixelLM.py; the same list caveat applies.
• Decoder gate. Alg. 1 downsamples the mask to feature resolution and multiplies the feature by σ(mask)+1 alg. 1. The code upsamples the image feature by 2× and gates it with the mean over channels of the previous mask mask_decoder_multi_scale.py.
• Mask output. Alg. 1 averages the mask over the output tokens alg. 1. The code takes the first of four SAM-style mask tokens when multimask_output is False mask_decoder_multi_scale.py.
• Algorithm indexing. As printed, Alg. 1 applies the feature update when l < L-1 at the start of each iteration, using mask from the previous iteration; for l = 0 that mask does not exist. The code gates only at level > 0 mask_decoder_multi_scale.py.
• Which layer goes to the LLM. §3.1 says the LLM gets the final-layer output §3.1; Tab. A.3 calls the baseline the penultimate layer tab. A.3. In the code the first scale is hidden_states[select_layer], but select_layer is not set in train_ds.py; the LLaVA default is -1 train.py. The second scale is hard-coded to hidden_states[-11] clip_encoder.py, which is the 14th layer (1-based), matching Tab. A.3. Confirm select_layer in the released checkpoint config.
• Tokens per group. Fig. 2's caption says two tokens per scale fig. 2; §5.5 and the README use three §5.5 README.md.
• Batch size. Paper: 16 tab. A.1. Code: 2 per device with 10 accumulation steps train_ds.py train_ds.py; GPU count not stated in the paper.
• Training data. The code's defaults list semantic-segmentation sets (ADE20K, COCO-Stuff, Pascal-Part, PACO-LVIS, Mapillary), RefCOCO-family sets, LLaVA-150k and ReasonSeg train_ds.py. The README run instead uses multi_reason_seg, which is MUSE, with sampling rates 2, 9, 2, 4 README.md. The paper's §5.1 names the sources but not the rates.
• Codebook as vocabulary. The code adds [SEG0]..[SEG5] to the tokenizer train_ds.py. The main text does not say this; only commented-out source text does.

No discrepancies found in the frozen-encoder setup, the input size (448 px with CLIP-ViT-L/14-336), the AdamW settings (learning rate, betas, weight decay), warmup of 100 steps, or the DICE weight of 0.5.

Pitfalls

• Refinement loss is ambiguous and not in the code. The paper's loss needs α and a binary M̂ §3.2. The code does not train it, so a faithful implementation needs an explicit choice of threshold and of how α applies.
• Two MUSE metric definitions. §4.3 says a GPT-3.5 score from 1 to 10 is multiplied with the IoU §4.3. App. A keeps the intersection only if the score is above 0.5 app. A. The text does not say how 1–10 scores are normalised to that threshold.
• The GPT scoring prompt is only in a figure (in App. A). A reimplementation needs the same prompt and model, which the paper does not fix in text.
• Matching uses BCE + DICE per mask pair app. A, which is not the same as the refinement loss.
• Text claims that tables do not support. §5.5 says token fusion gives "up to 2.4%" cIoU on MUSE val §5.5; Tab. 4 gives 48.0 → 50.1, which is 2.1 tab. 4. §5.5 says GPT-4V gains are "up to 7.0% gIoU and 8.6% cIoU" at equal data §5.5; the table gives 30.0 → 35.6 (+5.6) and 35.7 → 40.0 (+4.3) at 30k tab. 4. The MUSE gains on ReasonSeg (52.9 → 58.0 gIoU) appear only in the text §5.5.
• Dangling figure reference. The introduction and §4.1 refer to a figure label that does not exist in the source §1 §4.1, so the compiled paper shows a broken reference.
• Referring results are not uniformly best. LISA is higher on refCOCO val and testB tab. 3, though the text says PixelLM is best on most metrics §5.4.
• The chosen layer pair is not the best one in the table. Tab. A.3 gives MUSE val cIoU of 47.2 for the baseline, 50.7 for layers 23 and 14 (the default), 51.4 for 23, 14, 20 and 51.5 for 20, 17, 14 tab. A.3. The paper gives no reason for the default.
• Hidden costs. Each scale has its own deep-copied transformer mask_decoder_multi_scale.py, so parameters grow with L. The paper does not report the decoder size.
• The dataset needs GPT-4V. MUSE was generated with GPT-4V and filtered with GPT-4V and human checks §4.2 app. C.1. Rebuilding it needs API access; the generation prompt is shown only as an example figure fig. 3.
• TFLOPs are computed, not measured. The 512-token length is a convention app. A, so the 50% and 35% savings describe that convention only.

Lineage

• Builds on: LISA, arXiv:2308.00692 (reasoning segmentation with an LLM and SAM; PixelLM drops the SAM dependency) §1. LLaVA supplies the pre-trained multimodal model §5.1.
• Compared against: LISA (7B and 13B), LISA_rec (a two-step variant), LISA_aug (LISA trained on MUSE), SEEM, and referring-segmentation baselines MCN, VLT, CRIS, LAVT, ReLA and X-Decoder tab. 1 tab. 3.
• Superseded by: not known from content on this wiki.