Lavida-O: Elastic Large Masked Diffusion Models for Unified Multimodal Understanding and Generation

1 INTRODUCTION

• Unified models are motivated by the potential for understanding and generation capabilities to benefit one another, particularly in tasks such as image editing that require both.The introduction contrasts unified models with specialist detection, VQA, and diffusion models.Source: §1 · p. 2
• Existing unified MDMs lag state-of-the-art unified AR and AR+diffusion models in benchmark performance and task breadth, and MMaDa and Muddit cannot perform image editing.The cited passage also notes costly training, scarce masked-image-generation resources, and limited training and sampling literature.Source: §1 · pp. 2–3
• Lavida-O is a unified multimodal MDM for image understanding and generation, using Elastic-MoT, progressive upscaling, token compression, stratified sampling, and universal text conditioning.It also incorporates planning and self-reflection to use understanding capabilities in generation.Source: §1 · p. 3
• Lavida-O reports state-of-the-art results on RefCOCO grounding, GenEval text-to-image generation, and ImgEdit editing, outperforming Qwen2.5-VL and Flux .1 Kontext dev with up to a 6.8× speedup.The introduction identifies these as results across a wide range of benchmarks.Source: §1 · p. 3
• Table 1 summarizes multimodal MDM capabilities, with Lavida-O uniquely supporting localized understanding, high-resolution synthesis, image editing, and interleaved generation.The supplied caption names these capability areas but does not provide cell-level comparisons.Source: Table 1 · p. 3

2.1 MASKED DIFFUSION MODELS

• Masked diffusion models formalize masked generative modeling as forward and reverse diffusion in discrete space, with sequences progressively masked and then unmasked from an all-mask initialization.Earlier masked generative approaches include BERT for representation learning and MaskGIT for generation; image tokens can be produced by discrete tokenizers such as VQGAN.Source: §2.1 · pp. 3–4

2.2 UNIFIED MULTI-MODAL MODELS

• Multimodal MDMs extend masked diffusion to image understanding and text-to-image generation, while unified MDMs still trail leading unified AR and AR+diffusion models.Unified AR+diffusion systems use separate objectives; unified MDMs instead offer a single objective, while MoT architectures use different parameter sets for image and text with joint attention.Source: §2.1 · p. 4, §2.2 · p. 4

3.1 MODEL ARCHITECTURE

• For understanding, Lavida-O retains LaViDa’s semantic-image and text conditioning; for generation it represents target images as VQ tokens and compresses input-image VQ tokens by a factor of four for editing and interleaved tasks.The overall pipeline conditions the model on image semantic, image VQ, and prompt embeddings alongside a partially masked sequence to predict the unmasked sequence.Source: §3.1 · pp. 4–5, Figure 2 · p. 5

3.1.1 ELASTIC MIXTURE-OF-TRANSFORMERS (ELASTICMOT)

• Elastic-MoT uses a smaller generation branch and limits cross-modal interaction to the first M of N layers, allowing task-dependent partial parameter activation.In the final design, generation activates 6.4B parameters, text-to-image pretraining trains only the 2.4B generation branch, understanding uses 8B active parameters, and interleaved tasks use 10.4B.Source: §3.1.1 · p. 5, Figure 3 · p. 6

3.1.2 MODALITY-AWARE MASKING

• Modality-aware masking replaces fully masked image tokens with an [exp] text token at timestamp texp; when generated, that token expands into image mask tokens routed to the generation branch.This lets the model determine image-token routing during parallel decoding without requiring users to specify token locations.Source: §3.1.2 · p. 6

3.2 TASK-SPECIFIC DESIGNS

• Universal text conditioning appends resolution, crop, quality, luminance, and contrast conditions to prompts as plain text, improving image quality and giving users finer output control.The cited passage describes this as an alternative to specialized micro-conditioning embeddings.Source: §3.2 · p. 6
• Stratified random sampling recursively subdivides a 2 × 2 grid and unmasks one token per region, distributing unmasking across the image rather than clustering nearby high-confidence tokens.The stated motivation is to address adjacent-token correlation that conflicts with MDM independence assumptions.Source: §3.2 · p. 7, Figure 4 · p. 7
• Planning generates image layouts or identifies edit regions before image synthesis, while reflection evaluates whether an output meets the request and generates a correction when misalignment is detected.Both mechanisms explicitly use the model’s understanding capabilities to improve generation.Source: §3.2 · p. 7
• For grounding, Lavida-O normalizes box coordinates to [0, 1] and quantizes them into 1025 discrete tokens, enabling four-token boxes and parallel decoding of multiple boxes in as little as one diffusion step.The method uses multiple masked coordinate queries, such as for a dog and a cat, as input.Source: §3.2 · p. 7

4.1 SETUP

• Training proceeds from understanding-task continuation, to generation pretraining with resolution increasing from 256 to 512 and 1024, then joint end-to-end training.The final stage covers understanding, text-to-image generation, image editing, and interleaved tasks such as planning and self-reflection.Source: §4.1 · pp. 7–8

4.2 MAIN RESULTS

• Lavida-O outperforms MMaDa on MMMU, MME, and MMB, and improves over LaViDa on most benchmarks, including ChartQA, DocVQA, ScienceQA, and MathVista.The authors attribute the improvements over LaViDa to scaling the training data.Source: Table 2 · p. 8, §4.2 · p. 8
• Lavida-O significantly outperforms Meissonic, MMaDa, and Muddit on text-to-image generation; planning and reflection further enhance prompt-following performance.MJHQ FID comparisons exclude planning and reflection because FID does not reflect prompt-following capabilities.Source: Table 3 · p. 8, §4.2 · p. 8
• On RefCOCO referring expression comprehension, Lavida-O outperforms Qwen2.5VL-7B, InternVL3-8B, Grounding-DINO-L, and SegLLM-7B.The reported metric is Precision@0.5.Source: §4.2 · p. 8, Table 4 · p. 9
• Lavida-O outperforms BAGEL and FluxKontext-dev on ImgEdit, and surpasses GPT-4o on object replacement and removal.The authors link these localized-understanding edit tasks to the effectiveness of integrating object-grounding capabilities.Source: Table 5 · p. 9, §4.2 · p. 9

4.3 TRAINING AND INFERENCE SPEED

• Lavida-O achieves a 6.8× speedup over Qwen2.5-VL-7B on object grounding and improves training speed by 3.17× over BAGEL-style standard MoT.Reducing the generation branch size gives a 2.23× speedup, while decoupling attention in the last 16 layers adds 1.44×.Source: Figure 5 · p. 9, §4.3 · p. 9

5 CONCLUSION

• Lavida-O achieves state-of-the-art performance on text-to-image generation, image editing, and grounding, while remaining competitive with leading specialist and autoregressive unified models.The conclusion also identifies efficient training and inference techniques, including Elastic-MoT, universal text conditioning, and stratified random sampling.Source: §5 · p. 10
• Planning and self-reflection use the model’s understanding capabilities to improve generation through interleaved generation, critique, and iterative refinement.The conclusion presents this as a unified-model paradigm for generation.Source: §5 · p. 10

A.1 FORMULATION OF MASKED DIFFUSION MODELS

• MDMs gradually mask discrete tokens in a forward process, then iteratively sample reverse transitions from an all-mask sequence until it is mask-free.Training samples a timestep and a partially masked sequence; inference starts with mask tokens and repeatedly samples the reverse process.Source: §A.1 · p. 16

A.2 ELASTIC-MOT ARCHITECTURE

• Elastic-MoT pairs a smaller generation branch with an understanding branch and decouples their attention in later layers, while retaining joint attention in the first M layers.The design uses M = K = 16 across 32 layers; task settings activate different branches and parameter sets.Source: §A.2 · pp. 16–18, Figure 6 · p. 17

A.3 MODALITY-AWARE MASKING

• Modality-aware masking processes tokens through the understanding branch by default, then routes image tokens to the generation branch when an [exp] token is generated.For training, image-token sequences collapse into [exp] at t_exp, shortening the sequence and requiring a modified loss construction.Source: §A.3 · p. 18

A.4 UNIVERSAL TEXT CONDITIONING

• Universal text conditioning represents image controls as [KEY] : [VALUE] strings, allowing users to specify conditions such as resolution, crop, brightness, and contrast.Conditions are randomly dropped during training; at inference, users can provide all or only some, and high brightness and contrast produce highly stylized images.Source: §A.4 · pp. 19–20

A.5 STRATIFIED RANDOM SAMPLING

• Stratified sampling disperses early unmasked tokens across image regions, avoiding uniform sampling’s clusters and Halton sampling’s imperfect coverage.The first four tokens occupy four quadrants; the first sixteen occupy a 4 × 4 grid, and the appendix reports FID benefits in Section B.3.Source: §A.5 · p. 20, Figure 9 · p. 20

A.6 OBJECT GROUNDING WITH COORDINATE QUANTIZATION

• Coordinate quantization normalizes bounding boxes to [0,1] and encodes each with four tokens, enabling resolution-independent coordinates and parallel prediction of multiple boxes.The method pads each image to a square before normalization and discretizes coordinates into 1025 bins.Source: §A.6 · p. 21, Figure 10 · p. 22

A.7 REFLECTION AND PLANNING

• Planning generates object-and-box layouts before images, while reflection critiques and revises images with planning at each round, truncating history to at most three rounds.Unlike cited generation-only approaches requiring an external vision-language reward model, Lavida-O unifies planning, critique, and iterative improvement in one model.Source: §A.7 · pp. 21–22

B ADDITIONAL EXPERIMENT DETAILS AND RESULTS

• The appendix adds reproducibility details and experiments examining design choices including Elastic-MoT, stratified sampling, and the data pipeline.The passage describes the scope of the supplementary experiments, not their outcomes.Source: B ADDITIONAL EXPERIMENT DETAILS AND… · p. 23

B.1 SETUP

• Three-stage training first extends grounding, then scales text-to-image generation, and finally jointly trains understanding, generation, and interleaved tasks using datasets spanning those capabilities.The listed datasets include text-to-image, image-level and region-level understanding, image editing, and planning and reflection data.Source: §B.1 · p. 24

Size of Generation Branch

• 3.17× faster, the 2B generation branch achieves the highest GenEval and DPG scores, balancing performance and training efficiency.The comparison uses 50k Stage-2 training steps.Source: §B.2 · p. 24, Table 8 · p. 25

Number of Joint Attention Layers

• M = 16 joint-attention layers balances Stage-2 speed and text-to-image performance; M = 8 substantially reduces performance, while more layers add little.In Stage 3, increasing M does not significantly improve performance, including for interleaved editing; efficiency differences are less pronounced because all 10.4B parameters must be loaded.Source: §B.2 · pp. 25–26, Table 9 · p. 25

Weight Initialization

• Truncated initialization converges faster and yields lower validation loss than training from scratch during Stage 2.The comparison covers the first 20k Stage-2 steps.Source: §B.2 · p. 26, Figure 12 · p. 26

Sampling Strategy Comparison

• The stratified sampler achieves the best FID outcome on MJHQ-30k among the evaluated samplers.Evaluation uses 64 sampling steps and compares confidence-based, uniform random, and Halton sampling.Source: §B.3 · p. 26, Table 10 · p. 26

Effect of Task Scheduler

• Stage 1 addresses an imbalance of fewer than 1M grounding samples versus more than 10M image-level understanding samples with a changing New:Old sampling ratio.Each batch is drawn from a single dataset.Source: §B.4 · p. 27
• The dynamic scheduler achieves strong grounding and image-level understanding, unlike fixed ratios that under-train grounding or cause image-level forgetting.It starts at New:Old = 3:1 and gradually decreases to 1:3, and it outperforms fixed 3:1 on image-level understanding.Source: §B.4 · p. 27, Table 11 · p. 27

Effect of Grounding Data

• Grounding data enhances text-to-image generation and editing even without explicit planning, and planning compounds those benefits.The comparison removes grounding data during Stage 3 training.Source: §B.4 · p. 27, Table 12 · p. 27

Breakdown of Performance Improvements

• Planning improves GenEval object positioning by +0.19, while reflection additionally improves counting and attribution.On additional T2I-Compbench++ categories, planning also improves 3D positional constraints despite using only 2D bounding boxes.Source: §B.5 · p. 27, Table 13 · p. 28, Table 14 · p. 28

Image-Edit Results

• Action edits improve by +0.50 and hybrid edits by +0.16 with planning, while global style and background edits degrade slightly.Planning also improves adding/removing objects and subject actions.Source: §B.5 · pp. 27–28, Table 15 · p. 28

Inference-time Scaling

• Reflection improves performance with one step (N = 2), but gains saturate at N = 8 with little benefit beyond.Latency grows sublinearly with N because simple prompts can stop early; at N = 20, generation may terminate after two images.Source: §B.5 · p. 28, Table 16 · p. 28

B.6 ABLATION STUDIES ON UNIVERSAL TEXT CONDITIONING.

• Human evaluators strongly preferred images generated with universal text conditioning, which the authors report improves image quality.The study compared image pairs generated from randomly selected MJHQ-30k prompts and used 300 human responses.Source: §B.6 · pp. 28–29, Figure 13 · p. 29

B.7 SPEED–QUALITY TRADEOFF

• Lavida-O achieves faster inference and stronger quality on image generation and grounding; grounding reaches up to 6.8× speedup while surpassing Qwen2.5-VL-7B in precision.Figure 14 reports latency and task scores, using FID for MJHQ and Precision@0.5 for RefCOCO.Source: §B.7 · p. 29, Figure 14 · p. 30
• On MathVista, Lavida-O is much faster but less accurate than state-of-the-art autoregressive models, and its performance exceeds LaViDa (56.9 vs. 44.8).It remains stronger than the Open-LLaVA-Next-8B baseline.Source: §B.7 · p. 29, Figure 14 · p. 30

B.8 ADDITIONAL QUALITATIVE RESULTS

• Additional qualitative examples show Lavida-O text-to-image generations and image-editing results across diverse prompts and instructions.The examples are presented in Figures 15 and 16, respectively.Source: §B.8 · p. 30

• Training ran on 8 nodes with 8 A100 GPUs each and took 34.2 days of wall-clock time, or 53k GPU hours.Source: C COMPUTE COST · p. 30

D LIMITATIONS

• Text rendering and editing are very limited because training did not explicitly include text-rendering data and the VQ tokenizer cannot faithfully reconstruct small text.The authors propose adding text-rendering data and finetuning the tokenizer on document screenshots.Source: D LIMITATIONS · p. 30
• Lavida-O inherits pixel shift from editing data distilled from generative models, so unedited regions may undergo small but noticeable changes.The authors aim to mitigate this by obtaining cleaner, higher-quality editing data.Source: D LIMITATIONS · p. 30
• Math reasoning improves over LaViDa with additional training but remains considerably behind state-of-the-art models.Source: D LIMITATIONS · p. 30
• Like other generative models, Lavida-O may produce inaccurate or fabricated information; outputs should be validated when accuracy is critical.Source: D LIMITATIONS · p. 30

E BOARDER IMPACT

• Lavida-O’s generation and editing may be abused, and it may inherit biases from LaViDa and its training data; the authors intend it for research use.Source: E BOARDER IMPACT · p. 31