Update README.md

README.md

@@ -8,22 +8,20 @@ tags:
 - uniform-diffusion
 ---
 
+
+
 # Sumi-7B
 
-Sumi is a **7B uniform (discrete) diffusion language model**. Unlike an autoregressive
-LM, it does not predict the next token left-to-right. Instead it runs **full
-bidirectional attention** over the whole sequence and is trained to *denoise*: given a
-sequence whose tokens have been randomly corrupted to uniformly-sampled vocabulary
-tokens, it predicts the clean tokens at every position at once. Generation starts from a
-canvas of random tokens appended to the prompt and **iteratively denoises** it.
+<p align="left">
+  <a href="https://www.nlp.ecei.tohoku.ac.jp/projects/sumi/"><img src="https://img.shields.io/badge/%F0%9F%8C%90%20Project%20HP-1f72b8" alt="Project Page"></a>
+  <a href="https://arxiv.org/abs/2606.19005"><img src="https://img.shields.io/badge/arXiv-2606.19005-b31b1b?logo=arxiv&logoColor=white" alt="arXiv"></a>
+</p>
+Sumi is a native uniform diffusion language model trained from scratch, so it runs full bidirectional attention and denoises a canvas of randomly corrupted tokens.
+We provide Sumi in a custom model class, therefore you need to set `trust_remote_code=True` to use it in transformers.
 
-- Architecture: Llama-style transformer (RoPE, SwiGLU, RMSNorm, GQA) run **non-causally**, with an **off-by-one ("attention-sink") softmax**.
-- 36 layers, hidden 4096, 32 heads / 8 KV heads, vocab 100,278, context 4,864.
-- `model_type: "sumi"`, loaded via `trust_remote_code=True`.
+We recommend `transformers==5.8.1`.
 
-> **Requirements:** `transformers >= 5.8` and `torch >= 2.4` (any recent GPU). The model
-> ships its own modeling code; pass `trust_remote_code=True`. The tokenizer is bundled in
-> this repo.
+For more details, please refer to [our project page](https://www.nlp.ecei.tohoku.ac.jp/projects/sumi/) and [technical report](https://arxiv.org/abs/2606.19005).
 
 ## Quickstart
 
@@ -31,104 +29,37 @@ canvas of random tokens appended to the prompt and **iteratively denoises** it.
 import torch
 from transformers import AutoModelForMaskGeneration, AutoTokenizer
 
-model_id = "tohoku-nlp/sumi-7b"  # <-- this repo
-
+model_id = "tohoku-nlp/sumi-7b"
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForMaskGeneration.from_pretrained(
     model_id, trust_remote_code=True, dtype=torch.bfloat16
 ).to("cuda").eval()
 
-prompt = "The capital of France is"
+prompt = "Our journey into exploring diffusion language model begins,"
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
 out = model.generate(
     **inputs,
-    max_new_tokens=64,        # how many tokens to generate (denoised in place)
-    num_denoising_steps=128,  # refinement iterations — set higher for better quality
-    temperature=0.7,          # sampling temperature
+    max_new_tokens=256,       # content budget; the EOS/BOS delimiter is anchored here
+    num_denoising_steps=64,   # refinement iterations — the main quality/compute dial
+    sampler="ancestral",      # "ancestral" (default) or "adaptive" (sharper, for code/math)
+    temperature=0.7,
 )
 print(tokenizer.decode(out.sequences[0], skip_special_tokens=True))
 ```
 
-`generate` returns a `SumiGenerationOutput` whose `.sequences` is `[batch, prompt_len +
-max_new_tokens]` (prompt followed by the generated completion). It is a drop-in custom
-generation loop — `AutoModelForMaskGeneration` also works and is the semantically exact
-auto-class.
-
-## Generation parameters
-
-The diffusion generation loop replaces the usual autoregressive decoding knobs. The
-parameters you will normally set:
-
-| parameter | default | description |
-|---|---|---|
-| `max_new_tokens` | — (required) | Length of the completion. These tokens are initialised at random and denoised; the prompt stays fixed. |
-| `num_denoising_steps` | `128` | Number of denoising iterations. More steps → better quality at higher cost; quality saturates after a couple hundred. **This is the main quality/compute dial.** |
-| `sampler` | `"ancestral"` | `"ancestral"` — stochastic sampling from the analytic diffusion posterior (good general-purpose default). `"adaptive"` — confidence-based: each step commits only the most-confident position(s) and revises the rest, which gives sharper, more coherent outputs on structured tasks (e.g. code, math). |
-| `temperature` | `1.0` | Sampling temperature. **`0.7` is the recommended canonical setting.** Lower → more decisive denoising. |
-| `seed` / `generator` | `None` | Set a `seed` (int) or pass a `torch.Generator` for reproducible samples. |
-
-```python
-# More coherent / "decisive" decoding, e.g. for code or math:
-out = model.generate(**inputs, max_new_tokens=256,
-                     num_denoising_steps=256, sampler="adaptive", temperature=0.7)
-```
-
-### Advanced knobs
-
-You rarely need these; defaults are fine for normal use.
-
-| parameter | default | description |
-|---|---|---|
-| `tokens_per_step` | `1` | `adaptive` sampler only: how many positions to commit per step. Keep at `1`. |
-| `schedule` | `"linear"` | `ancestral` only: log-SNR noise schedule, `"linear"` or `"cosine"`. |
-| `min_log_snr` / `max_log_snr` | `-9.0` / `9.0` | `ancestral` only: bounds of the SNR schedule. |
-| `frozen` | `None` | Pin `(position, token_id)` pairs at fixed canvas positions (kept fixed during denoising). Used for in-filling and for the `[EOS][BOS]` document-boundary **anchor** used in evaluation. |
-| `denoise_end` | `None` | Per-row exclusive upper bound; freezes the tail of the canvas so the step budget concentrates on the content window. |
-
-## How it works
-
-1. The prompt is placed at the front of a canvas of length `prompt_len + max_new_tokens`.
-2. The completion region is filled with **uniformly random** vocabulary tokens (there is
-   no `[MASK]` token — this is *uniform* diffusion).
-3. For `num_denoising_steps`, the model does one bidirectional forward over the whole
-   canvas and updates the completion region (the prompt is frozen). `ancestral` resamples
-   every completion position from the analytic posterior each step; `adaptive` commits the
-   highest-confidence position(s) and leaves the rest for later steps.
-4. The final denoised canvas is returned.
-
-There is no KV cache (every step is a full bidirectional pass), so cost scales with
-`num_denoising_steps × sequence_length²`.
-
-## Evaluation
-
-Sumi is scored with diffusion-native metrics: a Monte-Carlo estimate of the diffusion
-**NELBO** for multiple-choice / log-likelihood tasks, and iterative denoising for
-generative tasks. A companion [`lm-evaluation-harness`](https://github.com/EleutherAI/lm-evaluation-harness)
-plugin (`sumi-eval`, `--model sumi`) implements this end-to-end.
-
-Representative 0-shot results (the HF conversion reproduces the original training-stack
-model — ARC-Easy acc_norm agrees within sampling CI and the diffusion NELBO is a byte-exact
-port):
-
-| task | metric | score |
-|---|---|---|
-| ARC-Easy | acc_norm | ≈ 0.69 |
-| HumanEval | pass@1 | 0.226 |
+`generate()` returns the trimmed completion in `out.sequences` and the full untrimmed canvas in `out.canvas`.
 
-(HumanEval generated with `sampler="adaptive"`, `num_denoising_steps=256`,
-`max_new_tokens=256`; scored on the raw model output, with no answer-extraction.)
 
 ## Citation
 
 ```bibtex
-@misc{ye2025sumi,
-      title={Sumi : Open Uniform Diffusion Language Model from Scratch}, 
-      author={Mengyu Ye and Keito Kudo and Wataru Ikeda and Ryosuke Matsuda and Keisuke Sakaguchi and Jun Suziki},
+@misc{ye2026sumi,
+      title={Sumi: Open Uniform Diffusion Language Model from Scratch}, 
+      author={Mengyu Ye and Keito Kudo and Wataru Ikeda and Ryosuke Matsuda and Keisuke Sakaguchi and Jun Suzuki},
       year={2026},
-      eprint={2606.xxxxx},
+      eprint={2606.19005},
       archivePrefix={arXiv},
-      primaryClass={cs.LG},
-      url={https://arxiv.org/abs/2606.xxxxx}, 
-}
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2606.19005}, 
+} 
 ```