{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/simplified-and-generalized-masked-diffusion","title":"Simplified and Generalized Masked Diffusion for Discrete Data","arxiv_id":"2406.04329","date":"2024-06-06","proceeding":null,"authors":["Jiaxin Shi","Kehang Han","Zhe Wang","Arnaud Doucet","Michalis K. Titsias"],"abstract":"Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes. Our code is available at https://github.com/google-deepmind/md4.","url_abs":"https://arxiv.org/abs/2406.04329v4","url_pdf":"https://arxiv.org/pdf/2406.04329v4.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"simplified-and-generalized-masked-diffusion","repo_url":"https://github.com/google-deepmind/md4","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"jax","reach":null}],"tasks":[{"task_slug":"language-modeling","task_name":"Language Modeling"},{"task_slug":"language-modelling","task_name":"Language Modelling"}],"methods":[{"method_slug":"adam","method_name":"Adam"},{"method_slug":"attention","method_name":"Attention"},{"method_slug":"attention-dropout","method_name":"Attention Dropout"},{"method_slug":"bpe","method_name":"BPE"},{"method_slug":"cosine-annealing","method_name":"Cosine Annealing"},{"method_slug":"dense-connections","method_name":"Dense Connections"},{"method_slug":"diffusion","method_name":"Diffusion"},{"method_slug":"discriminative-fine-tuning","method_name":"Discriminative Fine-Tuning"},{"method_slug":"dropout","method_name":"Dropout"},{"method_slug":"gpt-2","method_name":"GPT-2"},{"method_slug":"hoc","method_name":"HOC"},{"method_slug":"layer-normalization","method_name":"Layer Normalization"},{"method_slug":"linear-layer","method_name":"Linear Layer"},{"method_slug":"linear-warmup-with-cosine-annealing","method_name":"Linear Warmup With Cosine Annealing"},{"method_slug":"multi-head-attention","method_name":"Multi-Head Attention"},{"method_slug":"residual-connection","method_name":"Residual Connection"},{"method_slug":"softmax","method_name":"Softmax"},{"method_slug":"weight-decay","method_name":"Weight Decay"}],"datasets_introduced":[],"methods_introduced":[],"results":[{"leaderboard":"/sota/language-modelling-on-openwebtext","task":"Language Modelling","dataset":"OpenWebText","model":"GenMD4","rank_in_archive_order":10,"of":12,"metrics":{"eval_perplexity":"21.80","parameters":"131M"},"uses_additional_data":false}],"syntology":{"syntology_url":null,"atlas_url":"https://app.syntology.ai/?focus=2406.04329","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2406.04329"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"deterministic:regex_extraction","url":"https://github.com/google-deepmind/md4","reach":null}],"summary":{"ran":1,"ran_fixture":3,"ran_draft_wrong":1,"ran_honours":1,"unverified":15},"by_repo_kind":{"official":{"samples":21,"ran":6,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"66a81c44aa3c341b","entry":"MaskingSchedule","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"66a81c44aa3c341b"}},{"code_sha256_prefix":"84db74ee09e0bda4","entry":"get_timestep_embedding","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"84db74ee09e0bda4"}},{"code_sha256_prefix":"b6960d9da81c946a","entry":"jax_unstack","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"b6960d9da81c946a"}},{"code_sha256_prefix":"8873a5201399b10b","entry":"precompute_freqs_cis","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran_honours","verification_level":1,"contract_check":"HONOURS","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8873a5201399b10b"}},{"code_sha256_prefix":"4a1c75a9a5f46c4f","entry":"repeat_kv","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4a1c75a9a5f46c4f"}},{"code_sha256_prefix":"0dc13509bbb03398","entry":"reshape_for_broadcast","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"0dc13509bbb03398"}},{"code_sha256_prefix":"b5d57c4d33c17740","entry":"Attention","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"b5d57c4d33c17740"}},{"code_sha256_prefix":"e341090d86d921b9","entry":"CondEmbedding","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"e341090d86d921b9"}},{"code_sha256_prefix":"b74decb1e729d273","entry":"DiscreteClassifier","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"b74decb1e729d273"}},{"code_sha256_prefix":"d8f8ff9226830df3","entry":"Dropout1d","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"d8f8ff9226830df3"}},{"code_sha256_prefix":"170abc34ff9d4f81","entry":"FeedForward","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"170abc34ff9d4f81"}},{"code_sha256_prefix":"4b01cbf6a3ec21a6","entry":"MD4","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4b01cbf6a3ec21a6"}},{"code_sha256_prefix":"be5a00527b2e82e0","entry":"ModelArgs","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"be5a00527b2e82e0"}},{"code_sha256_prefix":"a75e172eea6d85f1","entry":"RMSNorm","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a75e172eea6d85f1"}},{"code_sha256_prefix":"2db631aa126ebf6e","entry":"ResBlock","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2db631aa126ebf6e"}},{"code_sha256_prefix":"daf4ef755a1978a3","entry":"SelfAttention","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"daf4ef755a1978a3"}},{"code_sha256_prefix":"87a69858dd77c308","entry":"Transformer","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"87a69858dd77c308"}},{"code_sha256_prefix":"7a351fb6b29e26ac","entry":"TransformerBlock","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"7a351fb6b29e26ac"}},{"code_sha256_prefix":"35c08ff52d7e2f6b","entry":"UNet","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"35c08ff52d7e2f6b"}},{"code_sha256_prefix":"a2d820814fa59452","entry":"UNet5DWrapper","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a2d820814fa59452"}},{"code_sha256_prefix":"582256db3844c9de","entry":"apply_rotary_emb","repo":"google-deepmind/md4","repo_kind":"official","path":"md4/models/diffusion/md4.py","file_url":"https://github.com/google-deepmind/md4/blob/HEAD/md4/models/diffusion/md4.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"582256db3844c9de"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}