--- abstract: | We propose `\ours`{=latex}, a unified Masked Diffusion Model (MDM) for multimodal understanding and generation. Unlike existing multimodal MDMs such as MMaDa and Muddit which only support simple image-level understanding tasks and low-resolution image generation, `\ours`{=latex}  presents a single framework that enables image-level understanding, object grounding, image editing, and high-resolution (1024px) text-to-image synthesis. `\ours`{=latex} incorporates a novel Elastic Mixture-of-Transformers (Elastic-MoT) architecture that couples a lightweight generation branch with a larger understanding branch, supported by token compression, universal text conditioning and stratified sampling for efficient and high-quality generation. Lavida-O further incorporates planning and iterative self-reflection in image generation and editing tasks, seamlessly boosting generation quality with its understanding capabilities. `\ours`{=latex} achieves state-of-the-art performance on a wide range of benchmarks including RefCOCO object grounding, GenEval text-to-image generation, and ImgEdit image editing, outperforming existing autoregressive models and continuous diffusion models such as Qwen2.5-VL and FluxKontext-dev, while offering substantial speedup at inference. These advances establish Lavida-O as a new paradigm for scalable multimodal reasoning and generation. Code and checkpoints are available at . author: - | Shufan Li$^{1,2,*}$, Jiuxiang Gu$^{1}$, Kangning Liu$^{1}$, Zhe Lin$^{1}$, Zijun Wei$^{1}$\ **Aditya Grover$^{2}$, Jason Kuen$^{1}$**\ $^1$Adobe  $^2$UCLA\ \* Work done primarily during internship at Adobe Research\ bibliography: - iclr2026_conference.bib title: "Lavida-O: Elastic Large Masked Diffusion Models for Unified Multimodal Understanding and Generation " --- \newif\ifarxiv \newcommand{\figleft}{{\em (Left)}} \newcommand{\figcenter}{{\em (Center)}} \newcommand{\figright}{{\em (Right)}} \newcommand{\figtop}{{\em (Top)}} \newcommand{\figbottom}{{\em (Bottom)}} \newcommand{\captiona}{{\em (a)}} \newcommand{\captionb}{{\em (b)}} \newcommand{\captionc}{{\em (c)}} \newcommand{\captiond}{{\em (d)}} \newcommand{\newterm}[1]{{\bf #1}} \def\figref#1{figure~\ref{#1}} \def\Figref#1{Figure~\ref{#1}} \def\twofigref#1#2{figures \ref{#1} and \ref{#2}} \def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}} \def\secref#1{section~\ref{#1}} \def\Secref#1{Section~\ref{#1}} \def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}} \def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}} \def\eqref#1{equation~\ref{#1}} \def\Eqref#1{Equation~\ref{#1}} \def\plaineqref#1{\ref{#1}} \def\chapref#1{chapter~\ref{#1}} \def\Chapref#1{Chapter~\ref{#1}} \def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}} \def\algref#1{algorithm~\ref{#1}} \def\Algref#1{Algorithm~\ref{#1}} \def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}} \def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}} \def\partref#1{part~\ref{#1}} \def\Partref#1{Part~\ref{#1}} \def\twopartref#1#2{parts \ref{#1} and \ref{#2}} \def\ceil#1{\lceil #1 \rceil} \def\floor#1{\lfloor #1 \rfloor} \def\1{\bm{1}} \newcommand{\train}{\mathcal{D}} \newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}} \newcommand{\test}{\mathcal{D_{\mathrm{test}}}} \def\eps{{\epsilon}} \def\reta{{\textnormal{$\eta$}}} \def\ra{{\textnormal{a}}} \def\rb{{\textnormal{b}}} \def\rc{{\textnormal{c}}} \def\rd{{\textnormal{d}}} \def\re{{\textnormal{e}}} \def\rf{{\textnormal{f}}} \def\rg{{\textnormal{g}}} \def\rh{{\textnormal{h}}} \def\ri{{\textnormal{i}}} \def\rj{{\textnormal{j}}} \def\rk{{\textnormal{k}}} \def\rl{{\textnormal{l}}} \def\rn{{\textnormal{n}}} \def\ro{{\textnormal{o}}} \def\rp{{\textnormal{p}}} \def\rq{{\textnormal{q}}} \def\rr{{\textnormal{r}}} \def\rs{{\textnormal{s}}} \def\rt{{\textnormal{t}}} \def\ru{{\textnormal{u}}} \def\rv{{\textnormal{v}}} \def\rw{{\textnormal{w}}} \def\rx{{\textnormal{x}}} \def\ry{{\textnormal{y}}} \def\rz{{\textnormal{z}}} \def\rvepsilon{{\mathbf{\epsilon}}} \def\rvtheta{{\mathbf{\theta}}} \def\rva{{\mathbf{a}}} \def\rvb{{\mathbf{b}}} \def\rvc{{\mathbf{c}}} \def\rvd{{\mathbf{d}}} \def\rve{{\mathbf{e}}} \def\rvf{{\mathbf{f}}} \def\rvg{{\mathbf{g}}} \def\rvh{{\mathbf{h}}} \def\rvu{{\mathbf{i}}} \def\rvj{{\mathbf{j}}} \def\rvk{{\mathbf{k}}} \def\rvl{{\mathbf{l}}} \def\rvm{{\mathbf{m}}} \def\rvn{{\mathbf{n}}} \def\rvo{{\mathbf{o}}} \def\rvp{{\mathbf{p}}} \def\rvq{{\mathbf{q}}} \def\rvr{{\mathbf{r}}} \def\rvs{{\mathbf{s}}} \def\rvt{{\mathbf{t}}} \def\rvu{{\mathbf{u}}} \def\rvv{{\mathbf{v}}} \def\rvw{{\mathbf{w}}} \def\rvx{{\mathbf{x}}} \def\rvy{{\mathbf{y}}} \def\rvz{{\mathbf{z}}} \def\erva{{\textnormal{a}}} \def\ervb{{\textnormal{b}}} \def\ervc{{\textnormal{c}}} \def\ervd{{\textnormal{d}}} \def\erve{{\textnormal{e}}} \def\ervf{{\textnormal{f}}} \def\ervg{{\textnormal{g}}} \def\ervh{{\textnormal{h}}} \def\ervi{{\textnormal{i}}} \def\ervj{{\textnormal{j}}} \def\ervk{{\textnormal{k}}} \def\ervl{{\textnormal{l}}} \def\ervm{{\textnormal{m}}} \def\ervn{{\textnormal{n}}} \def\ervo{{\textnormal{o}}} \def\ervp{{\textnormal{p}}} \def\ervq{{\textnormal{q}}} \def\ervr{{\textnormal{r}}} \def\ervs{{\textnormal{s}}} \def\ervt{{\textnormal{t}}} \def\ervu{{\textnormal{u}}} \def\ervv{{\textnormal{v}}} \def\ervw{{\textnormal{w}}} \def\ervx{{\textnormal{x}}} \def\ervy{{\textnormal{y}}} \def\ervz{{\textnormal{z}}} \def\rmA{{\mathbf{A}}} \def\rmB{{\mathbf{B}}} \def\rmC{{\mathbf{C}}} \def\rmD{{\mathbf{D}}} \def\rmE{{\mathbf{E}}} \def\rmF{{\mathbf{F}}} \def\rmG{{\mathbf{G}}} \def\rmH{{\mathbf{H}}} \def\rmI{{\mathbf{I}}} \def\rmJ{{\mathbf{J}}} \def\rmK{{\mathbf{K}}} \def\rmL{{\mathbf{L}}} \def\rmM{{\mathbf{M}}} \def\rmN{{\mathbf{N}}} \def\rmO{{\mathbf{O}}} \def\rmP{{\mathbf{P}}} \def\rmQ{{\mathbf{Q}}} \def\rmR{{\mathbf{R}}} \def\rmS{{\mathbf{S}}} \def\rmT{{\mathbf{T}}} \def\rmU{{\mathbf{U}}} \def\rmV{{\mathbf{V}}} \def\rmW{{\mathbf{W}}} \def\rmX{{\mathbf{X}}} \def\rmY{{\mathbf{Y}}} \def\rmZ{{\mathbf{Z}}} \def\ermA{{\textnormal{A}}} \def\ermB{{\textnormal{B}}} \def\ermC{{\textnormal{C}}} \def\ermD{{\textnormal{D}}} \def\ermE{{\textnormal{E}}} \def\ermF{{\textnormal{F}}} \def\ermG{{\textnormal{G}}} \def\ermH{{\textnormal{H}}} \def\ermI{{\textnormal{I}}} \def\ermJ{{\textnormal{J}}} \def\ermK{{\textnormal{K}}} \def\ermL{{\textnormal{L}}} \def\ermM{{\textnormal{M}}} \def\ermN{{\textnormal{N}}} \def\ermO{{\textnormal{O}}} \def\ermP{{\textnormal{P}}} \def\ermQ{{\textnormal{Q}}} \def\ermR{{\textnormal{R}}} \def\ermS{{\textnormal{S}}} \def\ermT{{\textnormal{T}}} \def\ermU{{\textnormal{U}}} \def\ermV{{\textnormal{V}}} \def\ermW{{\textnormal{W}}} \def\ermX{{\textnormal{X}}} \def\ermY{{\textnormal{Y}}} \def\ermZ{{\textnormal{Z}}} \def\vzero{{\bm{0}}} \def\vone{{\bm{1}}} \def\vmu{{\bm{\mu}}} \def\vtheta{{\bm{\theta}}} \def\va{{\bm{a}}} \def\vb{{\bm{b}}} \def\vc{{\bm{c}}} \def\vd{{\bm{d}}} \def\ve{{\bm{e}}} \def\vf{{\bm{f}}} \def\vg{{\bm{g}}} \def\vh{{\bm{h}}} \def\vi{{\bm{i}}} \def\vj{{\bm{j}}} \def\vk{{\bm{k}}} \def\vl{{\bm{l}}} \def\vm{{\bm{m}}} \def\vn{{\bm{n}}} \def\vo{{\bm{o}}} \def\vp{{\bm{p}}} \def\vq{{\bm{q}}} \def\vr{{\bm{r}}} \def\vs{{\bm{s}}} \def\vt{{\bm{t}}} \def\vu{{\bm{u}}} \def\vv{{\bm{v}}} \def\vw{{\bm{w}}} \def\vx{{\bm{x}}} \def\vy{{\bm{y}}} \def\vz{{\bm{z}}} \def\evalpha{{\alpha}} \def\evbeta{{\beta}} \def\evepsilon{{\epsilon}} \def\evlambda{{\lambda}} \def\evomega{{\omega}} \def\evmu{{\mu}} \def\evpsi{{\psi}} \def\evsigma{{\sigma}} \def\evtheta{{\theta}} \def\eva{{a}} \def\evb{{b}} \def\evc{{c}} \def\evd{{d}} \def\eve{{e}} \def\evf{{f}} \def\evg{{g}} \def\evh{{h}} \def\evi{{i}} \def\evj{{j}} \def\evk{{k}} \def\evl{{l}} \def\evm{{m}} \def\evn{{n}} \def\evo{{o}} \def\evp{{p}} \def\evq{{q}} \def\evr{{r}} \def\evs{{s}} \def\evt{{t}} \def\evu{{u}} \def\evv{{v}} \def\evw{{w}} \def\evx{{x}} \def\evy{{y}} \def\evz{{z}} \def\mA{{\bm{A}}} \def\mB{{\bm{B}}} \def\mC{{\bm{C}}} \def\mD{{\bm{D}}} \def\mE{{\bm{E}}} \def\mF{{\bm{F}}} \def\mG{{\bm{G}}} \def\mH{{\bm{H}}} \def\mI{{\bm{I}}} \def\mJ{{\bm{J}}} \def\mK{{\bm{K}}} \def\mL{{\bm{L}}} \def\mM{{\bm{M}}} \def\mN{{\bm{N}}} \def\mO{{\bm{O}}} \def\mP{{\bm{P}}} \def\mQ{{\bm{Q}}} \def\mR{{\bm{R}}} \def\mS{{\bm{S}}} \def\mT{{\bm{T}}} \def\mU{{\bm{U}}} \def\mV{{\bm{V}}} \def\mW{{\bm{W}}} \def\mX{{\bm{X}}} \def\mY{{\bm{Y}}} \def\mZ{{\bm{Z}}} \def\mBeta{{\bm{\beta}}} \def\mPhi{{\bm{\Phi}}} \def\mLambda{{\bm{\Lambda}}} \def\mSigma{{\bm{\Sigma}}} \newcommand{\tens}[1]{\bm{\mathsfit{#1}}} \def\tA{{\tens{A}}} \def\tB{{\tens{B}}} \def\tC{{\tens{C}}} \def\tD{{\tens{D}}} \def\tE{{\tens{E}}} \def\tF{{\tens{F}}} \def\tG{{\tens{G}}} \def\tH{{\tens{H}}} \def\tI{{\tens{I}}} \def\tJ{{\tens{J}}} \def\tK{{\tens{K}}} \def\tL{{\tens{L}}} \def\tM{{\tens{M}}} \def\tN{{\tens{N}}} \def\tO{{\tens{O}}} \def\tP{{\tens{P}}} \def\tQ{{\tens{Q}}} \def\tR{{\tens{R}}} \def\tS{{\tens{S}}} \def\tT{{\tens{T}}} \def\tU{{\tens{U}}} \def\tV{{\tens{V}}} \def\tW{{\tens{W}}} \def\tX{{\tens{X}}} \def\tY{{\tens{Y}}} \def\tZ{{\tens{Z}}} \def\gA{{\mathcal{A}}} \def\gB{{\mathcal{B}}} \def\gC{{\mathcal{C}}} \def\gD{{\mathcal{D}}} \def\gE{{\mathcal{E}}} \def\gF{{\mathcal{F}}} \def\gG{{\mathcal{G}}} \def\gH{{\mathcal{H}}} \def\gI{{\mathcal{I}}} \def\gJ{{\mathcal{J}}} \def\gK{{\mathcal{K}}} \def\gL{{\mathcal{L}}} \def\gM{{\mathcal{M}}} \def\gN{{\mathcal{N}}} \def\gO{{\mathcal{O}}} \def\gP{{\mathcal{P}}} \def\gQ{{\mathcal{Q}}} \def\gR{{\mathcal{R}}} \def\gS{{\mathcal{S}}} \def\gT{{\mathcal{T}}} \def\gU{{\mathcal{U}}} \def\gV{{\mathcal{V}}} \def\gW{{\mathcal{W}}} \def\gX{{\mathcal{X}}} \def\gY{{\mathcal{Y}}} \def\gZ{{\mathcal{Z}}} \def\sA{{\mathbb{A}}} \def\sB{{\mathbb{B}}} \def\sC{{\mathbb{C}}} \def\sD{{\mathbb{D}}} \def\sF{{\mathbb{F}}} \def\sG{{\mathbb{G}}} \def\sH{{\mathbb{H}}} \def\sI{{\mathbb{I}}} \def\sJ{{\mathbb{J}}} \def\sK{{\mathbb{K}}} \def\sL{{\mathbb{L}}} \def\sM{{\mathbb{M}}} \def\sN{{\mathbb{N}}} \def\sO{{\mathbb{O}}} \def\sP{{\mathbb{P}}} \def\sQ{{\mathbb{Q}}} \def\sR{{\mathbb{R}}} \def\sS{{\mathbb{S}}} \def\sT{{\mathbb{T}}} \def\sU{{\mathbb{U}}} \def\sV{{\mathbb{V}}} \def\sW{{\mathbb{W}}} \def\sX{{\mathbb{X}}} \def\sY{{\mathbb{Y}}} \def\sZ{{\mathbb{Z}}} \def\emLambda{{\Lambda}} \def\emA{{A}} \def\emB{{B}} \def\emC{{C}} \def\emD{{D}} \def\emE{{E}} \def\emF{{F}} \def\emG{{G}} \def\emH{{H}} \def\emI{{I}} \def\emJ{{J}} \def\emK{{K}} \def\emL{{L}} \def\emM{{M}} \def\emN{{N}} \def\emO{{O}} \def\emP{{P}} \def\emQ{{Q}} \def\emR{{R}} \def\emS{{S}} \def\emT{{T}} \def\emU{{U}} \def\emV{{V}} \def\emW{{W}} \def\emX{{X}} \def\emY{{Y}} \def\emZ{{Z}} \def\emSigma{{\Sigma}} \newcommand{\etens}[1]{\mathsfit{#1}} \def\etLambda{{\etens{\Lambda}}} \def\etA{{\etens{A}}} \def\etB{{\etens{B}}} \def\etC{{\etens{C}}} \def\etD{{\etens{D}}} \def\etE{{\etens{E}}} \def\etF{{\etens{F}}} \def\etG{{\etens{G}}} \def\etH{{\etens{H}}} \def\etI{{\etens{I}}} \def\etJ{{\etens{J}}} \def\etK{{\etens{K}}} \def\etL{{\etens{L}}} \def\etM{{\etens{M}}} \def\etN{{\etens{N}}} \def\etO{{\etens{O}}} \def\etP{{\etens{P}}} \def\etQ{{\etens{Q}}} \def\etR{{\etens{R}}} \def\etS{{\etens{S}}} \def\etT{{\etens{T}}} \def\etU{{\etens{U}}} \def\etV{{\etens{V}}} \def\etW{{\etens{W}}} \def\etX{{\etens{X}}} \def\etY{{\etens{Y}}} \def\etZ{{\etens{Z}}} \newcommand{\pdata}{p_{\rm{data}}} \newcommand{\ptrain}{\hat{p}_{\rm{data}}} \newcommand{\Ptrain}{\hat{P}_{\rm{data}}} \newcommand{\pmodel}{p_{\rm{model}}} \newcommand{\Pmodel}{P_{\rm{model}}} \newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}} \newcommand{\pencode}{p_{\rm{encoder}}} \newcommand{\pdecode}{p_{\rm{decoder}}} \newcommand{\precons}{p_{\rm{reconstruct}}} \newcommand{\laplace}{\mathrm{Laplace}} \newcommand{\E}{\mathbb{E}} \newcommand{\Ls}{\mathcal{L}} \newcommand{\R}{\mathbb{R}} \newcommand{\emp}{\tilde{p}} \newcommand{\lr}{\alpha} \newcommand{\reg}{\lambda} \newcommand{\rect}{\mathrm{rectifier}} \newcommand{\softmax}{\mathrm{softmax}} \newcommand{\sigmoid}{\sigma} \newcommand{\softplus}{\zeta} \newcommand{\KL}{D_{\mathrm{KL}}} \newcommand{\Var}{\mathrm{Var}} \newcommand{\standarderror}{\mathrm{SE}} \newcommand{\Cov}{\mathrm{Cov}} \newcommand{\normlzero}{L^0} \newcommand{\normlone}{L^1} \newcommand{\normltwo}{L^2} \newcommand{\normlp}{L^p} \newcommand{\normmax}{L^\infty} \newcommand{\parents}{Pa} \DeclareMathOperator*{\argmax}{arg\,max} \DeclareMathOperator*{\argmin}{arg\,min} \DeclareMathOperator{\sign}{sign} \DeclareMathOperator{\Tr}{Tr} \let\ab\allowbreak \newcommand{\fix}{\marginpar{FIX}} \newcommand{\new}{\marginpar{NEW}} \newcommand{\ours}{Lavida-O} \newcommand{\textgreen}[1]{\textbf{\textcolor{green!50!black}{#1}}} \newcommand{\textred}[1]{\textbf{\textcolor{red}{#1}}} \maketitle
We propose , a unified large masked diffusion model capable of multi-modal understanding and generation.
# Introduction The abilities to understand and generate images have been two essential objectives of image modeling research. Traditionally, these tasks are handled by a diverse set of specialist models, such as detection models for object localization [@liu2024grounding; @li2023mask], Visual Question Answering (VQA) models for question-answering [@li2022mplug], and diffusion models for text-to-image generation [@esser2024scaling-sd3; @podell2023sdxl; @rombach2022high]. Recently, the rise of unified multi-modal models such as GPT-4o [@openai2024gpt4o] has introduced a new paradigm: using a single generalist model to perform a wide range of image understanding and generation tasks. Not only is this unified approach more aligned with the goal of developing versatile multi-task Artificial General Intelligence (AGI), but it also demonstrates strong empirical performance by allowing understanding and generation capabilities to mutually benefit each other [@deng2025emerging]. This is especially notable in tasks requiring both understanding and generation capabilities, such as image editing, where unified models show unparalleled advantages over generation specialists. Most current unified models are built on Autoregressive (AR) large language models. Some works, such as BLIP3o [@chen2025blip3] and BAGEL [@deng2025emerging], employ AR modeling for text generation and continuous diffusion modeling for image generation (AR+diff), while others, such as Janus [@chen2025janus], first tokenize images into sequences of discrete tokens and then employ a unified AR next-token prediction objective for both image and text modalities. Recently, Masked Diffusion Models (MDMs) [@lou2023discrete-sedd; @sahoo2024simple] have emerged as a competitive alternative to AR models. Unlike AR models, MDMs treat token generation as a diffusion process over discrete tokens. In the forward process, the tokens of a sequence are gradually masked. At inference, we start with a sequence of mask tokens and gradually unmask them to obtain a sequence of meaningful tokens. Large-scale experiments in language modeling [@nie2025large; @dream2025] show that MDMs can achieve comparable performance to AR language models while offering many advantages, such as better speed-quality tradeoffs, controllability, and bidirectional context. Several recent works extend MDMs to multi-modal understanding and generation tasks [@li2025lavida; @yu2025dimple; @yang2025mmada; @shi2025muddit]. Compared with the AR+diff setup, unified MDMs avoid the need to carefully tune the balance between AR and diffusion losses by offering a unified objective, resulting in greater simplicity and scalability. Compared with unified AR modeling, unified MDMs offer significantly faster sampling speeds by allowing parallel decoding of multiple tokens. Despite these advantages, the latest unified MDMs---such as MMaDa [@yang2025mmada] and Muddit [@shi2025muddit]---still lag behind state-of-the-art unified AR and AR+diffusion models, both in the breadth of tasks they support and in benchmark performance. There are three main challenges in developing high-performing unified MDMs. First, unified models are expensive to train due to the large size of their language backbones. For example, to build a unified MDM with image generation capability, MMaDa pretrains an 8B model jointly on text and image generation, which is costly. This challenge is further exacerbated by the limited literature on training large-scale masked image generative models. In contrast, many open-source large-scale continuous diffusion models such as Flux [@flux2024] are readily available. Second, open-source resources for masked image generative models (MIGMs) are scarce, and the literature on their training techniques and sampling processes is less developed than that for continuous diffusion models. Even the best open-source MIGM, Meissonic-1B [@bai2024meissonic], significantly underperforms continuous diffusion models of comparable size [@xie2025sana1]. Lastly, while these models can perform both understanding and generation tasks, they lack explicit mechanisms to leverage image understanding capabilities to improve generation quality. In fact, MMaDa and Muddit cannot even perform image editing tasks, which require both understanding and generation capabilities. These models simply concatenate text-to-image data and image understanding data during training. To bridge this gap, we propose `\ours`{=latex}, a unified multi-modal Masked Diffusion Model (MDM) capable of both image understanding and generation tasks. To mitigate the cost of training large diffusion models, `\ours`{=latex} introduces several techniques such as Elastic Mixture-of-Transformers (Elastic-MoT), progressive upscaling (gradually increasing the image resolution during training), and token compression that enable efficient scaling. To improve generation quality, `\ours`{=latex} employs stratified sampling and universal text conditioning. To fully leverage the potential of a unified multi-modal model, `\ours`{=latex} incorporates planning and self-reflection mechanisms that explicitly utilize its understanding capabilities to enhance generation outputs. We highlight `\ours`{=latex}'s capabilities compared with previous multi-modal MDMs in Table `\ref{tab:task-capability}`{=latex}. Through extensive experiments, we show that `\ours`{=latex} achieves state-of-the-art performance on a wide range of benchmarks such as RefCOCO object grounding [@kazemzadeh2014referitgame], GenEval text-to-image generation [@ghosh2023geneval], and ImgEdit [@ye2025imgedit] image editing, outperforming existing autoregressive and continuous diffusion models such as Qwen2.5-VL [@bai2025qwen25-vl] and Flux .1 Kontext dev [@labs2025flux1kontextflowmatching], while offering up to a 6.8$\times$ speedup. Overall, our contributions can be summarized as follows: - We propose the first multi-modal MDM that achieves state-of-the-art performance on text-to-image generation, image editing, and grounding tasks, outperforming existing MDMs, AR models, and continuous diffusion models. - We propose several efficient and effective training and inference techniques for large-scale masked image generative models and unified multi-modal models, such as the Elastic-MoT architecture, universal text conditioning, and stratified sampling, significantly advancing the literature. - We introduce a novel paradigm that explicitly leverages the understanding capabilities of a unified model to improve its generation through planning and self-reflection. \resizebox{1.0\linewidth}{!}{ \begin{tabular}{lccccc} \toprule & \textbf{Understanding} & \textbf{Generation} \\ \cmidrule(lr){2-3} \cmidrule(lr){4-6} \textbf{Model} & \textbf{Image-level} & \textbf{Object-level} & \textbf{Text-to-image} & \textbf{Image-editing} & \textbf{Interleaved} \\ \midrule LaViDa, Dimple, LLaDa-V & $\checkmark$ & $\times$ & $\times$ & $\times$ & $\times$ \\ Muddit & $\checkmark$ & $\times$ & $ 512^2 $ & *\footnotemark[1] & $\times$ \\ MMaDa & $\checkmark$ & $\times$ & $ 512^2 $& $\times$ & $\times$ \\ LaViDa-O & $\checkmark$ & $\checkmark$ & $1024^2$ & $\checkmark$ & $\checkmark$ \\ \bottomrule \end{tabular} } # Background and Related Works ## Masked Diffusion Models {#sec:related-mdm} Masked Generative Modeling (MGM) has emerged as an alternative to AR models for modeling sequences of discrete tokens. Early works such as BERT [@devlin2019bert] used MGM as a representation learning objective. Later works [@chang2022maskgit; @chang2023muse] such as MaskGIT explored using MGM for generative modeling. In this setup, a sequence is initialized with only mask tokens, which are then gradually unmasked to generate the desired output. In these works, discrete tokenizers like VQGAN [@esser2021taming] are used to convert images into discrete tokens. More recently,MDMs [@austin2021structured-d3pm; @sahoo2024simple; @lou2023discrete-sedd; @shi2024simplified] have further developed the theory of MGM by formalizing the masking and unmasking process as the forward and reverse diffusion processes in discrete space. This provides a principled framework for training and sampling from these models. MDMs have renewed interest in masked modeling for language generation, offering theoretical advantages over AR models, such as better speed-quality tradeoffs and improved controllability. Notably, LLaDa-8B and Dream-8B [@nie2025large; @dream2025] demonstrated that MDMs can achieve competitive performance compared to AR models at scale. Several follow-up works [@li2025lavida; @yu2025dimple; @you2025lladav; @yang2025mmada] such as LaViDa extend MDMs to multi-modal tasks such as image understanding and text-to-image generation. Their capabilities are summarized in Table `\ref{tab:task-capability}`{=latex}. Formally, given a sequence of $L$ discrete tokens $X_0 = [X_0^1, X_0^2, \dots, X_0^L]$, the forward process $q(X_t|X_s)$ gradually masks the tokens over the time interval \[0,1\], with $1 \ge t \ge s \ge 0$. At $t = 1$, the sequence $X_1$ consists entirely of masked tokens, denoted by $[M]$. A neural network $p_\theta$ is used to model the reverse process $p(X_s|X_t)$. The masked diffusion objective is defined as: $$\begin{equation} \mathcal{L}_{\text{MDM}} = -\mathbb{E}_{t,X_0,X_t}\left[\frac{1}{t} \log p_\theta(X_0|X_t)\right] \label{eq:dlm-obj-ref} \end{equation}$$ where $p_\theta(X_0|X_t)$ is factorized into $\prod_{i=1}^L p_\theta(X_0^i|X_t)$ based on independence assumptions [@sahoo2024simple]. At inference time, the model starts from a fully masked sequence $X_1 = [M, M, \dots, M]$ and progressively applies the learned reverse process $\log p_\theta(X_0|X_t)$ to recover the original tokens. We provide a more detailed formulation of MDMs in Appendix `\ref{sec:appendix-formulation}`{=latex}. ## Unified Multi-modal Models {#sec:related_unified} Unified multi-modal models such as GPT-4o [@openai2024gpt4o] are capable of both image understanding and generation tasks, leading to strong performance on tasks requiring both capabilities, such as image editing. Generally, there are two dominant types of unified models based on their modeling objectives. The first type, such as BAGEL [@deng2025emerging], employs an AR objective for text generation and a diffusion objective for image generation (AR+diff). However, this design involves two different training objectives with distinct numerical scales and training dynamics, often requiring careful tuning of loss weighting and data mixtures. In contrast, the second type of models employ a unified objectives for both image and texts. Early works like Janus-Pro [@chen2025janus] employ a unified AR modeling objective. Recent works like [UMDD]{style="color: red"} and MMaDa [@yang2025mmada; @wang2008umdd] explore a unified MDM objective. Despite some success, a significant performance gap remains between these unified MDMs and state-of-the-art unified models in the AR and AR+diff categories. Architecturally, unified models also fall into two main categories. The first type, such as Janus and MMaDa, uses a single dense transformer to output both image and text tokens. The second type, such as BAGEL and MetaQueries [@pan2025transfer], employs separate parameter sets for handling image and text modalities. A common design in this category is the mixture-of-transformers (MoT) architecture [@liang2024mixture], where image and text inputs are processed by different parameter sets but can interact through joint attention mechanisms.[Under this paradigm, several works such as X-Fusion and LM-Fusion [@mo2025x; @shi2024lmfusion] further investigated architecture designs and training recipes of MoT.]{style="color: red"} These designs are illustrated in Figure `\ref{fig:elasticmot}`{=latex}. While being more flexible, training MoT experts can be expensive due to their large parameter counts. # Method
Overall Pipeline of . Given an input image and text prompt, we first concatenate the image semantic embedding Ci, image VQ embedding Cv, and text prompt embedding Cp to form the conditioning embedding C. The combined embedding is then passed to the model alongside the partially masked sequence Xt. The model then predicts the fully-unmasked sequence X0.
## Model Architecture `\ours`{=latex}'s model architecture is built on LaViDa [@li2025lavida], a diffusion model capable of only image understanding tasks. LaViDa uses a SigLIP [@zhai2023sigmoid] vision encoder to convert input images into continuous semantic embeddings $C_i$, which are concatenated with token embeddings of text prompts $C_t$ to form the final conditional embeddings $C=\text{Concat}(C_i,C_p)$ for visual understanding tasks. At each inference step, the diffusion model uses the partially unmasked answer $X_t$ and the conditional embedding $C$ to predict the clean text answer $X_0$. For image understanding tasks, `\ours`{=latex} maintains this exact setup of LaViDa. To incorporate visual generation tasks, we extend LaViDa's design by representing target images as sequences of discrete tokens using a VQ-Encoder [@esser2021taming]. When performing these tasks, $X_0$ and $X_t$ contain not only text tokens, but also VQ tokens that represent images. For image editing and interleaved generation tasks, we additionally incorporate VQ tokens of input images $C_{v}$ as part of the conditional embedding $C=\text{Concat}(C_i,C_v,C_p)$, since using semantic embeddings $C_i$ alone can degrade the low-level details needed for editing. To reduce the number of tokens and improve computational efficiency, we introduce a token compression module that reduces the number of VQ tokens by a factor of 4. The overall pipeline is illustrated in Figure `\ref{fig:explainer}`{=latex}. ### Elastic Mixture-of-Transformers (ElasticMoT) {#sec:elastic-mot} Our goal is to find an efficient method that can equip an understanding-only diffusion model with visual generation capabilities. However, both of the existing common choices described in Section `\ref{sec:related_unified}`{=latex}---dense models and MoT---are very expensive. Dense models use the same set of parameters for all tasks, requiring a mix of understanding and generation data during training to prevent catastrophic forgetting, which is not data-efficient. While the MoT setup allows freezing the understanding branch and training only the generation branch for image generation, its architecture doubles the total parameter count, leading to considerable computational overhead. Moreover, given an 8B base understanding model, both setups require training at least 8B parameters for generation tasks from scratch, which is prohibitively expensive. To address these limitations, we propose Elastic-MoT, a novel architecture design that efficiently adapts an understanding-only model for image generation tasks. Compared with the vanilla MoT architecture, Elastic-MoT introduces two major modifications. First, instead of using equally sized branches, the generation branch has a smaller hidden size. This reduces the parameter count and enables efficient training. We make this design choice based on the observation that many text-to-image models can generate high-quality images with only 2--4B parameters, suggesting that generation tasks may not require as much capacity as understanding tasks [@xie2025sana1; @xie2025sana]. Second, given an $N$-layer model, instead of having joint attention at all layers, we only allow text and image modalities to interact in the first $M$ layers. In the remaining $K=N-M$ layers, text and image tokens interact only within their modality through self-attention. This design activates only partial parameters for different tasks. For example, in `\ours`{=latex}'s final design, the generation branch has 2.4B new parameters and the understanding branch 8B parameters from LaViDa. With $N=32$ layers and $M=K=16$, image generation activates only 6.4B parameters (2.4B from generation + 4B from the first 16 understanding layers). During text-to-image pretraining, only the 2.4B generation branch is trainable, further improving the efficiency. Similarly, understanding tasks use 8B active parameters, while interleaved tasks requiring both branches use 10.4B. The full Elastic-MoT design is shown in Figure `\ref{fig:elasticmot}`{=latex}, with further details in Appendix `\ref{sec:appendix-elastic-mot-architecture}`{=latex} and `\ref{sec:appendix-ablation-mot}`{=latex}.
Design of Elastic MoT. Elastic-MoT introduces two major modifications to standard MoT. First, the generation branch has a smaller hidden size. Second, given an N-layer model, we only allow text and image modalities to interact in the first M layers. These two designs allow us to flexibly load only a portion of parameters depending on tasks, improving the efficiency.
### Modality-aware Masking {#sec:modality-aware-masking-main} One of the challenges in adapting MoT architecture for MDMs is routing---the mechanism to determine which branch should be activated for each token. This is trivial for unified AR MoT models, where the model can simply learn to generate a special token (e.g., \[img_start\]) to indicate that the next token should use the generation branch. However, MDMs decode tokens in parallel and must decide in advance which mask tokens should be routed to the understanding branch and which to the generation branch. A naive solution is to let the user specify the number and location of text and image tokens, but this is difficult for interleaved generation, such as image generation with self-reflection. To address this issue, we design a modality-aware masking process. Given a sequence of $M$ text tokens and $N$ image VQ tokens, the vanilla forward diffusion process gradually converts it into $M+N$ mask tokens during the time interval $[0,1]$. By contrast, our modality-aware forward process introduces a special timestamp $t_{\text{exp}}\in[0,1]$, at which fully masked image VQ tokens are collapsed into a special $[\text{exp}]$ text token. This process is illustrated in Figure `\ref{fig:modality-aware-masking}`{=latex} (Bottom-up). At inference, we assume all mask tokens are text tokens at the beginning. When a $[\text{exp}]$ token is generated, we replace it with a sequence of $L_{\text{img}}$ mask tokens, and specify that these tokens will be processed by the generation branch for image synthesis in subsequent forward calls. This process is also illustrated in Figure `\ref{fig:modality-aware-masking}`{=latex} (Top-down). We provide additional details in Appendix `\ref{sec:modality-aware-masking-appendix}`{=latex}. ## Task-Specific Designs In this section, we describe several additional technical innovations that improve the effectiveness and efficiency on newly incorporated tasks such as image generation, image editing and grounding. `\label{sec:method-image-gen}`{=latex} **Universal Text Conditioning**. A common approach to improving the quality of text-to-image models is micro-conditioning [@podell2023sdxl], which conditions the image generation process on extra parameters such as original image resolution, crop coordinates, and image quality scores. This is typically achieved via specialized embeddings. However, since a unified model has strong language understanding and reasoning capabilities, we can simply append these conditions as plain text (e.g., \`\`SCORE: 5.40") to the end of user prompts. In addition to common conditions, we also incorporate image luminance and contrast as micro-conditions. This simple and effective design not only improves image quality by biasing generation toward high-scoring distributions, but also gives users more refined control over outputs. We provide additional details in Appendix `\ref{sec:appendixuniversal-text-conditioning}`{=latex}.
Design choices of . (a) Forward diffusion process with modality-aware masking. (b) Visualization of the unmasking order in the proposed stratified random sampling process.
**Stratified Random Sampling.** Most MDMs use confidence-based sampling, unmasking high-confidence tokens first. In image generation, high-confidence tokens tend to cluster around already unmasked tokens. This negatively affecting image quality because adjacent tokens are highly correlated, which contradicts the independence assumption of MDMs. To mitigate this, we introudced a stratified sampling process. Starting with a $2\times2$ grid, we unmask one token per region to ensure broad spatial coverage. Each region is then recursively subdivided into four smaller subregions, and we continue unmasking one token from each new region. This process repeats until all tokens are revealed, producing a balanced, evenly distributed unmasking pattern across the entire image. This is illustrated in Figure `\ref{fig:stratified-sampling}`{=latex}. More details and analysis are provided in Appendix `\ref{sec:stratified-sampling}`{=latex} and `\ref{sec:ablation-sampling}`{=latex}. **Planning and Reflection.** While existing unified MDMs integrate image understanding and generation tasks with a single objective, they do not incorporate mechanisms that use understanding to improve generation, except for the assumption that joint training benefits both tasks. To address this, we introduce two explicit mechanisms that leverage understanding to improve generation: *planning* and *reflection*. With planning, the model first generates a layout of the image represented by bounding boxes, then creates the actual image accordingly. For image editing tasks, it first identifies the desired edit region before generating the edited image. With reflection, the model evaluates its own generation using its understanding capability and determines whether it satisfies the user's request. If misalignment is detected, the model generates a new image correcting the error. Examples are shown in Figure `\ref{fig:teaser}`{=latex}, with additional technical results and analysis in Appendix `\ref{sec:appendix-refl-detail}`{=latex} and `\ref{sec:appendix-ablation-reflection}`{=latex}. **Object Grounding with Coordinate Quantization**. The bi-directional context of MDMs naturally allows parallel decoding of bounding box coordinates. While `\ours`{=latex} can represent numbers as plain text, we adopt a specialized scheme that normalizes all bounding box coordinates to $[0,1]$ and quantizes them into 1025 discrete tokens representing $\frac{0}{1024},\frac{1}{1024},...,\frac{1024}{1024}$. This ensures each bounding box is represented by exactly four tokens. At inference, we construct a multiple query input with masked tokens such as \`\`A dog \[m\]\[m\]\[m\]\[m\]; A cat \[m\]\[m\]\[m\]\[m\]\", and unmask all coordinates in parallel. This design allow us to decode multiple bounding boxes in as low as a single diffusion step, greatly boosting the efficiency. We provide further details in Appendix `\ref{sec:appendix-grounding}`{=latex} # Experiment ## Setup We start with LaViDa [@li2025lavida] and extend it with a 2.4B image generation branch using the ElasticMoT architecture described in Section `\ref{sec:elastic-mot}`{=latex}. The training consists of three stages: **Stage 1:** We continue training the base model on object grounding and image-level understanding tasks. **Stage 2:** We incroprate an 2.4B image generation and pretrain for text-to-image generation. We start with a resolution of 256 and progressively increase it to 512 and 1024 during training. **Stage 3:** In the final stage, we jointly train the entire 2.4B + 8B model end-to-end on image understanding, text-to-image generation, image editing, and interleaved generation tasks such as planning and self-reflection. More details on the training data and process are provided in Appendix `\ref{sec:appendix-dataset}`{=latex}. ## Main Results **Image Understanding.** We report the performance of image understanding tasks in Table `\ref{tab:und-perf}`{=latex}. `\ours`{=latex} outperforms the previous state-of-the-art unified diffusion model, MMaDa, by a considerable margin on MMMU [@yue2023mmmu], MME [@fu2023mme], and MMB [@MMBench]. Compared with the base model LaViDa, `\ours`{=latex} achieves substantial improvements on most benchmarks such as ChartQA [@masry-etal-2022-chartqa], DocVQA [@mathew2021docvqa], ScienceQA [@lu2022learn], and MathVista [@lu2023mathvista], due to the scaling of the training data. \scriptsize `\setlength{\tabcolsep}{1.5pt}`{=latex}
Quantitative results on image-level understanding tasks.*Evaluated by us.
Model MMMU MME-P MME-C MMB ChartQA DocVQA InfoVQA Sci.QA AI2D M.Vista M.Verse
AR Und. Only
LLaVa-1.6-7B 35.1 1519.3 323 54.6 64.9 74.4 37.1 73.2 66.6 34.4 14.3
Qwen2.5-VL-7B 58.6 - - 83.5 84.9 82.6 - - 83.9 68.2 49.2
Intern-VL-3-8B 65.6 - - 83.4 86.6 92.7 76.8 - 85.2 75.2 39.8
AR Unified Und. and Gen.
BAGEL 55.3 1687 701 85 - - - - - 73.1 -
Janus-Pro-1B 36.3 1444 - 75.5 - - - - - - -
UniGen-1.5B 32.3 - - - - - - 79.4 67.4 44.6 -
Show-O 27.4 1233 - - - - - - - - - -
Masked Und. Only
Dimple (, ) 45.2 1514 432 74.6 63.4 - - 77.1 74.4 42.3 -
LaViDa (, ) 43.6 1366 341 70.5 64.6 59.0 34.2 80.2 70.0 44.8 27.2
Masked Unified Und. and Gen.
Muddit - 1104 - - - - - - - - -
MMaDa (, ) 30.2 1410 242* 68.5 9.8* 10.9* 14.9* 55.8* 66.6* 33.7* 13.5*

LaViDa-O

 45.1 1431 488 76.4 80.0 73.7 44.6 84.6 76.7 56.9 36.9
**Text-to-Image Generation.** We report text-to-image generation results on the GenEval [@ghosh2023geneval] and DPG [@hu2024equipdpg] benchmarks, and FID scores on 30k prompts from the MJHQ [@li2024playground] dataset. We compare against text-to-image models including Flux-dev [@flux2024], SD3-Medium [@esser2024scaling-sd3], Meissonic [@bai2024meissonic] and DALLE-3 [@openai_dalle3], unified models such as BAGEL [@deng2025emerging], MMaDa [@yang2025mmada] and Muddit [@shi2025muddit]. `\ours`{=latex} significantly outperforms the state-of-the-art Meissonic masked image generation model, as well as unified models such as MMaDa and Muddit. Planning and reflection further enhance prompt-following performance. We did not activate planning and reflection on MJHQ due to its large size and that FID does not reflect prompt-following capabilities. \scriptsize `\setlength{\tabcolsep}{12pt}`{=latex} \begin{tabular}{HHcccccc} \hline \textbf{Visual Gen} & \textbf{Language Gen} & \textbf{Method} & \textbf{Parms.} & \textbf{Type} & \textbf{GenEval} $\uparrow$ & \textbf{DPG-Bench}$\uparrow$ & \textbf{FID-30k}$\downarrow$\\ \hline \textit{Gen. Only} \\ & & Flux-dev \citep{flux2024}& 12B & Continuous & 0.68 & 84.0 & 10.15 \\ & & SD3-Medium \citep{esser2024scaling-sd3} & 2B & Continuous & 0.74 & 84.1 & 11.92 \\ && DALLE-3 \citep{openai_dalle3} & - & Continuous & 0.67 & 83.5 & - \\ Masked Diffusion & - & Meissonic \citep{bai2024meissonic} & 1B & Masked & 0.54 & - & -\\ \midrule \textit{Unified Und. and Gen.} \\ & & BAGEL \citep{deng2025emerging} & 7B+7B & Continuous & 0.82 & - & - \\ & & \textcolor{black}{Janus-Pro-1B}\citep{chen2025janus} & 1B & AR & 0.73 & 82.6 & - \\ & & \textcolor{black}{UniGen-1.5B} \citep{tian2025unigen} & 1B & AR & 0.63$\dagger$ & 82.8 $\dagger$ & - \\ & & OmniFlow \citep{li2024omniflow} & 3.4B & Continuous & 0.62 & - & -\\ Masked Diffusion & AR & Show-o \citep{xie2024show} & 1.3B & Masked & 0.67 & - & 15.18 \\ % Continuous Diffusion & AR & BLIP-3O & 0.84 & 0.81 \\ Masked Diffusion & Masked Diffusion & Muddit \citep{shi2025muddit}& 1B & Masked & 0.61 & - & -\\ Masked Diffusion & Masked Diffusion & MMaDA \citep{yang2025mmada} & 8B & Masked & 0.63 & 53.4* & 32.85* \\ \rowcolor{gray!20} Masked Diffusion & Masked Diffusion & LaViDa-O & 4B+2.4B & Masked & 0.77 & 81.8 & 6.68 \\ \rowcolor{gray!20} Masked Diffusion & Masked Diffusion & +Planning & 8B+2.4B& Masked& 0.85 & 82.9 & - \\ \rowcolor{gray!20} Masked Diffusion & Masked Diffusion & +Reflection & 8B+2.4B& Masked& 0.89 & 83.2 & -\\ \hline \end{tabular} **Object Grounding.** We evaluate the object grounding capabilities of `\ours`{=latex} on RefCOCO Referring Expression Comprehension (REC) tasks [@yu2016modeling; @mao2016generation], reporting the Precision@0.5 metric. `\ours`{=latex} outperforms autoregressive vision-language models such as Qwen2.5-VL-7B [@bai2025qwen25-vl] and InternVL3-8B [@zhu2025internvl3], as well as specialist models such as Grounding-DINO-L [@liu2024grounding] and SegLLM-7B [@wang2025segllm]. \scriptsize `\setlength{\tabcolsep}{10pt}`{=latex} \begin{tabular}{lccccccccH} \hline \textbf{Model} & \textbf{RefCOCO}$\uparrow$ & \textbf{RefCOCO+}$\uparrow$ & \textbf{RefCOCOg}$\uparrow$ & \textbf{Latency}$\downarrow$ \\ \cmidrule(lr){2-4} \cmidrule(lr){5-7} \cmidrule(lr){8-9} & val & testA & testB & val & testA & testB & val & test & (s/image) \\ \hline SegLLM-7B\citep{wang2025segllm} & 90.0 & 92.1 & 86.2 & 82.2& 85.5 &76.1 & 83.9 & 85.9 \\ Qwen2.5-VL-7B \citep{bai2025qwen25-vl} & 90.0 & 92.5 & 85.4 & 84.2 & 89.1 & 76.9 & 87.2 & 87.2 \\ GroundingDINO \citep{liu2024grounding}& 90.6 & 93.2 & 88.2 & 88.2 & 89.0 & 75.9 & 86.1 & 87.0 \\ InternVL3-8B \citep{zhu2025internvl3} & 92.5 & 94.6 & 88.0 & 88.2 & 92.5 & 81.8 & 89.6 & 90.0 \\ % LaViDa-Gnd-Specialist & 90.9 & 93.3 & 87.8 & 84.9 & 90.2 & 78.5 & 87.7 & 87.7 \\ % \hline \rowcolor{gray!20} LaViDa-O (4-step) & 92.3 & 94.8 & 89.0 & 88.7 & 92.5 & 83.3 & 90.0 & 90.6 \\ \rowcolor{gray!20} LaViDa-O (1-step) & 91.9& 94.6 &88.4 &87.4 & 91.7 & 82.2 &89.5 &89.8 \\ \hline \end{tabular} **Image Editing.** We report image editing results on ImgEdit benchmark [@ye2025imgedit] in Table `\ref{tab:image-edit}`{=latex}. `\ours`{=latex} outperforms state-of-the-art unified models such as BAGEL and specialized models like FluxKontext-dev. Most notably, `\ours`{=latex} even outperforms the state-of-the-art closed-source model GPT4-o[@openai2024gpt4o] on replacing and removing objects, which requires localized understanding. This underscores the effectiveness of `\ours`{=latex}'s design in integrating object-grounding capabilities. \scriptsize `\setlength{\tabcolsep}{3pt}`{=latex}
Per-Category and overall scores on ImgEdit benchmark.
Model Add Adjust Extract Replace Remove Background Style Hybrid Action Overall
GPT-4o 4.61 4.33 2.90 4.35 3.66 4.57 4.93 3.96 4.89 4.20
Qwen2.5VL+Flux 4.07 3.79 2.04 4.13 3.89 3.90 4.84 3.04 4.52 3.80
FluxKontext dev 3.76 3.45 2.15 3.98 2.94 3.78 4.38 2.96 4.26 3.52
OmniGen2 3.57 3.06 1.77 3.74 3.20 3.57 4.81 2.52 4.68 3.44
UniWorld-V1 3.82 3.64 2.27 3.47 3.24 2.99 4.21 2.96 2.74 3.26
BAGEL 3.56 3.31 1.70 3.30 2.62 3.24 4.49 2.38 4.17 3.20
Step1X-Edit 3.88 3.14 1.76 3.40 2.41 3.16 4.63 2.64 2.52 3.06
OmniGen 3.47 3.04 1.71 2.94 2.43 3.21 4.19 2.24 3.38 2.96
UltraEdit 3.44 2.81 2.13 2.96 1.45 2.83 3.76 1.91 2.98 2.70
AnyEdit 3.18 2.95 1.88 2.47 2.23 2.24 2.85 1.56 2.65 2.45
InstructAny2Pix 2.55 1.83 2.10 2.54 1.17 2.01 3.51 1.42 1.98 2.12
MagicBrush 2.84 1.58 1.51 1.97 1.58 1.75 2.38 1.62 1.22 1.90
Instruct-Pix2Pix 2.45 1.83 1.44 2.01 1.50 1.44 3.55 1.20 1.46 1.88

LaViDa-O

 4.04 3.62 2.01 4.39 3.98 4.06 4.82 2.94 3.54 3.71

+ Planning

 4.11 3.67 2.04 4.40 4.05 4.00 4.75 3.10 4.04 3.80
## Training and Inference Speed In Figure `\ref{fig:speed-main}`{=latex}, we benchmark the inference efficiency of `\ours`{=latex} across three tasks: text-to-image generation, object grounding, and math reasoning. We measure end-to-end latency in seconds per image. `\ours`{=latex} is significantly faster than autoregressive models. Notably, we achieve a $6.8\times$ speedup on object grounding tasks compared to Qwen2.5-VL-7B [@bai2025qwen25-vl]. We also report the training efficiency measured by per-step latency and compare our Elastic-MoT design with BAGEL-style standard MoT design, Elastic-MoT improves the training speed by $3.17\times$. Specifically, reducing the size of generation branch leads to a speedup of $2.23\times$, and decoupling the attention operation in the last 16 layers lead to an additional speedup of $1.44\times$, We provide additional analysis on the speed-quality tradeoff at inference time in Appendix `\ref{sec:appendix-speed-quality-tradeoff}`{=latex} and analysis on the training efficiency of Elastic-MoT design in `\ref{sec:appendix-ablation-mot}`{=latex}.
Training and Inference Speed of . We compare the end-to-end inference latency of  on three tasks, as well as pretraining efficiency measured by per-step latency.
\ifarxiv ## Additional Qualitative Results Finally, we provide additional qualitative examples demonstrating `\ours`{=latex}'s capabilities on diverse prompts and editing instructions. Figure `\ref{fig:demo-t2i}`{=latex} shows text-to-image generation, and Figure `\ref{fig:demo-editing}`{=latex} shows image editing results. `\fi`{=latex} `\ifarxiv`{=latex} ![**Qualitative examples of text-to-image generation.** We provide additional examples of text-to-image generation outputs on diverse prompts. ](figures/T2iDemo-5.png){#fig:demo-t2i width="100%"} ![**Qualitative examples of image editing.** We provide additional examples of image editing outputs on diverse instructions.](figures/EditingDemo-5.png){#fig:demo-editing width="100%"} \fi # Conclusion In summary, we proposed `\ours`{=latex}, the first multi-modal masked diffusion model that achieves state-of-the-art performance on text-to-image generation, image editing, and grounding tasks---competitive with the best specialist models and autoregressive unified models. We also introduced a novel paradigm of interleaved generation, which explicitly leverages understanding capabilities to improve generation results in a unified multi-modal model through planning and self-reflection. In developing `\ours`{=latex}, we proposed several efficient training and inference techniques, including the ElasticMoT architecture, universal text conditioning, and stratified random sampling, providing valuable insights for future work in masked diffusion models and unified multi-modal systems. \bibliographystyle{iclr2026_conference} \appendix \clearpage # Additional Technical Details ## Formulation of Masked Diffusion Models {#sec:appendix-formulation} \newcommand{\cat}[0]{\text{Cat}} \newcommand{\alphats}[0]{\frac{1-t}{1-s}} \newcommand{\oneminusalphats}[0]{\frac{t-s}{1-s}} Masked Diffusion Models (MDMs) model the generation process of discrete token sequences through a continuous-time Markov chain (CMTC). Formally, given a sequence of discrete tokens $X_0=[X_0^1,X_0^2,\ldots,X_0^L]$ of length $L$, the forward process $q(X_t|X_s)$ gradually converts it into a sequence of mask tokens $[M]$, denoted by $X_1=[X_1^1,X_1^2,\ldots,X_1^L]$, over the continuous time interval $[0,1]$, with $1 \ge t \ge s \ge 0$. Each token $X_t^i$ belongs to a fixed-size vocabulary set $V$. In our setup, $V$ consists of text tokens, image VQ tokens, and the special mask token $[M]$. This forward process is formally defined as $$\begin{equation} q(X_t^i|X_s^i) = \begin{cases} \cat(X_t^i;\textbf{M}), & \text{if } X_s^i=[M] \\ \cat(X_t^i;\alphats \mathbf{X_s^i}+\oneminusalphats \textbf{M}), & \text{if } X_s^i \ne [M], \end{cases} \end{equation}$$ where $\cat(\cdot)$ denotes a categorical distribution, and $\textbf{M}, \mathbf{X_0^i}, \mathbf{X_s^i} \in \mathbb{R}^{|V|}$ are probability vectors, with $|V|$ denoting the vocabulary size. In particular, $\textbf{M}$ is a one-hot vector representing the mask token $[M]$. This forward process yields the following marginal distribution: $$\begin{equation} q(X_t^i|X_0^i) = \cat(X_t^i;(1-t) \mathbf{X_0^i}+t \textbf{M}). \label{eq:q_process} \end{equation}$$ MDLM [@sahoo2024simple] demonstrated that the posterior of the reverse process $p(X_s|X_t,X_0)$ has the following form: $$\begin{equation} p(X_s^i|X_t^i,X_0^i) = \begin{cases} \cat(X_s^i;\mathbf{X_t^i}), & \text{if } X_s^i \ne [M] \\ \cat(X_s^i;\tfrac{t-s}{t} \mathbf{X_0^i}+\tfrac{s}{t} \textbf{M}), & \text{if } X_s^i = [M]. \end{cases} \label{eq:appendix-eq-p} \end{equation}$$ In practice, we replace $\mathbf{X_0^i}$ with the neural network prediction $p_\theta(X_0^i|X_t)$ when sampling from the reverse process, which gives the following transition: $$\begin{equation} p_\theta(X_s^i|X_t) = \begin{cases} \cat(X_s^i;\mathbf{X_t^i}), & \text{if } X_s^i \ne [M] \\ \cat(X_s^i;\tfrac{t-s}{t} p_\theta(X_0^i|X_t)+\tfrac{s}{t} \textbf{M}), & \text{if } X_s^i = [M]. \end{cases} \label{eq:appendix-inference} \end{equation}$$ **Sampling process.** At inference time, we initialize $X_1$ as a sequence of mask tokens, with $X_1^1=X_1^2=\cdots=X_1^L=[M]$. We discretize the continuous time interval $[0,1]$ into discrete timesteps $0=t_0
Activated parameters of  under different task settings. Elastic-MoT Design allow  to dynamically loads its parameters depending on the tasks. For understanding-only tasks, we only load the 8B generation branch. For text-to-image generation tasks, we load the first M = 16 layers of the understanding branch, which consists of 4B parameters, and the full 2.4B generation branch. For interleaved tasks, we load all 2.4B+8B parameters.
**Decoupled attention.** In standard MoT, understanding and generation tokens can interact with each other in all $N$ transformer layers through the joint attention mechanism. We decouple attention in the last $K$ layers and only allow tokens of the same type to interact with each other. In the first $M=N-K$ layers, all tokens can still interact with each other as in the standard MoT architecture. This design is motivated by two factors. First, it prevents text and image tokens from interfering with each other's representations in the later stages of generation. Second, and more importantly, it allows us to load only 4B out of 8B parameters for text-to-image generation tasks, greatly improving the scalability of pretraining while also reducing compute cost at inference time. We visualize the activated parameters for different tasks in Figure `\ref{fig:elastic-mot-appendix}`{=latex}. For understanding-only tasks, we activate only the understanding branch in all $N=M+K$ layers. For generation-only tasks, we activate the understanding branch in the first $M$ layers and the generation branch in all $N=M+K$ layers. For interleaved tasks with both text and image outputs, we activate all parameters. In our setup, we choose $M=K=16$, which yields $N=32$ layers in total. ## Modality-Aware Masking {#sec:modality-aware-masking-appendix} In this section, we provide details of the changes to the training and sampling process introduced by modality-aware masking, as described in Section `\ref{sec:modality-aware-masking-main}`{=latex}. Recall that in adapting the MoT architecture for MDMs, one of the main challenges is routing tokens. In particular, while we can easily decide which branch should process unmasked tokens based on whether they are image VQ tokens or text tokens, it is difficult to make such decisions for masked tokens, especially in interleaved generation tasks where the final output contains both images and text. Modality-aware masking addresses this problem by processing all tokens with the understanding branch by default and dynamically deciding when and where to invoke the generation branch during the sampling process. **Sampling Process.** For convenience, we denote masked tokens that will be processed by the understanding branch as $M_{und}$ and masked tokens that will be processed by the generation branch as $M_{gen}$. With this distinction, the routing policy becomes simple: all text tokens plus $M_{und}$ are processed by the understanding branch, while all image VQ tokens plus $M_{gen}$ are processed by the generation branch. We introduce a special text token `[exp]` to indicate when an image should be generated. When a `[exp]` token is generated in the unmasking process, it is automatically replaced with a sequence of $M_{gen}$ tokens. The number of $M_{gen}$ tokens representing each image is determined by prespecified output resolution. These tokens are then processed by the generation branch in subsequent rounds. For example, each $1024 \times1024$ image is represented by 1024 VQ tokens. This process is documented in Algorithm `\ref{alg:masked_inference}`{=latex} and illustrated in Figure `\ref{fig:masking-modality-appendix}`{=latex} (Left). \begin{algorithm}[h]\caption{Interleaved Generation with Modality-Aware Masking} \label{alg:masked_inference} \begin{algorithmic}[1] \Require Initial Generation Length $L$, discrete timestamps $0=t_0
Visualization of different sampling processes. We compare the unmasking order of the stratified sampler, Halton sampler, and uniform random sampler. Uniform random sampler produces the least desirable spatial pattern, with many unmasked tokens clustered together. Halton sampler is less ideal than stratified sampler because it does not guarantee perfectly stratified coverage. For example, when the number of unmasked tokens is 4, the upper-right quadrant remains unoccupied.
In this section, we provide detailed descriptions of the stratified random sampling process introduced in Section `\ref{sec:method-image-gen}`{=latex}. In the vanilla sampling process described in Equation `\ref{eq:appendix-inference}`{=latex}, each token is unmasked independently. In practice, this often leads to suboptimal generation quality. Instead of unmasking tokens randomly, several works adopt alternative sampling strategies in which the unmasking order of tokens is determined by heuristics such as the model's confidence at each token position [@nie2025large; @dream2025; @chang2022maskgit]. In image generation, tokens with high confidence are frequently adjacent to one another. As a result, confidence-based unmasking tends to reveal many adjacent tokens in a single step. Since tokens that are spatially adjacent are often highly correlated, this violates the independence assumption $p_\theta(X_{t_{k-1}} | X_{t_k})=\prod _{i=1}^L p_\theta(X_{t_{k-1}}^i | X_{t_k})$ stated in Section `\ref{sec:appendix-formulation}`{=latex}. To address this, we design a stratified sampling process that ensures the unmasked tokens are spatially dispersed. Specifically, we enforce that the first 4 unmasked tokens occupy the four quadrants of the image; the first 16 unmasked tokens occupy all 16 subregions obtained by dividing the image into a $4\times4$ grid; and so forth. The algorithm is formally described below: \begin{algorithm}[h]\caption{Stratified Unmasking Order} \label{alg:quadrant-unmask} \begin{algorithmic}[1] \Require Image size $N \times N$ \Ensure a list $\mathcal{O}$ of coordinates $(i,j)$ indicating unmasking order \State Initialize an empty list $\mathcal{O}$ \For{$d = 1, 2, \dots, \log_2 N$} \State Partition the image into $2^d \times 2^d$ grid cells \For{each grid cell $g$ in random order} \If{$\mathcal{O} \cap g = \text{\O}$} \State Sample $(i_g, j_g)$ uniformly within cell $g$ \State Append $(i_g, j_g)$ to $\mathcal{O}$ \EndIf \EndFor \EndFor \State \Return $\mathcal{O}$ \end{algorithmic} \end{algorithm} Our design is inspired by the stratified sampling process commonly used in numerical integration and computer graphics. It also follows a similar motivation to the recent Halton mask scheduler, which uses the low-discrepancy Halton sequence to ensure that unmasked tokens are spatially dispersed [@besnier2025halton]. We illustrate the differences among stratified sampling, Halton sampling, and uniform random sampling in Figure `\ref{fig:stratified-sampling-appendix}`{=latex}. As shown in the figure, uniform random sampling produces the least desirable spatial pattern, with many unmasked tokens clustered together. Compared with our proposed stratified sampling process, Halton sampling is less ideal because it does not guarantee perfectly stratified coverage. For example, when the number of unmasked tokens is 4, the upper-right quadrant remains unoccupied. The benefits of stratified sampling are also reflected in FID scores, which we document in Section `\ref{sec:ablation-sampling}`{=latex}. ## Object Grounding with Coordinate Quantization {#sec:appendix-grounding} In this section, we provide detailed descriptions of `\ours`{=latex}'s design for object grounding tasks. Given an image and a referring expression describing an object, the grounding task requires locating the described object in the image by predicting its bounding box coordinates. In autoregressive vision-language models such as Qwen2.5-VL [@bai2025qwen25-vl], bounding boxes are represented as plain text strings, such as \`\`\[123, 232, 300, 1021\]\". At inference, the coordinates are generated sequentially from left to right. This design has several limitations. First, since the model only sees a padded and resized image, it is difficult for the model to predict absolute pixel coordinates that depend on the original resolution of the input image. Second, the sequential generation order is slow and inefficient. To address these issues, we normalize the bounding box coordinates and quantize them into discrete bins. Specifically, given an image of size $H\times W$, we first pad it to a square image of size $D \times D$, where $D=\max(H,W)$, and normalize the bounding boxes in the padded image to the range \[0,1\] by dividing the raw pixel coordinates by $D$. This step makes the coordinates independent of the original input resolution. We then round each coordinate into 1025 bins representing $\frac{0}{1024},\frac{1}{1024},\frac{2}{1024},...,\frac{1024}{1024}$ and represent them with special tokens. This reduces the number of tokens needed to represent each bounding box to exactly 4. Finally, since `\ours`{=latex} is a masked diffusion model with a bi-directional attention mask and parallel decoding capabilities, we can predict multiple bounding boxes simultaneously. For example, if we want to obtain the bounding boxes of both \`\`a cute dog\" and \`\`a boy,\" we can initialize a text sequence \`\`a cute dog \[m\]\[m\]\[m\]\[m\]; a boy \[m\]\[m\]\[m\]\[m\]\" and perform parallel unmasking of multiple bounding box coordinates. This design is illustrated in Figure `\ref{fig:grounding}`{=latex}.
Coordinate Quantization. We normalize bounding box coordinates into the range [0,1] and discretize them into 1025 bins. This ensures that each bounding box is represented by exactly 4 tokens, allowing efficient parallel decoding of multiple bounding boxes in a single step.
## Reflection and Planning {#sec:appendix-refl-detail} The unique advantage of unified understanding and generation models is that they can leverage their understanding capabilities to improve generation results. Several works on unified models show that simple joint training on a combination of understanding and generation tasks improves performance on generation tasks [@xie2024show; @deng2025emerging], particularly in instruction-following capabilities. `\ours`{=latex} pushes this paradigm further by introducing two explicit mechanisms to exploit understanding capabilities: planning and reflection. At inference, these capabilities are invoked via specialized prompts, such as \`\`please generate a layout design before creating the final image\".
Interleaved Generation with Planning and Reflection. We provide visual examples of interleaved generation, including text-to-image generation with planning (Top), text-to-image generation with reflection (Middle), and image editing with planning (Bottom). We always enable planning during the reflection process. The layout traces is omitted in the middle figure for clarity and better presentation.
**Planning.** To improve prompt-following capabilities in text-to-image generation, we ask the model to first generate a layout design of objects, which consists of (object, bounding box) pairs, before generating the final image. Such interleaved generation is achieved through the modality-aware masking process described in Section `\ref{sec:modality-aware-masking-appendix}`{=latex}. We illustrate this process in Figure `\ref{fig:appendix-refl}`{=latex} (Top). As shown, planning enables `\ours`{=latex} to follow challenging and unintuitive prompts, such as \`\`a horse *above* an astronaut.\" Similarly, we can adopt planning for image editing tasks. Given an input image and an edit instruction, the model can first leverage its grounding capabilities to identify the regions that need to be edited before generating the edited image. This process is illustrated in Figure `\ref{fig:appendix-refl}`{=latex} (Bottom). **Reflection.** We can improve text-to-image generation performance by leveraging `\ours`{=latex}'s understanding capability to achieve self-critique and iterative self-improvement. Given an input prompt, the model first generates an image, then performs a self-critique step to evaluate whether the generated image matches the prompt. If it does, the generation process terminates. Otherwise, the model generates a revised image and attempts to fix the identified issues. This cycle is repeated until an image passes the self-critique process or the maximum number of rounds is reached. At each round, we also invoke the planning capability. Since `\ours`{=latex}'s context length is limited to 8192 tokens, we truncate the history when necessary to include at most three rounds. This process is illustrated in Figure `\ref{fig:appendix-refl}`{=latex} (Middle). Formally, the reflection process is defined by the following algorithm \begin{algorithm}[h]\caption{Iterative Image Refinement with Self-Reflection Loop} \label{alg:inference} \begin{algorithmic}[1] \Require Text prompt $P$, Unified Model $\Theta$, Max Iterations $N$ \Ensure Output image $I$ \State Initialize $I_1 \gets \text{GenerateWithPlanning}(\Theta,P)$ \Comment{Generate an initial image} \State $F_1 \gets \text{GetTextFeedback}(\Theta,P, I_1)$ \Comment{Obtain initial feedback} \For{$i = 2$ to $N$} \State $\mathcal{H}_i \gets \{(I_j, F_j) \mid j = 1, 2, ..., i-1\}$ \Comment{Construct History Context} \If{Run out of context limit of model $\Theta$} \State Truncate $\mathcal{H}_i$ by removing early rounds \EndIf \State $I_i \gets \text{GenerateWithPlanning}(\Theta,P, \mathcal{H}_i)$ \Comment{Generate a new image} \State $F_i \gets \text{GetTextFeedback}(P, I_i)$ \Comment{Obtain new feedback} \If{$F_i = \texttt{"{}"{}}$} \Comment{Stop if no more improvements} \State \Return $I_i$ \EndIf \EndFor \State \Return $I_N$ \end{algorithmic} \end{algorithm} Similar designs and algorithms have been explored for generation-only models in the context of inference-time scaling, such as Reflect-DiT [@li2025reflect] and ReflectionFlow [@zhuo2025reflection]. However, unlike these works, which require an external vision-language model as a reward model, `\ours`{=latex} uniquely unifies layout planning, self-critique, and iterative self-improvement in a single model through a unified generation process. # Additional Experiment Details and Results In this section, we document the details of the experiments for better reproducibility, including data pipeline, training hyperparameters, and compute cost. In addition, we also provide additional experimental results on the effectiveness of various design choices used by `\ours`{=latex}, such as the Elastic-MoT design, stratified sampling, and the data pipeline. ## Setup {#sec:appendix-dataset} **Pretrained Weights.** We use LaViDa [@li2025lavida] to initialize the understanding branch and semantic encoder. For the VQ encoder, we adopt the Meissonic encoder [@bai2024meissonic]. The image generation branch is initialized from the truncated weights of the understanding branch, as described in Section `\ref{sec:appendix-elastic-mot-architecture}`{=latex}. **Data Pipeline.** Unlike many frontier models, our model does not make use of proprietary images or documents. Our training data consists of the following components: - *A: Text-to-Image Pairs.* We source data from LAION-2B [@schuhmann2022laion] and COYO-700M [@kakaobrain2022coyo-700m]. We additionally include SA-1B [@kirillov2023segment], JourneyDB [@sun2023journeydb], BLIP3o-60k [@chen2025blip3], and ShareGPT4o-Image [@chen2025sharegpt]. Each dataset is heavily filtered to remove NSFW prompts, low CLIP scores [@radford2021learning], low aesthetic scores [@laion-aesthetics], and low-resolution images. This results in 200M images in our final mix. Where available, we use captions generated by VLMs instead of raw alt-texts. These captions are sourced from existing work including Recap-LAION, Recap-COYO [@LoTLIP], and BLIP-3o [@chen2025blip3]. - *B: Image-level Understanding Data.* We include LLaVA-OneVision [@li2024llava], Open-LLaVA-Next [@chen2024open], MAmmoth-VL [@guo2024mammoth], and VisualWebInstruct [@visualwebinstruct]. - *C: Region-level Understanding Data.* We include GranD [@hanoona2023GLaMM] and RefCOCO [@kazemzadeh2014referitgame]. - *D: Image Editing Data.* We include ShareGPT4o-Image [@chen2025sharegpt], GPT-Edit-1.5M [@wang2025gpt], and the image editing subset of UniWorld-V1 [@hu2022unified]. - *E: Interleaved Planning and Reflection Data.* For planning data, we manually construct a layout dataset by running an open-vocabulary object detector, GroundingDino-L [@liu2024grounding], on the outputs of image generation and editing datasets, including BLIP-3o [@chen2025blip3], ShareGPT4o-Image [@chen2025sharegpt], and GPT-Edit-1.5M [@wang2025gpt]. For reflection data, we leverage existing datasets including ReflectDiT [@li2025reflect] and ReflectionFlow [@zhuo2025reflection]. **Training Setup.** Training consists of three stages. In the first stage, we extend LaViDa to region-level tasks such as grounding. In the second stage, we perform large-scale pretraining on text-to-image generation tasks. In the final stage, we jointly train the model on a mix of understanding, generation, and interleaved tasks. We document the training hyperparameters, the datasets used, the active parameter count, and other relevant details in Table `\ref{tab:training-stages}`{=latex}. In addition, we implement a dataset mix scheduler that dynamically adjusts the sampling weight of each dataset throughout training to address data imbalance. Specifically, we assign a high weight to new capabilities at the beginning of each training stage and gradually decay the weight over time. For example, in Stage 1 we have fewer than 1M grounding samples but more than 10M image-level understanding samples. To enable efficient acquisition of grounding capability while preventing overfitting, we initially set the grounding-to-understanding ratio to 3:1, which is gradually decreased to 1:3. We provide further analysis of the scheduler in Section `\ref{sec:ablation-schedular}`{=latex}. **Stage 1** **Stage 2** **Stage 3** ---------------------- ------------------------------ ------------------------------------------ ------------------------------ Learning Rate $5 \times 10^{-6}$ $1 \times 10^{-4}$ $2 \times 10^{-5}$ Steps 80k 400k 100k $\beta_1$ 0.99 0.99 0.99 $\beta_2$ 0.999 0.999 0.999 optimizer AdamW AdamW AdamW Dataset Used B,C A A,B,C,D,E Loaded Parameters 8B 6.4B 10.4B Trainable Parameters 8B 2.4B 10.4B Und. resolution 384 $\times \{(1,3),(2,2)\}$ 384 $\times \{(1,3),(2,2)\}$ 384 $\times \{(1,3),(2,2)\}$ Gen. resolution \- 256 $\rightarrow$ 512 $\rightarrow$ 1024 1024 Semantic Encoder Trainable Not Loaded Trainable VQ Encoder Not Loaded Loaded Loaded Gen. Branch Not Loaded Trainable Trainable Und. Branch Trainable Partially Loaded Trainable : **Training configurations across three stages.** We use letters A-E to represent different dataset following Section `\ref{sec:appendix-dataset}`{=latex}. {#tab:training-stages} ## Ablation Studies on Elastic-MoT Design {#sec:appendix-ablation-mot} In this section, we report ablation results of the Elastic-MoT design, including the size of the generation branch and the number of joint attention layers. **Size of Generation Branch.** We report the performance of `\ours`{=latex} with different sizes of the generation branch during text-to-image pretraining (Stage 2) in Table `\ref{tab:model_scaling}`{=latex}. The results are obtained after 50k training steps with a global batch size of 1024. We also document the maximum per-GPU batch size, the gradient accumulation steps, and the training latency to measure efficiency. The results show that models of different sizes achieve comparable performance after 50k steps. Smaller models (1B, 2B) converge slightly faster and achieve marginally higher performance than larger models (4B, 8B). Larger models (4B, 8B) are harder to optimize as they need more steps, data, and tuning to realize their full capacity. In terms of latency, smaller models are considerably faster. The 2B model achieves the best balance between performance and efficiency, attaining the highest GenEval and DPG scores while being $3.17\times$ faster. **Architecture** **Performance** **Efficiency** ----------------------------------------------------------------- ----------------- -------------------- ---------------- ------------ ------------- ---------------------------- 0-1 `\cmidrule`{=latex}(lr)3-4 `\cmidrule`{=latex}(lr)5-7 Parm. Hidden Size GenEval $\uparrow$ DPG $\uparrow$ Batch Size Accum. Step Latency (s/it)$\downarrow$ 4B+1B 1536 0.56 60.8 16 1 **1.98** 4B+2B 2048 **0.57** **63.1** 16 1 3.67 4B+4B 3072 0.48 55.3 8 2 8.42 4B+8B 4096 0.55 58.6 8 2 11.64 : **Comparison of different model sizes on GenEval, DPG, and training efficiency.** We report the performance of `\ours`{=latex} with different sizes of the generation branch during the text-to-image pretraining (Stage-2) after 50k training steps. We also report the per-GPU batch size and training latency. {#tab:model_scaling} **Number of Joint Attention Layers.** To study the effect of varying the number of joint attention layers, we conducted two ablation experiments. The first experiment was performed during Stage 2 pretraining. We started with the Stage 1 checkpoint with $N=32$ layers in the understanding branch and fixed the generation branch size to 4B. We then varied $M$, the number of layers with joint attention, among {8,16,24,32}. The number of non-joint layers, $K$, is automatically determined by $K=N-M$. The results after 100k training steps are shown in Table `\ref{tab:ablation-joint-attention}`{=latex}. Among the four choices, $M=\{16,24,32\}$ yields a comparable performance, while $M=8$ shows a substantial drop. This suggests that a sufficient number of joint attention layers is necessary for strong text-to-image performance, but additional layers beyond a threshold provide little benefit. Training latency also decreases when $M$ is smaller (i.e., larger $K$), as fewer joint layers must be loaded. $M=16$ achieves the best balance of speed and performance. \begin{tabular}{ccccHcccc} \toprule \textbf{M} & \textbf{K} & \textbf{Pretraining (Stage-2)} & \textbf{SFT (Stage-3)} & \textbf{Efficiency}\\ \cmidrule(lr){3-4} \cmidrule(lr){6-8} \cmidrule(lr){9-9} & & GenEval $\uparrow$ & DPG $\uparrow$ & \textbf{FID} & GenEval $\uparrow$ & DPG$\uparrow$ & ImageEdit$\uparrow$ & Latency (s/it)$\downarrow$ \\ \midrule 8 & 24 & 0.57 & 69.3 & - & - & - & - & \textbf{2.45} \\ 16 & 16 & \textbf{0.63} & \textbf{75.0} & - & \textbf{0.89} & 83.2 & \textbf{3.66} & 3.67 \\ 24 & 8 & 0.63 & 73.3 & - & 0.81 & 83.0 & 3.60 & 4.12 \\ 32 & 0 & 0.61 & 71.2 & - & 0.85 & \textbf{83.2} & 3.55 & 5.20 \\ \bottomrule \end{tabular} We conducted a second experiment in Stage 3, where interleaved generation and editing tasks may benefit more from joint attention. Starting from a Stage 2 checkpoint pretrained with $M=16$ layers for 400k steps, we trained for 50k steps under $M=\{16,24,32\}$. The results show two key observations: (1) text-to-image tasks converge faster than image-editing tasks, reaching near-final performance after 50k steps, while editing tasks lag behind; (2) increasing $M$ does not significantly improve performance, even for interleaved editing. This may be due to token interference at later layers or the Stage 2 model being optimized with only 16 joint layers. Due to compute constraints, we were unable to retrain Stage 2 with alternative values of $M$. Nevertheless, keeping $M=16$ is a reasonable choice given our setup. Finally, in Stage 3 the efficiency difference is less pronounced, since all 10.4B parameters must be loaded for interleaved training and inference. **Weight Initialization.** We initialized the 2.4B generation branch with truncated weights from the understanding branch (Section `\ref{sec:appendix-elastic-mot-architecture}`{=latex}). We also explored initializing from scratch. Figure `\ref{fig:init-vs-scratch}`{=latex} shows the validation loss during the first 20k steps of Stage 2. Truncated initialization converges faster and yields lower loss.
Effect of truncated initialization. Validation loss comparison of truncated initialization vs. training from scratch during Stage 2. Truncated initialization converges faster and achieves lower loss.
## Ablation Studies on Stratified Sampling {#sec:ablation-sampling} We compared image generation quality under different sampling strategies on the MJHQ-30k dataset [@li2024playground] with 64 sampling steps. We evaluate the proposed stratified sampler against confidence-based sampling [@chang2022maskgit], uniform random sampling, and Halton sampling [@besnier2025halton]. The results are reported in Table `\ref{tab:appendix-sampling}`{=latex}. The stratified sampler achieves the best performance. **Method** **FID-30k** $\downarrow$ ------------ -------------------------- Confidence 11.42 Uniform 8.22 Halton 7.38 Stratified **6.68** : **Performance of Different Samplers in Text-to-Image Generation Tasks.** We report the FID scores on MJHQ-30K dataset using different samplers. The proposed stratified sampler achieves the best outcome. {#tab:appendix-sampling} ## Ablation Studies on Data Pipeline {#sec:ablation-schedular} **Effect of Task Scheduler.** To study the effect of the dataset scheduler described in Section `\ref{sec:appendix-dataset}`{=latex}, we compare three dataset mixing strategies in Stage 1 training. The goal of Stage 1 is to equip LaViDa with region-level understanding capabilities such as grounding. At this stage, training data includes fewer than 1M grounding samples but over 10M image-level understanding samples. To mitigate imbalance, we employ a scheduler that dynamically adjusts the sampling weights for new (grounding) and existing (image-level) capabilities. Each batch is drawn from a single dataset. For example, when New:Old=1:3, on average $\tfrac{1}{4}$ of batches contain grounding data and $\tfrac{3}{4}$ contain image-level data. We initialize the ratio as New:Old=3:1 and gradually reduce it to 1:3. We compare against fixed ratios of 1:3 and 3:1, reporting results after 20k steps in Table `\ref{tab:ablation-schedule-data}`{=latex}. Fixing New:Old=1:3 under-trains grounding, while fixing New:Old=3:1 improves grounding but causes forgetting on image-level understanding. In contrast, the dynamic scheduler achieves strong performance on both. Notably, it even outperforms the fixed 3:1 setup on image-level understanding, suggesting it also mitigates overfitting caused by the small grounding dataset. **Method** **New Capabilities** **Existing Capabilities** -------------------------------- ---------------------- --------------------------- ---------- --------- ---------- ----------- 2-4 `\cmidrule`{=latex}(lr)5-7 RefCOCO RefCOCO+ RefCOCOg MME ChartQA ScienceQA New:Old = 1:3 83.2 74.6 78.3 **449** 72.6 84.3 New:Old = 3:1 88.8 82.4 85.7 349 65.0 75.8 Dynamic **92.0** **86.9** **89.3** 436 **73.4** **86.4** : **Comparison of different task scheduling during Stage 1 Training.** We compare the performance under different dataset sampling weights of new capabilities (grounding) and old capabilities (image-level understanding). We explored two fixed sampling ratio 1:3 and 3:1 for New:Old. For the dynamic scheduler, the New:Old ratio is initialized as 3:1 and gradually decreased to 1:3. {#tab:ablation-schedule-data} **Does understanding data help generation tasks?** To examine whether incorporating understanding data benefits generation, we experimented with removing all grounding data from Stage 3. The results are shown in Table `\ref{tab:appendix-grounding-abl}`{=latex}. Even without explicit planning, incorporating grounding data enhances both text-to-image generation and editing, highlighting an inherent synergy between the tasks. When planning is enabled, these benefits compound, leading to even greater improvements. **Method** **GenEval** **DPG** **ImgEdit** -------------------- ------------- ---------- ------------- w/o grounding data 0.74 82.0 3.60 w/ grounding data 0.77 81.8 3.71 \+ planning **0.85** **82.9** **3.80** : **Effect of Grounding Data in Stage 3 Training.** To analyze the impact of the synergy between understanding and generation tasks, we explored removing object grounding in Stage 3 Training. This leads to worse overall performance. This demonstrates that jointly training on both understanding (grounding) and generation tasks is helpful for generation. {#tab:appendix-grounding-abl} ## Ablation Studies on Reflection and Planning {#sec:appendix-ablation-reflection} **Breakdown of Performance Improvements.** We provide a detailed breakdown of the gains introduced by planning and reflection. Table `\ref{tab:t2i-planning-breakdown}`{=latex} shows results on GenEval. Planning yields large improvements in object positioning (+0.19), while reflection additionally improves counting and attribution. To further examine the behavior of planning, we conducted additional text-to-image evaluations on categories of T2I-Compbench++ [@huang2025t2i] that are not covered by GenEval benchmark, including 3D spatial constraints and object texture attribution. We report these results in Table `\ref{tab:t2i-comp-bench}`{=latex}. We observe that `\ours`{=latex}  consistently demonstrate strong performance, with planning mechanism offering a significant performance boost. Notably, while our planning process use only 2D bounding boxes, we observe that it also improves satisfaction of 3D positional constraints by properly designing the size of relevant objects to reflect the distance. On Image-Edit (Table `\ref{tab:image-edit-breakdown}`{=latex}), planning improves adding/removing objects, subject actions, and hybrid instructions. The largest gains are in action (+0.50) and hybrid (+0.16). However, global edits (e.g., style, background) degrade slightly, as these tasks are less aligned with grounding. A promising direction for future is to let the model dynamically decide whether to invoke planning. **Single** **Two** **Position** **Counting** **Color** **Attribution** **Overall** ----------------------- ----------------------------- ----------------------------- ----------------------------- ----------------------------- ----------------------------- ----------------------------- ----------------------------- Baseline 0.99 0.85 0.65 0.71 0.86 0.58 0.77 +Planning 0.99 0.94 0.84 0.75 0.90 0.68 0.85 $\Delta$ vs. Baseline `\textgreen{=}`{=latex} `\textgreen{+0.09}`{=latex} `\textgreen{+0.19}`{=latex} `\textgreen{+0.04}`{=latex} `\textgreen{+0.04}`{=latex} `\textgreen{+0.10}`{=latex} `\textgreen{+0.08}`{=latex} +Reflection 1.00 0.95 0.89 0.85 0.90 0.74 0.89 $\Delta$ vs. Baseline `\textgreen{+0.01}`{=latex} `\textgreen{+0.10}`{=latex} `\textgreen{+0.24}`{=latex} `\textgreen{+0.14}`{=latex} `\textgreen{+0.04}`{=latex} `\textgreen{+0.16}`{=latex} `\textgreen{+0.12}`{=latex} : **Breakdown of performance improvements on GenEval Dataset.** We report the improvements of the planning and reflection mechanism on each category of the text-to-image generation tasks from GenEval Dataset. {#tab:t2i-planning-breakdown} **Model** **3D** **2D** **Texture** ------------------------------------------------------- ----------- ----------- ------------- Stable Diffusion 2 [@StabilityAI2022StableDiffusion2] 0.323 0.134 0.492 Janus-Pro-7B[@chen2025janus] 0.323 0.157 0.407 FLUX.1 Dev[@flux2024] 0.387 0.286 0.692 LaViDa-O 0.414 0.388 0.613 +Planning **0.442** **0.390** **0.715** : **Additional Text-to-Generation results on T2I-Compbench-++ Benchmark**. We report results on categories not included in GenEval benchmark, such as 3D spatial constraints and texture attribution. {#tab:t2i-comp-bench} \resizebox{1.0\linewidth}{!}{ \begin{tabular}{lcccccccccc} \toprule \textbf{Model} & \textbf{Add} & \textbf{Adjust} & \textbf{Extract} & \textbf{Replace} & \textbf{Remove} & \textbf{Background} & \textbf{Style} & \textbf{Hybrid} & \textbf{Action} & \textbf{Overall} \\ \midrule % \rowcolor{gray!20} Baseline & 4.04 &3.62 &2.01& 4.39 &3.98 &4.06 &4.82& 2.94 & 3.54& 3.71 \\ \midrule % \rowcolor{gray!20} + Planning & 4.11 & 3.67& 2.04 & 4.40& 4.05 &4.00& 4.75 & 3.10 &4.04& 3.80 \\ $\Delta$ vs. Baseline & \textgreen{+0.07} & \textgreen{+0.05} & \textgreen{+0.03} & \textgreen{+0.01} & \textgreen{+0.07} & \textred{-0.06} & \textred{-0.07} & \textgreen{+0.16} & \textgreen{+0.50} & \textgreen{+0.09} \\ \hline \end{tabular} } \begin{tabular}{ccccHcHcHcHc} \toprule Num. of Reflection Rrounds& \textbf{N=1} & \textbf{N=2} &\textbf{ N=4} & \textbf{N=6} & \textbf{N=8} & \textbf{N=10} & \textbf{N=12} & \textbf{N=14} & \textbf{N=16} & \textbf{N=18} & \textbf{N=20 }\\ \midrule GenEval Score $\uparrow$ &0.848 & 0.864 & 0.875 & 0.879 & 0.882 & 0.882 & 0.890 & 0.884 & 0.886 & 0.886 & 0.886 \\ \midrule Latency (s/image) $\downarrow$ & 27.2 & 32.6 & 39.3 & 43.6 & 47.1 & 50.4 & 53.4 & 56.0 & 58.3 & 60.4 & 62.2\\ \bottomrule \end{tabular} **Effect of Inference-time Scaling.** We evaluate reflection scaling by varying $N$, the maximum number of images generated per prompt. Table `\ref{tab:reflection-scaling}`{=latex} shows results. Even one reflection step ($N=2$) improves performance. Gains saturate at $N=8$, with little benefit beyond. Latency grows sublinearly with $N$ since simple prompts often trigger early stopping. For example, when $N=20$, the model may obtain a satisfactory output and terminate the generation process after generating just two images. ## Ablation Studies on Universal Text Conditioning. To examine the effectiveness of universal conditioning, we perform a user study and ask human evaluators to compare image generation with and without universal text conditioning. We curated a total of 300 human response on image pairs generated with randomly selected prompts from MJHQ-30k dataset. The human evaluators are provided with the following instruction: ::: center \resizebox{0.99\columnwidth}{!}{ \begin{tcolorbox} \textbf{Instructions} Both of these images were generated by AI models trained to create an image from a text prompt. Which image do you prefer given the associated text? Example criteria could include: detail, art quality, aesthetics, how well the text prompt is reflected, lack of distortions/irregularities (e.g. extra limbs, objects). In general, choose which image you think you would consider to be "better". \end{tcolorbox} } ::: We report the results in Figure `\ref{fig:human_eval_univ}`{=latex}. Results show that human evaluators exhibit a strong preference towards images generated with universal text conditioning, suggesting that conditioning the image generation with quality scores through our proposed universal text conditions method can effectively improve image quality. ![**Human Evaluation of Image Quality.** We conduct user study on image quality and compare text-to-image generations with and without universal text conditioning. Results show that universal text conditioning with aesthetic scores greatly improves image quality. ](figures/human_eval_abl-2.png){#fig:human_eval_univ width="50%"} ## Speed--Quality Tradeoff {#sec:appendix-speed-quality-tradeoff} A key advantage of masked diffusion models over autoregressive models is the speed--quality tradeoff enabled by parallel decoding. We study this in the unified setting by evaluating `\ours`{=latex} on MJHQ-30k text-to-image generation [@li2024playground], RefCOCO grounding [@kazemzadeh2014referitgame], and MathVista reasoning [@lu2023mathvista]. For MJHQ and RefCOCO, we vary the number of diffusion steps. For MathVista, we employ Fast-DLLM [@wu2025fast], which adaptively unmasks multiple tokens per step. The tradeoff is controlled via its threshold hyperparameter. Results are shown in Figure `\ref{fig:spped-appendix-2}`{=latex}. For MJHQ we report FID (lower is better), for RefCOCO Precision@0.5 (higher is better), and for MathVista accuracy (higher is better). We compare against several baselines: Flux [@flux2024] on T2I, Qwen2.5-VL-7B [@bai2025qwen25-vl] on grounding, and Qwen2.5-VL/Open-LLaVA-Next-8B [@chen2024open] on reasoning. `\ours`{=latex} achieves faster inference and stronger quality on image generation and grounding. For grounding, it reaches up to $6.8\times$ speedup while surpassing Qwen2.5-VL-7B in precision. On MathVista, while less accurate than state-of-the-art AR models, `\ours`{=latex} is much faster, and still stronger than popular AR baselines such as Open-LLaVA-Next-8B. Performance also exceeds the base LaViDa (56.9 vs. 44.8).
Speed–quality tradeoff on generation, grounding, and reasoning. Latency (s/sample) and benchmark scores are shown. For MJHQ: FID (lower is better). For RefCOCO: Precision@0.5 (higher is better). For MathVista: accuracy (higher is better). On MathVista, the maximum generation length is capped at 256 tokens.
\ifarxiv \else ## Additional Qualitative Results Finally, we provide additional qualitative examples demonstrating `\ours`{=latex}'s capabilities on diverse prompts and editing instructions. Figure `\ref{fig:demo-t2i}`{=latex} shows text-to-image generation, and Figure `\ref{fig:demo-editing}`{=latex} shows image editing results. `\fi`{=latex} # Compute Cost All experiments are conducted on 8 nodes, each equipped with 8 A100 GPUs. The total training amounts to 34.2 days measured by wall clock time, or 53k GPU hours. # Limitations In this section, we discuss several limitations of `\ours`{=latex}. **Text Rendering.** Since the image generation branch is trained from scratch and we did not explicitly include datasets for text rendering, `\ours`{=latex} 's capability to render and edit text is very limited. We also find that the VQ image tokenizer we use cannot faithfully reconstruct small texts. We aim to address this issue in future work by incorporating additional text rendering data and finetune the VQ image tokenizer on screenshots of documents. **Pixel Shift.** Our image editing datasets, such as GPT-Image-Edit-1.5M [@wang2025gpt] contains images distilled from generative models like GPT-4o, which is known to have \`\`pixel shift\" problems. Specifically, even if the instruction only requires editing a specific region, the other regions may still experience small but noticeable changes. As a consequence, `\ours`{=latex} inherit this problem. We aim to mitigate this by obtaining more clean and high-quality image-editing data. **Math Reasoning.** The focus of `\ours`{=latex} is to build a unified multi-modal MDMs capable of both understanding and generation tasks. Although its math reasoning capabilities has improved from the base model LaViDa thanks to additional training, there remains a considerable gap when compared against state-of-the-art models. We leave further improvements on math reasoning tasks to future work. **Hallucination.** Like all generative models, ours may occasionally produce inaccurate or fabricated information. We recommend using model outputs as guidance rather than unquestioned truth, and validating them where accuracy is critical. # Boarder Impact `\ours`{=latex} has strong text-to-image generation capabilities and image-editing capabilities, which may be abused to create various harmful and offensive content. We strongly caution the community against such use cases. Additonally, our model may inherit the biases embedded in the base model LaViDa, as well as biases incorporated in the images and texts of the training data. Our model is intended to be used by researchers to build a strong diffusion model for multi-modal applications and explore methods of building future multi-modal foundational models. We do not recommend that it be used for any other purposes. \ifarxiv \else ![**Qualitative examples of text-to-image generation.** We provide additional examples of text-to-image generation outputs on diverse prompts. ](figures/T2iDemo-5.png){#fig:demo-t2i width="100%"} ![**Qualitative examples of image editing.** We provide additional examples of image editing outputs on diverse instructions.](figures/EditingDemo-5.png){#fig:demo-editing width="100%"} \fi # LLM Usage We use LLM to correct typos and grammatical errors only.