Update README.md

README.md

@@ -1,3 +1,108 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+datasets:
+- Zigeng/DMax-LLaDA-2.0-Mini-Math-Trajectories
+base_model:
+- inclusionAI/LLaDA2.0-mini
+---
+
+<div align="center">
+<h1>🚀 DMax: Aggressive Parallel Decoding for dLLMs</h1>
+  <div align="center">
+  <a href="https://github.com/czg1225/DMax/blob/main/LICENSE">
+    <img alt="Apache" src="https://img.shields.io/badge/License-Apache-4E94CE.svg">
+  </a>
+  <a href="https://github.com/czg1225/DMax">
+    <img src="https://img.shields.io/badge/Paper-Arxiv-darkred.svg" alt="Paper">
+  </a>
+  <a href="https://github.com/czg1225/DMax">
+    <img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub">
+  </a>
+</div>
+</div>
+
+> **DMax: Aggressive Parallel Decoding for dLLMs**  
+> [Zigeng Chen](https://czg1225.github.io/chenzigeng99/), [Gongfan Fang](https://fangggf.github.io/), [Xinyin Ma](https://horseee.github.io/), [Ruonan Yu](https://scholar.google.com/citations?user=UHP95egAAAAJ&hl=en), [Xinchao Wang](https://sites.google.com/site/sitexinchaowang/)  
+> [xML Lab](https://sites.google.com/view/xml-nus), National University of Singapore
+
+
+## 💪 Highlights
+
+- **Aggressive Decoding Parallelism**: Achieves 6.0 TPF on math and reasoning tasks and 6.6 TPF on code tasks while preserving accuracy.
+- **Self-Revising dLLM**: Extends a pretrained MDLM into a UDLM with an intrinsic ability to revise its own erroneous predictions during decoding.
+- **Soft Parallel Decoding**: Uses interpolation between mask and token embeddings to propagate confidence priors from previous steps.
+
+<div align="center">
+  <img src="assets/tradeoff.png" width="90%" />
+  <br>
+  <em>Superior Parallelism-Accuracy Trade-off, Increased TPF with Maintained Accuracy.</em>
+</div>
+
+## 💡 Introduction
+
+We present DMax, a new paradigm for efficient dLLMs. It mitigates error accumulation in parallel decoding, enabling aggressive decoding parallelism while preserving generation quality. Unlike conventional masked dLLMs that decode through a binary mask-to-token transition, DMax reformulates decoding as a progressive self-refinement from mask embeddings to token embeddings. At the core of our approach is On-Policy Uniform Training, a novel training strategy that efficiently unifies masked and uniform dLLMs, equipping the model to recover clean tokens from both masked inputs and its own erroneous predictions. Building on this foundation, we further intoduce Soft Parallel Decoding. Extensive experiments across a variety of benchmarks demonstrate the effectiveness of DMax.
+
+<!-- ![figure](assets/intro.png) -->
+<div align="center">
+  <img src="assets/train.png" width="100%" />
+  <br>
+  <em>Overview of the On-Policy Uniform Training.</em>
+</div>
+
+## 💻 Model and Datasets
+
+| Model | Description | Source Model | Link |
+| --- | --- | --- | --- |
+| 🤖 DMax-Math-16B | Highly parallel dLLM for math and reasoning. | LLaDA-2.0-mini | [Hugging Face](https://huggingface.co/Zigeng/DMax-Math-16B) |
+| 🤖 DMax-Coder-16B | Highly parallel dLLM for code generation. | LLaDA-2.0-mini | [Hugging Face](https://huggingface.co/Zigeng/DMax-Coder-16B) |
+
+| Dataset | Description | Link |
+| --- | --- | --- |
+| 📊 DMax-Math-Training-Data | Trajectories on math problems generated by LLaDA-2.0-mini | [Hugging Face](https://huggingface.co/datasets/Zigeng/DMax-LLaDA-2.0-Mini-Math-Trajectories) |
+| 📊 DMax-Code-Training-Data | Trajectories on code problems generated by LLaDA-2.0-mini | [Hugging Face](https://huggingface.co/datasets/Zigeng/DMax-LLaDA-2.0-Mini-Code-Trajectories) |
+
+## 🚀 Quick Start
+
+```python
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "Zigeng/DMax-Math-16B", trust_remote_code=True, device_map="cuda:0"
+)
+model = model.to(torch.bfloat16)
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained("Zigeng/DMax-Math-16B", trust_remote_code=True)
+
+prompt = "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?" + "\nLet's think step by step\n"
+
+input_ids = tokenizer.apply_chat_template(
+    [{"role": "user", "content": prompt}],
+    add_generation_prompt=True,
+    tokenize=True,
+    return_tensors="pt",
+)
+
+nfe, generated_tokens = model.generate_spd(
+    inputs=input_ids,
+    gen_length=2048,
+    block_length=32,
+    threshold=0.0,
+)
+
+generated_answer = tokenizer.decode(
+    generated_tokens[0],
+    skip_special_tokens=True,
+)
+
+print(generated_answer)
+print("nfe:",nfe,"token length",len(generated_tokens[0]))
+```
+
+## 📖 Experimental Results
+
+![trade-off](assets/exp.png)
+
+---
+