Simplex Relaxation for Discrete Diffusion

Jinya Sakurai1,4, Patrick Pynadath2, Satoshi Hayakawa3,
Jaehong Yoon1, Xulei Yang4, Nancy F. Chen4,5, Xun Xu4,5
1NTU Singapore, 2Purdue University, 3The University of Tokyo, 4IAIC, A*STAR, 5CFAR, A*STAR
Paper arXiv GitHub Hugging Face

TL;DR

Simplax enriches uniform discrete diffusion with an exact simplex-valued auxiliary state, enabling Rao–Blackwellized reverse-bridge training and sampling while preserving categorical corruption and token-based denoising.

Standard discrete diffusion matches one bridge, while Simplax matches all K virtual bridges through a simplex state.
Standard discrete diffusion matches one realized bridge. Simplax marginalizes matching over all K virtual bridge anchors from the same simplex state.

Abstract

Discrete diffusion models for categorical generation are defined by a corruption kernel, which determines the intermediate state space and the associated reverse prediction problem. We study whether uniform discrete diffusion can be augmented with an explicit continuous state while leaving its categorical corruption process unchanged.

We introduce Simplax, an exact Dirichlet–categorical augmentation that couples each corrupted categorical state with an auxiliary simplex-valued variable while preserving uniform diffusion as its categorical marginal. This yields a tractable Rao–Blackwellized reverse-bridge objective and a stochastic reverse sampler, while retaining the corrupted categorical state as the denoiser input.

Simplax improves the generative perplexity–entropy tradeoff on unconditional OpenWebText generation. On Sudoku, a model trained exclusively on 30-clue puzzles achieves the highest accuracy among the compared methods across all evaluated clue densities, including the minimum uniquely solvable 17-clue regime, and the highest validity in unconditional generation.

Simplex Relaxation

Simplax couples a corrupted token zt to an auxiliary simplex state wt through a shifted Dirichlet conditional. Marginalizing wt exactly recovers the original uniform categorical process, and decoding from the simplex is ordinary categorical sampling.

During training, the denoiser still receives the sampled token. The auxiliary simplex state defines a distribution over all categorical bridge anchors, whose reverse-bridge losses can be marginalized analytically. The same hierarchy gives a reverse sampler that carries simplex information between steps without requiring dense simplex vectors as network inputs.

OpenWebText Generation

Simplax lies on or near the lower generative perplexity–entropy envelope across inference budgets. With UDLM initialization, it achieves the strongest matched-entropy perplexity among the compared diffusion methods at low and high NFE under both reported evaluators.

Perplexity evaluator Llama-2 7B
Llama-2 7B OpenWebText generative perplexity versus entropy frontiers from 8 to 1024 function evaluations.

Click or drag to magnify · Double-click to close

Llama-2 7B generative frontiers on OpenWebText. Lower perplexity is better at comparable entropy; the dotted vertical line marks the data entropy.

Generated Text Samples

Representative unconditional generations selected at the temperature whose unigram entropy is closest to the OpenWebText entropy of 5.44. Text is shown without manual rewriting and truncated at […].

Sampling budget NFE = 16
CANDIEntropy 5.44

July on our records and the song in July, and we're signing people just for the next album, but it didn't couple with laser or anything. I had just a video. I thought they were crazy, because there were a lot of people out there who did it right. They felt like it was completely under my radar. […]

UDLMEntropy 5.53

What with talk of a new bus un the altitude in the first place. In South Bend in November came the focus upon new cyclists -- with an being builder. The other duck, Alabama. In total at least 82 bikes were donated to the department. […]

MDLMEntropy 5.43

Select Rally and Comm Rally was a great game with strong designed scenarios, but I did think about no-one should play it as middle of the road. You've already started to love my opinion on Two vs Two. You should check it, and stick and play this once again! We again stack you. This is the final year for the year. […]

DuoEntropy 5.45

I ever thought that it would be torrent be progression to try something new like my team base's football department. While I would have been optimistic about both the br of young players that hockey at the University at which went down in international football over the last of years and the way paths kept me focused so well I was pretty. […]

FLMEntropy 5.58

Dec 16, 2015 Edit: I cant let my reader now clear that this painting looks like a collegecommunication.'' On Z. Thanks the box. that by the way I am out over to the river to buy your best chance to date for such a horrible year. But at which point after year they release the windows of 3,100 by 11 inches. St. […]

LangFlowEntropy 5.42

God, will thrive. No others, and not at all, will determine our faith. We live in transient despair. But we don't know how to do it tomorrow. We want to live. We are sinners, we do not have God. God. we have identified the freedom of the one Being, is made to the only kind that governs our nation. […]

S-FLMEntropy 5.45

I feel like, how much does that have to work out there? Are you going into a different development company? I think speaking about the timing and the reputation of what I mean as something that have really enjoyed my career doing. I kind of think the company can change things. […]

Simplax (UDLM init)Entropy 5.46

It was basically a percentage league, is there going to have to be one a year? They always were because they had them. If they knew that rule, they got to play one for their bodies. Then, when they voted for a “bolt 1” rule and what they got was, they want the fifth best defender to attack you, so they added […]

Sudoku Generation

All models are trained only on 30-clue puzzles. Simplax obtains the highest conditional accuracy at every tested clue density and reaches 95.30% validity in unconditional generation.

NFE–Accuracy Comparison

All models are trained on 30-clue puzzles. Across inference-time clue counts, Simplax maintains strong performance even at low NFE.

Inference-time clues 30
Sudoku generation accuracy versus NFE with 30 inference-time clues.
Accuracy versus NFE with 30 clues for both training and inference.

Qualitative Samples

Sudoku generations from baseline models and Simplax for the same 25-clue puzzle; only Simplax is valid.
A representative 25-clue puzzle. Blue entries are given clues, black entries match the solution, and red entries are errors. Simplax produces the valid solution.

BibTeX

@misc{sakurai2026simplex,
  title         = {Simplex Relaxation for Discrete Diffusion},
  author        = {Jinya Sakurai and Patrick Pynadath and Satoshi Hayakawa
                   and Jaehong Yoon and Xulei Yang and Nancy F. Chen and Xun Xu},
  year          = {2026},
  eprint        = {2608.10615},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL}
}