• Overview
• Getting Started
  • Environment Setup
  • Base Docker Image
  • Dependencies
  • Setup
• Training and evaluation
• Implementation Credits
• Checkpoints

This repository contains the official implementation of Multiplex Thinking: Reasoning via Token-wise Branch-and-Merge.

Multiplex Thinking proposes a token-wise branch-and-merge reasoning mechanism, enabling efficient and expressive multi-pat reasoning while maintaining a compact token representation.

The codebase is built upon several high-quality open-source projects. We sincerely thank the original authors and contributors for their outstanding work.

We recommend using Docker to ensure a consistent and reproducible environment. If you prefer Conda, we also provide an environment specification in conda_env.yaml.

We suggest starting from the official verl SGLang worker Docker image:

• https://github.com/volcengine/verl/blob/325cbc770bfe32ef022f1cd67feab1a23bba9e42/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.9.post6.mcore0.13

For general system configuration, please refer to the official documentation of verl:

• https://verl.readthedocs.io/en/latest/workers/sglang_worker.html

Please ensure the following package versions are installed:

• sglang == 0.4.9.post6
• transformers == 4.54.0

Run the setup script:

bash setup.sh

The `setup.sh` script handles the installation of required dependencies and ensures the correct versions of our customized libraries are active by running:
* `pip install sglang-0.4.9.post6`
* `pip install transformers-4.54.0`

Train and evaluate by running:

bash scripts/train.sh \
  --model deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \
  --exp_name your_exp_name \
  --enable_unweighting True \ # True for average embedding; False for weighted embedding
  --total_training_steps 300 \
  --train_batch_size 128 \
  --max_token_len_per_gpu 32768 \
  --loss_mode multiplex_thinking \
  --multiplex_width 3 \
  --n_gpus_per_node 8 \
  --max_response_length 4096 \
  --val_rollout_n 4 \
  --val_dataset math \
  --val_batch_size 1024

Or run evaluation:

bash scripts/eval.sh

This codebase is built upon and inspired by the exceptional work from the following projects:

• Training & RL Framework: verl & DeepScaleR
• Inference Engine: sglang
• Code Inspiration & Adaptations: Soft Thinking

Model weights are available on Hugging Face: 👉 Multiplex-Thinking-HF-Checkpoints

If you find this work useful for your research, please cite our paper as:

@article{tang2026multiplexthinking,
  title   = {Multiplex Thinking: Reasoning via Token-wise Branch-and-Merge},
  author  = {Tang, Yao and Dong, Li and Hao, Yaru and Dong, Qingxiu and Wei, Furu and Gu, Jiatao},
  journal = {arXiv preprint arXiv:2601.08808},
  year    = {2026}
}