Composable Decoding on the Probability Simplex: Theory and Implementation
Xiaotong Ji, Ahmed Khaled Khamis, Rasul Tutunov, Matthieu Zimmer, Haitham Bou-Ammar
cs.LG, cs.AI
2026-09-28
Casts LLM decoding as simplex optimisation; composing KL with Diversity keeps 64.4% pass@1 and 90% pass@16 on MATH500, while pure Diversity drops to 56%. Cost is 1.13 to 2.88x.
Every LLM pipeline ends in decoding, and decoding is still a shelf of disconnected knobs. Teams pick among greedy, temperature, Top-k, Top-p, Min-p, typical, and η-sampling by trial and error. Prior work mapped shared properties of these transforms and measured quality-diversity trade-offs. What it did not give is a single interface where those preferences become explicit objectives that can be added, configured, and compared.
Sequence-level methods take a different route. DAEMON targets text metrics, power sampling sharpens the full-sequence distribution, and Controlled Decoding uses learned prefix values. The shared bill is external rewards, critics, or extra rollouts. This paper, from Huawei Noah's Ark Lab and UCL, pulls the problem back to the next-token distribution at each step: first pick a distribution q on the probability simplex, then sample or take a mode. Decoders differ by the objective they implicitly optimise, not by the name of the sampler.
At step t the model emits a score vector st. The decoder maximises the inner product of q and st, minus λ times a regulariser Ω(q), subject to q lying in a feasible set Ct. Ct is typically a support constraint: Top-k, Top-p, or Min-p first selects a candidate set St and pins all other coordinates to zero. Ω decides how mass is spread on that face.
With no regulariser, the optimum puts all mass on the highest-scoring tokens (greedy). Negative Shannon entropy recovers softmax, with temperature equal to λ. The same entropy regulariser restricted to St is temperature-scaled Top-k or Top-p. An ℓ2 penalty yields Sparsemax; zeros come from the boundary condition, so the support need not be chosen in advance.
Composition adds regularisers with weights αi that are nonnegative and sum to one, then solves the same problem. Each regulariser contributes an additive term to the regularised score in the KKT stationarity condition. KL and entropy have closed forms. After composition the stacked derivatives usually cannot be inverted coordinate-wise, so the library runs mirror ascent on the simplex. Negative entropy as the potential turns the Bregman divergence into KL, and the update becomes the multiplicative map q ← q ⊙ exp(ρ∇f) followed by normalisation. Nonnegativity and the unit-sum constraint hold by construction. Projected gradient ascent with Euclidean distance is the wrong geometry for distributions, so it is not used.
Best-of-K is one concrete mix. KL anchors q to a reference pt, the softmax of model logits on St. The coverage term is an occupancy probability: token v appears at least once in K independent draws with probability 1-(1-q(v))^K. Coverage spreads equal positive weights over the top-r tokens of pt. Diversity weights by the gap from the largest logit, favouring alternatives that are close but not already dominant. For K>1 the bonus shrinks as q(v) grows, so tokens that are already likely to show up get diminishing returns. Temperature scaling stretches all score gaps by one factor; it cannot apply this nonuniform bonus.
CompoSimplex exposes support rules, regulariser primitives, and solvers as configuration, wired into Transformers and vLLM as a logits processor. No extra forward pass. Main runs use closed forms for KL and entropy, and 10 mirror-ascent steps with step size 0.1 otherwise.
Four pairs: Qwen2.5-7B on MATH500, Qwen3-4B-Base on GPQA Diamond, LFM2.5-1.2B-Base on IFEval, Gemma-4-26B-A4B-IT on LiveCodeBench v6. Default: 16 samples, temperature 0.5, λ=1, equal mix weights.
MATH500, Top-k (k=200):
| Method | pass@1 | pass@16 | SC@16 | SemDiv |
| Top-k | 59.8 | 90.0 | 78.0 | 0.150 |
| KL | 64.4 | 88.4 | 76.8 | 0.132 |
| Diversity | 56.0 | 91.0 | 76.2 | 0.159 |
| KL+Diversity | 64.4 | 90.0 | 77.2 | 0.149 |
KL+Diversity matches KL at 64.4% pass@1 and lifts pass@16 from 88.4% to 90.0%, near Diversity's 91.0%. Under Min-p, JS+Coverage reaches 67.2% pass@1, above the Min-p baseline 64.4% and above JS at 66.6% and Coverage at 65.4%.
The largest gain over the matching base sampler is +10.6 points of pass@1 (JS+Diversity 26.3 vs Top-k 15.7 on GPQA) and +5.1 points of pass@16. Regularisation is better at concentrating mass on correct completions when the sample budget is small.
Composition does not dominate every axis. On GPQA Top-k the same JS+Diversity mix records 79.8% pass@16, below Diversity at 85.4% and JS at 80.3%. On IFEval with η-sampling, KL+Coverage raises pass@16 from 80.6% to 83.4% while all-pass@16 falls from 17.2% to 16.1%. On LiveCodeBench Top-p, JS+Coverage hits 66.9% pass@4, above both parts, then drops to 67.4% pass@16 against 69.1% on each side.
Raising λ from 0.5 to 2 increases the declared utility and cuts KL or JS. Pure Coverage and Diversity collapse to 41.0% and 39.8% pass@1 at λ=2; the four equal-weight mixes stay between 57.0% and 62.6%. Raising the utility weight α monotonically improves the utility and lowers KL, with no stable direction on pass@1 or pass@16.
Closed-form KL is nearly free: 2.238 vs 2.258 ms/token on MATH500. Iterative single regularisers and mixes cost about 1.13x to 2.88x. With 10 steps at step size 0.1, mean L1 to a reference optimum sits between 0.033 and 0.062. These runs are finite-step approximations.
This is a decode-time configuration. Weights stay frozen, no critic is trained, and the step-wise solve does not query future rollouts. Stacks that already plug in a logits processor can swap support and Ω through config. The intended use is reasoning with a K-sample budget that needs both single-draw quality and coverage; the Best-of-K coverage term is written for that budget.
Most cells move by one to three points, and some metrics fall below a single regulariser. What you get is a composable objective space plus a KKT derivation that recovers greedy through Sparsemax. Latency grows by roughly 1x to 2x depending on the mix; closed-form KL and entropy add almost nothing.
There is no dedicated Limitations section. The body already states that λ is hard to set, that raising α does not stably lift task scores, that strong utility regularisation can wreck accuracy, that some metrics fall below the parts, and that the main runs use a 10-step approximation.
Coverage is a prefix-level token occupancy, not sequence-level Best-of-K. Any link to self-consistency or rerank candidate quality is empirical. Evaluation stays on math, science QA, instruction following, and code; open-ended generation is absent. The main tables do not compare against power sampling or Controlled Decoding under the same protocol. GPQA has 198 items, LiveCodeBench v6 adds 175, and seed standard deviations are reported for three seeds on MATH500 Top-k only. vLLM is wired up; the reported numbers come from Transformers.