DiBO-AntMorphology

Final task-specific DiBO model for AntMorphology-Exact-v0, released with Training Diffusion Language Models for Black-Box Optimization (ICML 2026 Spotlight). See also the Hugging Face paper page and the DiBO code repository.

This model completed domain adaptation (DA), supervised fine-tuning (SFT), and reinforcement learning (RL).

Available model formats

This repository provides the same final task-specific DiBO model in two formats.

Original PyTorch checkpoint. dibo_ant_morphology_final.pt is the canonical paper-faithful checkpoint produced by the DiBO training pipeline. It stores the state dictionary under the model key and is loaded through the DiBO codebase on top of the pinned LLaDA base revision.
Transformers/safetensors export. The root-level config, tokenizer, custom modeling code, and sharded safetensors files are a validated convenience export derived deterministically from the original checkpoint. They load directly with AutoModel.from_pretrained(...).

The safetensors model was not trained separately. The LLaDA base weights are not duplicated in this repository.

A. Load the standard Transformers export

from transformers import AutoModel, AutoTokenizer

repo_id = "zpointsun/DiBO-AntMorphology"
tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    revision="v1.0.7",
    trust_remote_code=True,
)
model = AutoModel.from_pretrained(
    repo_id,
    revision="v1.0.7",
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype="auto",
)
model.eval()

The packaged tokenizer already includes the four DiBO delimiter tokens. Do not add them or resize embeddings again after loading this export.

The tokenizer configuration retains LLaDA's chat_template metadata, but DiBO does not call apply_chat_template during training or evaluation. DiBO directly tokenizes its rendered unified prompt-response corpus with the delimiter tokens above; do not insert chat headers when reproducing the released evaluation path.

B. Download and load the original checkpoint

The original artifact uses the released DiBO loader, which initializes the pinned LLaDA base, adds the four delimiter tokens, resizes the input embedding, and strictly loads checkpoint["model"].

hf download zpointsun/DiBO-AntMorphology dibo_ant_morphology_final.pt \
  --revision v1.0.7 --local-dir checkpoints/dibo-antmorphology

import torch
from huggingface_hub import hf_hub_download
from src.model.dllm import DEFAULT_MODEL_ID, LLADA_MODEL_REVISION, load_model_and_tokenizer

assert DEFAULT_MODEL_ID == "GSAI-ML/LLaDA-8B-Instruct"
assert LLADA_MODEL_REVISION == "08b83a6feb34df1a6011b80c3c00c7563e963b07"
checkpoint_path = hf_hub_download(
    "zpointsun/DiBO-AntMorphology",
    filename="dibo_ant_morphology_final.pt",
    revision="v1.0.7",
)
model, tokenizer = load_model_and_tokenizer(DEFAULT_MODEL_ID, device="cuda")
checkpoint = torch.load(checkpoint_path, map_location="cuda")
model.load_state_dict(checkpoint["model"], strict=True)
model.eval()

C. Evaluate either format

From a checkout of the released DiBO code and its oracle environment:

# Standard Transformers export
python eval.py --tasks AntMorphology-Exact-v0 \
  --model_name_or_path zpointsun/DiBO-AntMorphology --model_revision v1.0.7 \
  --seeds <SEEDS> --max_attempts 1000

# Canonical local .pt checkpoint
python eval.py --tasks AntMorphology-Exact-v0 \
  --checkpoint_path checkpoints/dibo-antmorphology/dibo_ant_morphology_final.pt \
  --seeds <SEEDS> --max_attempts 1000

Both choices share the same downstream DiBO evaluation path. Direct oracle evaluation requires the Design-Bench data cache and task dependencies described in the DiBO repository. For the exact Design-Bench snapshot used in the DiBO experiments, see DiBO-DesignBench-Snapshot.

Limitations

Practical inference requires a CUDA-capable PyTorch environment. These task-specific models are designed for DiBO's masked-response generation and evaluation workflow; this release does not claim generic text-generation pipeline support. Loading a released final model is for evaluation or use and does not reproduce the DA/SFT/RL training process.

Other DiBO task models

Citation

If you find DiBO helpful, please cite:

@article{sun2026training,
  title={Training diffusion language models for black-box optimization},
  author={Sun, Zipeng and Chen, Can and Yuan, Ye and Wu, Haolun and Gu, Jiayao and Pal, Christopher and Liu, Xue},
  journal={arXiv preprint arXiv:2603.17919},
  year={2026}
}