Discrete Markov Bridge

This is the code for the implementation of the Discrete Markov Bridge. For description and theory, refer to the paper by Hengli Li, Yuxuan Wang, Song-Chun Zhu, Ying Nian Wu, and Zilong Zheng.

Figure 1. Overview of the Discrete Markov Bridge framework.

Introduction

Discrete Markov Bridge (DMB) consists of two component: the Matrix-learning and the Score-learning. The Matrix-learning process is designed to learn an adaptive transition rate matrix, which facilitates the estimation of an adapted latent distribution. Concurrently, the score-learning process focuses on estimating the probability ratio necessary for constructing the inverse transition rate matrix, thereby enabling the reconstruction of the original data distribution.

Installation

conda create -n DMB python=3.10
conda activate DMB
pip3 install torch torchvision torchaudio
pip install transformers datasets tqdm accelerate
pip install wandb

Usage

cd src
sh scripts/example.sh

For understanding of the shell scripts, please check parse-file for description of the args.

The example shell scripts

CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" torchrun --nproc_per_node=8 --master_port=29501 main_ddp.py \
    --ngpus 8 \
    --sche_name "loglinear" \ # diffusion scheduler
    --sigma_min 1e-4 \
    --sigma_max 20 \
    --Q_lr 1e-3 \ # learning rate for matrix learning
    --Q_weight_decay 1e-2 \
    --Q_initialization 'gather' \
    --Q_epochs 15 \
    --vocab_size 27 \
    --score_epoch 10 \
    --hidden_size 768 \ # transformer args
    --time_hidden_size 128 \ # transformer args
    --dropout 0.1 \ # transformer args
    --n_blocks 12 \ # transformer args
    --n_heads 12 \ # transformer args
    --score_lr 3e-4 \ # score learning rate
    --score_warmup_steps 2500 \
    --mu_train_dataset_name 'text8' \ # training
    --mu_eval_dataset_name 'text8' \ # eval
    --mu_test_dataset_name 'text8' \ # test
    --cache_dir './cache/' \
    --seqlen 256 \
    --score_accum 1 \
    --score_train_batch_size 512 \
    --score_eval_batch_size 512 \
    --score_grad_clip 1.0 \
    --sample_batch_size 512 \
    --Q_accum 1 \
    --Q_train_batch_size 512 \
    --eval_times 1000 \
    --ema 0.9999 \
    --epoch 1000 \
    --run_name 'text8' \
    --random_seed 42 \

Citation

@misc{li2025discretemarkovbridge,
      title={Discrete Markov Bridge}, 
      author={Hengli Li and Yuxuan Wang and Song-Chun Zhu and Ying Nian Wu and Zilong Zheng},
      year={2025},
      eprint={2505.19752},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2505.19752}, 
}

Contact

• If you have any questions, please send me an email at: lihengli@stu.pku.edu.cn