STaRR: Spatial-Temporal Token-Dynamics-Aware Responsive Remasking for Diffusion Language Models

This repository contains STaRR inference and evaluation code for diffusion language models, including LLaDA and Dream.

Setup

The project uses uv to manage Python dependencies. Use Python 3.10 or newer and a CUDA-capable PyTorch environment for model inference.

Install uv if it is not already available:

curl -LsSf https://astral.sh/uv/install.sh | sh
# or
pip install uv

Create and activate the virtual environment, then install the locked dependencies:

uv venv .venv
source .venv/bin/activate
uv sync --locked

Always activate the virtual environment before running demo or evaluation scripts:

cd /path/to/script
source .venv/bin/activate

The lock file and pyproject.toml use the Tsinghua PyPI mirror. The first run will download model weights and datasets from Hugging Face, so make sure your Hugging Face access and cache are configured correctly.

Models

The included scripts use these default model checkpoints:

• LLaDA: GSAI-ML/LLaDA-8B-Instruct
• Dream: Dream-org/Dream-v0-Instruct-7B

The demo scripts currently set the device in code, for example cuda:1. If your machine uses a different GPU layout, update the device variable in llada/llada_starr.py or dream/dream_starr.py before running.

Run Demos

Run the LLaDA STaRR demo:

source .venv/bin/activate
python llada/llada_starr.py --starr

Run the Dream STaRR demo:

source .venv/bin/activate
python dream/dream_starr.py --starr

Evaluate on Datasets

Evaluation scripts run multiple lm-eval tasks by default:

• gsm8k
• minerva_math500
• humaneval
• mbpp

The scripts set code-evaluation environment variables internally and call uv run accelerate launch. Activate the virtual environment before running them.

LLaDA

Baseline evaluation:

source .venv/bin/activate
bash llada/eval/eval_baseline.sh

STaRR evaluation:

source .venv/bin/activate
bash llada/eval/eval_starr.sh

STaRR with dual cache:

source .venv/bin/activate
bash llada/eval/eval_starr_dual_cache.sh

Default LLaDA outputs are written under:

• llada/eval/baseline_results/
• llada/eval/starr_results/
• llada/eval/dual_cache_starr_results/

Dream

Baseline evaluation:

source .venv/bin/activate
bash dream/eval_baseline.sh

STaRR evaluation:

source .venv/bin/activate
bash dream/eval_starr.sh

STaRR with dual cache:

source .venv/bin/activate
bash dream/eval_starr_dual_cache.sh

Default Dream outputs are written under:

• dream/baseline/
• dream/starr/
• dream/starr_cache/

Notes

• The evaluation scripts default to CUDA_VISIBLE_DEVICES=0,1,2,3. Edit the device or CUDA_VISIBLE_DEVICES setting in the target script if you want to use a different GPU set.
• humaneval and mbpp run generated code as part of evaluation. Review the scripts and run them only in an environment where code execution is acceptable.
• To change generation length, diffusion steps, block length, or model paths, edit the variables near the top of the corresponding shell script.