PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse-LaViDa: Sparse Multimodal Discrete Diffusion Language Models

Shufan Li, Jiuxiang Gu, Kangning Liu, Zhe Lin, Zijun Wei, Aditya Grover, Jason Kuen

arXiv:2512.14008v2cs.CV

TL;DR

MDMs repeatedly process redundant masked tokens during sampling, limiting inference efficiency. Sparse-LaViDa truncates unnecessary tokens, uses register tokens as compressed representations, and applies a specialized attention mask. It reports speedups across generation, editing, and visual reasoning while maintaining quality comparable to LaViDa-O.

  • Problem

    MDMs repeatedly process full sequences, including many redundant masked tokens, at every sampling step.

    Source: §1 · p. 1, §2.1 · p. 3

  • Method

    Sparse-LaViDa truncates unnecessary masked tokens, represents truncated tokens with registers, and uses a step-causal attention mask for training and inference consistency.

    Source: §1 · p. 2, §3.1 · p. 4, §3.3 · p. 5

  • Results

    Sparse-LaViDa reports speedups across image generation, editing, and visual math reasoning, including 2.80× speedup on image editing, while maintaining quality comparable to LaViDa-O.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Sparse-LaViDa offers significant speedups for long-sequence visual generation and reasoning without compromising generation quality.

    Source: §5 · p. 8, §9 · p. 17

    Speedups are limited to long-sequence generation; short QA and object grounding do not benefit, and Sparse-LaViDa requires additional training.

    Source: §5 · p. 8, §9 · p. 17

Abstract

from arXiv · show

Masked Discrete Diffusion Models (MDMs) have achieved strong performance across a wide range of multimodal tasks, including image understanding, generation, and editing. However, their inference speed remains suboptimal due to the need to repeatedly process redundant masked tokens at every sampling step. In this work, we propose Sparse-LaViDa, a novel modeling framework that dynamically truncates unnecessary masked tokens at each inference step to accelerate MDM sampling. To preserve generation quality, we introduce specialized register tokens that serve as compact representations for the truncated tokens. Furthermore, to ensure consistency between training and inference, we design a specialized attention mask that faithfully matches the truncated sampling procedure during training. Built upon the state-of-the-art unified MDM LaViDa-O, Sparse-LaViDa achieves up to a 2x speedup across diverse tasks including text-to-image generation, image editing, and mathematical reasoning, while maintaining generation quality.

1. Introduction

Sparse-LaViDa targets redundant masked-token computation in multimodal MDM sampling while preserving flexible decoding and bidirectional context. It combines sparse token truncation, register representations, and a specialized attention mask, with reported speedups across generation, editing, and reasoning.

1. Introduction

  • Unified multimodal models combine understanding and generation, which can mutually benefit from shared representations, especially for tasks such as image editing.LaViDa-O is a unified MDM with strong performance across image understanding and generation tasks.Source: §1 · p. 1
  • MDMs repeatedly process full sequences, including redundant masked tokens, while Block Diffusion imposes left-to-right decoding that sacrifices bidirectional context.The left-to-right order is poorly suited to image generation and editing, where tokens lack a natural ordering.Source: §1 · pp. 1–2, Figure 2 · p. 2
  • Sparse-LaViDa represents partially masked sequences without materializing every mask and uses register tokens as compressed representations of truncated tokens.These components address the modeling capacity lost through truncation.Source: §1 · p. 2, Figure 2 · p. 2
  • Its step-causal attention mask supports KV caching while preserving the bidirectional context important for image generation, editing, and inpainting.Source: §1 · p. 2
  • 1.96× speedup on text-to-image generation, 2.80× on image editing, and 2.84× on visual math reasoning, with quality comparable to LaViDa-O.The comparison is against LaViDa-O.Source: Figure 1 · p. 1, §1 · p. 2

2. Related Work

Related work establishes MDMs as a discrete-diffusion approach to text and image modeling, then surveys acceleration methods that trade off decoding flexibility or bidirectional context. Sparse-LaViDa aims to combine token truncation and KV caching without those constraints.

2.1. Masked Diffusion Models

  • At inference, inputs include cached and newly decoded tokens, current masks, and registers; newly decoded tokens cannot attend to current masks or registers.Source: Figure 3 · p. 3
  • MDMs progressively replace clean tokens with masked tokens over time; at t = 1 the sequence is fully masked, and a learned reverse process iteratively unmasks tokens until a clean sequence is obtained.Source: §2.1 · p. 3
  • Dense MDMs process all L tokens at each step even when predictions are needed for only a small subset, creating computational inefficiency.Source: §2.1 · p. 3

2.2. Acceleration Methods for MDMs

  • Training-free accelerators Fast-dLLM [61], dKV-Cache [41], and Sparse-dLLM [56] use heuristic KV caching and can cause unpredictable performance degradation across tasks.Source: §2.2 · p. 3
  • Block Diffusion [1], SDAR [11], and D2F [58] use block-causal masks for KV caching and truncation, but assume left-to-right decoding and sacrifice bidirectional context.Source: §2.2 · pp. 3–4
  • Sparse-LaViDa supports KV caching and arbitrary-position truncation without left-to-right decoding or sacrificing bidirectional context; sparse sequences encode non-mask tokens, their locations, and total length.Source: §2.2 · p. 4, Figure 4 · p. 4

3. Method

Sparse-LaViDa represents masked sequences compactly, samples selected tokens with KV caching, and trains with a step-causal mask to match inference. Its sampling order varies by task while retaining bidirectional context for text infilling and constraint generation.

3.1. Sparse Parameterization of Masked Sequence

  • Masked positions can be inferred from the positions of clean tokens and the sequence length, avoiding separate input tokens for every mask.For example, “I have [m] dog [m] [m] [m]” is encoded using the clean tokens with positional embeddings and a length of 7.Source: §3.1 · p. 4
  • Sparse-LaViDa uses 64 register tokens because a single special token is insufficient, and the fixed register count does not grow with truncated masked tokens.Aggressive truncation can reduce model capacity; the paper notes that a 1024×1024 image has 4096 tokens, while fewer than 100 may be sampled in early steps.Source: §3.1 · p. 4

3.2. Sampling with Sparse Parameterization

  • At each sampling step, the model processes the prompt, previously decoded tokens, and tokens to decode; cached tokens are reused and logits are produced only for the decode tokens.The prompt is prefetched into the KV cache, and tokens decoded in the previous step are added to the cache before the next decode.Source: §3.2 · p. 4
  • For text-to-image generation and image editing, Sparse-LaViDa uses LaViDa-O’s stratified random sampler to set a pre-generated unmasking order without confidence scores.Prior work [4] [32] proposed pre-generated 2D sequences, which are reported to yield higher generation quality than confidence-based unmasking.Source: §3.2 · p. 4
  • For text generation and image understanding, blocks are sampled left-to-right, but the method supports bidirectional context and arbitrary block order for text infilling and constraint generation.Within a block, confidence scores determine which tokens to unmask after the forward pass.Source: §3.2 · pp. 4–5

3.3. Training Pipeline

  • The step-causal attention mask partitions prompt, clean, and masked tokens into blocks so clean tokens attend only to their own or earlier blocks, while masked blocks cannot attend to one another.Masked tokens may attend to prompt and clean tokens; each masked block receives corresponding register tokens with the same block number.Source: §3.3 · p. 5
  • Sparse-LaViDa trains with the standard MDM objective using sparse parameterization and step-causal masking, without an additional distillation stage.The mask simulates incremental token caching and truncation during inference, closing the training–inference gap.Source: §3.3 · pp. 5–6

4. Experiments

Experiments evaluate Sparse-LaViDa across generation, editing, understanding, and qualitative tasks, alongside ablations of caching, truncation, registers, and training strategy. Reported results include faster inference with task-specific quality comparisons and scope-dependent speedups.

4.1. Setup

  • Sparse-LaViDa starts from pretrained LaViDa-O weights, is fine-tuned on mixed understanding, generation, and editing data, and is evaluated across multimodal benchmarks.Unless otherwise stated, image-generation experiments use 1024 resolution on one A100 GPU; inference latency and speedup are measured relative to LaViDa-O.Source: §4.1 · p. 6

4.2. Text-to-Image Generation

  • 1.95× speedup on GenEval (21.27s vs. 10.86s), with performance comparable to LaViDa-O (+0.01).On DPG-bench, Sparse-LaViDa scores +0.6; under the same 20M-data training condition, its FID is 7.63 versus 8.11 for LaViDa-O*.Source: §4.2 · pp. 6–7, Table 2 · p. 6

4.3. Image Editing

  • 2.83× speedup on ImgEdit, with accuracy +0.08 and end-to-end latency reduced from 63.98s to 22.55s.ImgEdit evaluates visual quality and prompt compliance; Table 3 reports per-category and overall scores.Source: Table 3 · p. 7, §4.3 · p. 7

4.4. Image Understanding

  • 2.80× speedup on MathVista while matching LaViDa-O accuracy and outperforming Fast-dLLM on both accuracy and latency.The evaluation uses generation length L = 1024 tokens and block size S = 32; Table 4 compares MathVista accuracy and latency.Source: Table 4 · p. 7, §4.4 · p. 7
  • Image understanding performance is competitive across additional benchmarks, but speedups are minimal for QA outputs shorter than 32 tokens.For these short outputs, Sparse-LaViDa effectively reduces to prompt caching without truncation.Source: Table 5 · p. 7, §4.4 · p. 7

4.5. Qualitative Results

  • Figure 6 illustrates bidirectional-context tasks, including inpainting/outpainting, parallel grounding, and constrained captioning, unlike semi-AR Block Diffusion.Colored regions in text-generation examples mark masked tokens initialized for infilling.Source: §4.5 · p. 7, Figure 6 · p. 8

4.6. Ablation Studies

  • Each caching or truncation component improves efficiency, and combining them yields maximum speedup; removing registers has little GenEval impact but a larger DPG alignment cost.Register ablations also show measurable image-quality differences, suggesting stronger effects on low-level visual detail than high-level structure.Source: Table 6 · p. 8, §4.6 · p. 8, Table 7 · p. 8
  • Without fine-tuning, applying Sparse-LaViDa’s inference pipeline degrades GenEval and DPG performance; removing the step-causal attention mask also hurts performance.Table 8 reports the training-strategy comparison on GenEval and DPG-bench.Source: Table 8 · p. 8, §4.6 · p. 8

5. Conclusion and Future works.

Sparse-LaViDa is an efficient parameterization for multimodal MDMs, reporting speedups across visual and text-generation tasks while maintaining generation quality. The conclusion also notes additional training requirements and that experiments used post-training rather than training from scratch.

5. Conclusion and Future works.

  • Sparse-LaViDa reports speedups across text-to-image generation, image editing, and visual math reasoning without compromising generation quality.The conclusion describes it as a novel parameterization for multimodal MDMs.Source: §5 · p. 8
  • Sparse-LaViDa requires additional training, yet the conclusion reports faster speedup than the most aggressive KV-caching strategy.The cited strategy caches all possible tokens and is described as an upper bound for heuristic-based training-free methods.Source: §5 · p. 8
  • Experiments used a post-training setup because of compute costs; pretraining from scratch is left for future work.The conclusion says the parameterization could in principle be used to pre-train large models from scratch.Source: §5 · p. 8

7. Additional Technical Details

The appendix reviews standard masked diffusion, then details sparse sequence representation with register tokens and a step-causal attention mask aligned with cached inference and parallelizable training.

7.1. Formulation of Discrete Diffusion Models

  • Standard MDMs progressively mask clean tokens, train on partially masked sequences, and iteratively reverse-sample from a fully masked sequence; Sparse-LaViDa keeps the formulation but parameterizes it sparsely.The standard sequence contains clean and masked tokens at intermediate times, while Sparse-LaViDa introduces a sparse representation of that sequence.Source: §7.1 · p. 13

7.2. Register Tokens

  • Register tokens represent truncated masked positions, reducing the model input to |A| + |C| + m tokens when only a subset C of masked positions is decoded.When C = B, the sparse parameterization reduces to the standard MDM parameterization.Source: §7.2 · p. 14

7.3. Step-Causal Attention Masks

  • The step-causal mask supports KV-cached inference while preserving the attention pattern needed to simulate different sampling orders together during training.The training construction uses a block-diagonal mask to keep masked-token inputs from the separate sampling processes from attending to one another.Source: §7.3 · pp. 14–15

8. Additional Experiment Details and Results

The additional experiments describe the data and training setup, report object-grounding and larger-batch findings, and compare image-generation and editing outputs qualitatively.

8.1. Data pipeline

  • The training mix includes 20M text-to-image images plus image-level understanding, region-level understanding, and image-editing data.The listed sources include LAION-2B and COYO-700M for text-to-image pairs; the data are filtered for quality and safety.Source: §8.1 · p. 16

8.2. Training Setup

  • Sparse-parameterization SFT on LaViDa-O weights took five days on 64 A100 GPUs across eight nodes, 15% of LaViDa-O's from-scratch training budget.The setup largely followed LaViDa-O's SFT hyperparameters.Source: §8.2 · p. 16

8.3. Additional Results on Object Grounding

  • Object-grounding performance is reported as comparable to LaViDa-O, with identical inference speed because all coordinates are generated in one step.The authors attribute the lack of speedup to token truncation not benefiting this setup.Source: §8.3 · pp. 16–17, Table 10 · p. 17

8.4. Addtional Results with Larger Batch Size

  • At varying batch sizes, LaViDa-O consistently outperforms baselines on text-to-image throughput.Table 11 reports throughput in images per minute, with higher values better.Source: §8.4 · p. 17, Table 11 · p. 17

8.5. Additional Qualitative Comparisons

  • 1.95× speedup on text-to-image generation and 2.83× on image editing accompany comparable visual quality in the qualitative comparisons.The comparisons are between Sparse-LaViDa and the LaViDa-O baseline.Source: §8.5 · p. 17

9. Limitations

The reported speedup is concentrated in long-sequence generation, while several task and model-quality limitations remain and pretraining use is not yet demonstrated.

9. Limitations

  • Speedups benefit long-sequence generation, but not short QA or object grounding with very few output tokens.The limitations passage identifies text-to-image generation, image editing, and visual math problem-solving with long reasoning chains as benefiting tasks.Source: §9 · p. 17
  • The model inherits LaViDa-O's response hallucinations, and image edits may subtly shift regions that were not edited.The passage says the pixel-shift issue was also observed in LaViDa-O.Source: §9 · p. 19
  • Experiments demonstrate post-training use only; applying the method during pretraining remains a proposed extension.The passage says the method should in principle generalize to all training stages, including pretraining.Source: §9 · p. 19
Loading 2512.14008v2…