Hengchang-Liu/D3LM-from-nt

D3LM: A Discrete DNA Diffusion Language Model for Bidirectional DNA Understanding and Generation

This repository contains the model presented in D3LM: A Discrete DNA Diffusion Language Model for Bidirectional DNA Understanding and Generation.

A masked diffusion language model for unconditional mammalian DNA sequence generation, built on the ESM encoder with Rotary Positional Embeddings.

Initialization: pretrained from Nucleotide Transformer, then fine-tuned with masked diffusion objective on mammalian DNA.

Parameters	Hidden	Layers	Heads	Max Length	Vocab
~50M	512	12	16	2,048	4,107

Usage

pip install transformers torch tqdm

import sys
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

model_name = "Hengchang-Liu/D3LM-from-nt"
model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True).eval()
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Import MDMGenerationConfig from the model's auto-downloaded module
MDMGenerationConfig = getattr(sys.modules[type(model).__module__], "MDMGenerationConfig")

# Unconditional generation: create a fully-masked prompt of desired length
length = 200
input_ids = torch.full((1, length), tokenizer.mask_token_id, dtype=torch.long)

config = MDMGenerationConfig(
    mask_token_id=tokenizer.mask_token_id,
    max_length=length,
    steps=50,
    temperature=1.0,
    top_p=0.9,
    alg="random",
    num_return_sequences=4,
    return_dict_in_generate=True,
)

with torch.no_grad():
    outputs = model.diffusion_generate(inputs=input_ids, generation_config=config)

for i, seq in enumerate(outputs.sequences):
    print(f">{i}
{tokenizer.decode(seq, skip_special_tokens=True).replace(' ', '')}")

Generation Parameters

Parameter	Default	Description
steps	50	Diffusion denoising steps
temperature	1.0	Sampling temperature
top_p	0.9	Nucleus sampling cutoff
top_k	0	Top-k cutoff (0 = off)
alg	"random"	Unmasking order: random, entropy, maskgit_plus, topk_margin, origin, p2
alg_temp	0.9	Gumbel temperature for confidence ordering (0 = deterministic)

Citation

@misc{yang2026d3lmdiscretednadiffusion,
      title={D3LM: A Discrete DNA Diffusion Language Model for Bidirectional DNA Understanding and Generation}, 
      author={Zhao Yang and Hengchang Liu and Chuan Cao and Bing Su},
      year={2026},
      eprint={2603.01780},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.01780}, 
}

License

Apache 2.0