CDLM-0.5B

Model Description

CDLM-0.5B is a fine-tuned version of fredzzp/open-dcoder-0.5B, trained using error-aware training with the mixture objective proposed in our paper on Corrective Diffusion Language Models. This model is designed to improve error-aware confidence and targeted refinement capabilities in code generation tasks.

Key Features

Base Model: fredzzp/open-dcoder-0.5B (a masked diffusion language model based on Qwen2)
Training Method: Error-aware training with mixture objective that explicitly supervises visible incorrect tokens
Architecture: Masked Diffusion Language Model (MDLM)
Parameters: ~0.5B

Training Details

This model was fine-tuned from fredzzp/open-dcoder-0.5B using error-aware training with a mixture objective. For detailed information on the training methodology, please refer to our paper: Corrective Diffusion Language Models.

Usage

Installation

pip install torch transformers

Code Generation

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Shuibai12138/CDLM-0.5B"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
).to(device)

# Generate code
prompt = "def fibonacci(n):"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

# Use diffusion generation
outputs = model.diffusion_generate(
    inputs=input_ids,
    max_new_tokens=100,
    steps=16,
    temperature=0.8
)

prompt_len = input_ids.shape[1]
generated_text = tokenizer.decode(outputs.sequences[0][prompt_len:], skip_special_tokens=True)

print("Generated Code:")
print(generated_text)

Note: This model uses a custom diffusion_generate method, so trust_remote_code=True is required when loading the model.

Iterative Refinement

The model supports iterative refinement for code correction. See the CDLM repository for examples of using the model for code correction tasks.

Citation

If you use this model in your research, please cite:

@misc{zhang2025correctivediffusionlanguagemodels,
      title={Corrective Diffusion Language Models}, 
      author={Shuibai Zhang and Fred Zhangzhi Peng and Yiheng Zhang and Jin Pan and Grigorios G. Chrysos},
      year={2025},
      eprint={2512.15596},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2512.15596}, 
}