Excited to share what we have worked on for the past 6 months.
Can we really call masked diffusion an "any-order" model?
In our new work, we argue that off-the-shelf masked diffusion models have any-order predictive power but fail to reason in a 'genuinely any-order fashion'.
We introduce two models that do: LatentMDM and 7B-scale FlexMDM, extending our prior work. LatentMDM, our new latent-space masked diffusion, outperforms AR (with KV caching) on TinyGSM, at a matched inference budget. To the best of our knowledge, LatentMDM is the first model on the same setup to do so.
Joint work with co-first authors
@seunggeun_kr,
@tklll_tx, and Yuyuan Chen, and advisors
@du_yilun,
@ShamKakade6, and
@sitanch.