| # Riemannian Diffusion Langauge Model |
|
|
| Official code repository for the paper [Continuous Diffusion Model for Language Modeling](https://arxiv.org/abs/2502.11564) (NeurIPS 2025) |
|
|
| We provide an implementation for Riemannian Diffusion Language Model (RDLM) on language modeling tasks. |
|
|
|
|
| ## Dependencies |
| Create an environment with Python 3.9, and Pytorch 2.3.1. Install requirements with the following command: |
|
|
| ``` |
| pip install -r requirements.txt |
| ``` |
|
|
| ## Running Experiments |
|
|
| ### 1. Configurations |
| The configurations are provided in the `config/` directory in YAML format. |
| - To use new dataset, refer to `configs/exp`. |
| - To use new model architecture, refer to `configs/model`. |
| - To use new type of generative process, refer to `configs/sde`. |
|
|
| ### 2. Preparations |
| Datasetes are automatically downloaded when running the training script. |
| - Data cache directory is set to `data/`. This can be modified via `data.cache_dir` in the config file. |
|
|
| - To add new dataset or modify dataset/tokenizer setting, please refer to the `data.py`. |
|
|
| ### 3. Training |
|
|
| To run on Text8 dataset use the following command: |
| ```sh |
| CUDA_VISIBLE_DEVICES=0 python main.py \ |
| ngpus=1 \ |
| training.accum=1 \ |
| exp=text8 \ |
| sde=mixture \ |
| sde.step_thr=0.35 \ |
| scheduler=geometric \ |
| scheduler.weight_type=step \ |
| scheduler.left=0.3 \ |
| scheduler.right=0.6 |
| ``` |
| - `ngpus` is the number of GPUs used for training |
| - `training.accum` is the number of gradient accumulation steps. |
|
|
| Modify these two hyperparameters to fit your hardware. |
|
|
|
|
| Similarly, to run on One Billion Words (LM1B) dataset use the following command: |
| ```sh |
| CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py \ |
| ngpus=4 \ |
| training.accum=1 \ |
| exp=lm1b \ |
| tokens=3 \ |
| sde=mixture \ |
| sde.rho_scale=1.14 \ |
| sde.step_thr=0.38 \ |
| scheduler=geometric \ |
| scheduler.weight_type=step \ |
| scheduler.left=0.3 \ |
| scheduler.right=0.75 |
| ``` |
|
|
|
|
|
|
| ### 4. Generation and Evaluation |
|
|
| Run the following command to generate samples and evaluate: |
| ```sh |
| CUDA_VISIBLE_DEVICES=0 python main.py \ |
| ngpus=1 \ |
| run_mode=sample \ |
| server=sample \ |
| exp=sample_lm1b \ |
| "model_path='PATH_TO_MODEL_CHECKPOINT'" \ |
| seed=0 |
| ``` |
|
|
| ## Pretrained checkpoints |
| The checkpoints for the models trained on Text8 and LM1B datasets are available in this [Google Drive folder](https://drive.google.com/drive/folders/1aDTZtPIxAxQrkaRSahjuWbkxX1OYq9CC?usp=sharing). |
|
|
| - Download `checkpoint.pth` and pass the path to the downloaded file to `PATH_TO_MODEL_CHECKPOINT`. Use the command provided in the above section to generate and evaluate the samples. |
|
|
| - Additional files `sde.pkl` and `config.yaml` are provided for reproducibility and further analysis. |
|
|
|
|
| ## Citation |
|
|
| If you found the provided code with our paper useful in your work, we kindly request that you cite our work. |
|
|
| ```BibTex |
| @inproceedings{jo2025RDLM, |
| author = {Jaehyeong Jo and |
| Sung Ju Hwang}, |
| title = {Continuous Diffusion Model for Language Modeling}, |
| booktitle = {Advances in Neural Information Processing Systems}, |
| year = {2025}, |
| } |
| ``` |
