Composable Decoding on the Probability Simplex: Theory and Implementation

• The paper addresses the lack of a common practical interface for configuring, combining, and evaluating decoding objectives that have often been treated as isolated strategies.Source: §1 · p. 1
• Decoding can be framed as regularised optimisation over the probability simplex, with familiar methods arising from choices of regularisers and support constraints.Source: §1 · pp. 1–2
• Weighted regularisers combine distributional preferences within one optimisation problem, enabling decoder compositions without external rewards, learned critics, or model updates.Source: §1 · p. 2
• Best-of-K decoding combines KL regularisation with local token-coverage utility for a K-sample budget.Source: §1 · p. 3
• CompoSimplex implements configurable support constraints, regularisation primitives, and simplex solvers through a shared interface for Transformers and vLLM.Source: §1 · pp. 2–3
• Across four models and benchmarks, compositions retain complementary strengths in single-sample accuracy, multi-sample success, and diversity.Source: §1 · p. 3

2 DECODING ON THE PROBABILITY SIMPLEX

• Decoding first constructs a distribution over the vocabulary; the final token is sampled from it or selected as its mode.Source: §2 · p. 3

2.1 DECODING AS OPTIMISATION OVER DISTRIBUTIONS

• The objective balances expected model score against a regulariser, while the feasibility set can restrict probability mass to selected candidate tokens.Source: §2.1 · p. 3

2.2 OPTIMALITY CONDITIONS ON THE SIMPLEX

• At the optimum, positive-mass tokens share a regularised score, while boundary tokens cannot exceed it when the derivative at zero is finite.Source: §2.2 · p. 4
• On a support-constrained simplex, the optimality condition applies on the feasible face, with excluded tokens fixed to zero.Source: §2.2 · p. 4

3.1 COMPOSITION THROUGH THE REGULARISER

• Weighted component regularisers combine multiple distributional preferences in one objective, with additive gradient contributions balancing against the model score.Weights are non-negative and sum to one; a component can penalise an undesirable property or represent the negative of a quantity to encourage.Source: §3.1 · pp. 4–5

3.2 SOLVING THE COMPOSED OBJECTIVE

• Mirror ascent solves composed objectives whose summed regulariser derivatives resist a closed-form solution, using simplex geometry rather than Euclidean projection.With the entropy-based divergence, its multiplicative update preserves non-negativity and normalisation; support constraints restrict the update to the feasible simplex face.Source: §3.2 · p. 5

3.3 USE CASE: BEST-OF-K DECODING

• Best-of-K combines KL regularisation with weighted expected coverage, encouraging alternatives across K samples while keeping the distribution close to the reference model.For K > 1, coverage gives diminishing returns to tokens already likely to appear among samples; the weighted utility reflects the probability of observing a token at least once in K draws.Source: §3.3 · pp. 5–6

COMPOSIMPLEX: A LIBRARY FOR COMPOSABLE DECODING

• A decoder configuration selects support, regularisation primitives, and optimiser settings; λ sets overall regularisation strength and α weights set relative primitive contributions.At each generation step, these components define an optimisation problem whose solution is the next-token distribution.Source: §4 · p. 6
• Available support rules include full vocabulary, Top-k, Top-p, Min-p, η-sampling, and typical sampling, each defining the candidate-token set.Top-k uses a fixed candidate count, while Top-p uses cumulative probability mass; the other rules use their stated probability, entropy, or information-content criteria.Source: §4 · p. 6
• Regularisation primitives include KL and JS divergences, negative entropy, and multi-sample Coverage and Diversity variants with different token-weighting schemes.Coverage weights the top-r tokens under the reference distribution; Diversity favours tokens with moderate logit gaps.Source: §4 · pp. 6–7
• The optimiser uses closed-form solutions for supported single-objective cases and mirror ascent otherwise, reusing current logits without additional model forward passes.A small number of mirror-ascent steps is used to limit added computation.Source: §4 · p. 7
• Custom logits processors integrate CompoSimplex with Transformers and vLLM, allowing either backend to sample from or take the argmax of the computed distribution.Tokens outside the selected support are masked.Source: §4 · p. 7

5.1 PERFORMANCE EVALUATION

• On MATH500, KL+Diversity matches KL at 64.4% pass@1 and reaches 90.0% pass@16, near Diversity’s 91.0% and above KL’s 88.4%.Across GPQA Diamond, IFEval, and LiveCodeBench, compositions cover a larger area than either constituent in most settings, though some metrics fall below individual regularisers.Source: Figure 2 · p. 8, §5.1 · p. 8
• Maximum gains over the base sampler are larger for pass@1 (+10.6 percentage points) than pass@16 (+5.1 percentage points).The authors suggest regularisation can concentrate probability mass on correct completions, making them easier to obtain with fewer samples.Source: §5.1 · p. 8
• Approximate mirror-ascent decoding costs 1.13× to 2.88× the corresponding base method, while reusing model logits without additional forward passes.Closed-form KL and entropy primitives have generation costs similar to corresponding standard samplers.Source: §5.1 · p. 8

5.2 DISTRIBUTIONAL BEHAVIOUR ANALYSIS

• At λ = 2, three of four compositions are non-dominated among evaluated configurations in their divergence–utility planes.Figure 3 compares individual objectives and equally weighted compositions on MATH500 with Qwen2.5-7B and Top-k support; marker size encodes λ and dashed curves are Pareto guides.Source: Figure 3 · p. 9, §5.2 · p. 9
• As λ rises from 0.5 to 2, Diversity’s utility rises while pass@1 falls from 62.4% to 39.8%; KL + Diversity retains 57.0% pass@1 at λ = 2.At the largest tested strength, all four compositions achieve higher pass@1 than their corresponding pure utility primitives.Source: §5.2 · p. 9
• Across tested weights α ∈ {0, 0.25, 0.5, 1}, increasing the utility weight raises utility and reduces measured KL divergence, but task-performance effects vary by metric.No consistent task-performance improvement occurs as α increases.Source: §5.2 · p. 9

6 RELATED WORK

• Sampling transformations have shared properties, but their quality–diversity trade-offs depend on task and configuration (Nadeem et al., 2020; Wiher et al., 2022).Prior examples include Top-K (Fan et al., 2018), nucleus sampling (Holtzman et al., 2020), and entropy-based temperature selection for multi-sample answer aggregation (Du et al., 2025).Source: §6 · p. 9
• Unlike sequence-level methods such as DAEMON (Ji et al., 2024) or reward-supervised token control such as Controlled Decoding (Mudgal et al., 2024), this work optimizes next-token distributions using current model scores.Bregman decoding also optimizes next-token distributions, but combines divergence with an ℓ0 penalty for sparse support.Source: §6 · p. 9

7 CONCLUSION

• The framework recovers familiar decoders and composes distributional preferences; experiments find complementary strengths across single-sample accuracy, multi-sample success, and diversity.Under strong regularisation, compositions can achieve non-dominated distribution-space points while maintaining more robust task performance than single primitives.Source: §7 · p. 10

A ADDITIONAL RELATED WORK

• The related-work discussion motivates separating support selection from distribution shaping, including approaches where sparse distributions arise from regularisation (Martins & Astudillo, 2016; Peters et al., 2019).It distinguishes this work’s joint divergence, entropy, and token-coverage terms from locally optimized sparse distributions.Source: A ADDITIONAL RELATED WORK · p. 15
• Unlike sequence-level or reward-guided methods, the implemented objectives shape the current next-token distribution without complete trajectories or trained critics (Ji et al., 2024; Qin et al., 2022; Mudgal et al., 2024).The related-work discussion also contrasts methods that use future rollouts, learned predictors, or expert and anti-expert logits (Karan & Du, 2026; Ji et al., 2026; Yang & Klein, 2021).Source: A ADDITIONAL RELATED WORK · p. 15
• Best-of-K rewards coverage of weighted token alternatives across K independent draws at one prefix, with its effects on accuracy and semantic diversity tested empirically (Wang et al., 2023).COMPOSIMPLEX exposes support, weighted regularisers, and a simplex solver through a shared implementation interface (Kwon et al., 2023; Kruszewski et al., 2023).Source: A ADDITIONAL RELATED WORK · p. 15

B MIRROR ASCENT CLOSED-FORM EXPRESSION

• The mirror-ascent appendix derives a closed-form optimal primal update using the negative-entropy potential, whose Bregman divergence is KL.The derivation proceeds through first-order stationary conditions and normalization.Source: B MIRROR ASCENT CLOSED-FORM EXPRESSION · p. 16

C.1 STANDARD DECODERS AS SPECIAL CASES

• Greedy decoding places all mass on highest-scoring tokens, while negative-entropy regularisation recovers softmax sampling with temperature λ.Tied highest-scoring tokens permit deterministic greedy decoding through a tie-breaking rule.Source: §C.1 · pp. 16–17
• Top-K sampling is restricted to the K highest-scoring tokens, and negative entropy yields temperature-scaled probabilities on that support.Top-P uses the same restricted entropy objective but selects its nucleus support from the base distribution before optimization.Source: §C.1 · p. 17
• A quadratic regulariser yields sparsemax, with zero probabilities arising from the regulariser and boundary condition rather than a prescribed support set.Standard sparsemax is recovered at λ = 1.Source: §C.1 · p. 17

C.2 SINGLE AND COMPOSED REGULARISERS: RELATION TO TEMPERATURE SCALING

• With the same λ and support, KL and entropy solutions both temperature-scale model logits, but KL is more concentrated on high-scoring tokens.As λ →∞, entropy approaches the uniform distribution on the support, whereas KL approaches the reference distribution p_t.Source: §C.2 · p. 18
• Best-of-K’s utility bonus varies by token and need not scale score differences uniformly, so its composed optimum is not restricted to temperature scaling.The exact optimum is approximated by the finite-step solver; on the tested grid, BoK matches or improves on the Top-k baseline at each temperature.Source: §C.2 · p. 18

D ADDITIONAL EXPERIMENTAL RESULTS

• The supplementary experiments cover checkpoints, benchmarks, decoding settings, evaluation protocols, selected configurations, seed variation, solver convergence and cost, and regularisation-strength and weight effects.These analyses supplement the experimental results in Section 5.Source: D ADDITIONAL EXPERIMENTAL RESULTS · p. 18

Implementation and solver settings

• The experiments use COMPOSIMPLEX's shared evaluation interface, and non-closed-form regularised objectives use 10 mirror-ascent updates with step size 0.1.Closed-form solutions are used for single KL and entropy objectives; solver settings are varied in solver studies.Source: §D.1 · p. 19

Models and benchmarks

• Four model–benchmark pairs cover instruction following, scientific reasoning, mathematics, and code generation, using base and instruction-tuned models.The pairs are LFM2.5-1.2B-Base on IFEval, Qwen3-4B-Base on GPQA Diamond, Qwen2.5-7B on MATH500, and Gemma-4-26B-A4B-IT on LiveCodeBench v6.Source: §D.1 · p. 19

Generation and decoding settings

• Unless stated otherwise, experiments sample 16 completions at temperature T = 0.5, use λ = 1 and equal composition weights, and compare decoders with matched prompt and completion-index seeds.Top-k uses k = 200 and Top-p uses p = 0.9 by default; other support thresholds are specified for Min-p, typical, and η-sampling.Source: §D.1 · p. 19

Task-specific grading

• Task-specific grading checks extracted answers or outputs against references, official evaluators, or execution tests; semantic diversity is mean within-prompt pairwise embedding distance.IFEval uses strict prompt-level correctness, while LiveCodeBench treats compilation errors, runtime errors, and timeouts as failures.Source: §D.1 · p. 19

Distributional metrics

• Distributional statistics average KL, JS, entropy, coverage, and diversity-gap utility over generation steps and completions; coverage scores use a reporting normalisation distinct from the objective weights.Diversity-gap utility uses K = 16 and weights proportional to Δv exp(−Δv).Source: §D.1 · pp. 19–20

Qwen2.5-7B

• On MATH500, JS + Coverage under Min-p exceeds both primitives on pass@1 and pass@16 while matching Coverage on pass@4 and SC@16; under Top-k, KL + Diversity retains KL’s 64.4% pass@1 and raises pass@16 from 88.4% to 90.0%.For Top-k, 90.0% remains below Diversity’s 91.0%; SemDiv lies between the constituent values for both compositions.Source: §D.2 · p. 20, Table 2 · p. 20

Qwen3-4B-Base

• On GPQA Diamond, JS + Diversity under Top-k exceeds both constituents on pass@1 and pass@4 but has lower pass@16; under typical support, it retains JS’s 26.3% pass@1 and raises pass@16 from 80.8% to 82.8%.The typical-support result is closer to Diversity’s 83.3%, and SC@16 lies between the constituent values.Source: §D.2 · p. 21, Table 3 · p. 21

LFM2.5-1.2B-Base

• On IFEval, KL + Diversity under η-sampling reaches 81.9% pass@16 versus 80.0% for KL and 81.5% for Diversity, with higher pass@4 and SemDiv but lower pass@1 and all-pass@16 than either constituent.Success uses strict prompt-level grading.Source: §D.2 · p. 21, Table 4 · p. 22

Gemma-4-26B-A4B-IT

• On LiveCodeBench, JS + Coverage under Top-p matches Coverage’s 56.6% pass@1 and reaches 66.9% pass@4, above JS’s 65.7% and Coverage’s 64.6%, but pass@16 is 67.4% versus 69.1% for both.Hard pass@1 is measured on 80 Hard problems.Source: §D.2 · p. 21, Table 5 · p. 22

Standard deviation across random seeds

• Across seeds on MATH500 with Qwen2.5-7B and Top-k, most standard deviations for pass@1, pass@4, pass@16, and SC@16 are below one percentage point, within a 0.12–1.33-point range.The reported deviations use seeds 0, 42, and 1234.Source: §D.2 · p. 21, Table 6 · p. 23

Solver convergence

• The convergence study varies step sizes 0.05, 0.1, and 0.5 and update counts 5, 10, 25, and 50 against reference optima computed analytically or from KKT conditions.The study uses the same 128 cached prefixes for every configuration.Source: §D.3 · p. 23, Figure 4 · p. 24
• At step size 0.1, 50 mirror-ascent updates bring mean L1 distance below 0.009 for every displayed objective; task-level runs use 10 updates and remain finite-step approximations.Task-level mean distances range from 0.033 to 0.062; KL and entropy use closed-form solutions.Source: §D.3 · p. 23, Figure 4 · p. 24

Computational cost

• On MATH500, iterative JS, Coverage, and Diversity cost 4.329–5.751 ms/token versus 2.258 ms/token for baseline decoding; on LiveCodeBench, they cost 20.954–21.255 ms/token versus 18.560 ms/token.The relative overhead varies across recorded model and batching configurations; closed-form KL has a cost close to baseline.Source: §D.3 · p. 23, Table 7 · p. 24

Sweep design

• The regularisation sweep compares four KL/JS coverage-diversity families at λ = 0.5, 1, and 2; at λ = 1, α varies over 0, 0.25, 0.5, and 1 in KL + Coverage and KL + Diversity.Across evaluated α weights, utility increases and KL divergence decreases in both families, without consistent gains in pass@1 or pass@16.Source: §D.4 · p. 25

Regularisation-strength sweep

• At λ = 2, the four compositions retain pass@1 between 57.0% and 62.6%, while single Coverage and Diversity reach their highest utilities but fall to 41.0% and 39.8% pass@1, respectively.Across the four compositions, increasing λ raises corresponding utility and SemDiv while reducing KL or JS divergence.Source: §D.4 · p. 25, Table 8 · p. 26

Relative composition weights

• Across evaluated α weights in KL + Coverage and KL + Diversity, increasing α monotonically raises the corresponding utility and lowers KL divergence, but pass@1 and pass@16 do not improve consistently.Nearby weights yield broadly similar task performance, suggesting limited sensitivity to α within a small range.Source: §D.4 · p. 25, Figure 5 · p. 25