MDLM-en-1.7b-SFT

Instruction-tuned version of LumiOpen/mdlm-en-1.7b. Fine-tuned on instruction-following and conversational data using supervised fine-tuning (SFT).

• Architecture: Bidirectional transformer, 1.7B parameters, loglinear noise schedule (SUBS parameterization)
• Base model: LumiOpen/mdlm-en-1.7b (pretrained on 10B FineWeb tokens)
• Equivalent AR model: A 1.7B autoregressive transformer trained on the same pretraining data and fine-tuned on the same SFT corpus

---

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "LumiOpen/mdlm-en-1.7b-sft",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")
tok = AutoTokenizer.from_pretrained("LumiOpen/mdlm-en-1.7b-sft")
model.eval()

prompt = "User: What is artificial intelligence?\n\nAssistant: "
ids = tok(prompt, return_tensors="pt", add_special_tokens=False)
ids = {k: v.to("cuda") for k, v in ids.items()}

out = model.mdlm_generate(
    ids["input_ids"],
    attention_mask=ids["attention_mask"],
    max_new_tokens=128,
    num_steps=64,
    temperature=0.8,
    top_p=0.9,
)
print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))

Multi-turn conversation

def build_prompt(turns):
    """turns: list of (role, text) where role is 'user' or 'assistant'."""
    parts = []
    for i, (role, text) in enumerate(turns):
        prefix = "User" if role == "user" else "Assistant"
        sep = "" if i == 0 else "\n\n"
        parts.append(f"{sep}{prefix}: {text.strip()}")
    parts.append("\n\nAssistant: ")
    return "".join(parts)

turns = [
    ("user", "What is machine learning?"),
    ("assistant", "Machine learning is a branch of AI where models learn patterns from data."),
    ("user", "Can you give me an example?"),
]
prompt = build_prompt(turns)
ids = tok(prompt, return_tensors="pt", add_special_tokens=False)
out = model.mdlm_generate(ids.input_ids, attention_mask=ids.attention_mask,
                     max_new_tokens=128, num_steps=64)
print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))

---

mdlm_generate() arguments

Argument	Type	Default	Description
input_ids	LongTensor (1, L)	required	Prompt token ids
attention_mask	LongTensor (1, L)	None (all ones)	1 for real tokens, 0 for padding. Always pass this — the model was trained with right-padding
max_new_tokens	int	128	Number of answer tokens to generate
num_steps	int	64	Denoising steps. More steps = slower but more coherent
temperature	float	0.8	Sampling temperature. Lower = more conservative
top_p	float	0.9	Nucleus sampling threshold

Recommended defaults

Use case	max_new_tokens	num_steps	temperature	top_p
Conversational reply	128	64	0.8	0.9
Long-form answer	256–512	128	0.8	0.9

T=0.8, top_p=0.9, steps=64 is the recommended default across all use cases. Lowering temperature to 0.5–0.6 tends to increase repetition on instruction prompts and does not improve casual prompt quality. Increasing steps to 128 does not reliably improve output quality.

attention_mask note: This model was fine-tuned with right-padding (EOS tokens on the right, prompt content starting at position 0). Always pass attention_mask from the tokenizer so padding tokens are correctly ignored.

Casual prompt note: Short conversational inputs ("Hello!", "What's up?") are handled by a 500-step synthetic data injection. Output quality is good but stochastic — some samples are coherent, others may drift. Running with a fixed seed will give reproducible results.

---

Training details

Stage	Steps	Data	LR
SFT — Dolci	20 000	Dolci-Instruct-SFT, English-filtered, Chat domain 3× oversampled	2e-5
Casual injection	500	507 synthetic examples (232 unique prompts) covering greetings, identity, gibberish handling, short inputs	5e-6

The casual injection step teaches the model to respond naturally to short conversational inputs ("Hi", "Who are you?", "help") which are absent from the instruction-only Dolci corpus.

---

Comparison to equivalent AR model

Both models share the same 1.7B architecture template, tokenizer, pretraining corpus, and SFT data. The only difference is the generation mechanism: MDLM is bidirectional and generates by iterative denoising; AR generates left-to-right.

MCQ benchmarks (300 items, chain-rule masked scoring)

Benchmark	AR-SFT	MDLM-SFT
ARC-Challenge	0.243	0.273
HellaSwag	0.317	0.450
TruthfulQA	0.260	0.280

IFEval — instruction following (541 prompts, programmatic scoring)

Model	Prompt strict	Prompt loose	Instruction strict	Instruction loose
AR-SFT	14.0%	15.7%	28.3%	28.7%
MDLM-SFT	22.2%	25.0%	35.5%	38.1%

MDLM-SFT leads AR-SFT by +8.2 pp on strict prompt-level instruction following.

BERTScore F1 (roberta-large, vs human references, 10 diverse prompts)

Model	Mean F1
AR-SFT	0.8444
MDLM-SFT	0.8590

Generation quality (10 prompts, human evaluation)

Model	Success	Partial	Failure
AR-SFT	1 / 10	2 / 10	7 / 10
MDLM-SFT	5 / 10	3 / 10	2 / 10

AR-SFT shows a notable refusal pattern (refuses diet tips, creative writing, technical explanations) and a repetition loop problem. MDLM-SFT attempts all prompts directly.

Speed

Setting	AR	MDLM
Full 1024-token generation (64 steps)	175 tok/s	790 tok/s (4.5×)
Short chat reply ~150 new tokens	169 tok/s	64 tok/s

MDLM's parallel generation is fastest when generating many new tokens relative to prompt length. For short interactive chat replies into a long prompt, AR has lower latency.

---

Sampler

Uses the loglinear SUBS ancestral sampler from MDLM: Simple and Effective Masked Diffusion Language Models with Gumbel-max sampling and nucleus (top-p) filtering.

---

Citation

@article{sahoo2024simple,
  title={Simple and Effective Masked Diffusion Language Models},
  author={Sahoo, Subham Sekhar and Arriola, Marianne and Schiff, Yair and Gokaslan, Aaron and Marroquin, Edgar and Chiu, Justin T and Rush, Alexander and Kuleshov, Volodymyr},
  journal={arXiv preprint arXiv:2406.07524},
  year={2024}
}