---
license: apache-2.0
language:
  - en
tags:
  - latent-reasoning
  - chain-of-thought
  - continuous-thought
  - math
  - gsm8k
datasets:
  - whynlp/gsm8k-aug
pipeline_tag: text-generation
---

# MUX: Continuous Reasoning via Multiplexed Tokens

**[Ayhan Suleymanzade](https://misakitaro0414.github.io)**<sup>1</sup> &nbsp; **[Halil Alperen Gozeten](https://scholar.google.com/citations?user=aweeY7QAAAAJ&hl=en)**<sup>2</sup> &nbsp; **[Michael Bronstein](https://www.cs.ox.ac.uk/people/michael.bronstein/)**<sup>1,3</sup> &nbsp; **[Ismail Ilkan Ceylan](https://informatics.tuwien.ac.at/people/ismail-ceylan)**<sup>1,4†</sup> &nbsp; **[Jinwoo Kim](https://jw9730.github.io)**<sup>5†</sup>

<sup>1</sup>AITHYRA &nbsp; <sup>2</sup>University of Michigan &nbsp; <sup>3</sup>University of Oxford &nbsp; <sup>4</sup>TU Wien &nbsp; <sup>5</sup>KAIST &nbsp; <sup>†</sup>Equal advising

## Overview

**MUX** compresses discrete chain-of-thought reasoning into continuous *multiplexed tokens*. Each latent token is trained to represent a **weighted linear superposition** of a span of discrete reasoning subwords via KL-divergence distillation.

## Available Checkpoints

| File | Base Model | Dataset | Latent Tokens | GSM8K Accuracy |
|:---|:---|:---|:---:|:---:|
| `mux-gpt2-gsm8k.bin` | GPT-2 (124M) | GSM8K-AUG | 6 | 48.45% |
| `mux-llama1b-gsm8k.bin` | LLaMA 3.2 1B-Instruct | GSM8K-AUG | 6 | 57.01% |

## Usage

```python
# Clone the MUX codebase
git clone https://github.com/<your-org>/MUX.git
cd MUX

# Download checkpoint
# Option 1: huggingface-cli
huggingface-cli download MisakiTaro0414/MUX mux-llama1b-gsm8k.bin --local-dir .

# Option 2: direct download
wget https://huggingface.co/MisakiTaro0414/MUX/resolve/main/mux-llama1b-gsm8k.bin

# Evaluate
CKPT_DIR=. bash scripts/eval_llama1b.sh
```

## Training Details

- **Training data**: [whynlp/gsm8k-aug](https://huggingface.co/datasets/whynlp/gsm8k-aug) (385K augmented GSM8K samples)
- **Optimization**: LoRA (r=128, alpha=32) with cosine LR schedule
- **Loss**: Superposition KL-divergence distillation + reference loss
- **Hardware**: NVIDIA H100 GPUs

## Citation

```bibtex
@article{suleymanzade2025mux,
  title   = {MUX: Continuous Reasoning via Multiplexed Tokens},
  author  = {Suleymanzade, Ayhan and Gozeten, Halil Alperen and Bronstein, Michael and Ceylan, {\.I}smail {\.I}lkan and Kim, Jinwoo},
  journal = {arXiv preprint},
  year    = {2025}
}
```

## License

Apache 2.0