LumiOpen/mdlm-en-1.7b-sft

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

Recommended defaults

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

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)

IFEval — instruction following (541 prompts, programmatic scoring)

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)

Generation quality (10 prompts, human evaluation)

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

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}
}