Title: Foundations of Top-$k$ Decoding for Language Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models

Top-$k$ Decoding

Adaptive Decoding

Sparse Optimization

$\ell_0$-Regularization

Bregman Divergence
Abstract: Top-$k$ decoding is a widely used method for sampling from LLMs: at each token, only the largest $k$ next-token-probabilities are kept, and the next token is sampled after re-normalizing them to sum to unity. Top-$k$ and other sampling methods are motivated by the intuition that true next-token distributions are sparse, and the noisy LLM probabilities need to be truncated. However, to our knowledge, a precise theoretical motivation for the use of top-$k$ decoding is missing. In this work, we develop a theoretical framework that both explains and generalizes top-$k$ decoding. We view decoding at a fixed token as the recovery of a sparse probability distribution. We introduce *Bregman decoders* obtained by minimizing a separable Bregman divergence (for both the *primal* and *dual* cases) with a sparsity-inducing $\ell_0$-regularization; in particular, these decoders are *adaptive* in the sense that the sparsity parameter $k$ is chosen depending on the underlying token distribution. Despite the combinatorial nature of the sparse Bregman objective, we show how to optimize it efficiently for a large class of divergences. We prove that (i) the optimal decoding strategies are greedy, and further that (ii) the objective is discretely convex in $k$, such that the optimal $k$ can be identified in logarithmic time. We note that standard top-$k$ decoding arises as a special case for the KL divergence, and construct new decoding strategies with substantially different behaviors (e.g., non-linearly up-weighting larger probabilities after re-normalization).

Foundations of Topk Decoding for Language Models

Abstract

Topk decoding is a widely used method for sampling from LLMs: at each token, only the largest k next-token-probabilities are kept, and the next token is sampled after renormalizing them to sum to unity. Topk and other sampling methods are motivated by the intuition that true next-token distributions are sparse, and the noisy LLM probabilities need to be truncated. However, to our knowledge, a precise theoretical motivation for the use of topk decoding is missing. In this work, we develop a theoretical framework that both explains and generalizes topk decoding. We view decoding at a fixed token as the recovery of a sparse probability distribution. We introduce Bregman decoders obtained by minimizing a separable Bregman divergence (for both the primal and dual cases) with a sparsity-inducing ℓ 0 -regularization; in particular, these decoders are adaptive in the sense that the sparsity parameter k is chosen depending on the underlying token distribution. Despite the combinatorial nature of the sparse Bregman objective, we show how to optimize it efficiently for a large class of divergences. We prove that (i) the optimal decoding strategies are greedy, and further that (ii) the objective is discretely convex in k , such that the optimal k can be identified in logarithmic time. We note that standard topk decoding arises as a special case for the KL divergence, and construct new decoding strategies with substantially different behaviors (e.g., non-linearly up-weighting larger probabilities after renormalization).

1 Introduction

Large language models (LLMs) are powerful generative AI tools for producing text. When pre-trained on large text corpora and aligned according to human preferences, they can be used for a wide range of tasks. On a technical level, they are probability distributions over text: given any user text prompt x , an LLM samples an answer Y ∼ π ( ·| x ) from a probability distribution π ( ·| x ) over text. However, even after obtaining a pre-trained, fine-tuned, and human preference-aligned model π , it is uncommon to directly sample from the model. Instead, several sampling/decoding methods are commonly used, including topk [21] or topp sampling [32]. Due to their improved empirical performance compared to direct sampling, they are used by default or as an option in many popular LLMs, including the GPT series, Gemini, and Claude. From a broader perspective, per-token samplers/decoders belong to an expanding collection of post-hoc methods for improving LLM performance, which range from pre-sampling transforms (e.g. temperature decoding), to sequence-level decoding strategies (e.g. beam search), to post-hoc selection (e.g. best-ofN or self-consistency), to a variety of test-time scaling approaches; see e.g., [12, 21, 32].

In this paper, we focus on decoding methods that modify each next-token-probability distribution to induce sparsity , i.e., to keep only a small number of tokens with a nonzero probability. This includes the widely used topk [21] and topp [32] sampling methods, among others. These methods

are motivated by the intuition that the noisy LLM probabilities need to be truncated to denoise the 'unreliable tail' [32]. In particular, we focus on the popular topk decoding method, which keeps only the largest k next-token-probabilities at each decoding step. These are renormalized-via dividing by their sum-to a probability distribution from which the next token is sampled.

Despite the wide use and rich intuition behind topk decoding, to our knowledge, a precise theoretical understanding of topk decoding is not available (see Section 6 for a discussion of related work). Therefore, in this work, we develop a theoretical framework that flexibly generalizes and sheds light on the key properties of topk decoding. For a fixed token, we view decoding as recovering a sparse probability distribution from a given (generally non-sparse) LLM token distribution. We consider denoisers obtained by minimizing a Bregman divergence (such as KL divergence or Brier score) from the 'raw' LLM token distribution, with a sparsity-inducing ℓ 0 regularization. This approach is motivated by a rich literature of both Bregman divergences and sparsity, see Section 6 for details.

Bregman renormalizations. Our approach both generalizes topk decoding and opens up a rich field of efficient adaptive 'Bregman decoding' methods with a wide and tunable range of behaviors. As an example, we consider Bregman divergences generated by the α -entropies x ↦→ x α / [ α ( α -1)] [29, 51], and display in the adjacent figure how Bregman decoders modulate a token distribution for several values of α . For α → 1 , we obtain standard topk decoding (for k = 10 here). By contrast, for α = 0 . 25 , the decoder shifts most of the mass onto the top few tokens; while for α = 2 . 5 , the mass is spread much more uniformly across the topk tokens. This exemplifies how our framework enables the design of novel decoders eliciting a wide range of behaviors.

Provable adaptivity. An important feature of our framework is that it studies, and provides, provably adaptive decoding strategies. Namely, given any raw LLM token probability vector p , our Bregman decoders effectively perform (a generalization of) topk ∗ decoding of p for an optimal k ∗ = k ∗ ( p ) : the utilized k ∗ varies depending on the LLM token distribution p , and is chosen to minimize the decoder's ℓ 0 -regularized Bregman divergence from p . This rigorous sparse-objective-centric foundation of adaptivity in LLM decoding is, to our knowledge, new in the literature. Moreover, perhaps surprisingly, we are able to show in substantial generality that an optimal k ∗ can be found provably and efficiently without relying on grid search or other heuristics.

1.1 A roadmap of our contributions

In Section 2, we introduce our theoretical framework. We view topk decoding strategies as two-step: (i) select a number of tokens k , and (ii) renormalize the selected k tokens' entries to a probability distribution (Section 2.1). We introduce two rich classes of decoding strategies (Section 2.2): primal Bregman decoding and dual Bregman decoding . These correspond to ℓ 0 -regularized minimization of a Bregman divergence to the 'raw' LLM distribution over tokens, in its first vs. second argument. 2

In general, ℓ 0 -regularization leads to combinatorial optimization problems, for which there are no known polynomial-time algorithms [11, 42]. Our main contribution is to show that, despite this, the sparse Bregman decoding objective can be efficiently optimized under mild assumptions, by virtue of having two key structural properties: (1) Greedy selection : Choosing the k largest probabilities is optimal (Theorems 3.2 and 3.3 in Section 3.2); (2) k -convexity : Searching for the optimal k ∗ is a (discretely) convex problem in k (Theorem 3.4 in Section 3.3). While simple to state and desirable, these properties are non-trivial to establish, and require a range of novel structural insights into the sparse Bregman objective that could be of independent interest.

In Section 4, we illustrate our theory by introducing α -Bregman decoding strategies, generated by Tsallis α -entropies x ↦→ x α / [ α ( α -1)] . We study how their behavior depends on α , and highlight several closed-form cases of interest. One example of the optimization-theoretic elegance of α -decoders is their convergence to water-filling as α →∞ . Finally, in Section 5, we study the empirical performance of some of the novel decoding schemes on open-ended text generation and mathematical problem solving tasks with LLMs, and find that they perform competitively with topk decoding.

2 Regularized sparse Bregman decoding

2.1 Topk decoding preliminaries

Topk decoding. Given a probability distribution p = ( p 1 , . . . , p V ) (where V stands for 'vocabulary size'), and some 1 ⩽ k ⩽ V , topk decoding first selects the indices S k = ( i 1 , . . . , i k ) of the largest k probabilities, breaking ties arbitrarily. Setting all other coordinates to zero in p , one obtains the vector p [1 : k ] of the k largest entries. Then, it renormalizes this vector by dividing it by its sum. Letting ( p (1) , p (2) , . . . , p ( k ) ) = ( p i 1 , . . . , p i k ) be the largest k entries of p ,

One then draws a sample from the distribution topk ( p ) .

Decoding strategies. Next, we aim to generalize topk decoding. We will refer to any operator Dec on probability distributions as a decoding strategy ; formally Dec : ∆ V → ∆ V , where ∆ V = { x ∈ [0 , 1] V : ∑ V i =1 x i = 1 } is the simplex of V -dimensional probability distributions. Observe that topk decoding consists of two steps: selecting the largest coordinates and renormalizing them. The second step can be viewed as 're-distributing' the probability mass that has been thresholded away by selection among the remaining indices. This step can be performed in a lot of other meaningful ways besides division by the sum. For instance, we may put a larger weight on the larger remaining probabilities, if we consider them more reliable.

Renormalization. Motivated by this, we define the notion of a renormalization mapping, which takes as input a thresholded probability vector with k nonzero entries remaining. We consider renormalization maps that are permutation-equivariant , i.e., when their input is permuted, their output is permuted accordingly; which clearly holds for the sum-division used in topk . Therefore, since the sum of probabilities after selection can be less than unity, we can define them as maps from the sub-probability simplex ∆ sub ,k = { x ∈ [0 , 1] k : ∑ k i =1 x i ⩽ 1 } to the simplex ∆ k .

Definition 2.1 (Renormalization) . For a positive integer k , we call a permutation-equivariant map T : ∆ sub ,k → ∆ k a renormalization map .

A renormalization map can be extended to the full simplex ∆ V , by applying it only on the nonzero coordinates. 3 We can now define generalized topk decoding as renormalizing the topk entries via a general renormalization map.

Definition 2.2 (Generalized topk decoding) . For a fixed k , a generalized topk decoding strategy Dec k,T : ∆ V → ∆ V , parameterized by the choice of k and renormalization map T , takes as input any V -class probability vector p , thresholds it to the sub-vector p [1 : k ] consisting of its topk elements, and renormalizes it to T ( p [1 : k ]) ∈ ∆ V .

Adaptivity. Anatural extension is to choose k adaptively based on p . For this, we consider a k -selector map ˆ k : ∆ V → [ V ] := { 1 , . . . , V } , and a collection of renormalization maps T k : ∆ sub ,k → ∆ k , k = 1 , . . . , V . We define an adaptive generalized topk decoding strategy Dec T : ∆ V → ∆ V via p ↦→ T ˆ k ( p ) ( p [1 : ˆ k ( p )]) . Below, we will design specific renormalizers T and ways to choose k .

2.2 Regularized sparse Bregman decoding

Decoding via sparse divergence minimization. Consider a divergence Div( · , · ) : ∆ V × ∆ V → R between two distributions. Classical examples include the squared error Div( p, q ) = ∥ p -q ∥ 2 2 and the KL divergence Div( p, q ) = ∑ V j =1 p j ln( p j /q j ) . We define the decoding strategy Dec Div , via sparsity-regularized divergence minimization 4 under divergence Div , for any probability vector p as:

Here, the ℓ 0 -pseudonorm ∥ ˆ p ∥ 0 is the number of nonzero entries of ˆ p , and λ ⩾ 0 is a sparsity cost hyperparameter. As λ increases, the optimal solution ˆ p = p ∗ gets increasingly more sparse.

Separable Bregman divergences. In this work, we shall instantiate Div in Problem 2 with separable Bregman divergences [1, 10]. We will see that this class is expressive enough to induce topk decoding and many fruitful generalizations of it. For a convex domain Dom ⊆ R and a convex differentiable function ϕ : Dom → R , the one-dimensional Bregman ϕ -divergence d ϕ is defined as: d ϕ ( x, y ) = ϕ ( x ) -ϕ ( y ) -ϕ ′ ( y )( x -y ) , for x, y ∈ Dom . The separable V -dimensional Bregman ϕ -divergence D ϕ : Dom V → R is then defined as:

A well-known property of Bregman divergences is that D ϕ ( x, y ) ⩾ 0 for all x, y , with equality if x = y ; when ϕ is strictly convex, x = y in fact becomes the unique minimum.

0.07

Value

0.12

Value

0.10

0.06

0.05

Objective

0.08

Objective

0.04

0.06

0.03

0.04

0.02

0.02

1.00.8

0.6

0.4

0.20.0

0.8

0.6

0.4

p2

A

0.2

0.0

1.00.80.60.4

0.2

0.0

0.8

0.6

0.4

p2

A

0.2

0.0

Primal and dual Bregman decoding. Since Bregman divergences are generally non-symmetric in their arguments, we may instantiate the sparse Bregman decoding Problem 2 in two substantially distinct ways: by placing the estimand ˆ p in the first ( primal ) or second ( dual ) argument:

Both formulations possess a sound theoretical motivation. Bregman projections are commonly defined as minimization in the first argument, while Bregman-based proper scoring rules for mean elicitation correspond to minimization in the second argument [see e.g., 24, 39, etc].

The landscapes of primal and dual decoding are illustrated in Figure 1. The dual objective can be non-convex even in the interior of the simplex. However, crucially, the objectives are discontinuous at the edges of the simplex due to the ℓ 0 penalty. While in general these decoding objectives could be combinatorial problems that may be hard to solve, we will show in Section 3 that for separable Bregman divergences, both the primal and dual problems can be solved efficiently.

In both the primal and the dual Bregman case, when λ = 0 , the corresponding sparse decoding Problem 2 is solved at ˆ p = p (and uniquely so if ϕ is strictly convex), with the intuition that absent sparsity requirements the best guess is to preserve the original distribution p . Henceforth we focus on the sparse regime λ > 0 , which forces some entries of ˆ p to be zero at optimality. Our main results in Section 3 show for both primal and dual decoding that, under mild conditions on D ϕ , the optimal solution keeps exactly the topk ⋆ coordinates of p , for an objective-chosen k ⋆ = k ⋆ ( p ) . This yields a principled and broad generalization of topk decoding.

3 Efficient computation of primal and dual Bregman decoding

We now investigate the optimization of the sparse objectives that give rise to primal and dual Bregman decoding. Absent further structure in these objectives, for any fixed k one would have to search over all (combinatorially many) sizek sparsity patterns to decide which k probabilities to keep; and one would have to try all k ∈ [ V ] to determine the optimal k . Fortunately, we will now show that Bregman decoding objectives admit computationally efficient optimization, which rests on two pillar properties: (1) The greedy property : Given any k , it is optimal to select the top k tokens. (2) k -convexity : The sparse Bregman objective is (discretely) convex as a function of k .

First, in Section 3.1, we deal with the innermost optimization layer: the renormalization of the selected token probabilities (which is performed after the optimal k and the optimal sparsity pattern have been identified). Under certain conditions on the Bregman generator, we show that it reduces to scalar root-finding both in the primal and in the dual case (the dual case in fact necessitates nested root finding). We then proceed to show the greedy property in Section 3.2. Finally, for the outermost layer of our optimization problem, we demonstrate the k -convexity property in Section 3.3.

3.1 Renormalization for a fixed sparsity pattern

We first investigate the renormalization component of a Bregman decoding strategy. Once the optimal sparsity pattern S ⊆ [ V ] (of some size | S | = k ) has been identified, the vector x -which denotes the sub-vector of p restricted to indices in S -needs to be projected onto the simplex ∆ k . Since the ℓ 0 regularization term becomes fixed to λk , Problem 2 becomes equivalent to: arg min ˆ p ∈ ∆ k Div(ˆ p, x ) . This is a k -dimensional Bregman projection problem to the simplex (without sparsity regularization). We will now, for both primal and dual decoding, (i) derive conditions under which this problem is well defined, and (ii) show that it can be efficiently solved by reduction to scalar root finding .

Primal renormalization. We impose the following mild condition on the Bregman generator ϕ ; compared to a minimal set of assumptions for a Bregman divergence to be well-defined, it additionally requires first-order smoothness and strict convexity to hold on the entirety of the relevant interval.

Assumption 3.1 (Primal validity) . The map ϕ is convex and continuously differentiable on [0 , 1] as well as strictly convex on (0 , 1) .

Existing results [33, 34] then imply that for a primal valid potential ϕ , denoting f = ϕ ′ (and extending its inverse f -1 so that f -1 ( x ) = 0 for x < f (0) and f -1 ( x ) = 1 for x > f (1) , making it continuous and non-decreasing on all of R ), the primal renormalization map T ϕ is given for x ∈ ∆ sub ,k by:

Since ν ↦→ f -1 ( f ( x i ) + ν ) is non-decreasing 5 in ν , the solution can be found efficiently using off-the-shelf root-finding algorithms such as Brent's method.

Dual renormalization. In contrast to the primal case, dual Bregman projections have (to our knowledge) not been directly studied in prior literature. They also offer new challenges: even their uniqueness cannot be taken for granted due to the general nonconvexity of Bregman divergences in the second argument [3]. To pave the road towards dual Bregman projections, we will therefore rely on additional structure in ϕ and d ϕ , expressed as the following dual validity condition.

Assumption 3.2 (Dual validity) . The map ϕ is thrice differentiable on (0 , 1] with lim x → 0 + xϕ ′′ ( x ) = 0 . For x ∈ (0 , 1] ,y ↦→ d ϕ ( x,y ) is strictly convex for y ∈ [ x, 1] , and y ↦→ d ϕ (0 ,y ) is strictly convex for y ∈ (0 , 1] .

̸

We establish in Theorem A.1 (see Appendix A) that subject to dual validity, the dual renormalization map T ∗ ϕ is uniquely defined for any x ∈ ∆ sub ,k with x = 0 k by the following implicit equations:

This transformation is interpretable despite its implicit nature: For every index i ∈ [ k ] , Equation 5 has the effect of increasing the corresponding probability x i by a positive additive amount regulated by an auxiliary variable ν ∗ ; the latter is chosen to make the increased topk probabilities sum to 1 .

Assumption 3.2, short of requiring global convexity of d ϕ ( x, · ) on [0 , 1] , only enforces it for y ∈ [ x, 1] . To enable this relaxation, the proof of Theorem A.1 carefully excludes optimal solutions belonging to the region y ≤ x or to the simplex boundary. Rather than a mere curiosity, this refinement substantially expands the scope of dual decoding. In particular, in our later specialization, it is essential for ensuring that dual α -decoding is uniquely defined for all α > 1 , not just α ∈ (1 , 2] : as plots in Appendix G.4 demonstrate, α -Bregman divergences are nonconvex for y ≤ x for α > 2 .

0.0200

0.0200

Originalx

Originalx

0.0175

Rescaledx

0.0175

Rescaledx

以

lue

0.0150

α=1.5

0.0150

α=1.5

izec

0.0125

α=4

α=10

lizec

0.0125

α=4

α=10

0.0100

0.0100

w

0.0075

0.0075

o

0

0.0050

n

0.0050

R

R

0.0025

0.0025

0.0000

0.0000

0.0000.0020.0040.0060.0080.0100.012

0.0000.0020.0040.0060.0080.0100.012

Originalxi

Originalxi

See Section F for algorithmic details on computing the dual map, as well as pseudocode for our algorithms. Figure 2 illustrates the primal and dual renormalization maps for α -Bregman divergences (introduced in Section 4). In this concrete example, T ϕ and T ∗ ϕ appear similar; however, for different, e.g. more 'peaked', inputs x ∈ ∆ sub ,k , they are more distinct, as we illustrate in Appendix G.3.

3.2 Greedy property: Justifying topk selection

The viewpoint that lower-probability tokens can be considered as noisy [32] suggests that it would be natural and indeed desirable for a decoding strategy to be 'greedy'-dictating that it is optimal to renormalize over the topk -probability tokens, for some k ∈ [ V ] . We formalize this as follows.

Definition 3.1 (Greedy decoding) . A decoding strategy Dec : ∆ V → ∆ V is called greedy if for every p ∈ ∆ V , the set of nonzero entries of Dec( p ) is a set of topˆ k entries of p , for some ˆ k = ˆ k ( p ) .

While many popular decoding methods are greedy [12, 21, 32, 38], some are not [22, 36]; justifications for non-greediness, i.e., the ability to occasionally throw out some of the topk tokens, include that this can e.g. help generate more 'typical' text. As such, our assertion that Bregman decoding strategies are greedy is nontrivial and requires proof. First, we state our result for primal Bregman decoding.

Theorem 3.2 (Primal Bregman decoding is greedy) . The primal Bregman decoding strategy from (2) is greedy for any primal valid potential ϕ .

The proof is provided in Appendix B. It proceeds by decomposing the Bregman objective into several terms, see Lemma B.2, and bounding them with the help of the primal renormalization equations (4).

The dual case, owing i.a. to the implicit form of the dual renormalization formulas (5), is correspondingly more complex to handle. Unlike in Theorem 3.2, our next result requires further conditions, which we state as a menu of two options. The relationship between the extra assumptions is intricate; Assumption (A2) is implied by, but is strictly weaker than, log-convexity of ϕ ′ .

Theorem 3.3 (Dual Bregman decoding is greedy) . The dual Bregman decoding strategy from (2) is greedy for any dual-valid ϕ with ϕ ′ (0) = 0 that further satisfies either of the following conditions: (A1) ϕ ′ is convex;

(A2) The maps 6 u defined as u ( x ) := xϕ ′′ ( x ) /ϕ ′ ( x ) for x ∈ (0 , 1] and ϕ are nondecreasing.

The proof is provided in Appendix C. In it, we use two different proof techniques for both conditions: For Condition (A1), our proof in Appendix C.1 leverages the decomposition from the primal case along with the change of variables d ϕ ( x, y ) = d ϕ ∗ ( ϕ ′ ( y ) , ϕ ′ ( x )) , where ϕ ∗ is the convex conjugate of ϕ . For Condition (A2), we develop a saddle-point proof approach in Appendix C.2. For that, we perform a sensitivity analysis of both the renormalized values [ T ∗ ϕ ( p )] i and of the per-coordinate Bregman loss terms, relative to hypothetical changes in the dual Lagrange multiplier ν ∗ and in the entries p i of p ; we carry this out via implicit differentiation of the defining equations (5).

3.3 k -convexity: Speeding up the search for optimal adaptive k

We have seen that for fixed k , greedily selecting the top k tokens is optimal. However, without further structure, we would still have to search over all k ∈ [ V ] to determine the optimal k ∗ , which would be cost-prohibitive for large token vocabularies. Fortunately, as we will see, only logarithmically many values of k will need to be tried, as under greedy selection, the primal and dual Bregman decoding objectives both enjoy discrete convexity with respect to k .

To formally state our result, fix a divergence Div , probability vector p ∈ ∆ V , and hyperparameter λ . We denote the regularized cost of selecting the topk entries of p , as a function of k ∈ [ V ] , by:

Recall that a function h : [ V ] → R is discretely convex if for all k ∈ { 2 , . . . , V -1 } , its discrete second derivative ∆ 2 h ( k ) := ∆ h ( k +1) -∆ h ( k ) := { h ( k +1) -h ( k ) } - { h ( k ) -h ( k -1) } ⩾ 0 . Theorem 3.4 (Discrete primal and dual cost convexity) . cost( · ) is discretely convex in k ∈ [ V ] for:

2 . Div(ˆ p, p ) = D ( p, ˆ p ) , if ϕ is dual valid.

Figure 6 in Appendix G.5) illustrates the result of Theorem 3.4 by displaying the cost( · ) functions for primal and dual Bregman α -decoding (defined in Section 4 below) for assorted α .

Implications for efficient computation. As a corollary of Theorem 3.4, an optimal k ∗ is provably identifiable by searching for k for which ∆cost( k ) ⩽ 0 and ∆cost( k +1) ⩾ 0 , for which repeated bisection (binary search) over 1 ≤ k ≤ V suffices - and thus, only O (log V ) tries of k are necessary. However, even less computation can suffice if one leverages that the optimal k is typically small. First, if one heuristically sets a hard limit k u on k (e.g. k u = 50 ), then identifying an optimal k ∈ [ k u ] requires O ( k u ) tries. Secondly, one may use exponential search 7 instead of binary search over k : this requires only O (log k ∗ ) tries - very small for typical values of k ∗ -and has the added benefit that only renormalizations over at most O ( k ∗ ) tokens are performed at each step.

Proving Theorem 3.4. Our proof uses two distinct approaches for the primal and the dual cases:

Primal k -convexity. The proof is developed in Appendix D. As its cornerstone, we use the Legendre dual mapping ϕ ∗ of the generator ϕ to establish and leverage the following cost structure: for any k , cost( k ) can up to additional terms be represented as max ν ≥ 0 [ ν -∑ k i =1 ϕ ∗ ( ϕ ′ ( p i ) + ν ) ] . This expression is concave in ν , and its unique optimizer is ν k , the optimal Lagrange multiplier for renormalizing the top k probabilities of p from (4). Using this, we then establish ∆ 2 cost( k ) ≥ 0 .

Dual k -convexity. The proof is in Appendix E. The above dualization strategy does not directly apply. Instead, we lower bound ∆ 2 cost ∗ ( k ) by regrouping the loss contributions of the indices i ∈ [ k +1] , and -via intricate term rearrangement and bounding-reduce to proving the local concavity of a special transformation (Equation 20) that turns out to hold by our dual-validity assumption.

4 Example: Bregman α -decoding

We now consider, as an illustration, a single-parameter family of Bregman decoding strategies, which arises via the generators of the Havrda-Charvát-Tsallis α -entropies [8, 29, 45, 51, 52]:

̸

When α < 0 and x = 0 , we set x α := + ∞ so that ϕ α (0) = ∞ . For α = 1 , one defines ϕ 1 ( x ) = x log( x ) , which corresponds to the Shannon entropy, arising in the limit 8 as α → 1 . Observe that ϕ α is primal valid for all α = 0 , as ϕ ′′ α ( x ) = x α -2 . This yields the following primal family of renormalizations, which we will index by α rather than ϕ :

Definition 4.1 (Primal Bregman α -decoding) . Fix α ∈ J, k ∈ [ V ] . The renormalization map T α is given for p ∈ ∆ sub ,k as: [ T α ( p )] i =( p α -1 i + ν ) 1 α -1 for i ∈ [ k ] , with ν ∈ R chosen so that ∑ i ∈ [ k ] [ T α ( p )] i =1 .

Note that for α = 1 , we have ϕ ′ 1 ( x ) = log x +1 . Hence, (4) implies e ν ∑ k i =1 p i = 1 , and we obtain the 'standard' renormalization: [ T 1 ( p )] i = p i / ( ∑ k j =1 p j ) , for i ∈ [ k ] . Therefore, primal Bregman 1-decoding is topk decoding , showing how one recovers topk in our framework. It turns out that some further values of α also lead to renormalization maps of special interest. For any fixed p , we let T -∞ ( p ) = lim inf α →-∞ T α ( p ) and T ∞ ( p ) = lim inf α →∞ T α ( p ) , where the limits are entrywise.

Proposition 4.2 (Special primal α -renormalization maps) . We have the following special instances 9 of the primal Bregman α -renormalization map, defined for all i ∈ [ k ] as follows:

Along with the primal family, the dual α -decoding family can also be defined based on ϕ α . Unlike α -decoding, the dual Bregman sparse decoding Problem 2 can be non-convex, as displayed in Figure 1 above. Figure 5 in Appendix G.4 further demonstrates the nonconvexity of D ϕ α on the unit square for some α . Yet, we can still show that any dual α -decoding with α > 1 is valid, greedy and k -convex:

Lemma 4.3. All generator functions ϕ α , α > 1 , are dual-valid and satisfy Assumption (A2).

We give an illustration contrasting primal and dual α -decoding for various α > 1 in Appendix G.3.

5 Experiments

We now illustrate some of the decoding schemes described in our paper in the context of LLMs. Since our goal is to develop the theoretical foundations of topk decoding, our aim in this section is simply to illustrate that the performance of our novel decoding schemes can be competitive with standard topk decoding. In particular, we do not aim to compare or compete with other popular and established decoding methods, which is beyond the scope of our theory-focused paper.

5.1 Experimental Setup

Method. In addition to standard topk decoding, which coincides with the α = 1 case of our primal α -decoding family described in Section 4, we illustrate primal α -decoding strategies for α = 1 . 5 and α = 2 . These have closed-form renormalization maps that are as fast as standard renormalization.

Full and partial evaluation. Further, we perform two types of experiments: (1) For the evaluation of our full decoding strategy, we decode by adaptively selecting the optimal sparsity parameter k ∗ by optimizing our sparse Bregman objective. In this approach, we aim to observe the behavior when adaptively choosing k ∗ . Since practical choices of k ∗ are always upper bounded, we set a maximum k ∗ ⩽ k max := 50 . (2) In the partial evaluation approach, we instead directly evaluate-for each fixed choice of k in the grid k ∈ { 5 , 10 , . . . , 50 } -our proposed renormalization strategies along with standard topk renormalization.

Models and benchmarks. We conduct experiments using the GPT-2 Large [43] and Llama 3.1 8B [25] models. We evaluate on two benchmarks: (1) open-ended text generation using the WebText test set from the GPT-2 output dataset [40], and (2) grade school math reasoning using the GSM8K Chainof-Thought benchmark [13]. Additional experiments with larger models, Qwen2.5-14B-Instruct and Phi-3-Medium-4K-Instruct, as well as evaluations on the TriviaQA benchmark, are presented in Appendix H.

Evaluation metrics. For open-ended text generation, following Chen et al. [12], we use the first 35 tokens of each WebText test sample as a prompt and generate up to 256 tokens. We evaluate the following standard metrics [see e.g., 12, 32, 38, etc]:

(1) Perplexity difference , which measures the perplexity (according to base model p base ) of human text compared to that obtained from a decoding strategy p decoding derived from the base model, where lower is better. This equals E X ∼D [ E Y ∼D ( ·|X ) ( p base ( Y | X ) -1 / | Y | ) -E Y ∼ p decoding ( ·| X ) ( p base ( Y | X ) -1 / | Y | )] , where X ∼ D is a

prompt drawn from the dataset, Y ∼ D ( ·|X ) denotes a human-written continuation drawn from the dataset, and Y ∼ p decoding ( · | X ) denotes a model-generated continuation using a specific decoding strategy. Here, | Y | is the length of the continuation.

(2) Repetition difference : E X ∼D [ P Y ∼ p decoding ( ·| X ) (rep( Y )) -P Y ∼D ( ·| X ) (rep( Y )) ] , where rep( Y ) is the event that Y contains two contiguous and identical token spans of length ⩾ 2 ; lower is better.

5.2 Results

Open-ended text generation. Using the partial evaluation setup with temperature fixed at 1.0, Figure 3 reports the differences in perplexity and repetition frequency between model-generated and human-written text across a range of k values. Primal decoding strategies are competitive with topk in terms of both metrics. In particular, α = 2 . 0 has the smallest gaps in perplexity and repetition frequency. The marked decrease in repetitiveness for α ∈ { 1 . 5 , 2 } relative to standard topk ( α = 1 ) is consistent with the theoretical behavior of α -Bregman renormalization that we have derived. For fixed k , increasing α shifts the renormalization from upweighting the largest probabilities to comparatively upweighting the smaller probabilities within the selected topk set, thereby increasing sampling diversity.

GSM8K dataset. Using the full decoding strategy, we evaluate the LLaMA 3.1 8B model using 8-shot CoT prompting. We test various temperatures, regularization strengths λ ∈ { 0 . 01 , 0 . 0001 } and primal decoding parameters α ∈ { 1 . 5 , 2 . 0 } . Results for other settings are in Appendix H. To ensure a matched comparison, we run topk with k = k ∗ for the Bregman decoding run with the same temperature, λ , and α , rounded to the nearest integer, see Table 11 in Appendix H. As seen in Table 1, across all temperature settings, primal decoding with adaptive k ∗ achieves accuracy comparable to topk . At higher temperatures (such as 1 . 5 ), the performance of topk decoding degrades more rapidly than that of primal decoding.

10

20

30

40

50

k

4

6

8

10

12

Perplexity Diff

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

2

4

6

8

10

Repetition Freq Diff(%)

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

5

6

7

8

9

10

Perplexity Diff

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Repetition Freq Diff(%)

Top-k

Primal-1.5

Primal-2.0

6 Related work

Bregman projection. Michelot [37] considered the Brier score projection problem and derived an efficient algorithm. Later, Shalev-Shwartz et al. [48] revisited the properties of optimal Brier projection, and Duchi et al. [17] gave and analyzed the explicit algorithm that we discuss in what follows. Wang and Carreira-Perpinán [53] simplified and distilled the proof. [35] further studied the projection as a method for generating sparse probability predictions in multiclass prediction problems. [33, 34] developed methods for efficient Bregman projections to the simplex; for a fixed support, these results characterize our primal decoding. [44, 46] developed differentiable variants of topk decoding. In contrast to these works, we: (1) consider Bregman projections under ℓ 0 regularization, and (2) offer, to the best of our knowledge, novel analyses of dual Bregman projections.

ℓ 0 regularization. Regularization via the ℓ 0 -pseudonorm has been studied widely, with various approximate algorithms (based on surrogates, integer programming, branch-and-bound methods, etc.) developed for problems ranging from linear regression to more general learning tasks [see e.g., 2, 6, 9, 15, 18-20, 30, 41, 49, 50, 58, 61, etc]. In contrast, the algorithms we propose are exact within numerical precision for the specific class of problems we consider.

Bregman divergences. The properties of Bregman divergences [10] have been widely studied; see, e.g., [1, 3, 5, 8, 27, 39, 47, 55, 57], etc. In particular, there are a number of relations between Bregman divergences and their versions with reversed arguments, motivated by the fact that convexity in the first parameter allows for minimization, making it useful to switch the order of the variables, see e.g., [1, 26] etc. We both leverage some of these results in our work, and contribute some, to the best of our knowledge, novel proof techniques and insights into the (primal and dual) Bregman geometry.

LLMdecoding. There is a vast range of work on LLM sampling (or decoding), see e.g., [54] and references therein. Classical methods include greedy sampling and beam search. Sparse sampling methods such as topk sampling [21] are motivated by intuition that the 'unreliable tail' of lowprobability tokens is mis-estimated [32]. In particular, [32] propose topp sampling, and [38] propose minp sampling. Other sampling methods were proposed in [4, 22, 31, 36]. [12] propose the decoding game, a two-player game between an LLM and an adversary that distorts the true distribution. They show that certain sparse truncated sampling methods are approximately minimax optimal. For other approaches to make LLM output probabilities sparse, see e.g., [14, 59, 60]. In contrast, our goal is to develop a deep theoretical understanding of topk decoding, placing it into a broader framework.

General motivation. The motivation for our general approach is two-fold: (1) Without sparsity considerations, Bregman divergences closely correspond to proper scoring rules, and are minimized at the true probability distribution, see e.g., [10, 24]. This property is highly desirable in probabilistic forecasting and prediction, incentivizing a forecaster to predict the true distribution in order to minimize their loss. (2) The ℓ 0 -pseudonorm has been widely argued to both be a reasonable measure of sparsity, and to have good properties as a regularizer in certain sparse estimation problems such as sparse regression [see e.g., 7, 16, 23, 28, etc]. Combining these two lines of thought provides the motivation for studying ℓ 0 -regularized Bregman divergence minimization.

7 Discussion

This paper develops a theoretical foundation for topk decoding. We hope that our framework, which rests on the structural pillars of (i) greedy selection and (ii) k -convexity, will spur the development of novel theoretically motivated adaptive sparse decoding methods for LLMs and beyond.

LLM decoding beyond topk . Our analysis treats per-step decoding as jointly selecting a support size k and applying a renormalization via Bregman projection. A natural extension that future work may address is to develop analogous optimization-theoretic foundations for other popular truncation rules, such as topp and minp , which also follow the template 'truncate a tail, then renormalize'.

Quantifying and controlling adaptivity. A central feature of our framework is that the parameter k ⋆ ( p ) is chosen adaptively with p . A useful next step is to characterize how k ⋆ ( p ) varies with the hyperparameter λ and with properties of the token distribution p (e.g., concentration or tail decay). For α -decoding, our preliminary calculations in stylized settings suggest that there may be natural scaling laws relating k ⋆ to ( α, λ ) ; an interesting open problem is to rigorously derive such laws under distributional models that reflect LLM next-token probabilities (e.g., peaked or heavy-tailed regimes), and to translate them into practical guidance for targeting a desired average sparsity level.

Primal vs. dual decoding. While we established greedy selection and k -convexity for both primal and dual formulations, their relationship is not yet well understood. As we illustrate via several examples, primal and dual α -decoders are similar yet subtly different for corresponding α . Thus, more generally quantifying the primal-dual Bregman decoding gap would be of interest.

Broader evaluation. Our initial experiments hint that some performance gains over vanilla topk might be found even amongst α -decoders. It is promising and essential to further understand how Bregman decoders could be optimized towards downstream quality/factuality measures in practice.

Beyond LLM decoding. Finally, at a technical level, our main results identify a broad class of nonconvex sparse optimization problemsℓ 0 -regularized separable Bregman projection onto the simplex-that nevertheless admit exact, efficient solutions via greedy support selection and discrete convexity in the support size. We conjecture that the resulting general-purpose efficient optimization primitive for sparsifying probability vectors could be applicable well beyond LLM decoding.

References

Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry , volume 191. American Mathematical Soc., 2000.

Sohail Bahmani, Bhiksha Raj, and Petros T Boufounos. Greedy sparsity-constrained optimization. Journal of Machine Learning Research , 14(1):807-841, 2013.

Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, and Joydeep Ghosh. Clustering with Bregman divergences. Journal of Machine Learning Research , 6(Oct):1705-1749, 2005.

Sourya Basu, Govardana Sachitanandam Ramachandran, Nitish Shirish Keskar, and Lav R. Varshney. Mirostat: A neural text decoding algorithm that directly controls perplexity. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=W1G1JZEIy5_ .

Heinz H Bauschke and Patrick L Combettes. Iterating Bregman retractions. SIAM Journal on Optimization , 13(4):1159-1173, 2003.

Dimitris Bertsimas, Angela King, and Rahul Mazumder. Best subset selection via a modern optimization lens. The Annals of Statistics , 44(2):813, 2016.

Lucien Birgé and Pascal Massart. Gaussian model selection. Journal of the European Mathematical Society , 3(3):203-268, 2001.

Mathieu Blondel, André F.T. Martins, and Vlad Niculae. Learning with Fenchel-Young losses. Journal of Machine Learning Research , 21(35):1-69, 2020. URL http://jmlr.org/papers/ v21/19-021.html .

Thomas Blumensath and Mike E Davies. Iterative hard thresholding for compressed sensing. Applied and Computational Harmonic Analysis , 27(3):265-274, 2009.

Lev M Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR Computational Mathematics and Mathematical Physics , 7(3):200-217, 1967.

Emmanuel J Candes and Terence Tao. Decoding by linear programming. IEEE Transactions on Information Theory , 51(12):4203-4215, 2005.

Sijin Chen, Omar Hagrass, and Jason Matthew Klusowski. Decoding game: On minimax optimality of heuristic text generation strategies. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/forum?id=Wfw4ypsgRZ .

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021.

Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 2174-2184, 2019.

Antoine Dedieu, Hussein Hazimeh, and Rahul Mazumder. Learning sparse classifiers: Continuous and mixed integer optimization perspectives. Journal of Machine Learning Research , 22(135):1-47, 2021. URL http://jmlr.org/papers/v22/19-1049.html .

David L Donoho and Iain M Johnstone. Ideal spatial adaptation by wavelet shrinkage. Biometrika , 81(3):425-455, 1994.

John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto the ℓ 1 -ball for learning in high dimensions. In Proceedings of the 25th International Conference on Machine Learning , pages 272-279, 2008.

M'hamed Essafri, Luca Calatroni, and Emmanuel Soubies. Exact continuous relaxations of ℓ 0 -regularized criteria with non-quadratic data terms, 2024. URL https://arxiv.org/abs/ 2402.06483 .

M'hamed Essafri, Luca Calatroni, and Emmanuel Soubies. On ℓ 0 Bregman-relaxations for Kullback-Leibler sparse regression. In 2024 IEEE 34th International Workshop on Machine Learning for Signal Processing (MLSP) , pages 1-6. IEEE, 2024.

Mhamed Essafri, Luca Calatroni, and Emmanuel Soubies. Box-constrained ℓ 0 Bregmanrelaxations, 2025. URL https://arxiv.org/abs/2503.15083 .

Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) . Association for Computational Linguistics, 2018.

Matthew Finlayson, John Hewitt, Alexander Koller, Swabha Swayamdipta, and Ashish Sabharwal. Closing the curious case of neural text degeneration. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum? id=dONpC9GL1o .

Dean P Foster and Edward I George. The risk inflation criterion for multiple regression. The Annals of Statistics , 22(4):1947-1975, 1994.

Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association , 102(477):359-378, 2007.

Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Llama 3 herd of models. arXiv preprint arXiv:2407.21783 , 2024.

Sebastian Gruber and Florian Buettner. Uncertainty estimates of predictions via a general biasvariance decomposition. In International Conference on Artificial Intelligence and Statistics , pages 11331-11354. PMLR, 2023.

Peter D Grünwald and A Philip Dawid. Game theory, maximum entropy, minimum discrepancy and robust bayesian decision theory. Ann. Statist. , 32(1):1367-1433, 2004.

Trevor Hastie, Robert Tibshirani, and Martin Wainwright. Statistical Learning with Sparsity: The Lasso and Generalizations . CRC Press, 2015.

Jan Havrda and František Charvát. Quantification method of classification processes. concept of structural a -entropy. Kybernetika , 3(1):30-35, 1967.

Hussein Hazimeh, Rahul Mazumder, and Tim Nonet. L0learn: A scalable package for sparse learning using ℓ 0 regularization. Journal of Machine Learning Research , 24(205):1-8, 2023.

John Hewitt, Christopher D Manning, and Percy Liang. Truncation sampling as language model desmoothing. In Findings of the Association for Computational Linguistics: EMNLP 2022 , pages 3414-3427, 2022.

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=rygGQyrFvH .

Walid Krichene, Syrine Krichene, and Alexandre Bayen. Efficient Bregman projections onto the simplex. In 2015 54th IEEE Conference on Decision and Control (CDC) , pages 3291-3298. IEEE, 2015.

Cong Han Lim and Stephen J. Wright. Efficient Bregman projections onto the permutahedron and related polytopes. In Arthur Gretton and Christian C. Robert, editors, Proceedings of the 19th International Conference on Artificial Intelligence and Statistics , volume 51 of Proceedings of Machine Learning Research , pages 1205-1213, Cadiz, Spain, 09-11 May 2016. PMLR.

Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. In International Conference on Machine Learning , pages 16141623. PMLR, 2016.

Clara Meister, Tiago Pimentel, Gian Wiher, and Ryan Cotterell. Locally typical sampling. Transactions of the Association for Computational Linguistics , 11:102-121, 2023.

Christian Michelot. A finite algorithm for finding the projection of a point onto the canonical simplex of R n . Journal of Optimization Theory and Applications , 50:195-200, 1986.

Nguyen Nhat Minh, Andrew Baker, Clement Neo, Allen G Roush, Andreas Kirsch, and Ravid Shwartz-Ziv. Turning up the heat: Minp sampling for creative and coherent LLM outputs. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/forum?id=FBkpCyujtS .

Frank Nielsen. An elementary introduction to information geometry. Entropy , 22(10):1100, 2020.

OpenAI. GPT-2 output dataset. https://github.com/openai/gpt-2-output-dataset , 2019. URL https://github.com/openai/gpt-2-output-dataset .

Jianting Pan and Ming Yan. Efficient sparse probability measures recovery via Bregman gradient. Journal of Scientific Computing , 102(3):66, 2025.

Christos H Papadimitriou and Kenneth Steiglitz. Combinatorial optimization: Algorithms and complexity . Courier Corporation, 1998.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Luca Ragazzi, Paolo Italiani, Gianluca Moro, and Mattia Panni. What are you token about? Differentiable perturbed topk token selection for scientific document summarization. In Findings of the Association for Computational Linguistics ACL 2024 , pages 9427-9440, 2024.

Daniel Reem, Simeon Reich, and Alvaro De Pierro. Re-examination of Bregman functions and new properties of their divergences. Optimization , 68(1):279-348, 2019.

Michael Eli Sander, Joan Puigcerver, Josip Djolonga, Gabriel Peyré, and Mathieu Blondel. Fast, differentiable and sparse topk : A convex analysis perspective. In International Conference on Machine Learning , pages 29919-29936. PMLR, 2023.

Leonard J Savage. Elicitation of personal probabilities and expectations. Journal of the American Statistical Association , 66(336):783-801, 1971.

Shai Shalev-Shwartz, Yoram Singer, Kristin P Bennett, and Emilio Parrado-Hernández. Efficient learning of label ranking by soft projections onto polyhedra. Journal of Machine Learning Research , 7(7), 2006.

Yiyuan She, Zhifeng Wang, and Jiuwu Jin. Analysis of generalized Bregman surrogate algorithms for nonsmooth nonconvex statistical learning. The Annals of Statistics , 49(6): 3434-3459, 2021.

Emmanuel Soubies, Laure Blanc-Féraud, and Gilles Aubert. A continuous exact ℓ 0 penalty (cel0) for least squares regularized problem. SIAM Journal on Imaging Sciences , 8(3):16071639, 2015.

Constantino Tsallis. Possible generalization of Boltzmann-Gibbs statistics. Journal of Statistical Physics , 52:479-487, 1988.

Constantino Tsallis. Introduction to nonextensive statistical mechanics: Approaching a complex world . Springer, 2009.

Weiran Wang and Miguel A Carreira-Perpinán. Projection onto the probability simplex: An efficient algorithm with a simple proof, and an application. arXiv preprint arXiv:1309.1541 , 2013.

Sean Welleck, Amanda Bertsch, Matthew Finlayson, Hailey Schoelkopf, Alex Xie, Graham Neubig, Ilia Kulikov, and Zaid Harchaoui. From decoding to meta-generation: Inference-time algorithms for large language models. Transactions on Machine Learning Research , 2024. ISSN 2835-8856. URL https://openreview.net/forum?id=eskQMcIbMS . Survey Certification.

Robert C Williamson, Elodie Vernet, and Mark D Reid. Composite multiclass losses. Journal of Machine Learning Research , 17(222):1-52, 2016.

Wotao Yin, Stanley Osher, Donald Goldfarb, and Jerome Darbon. Bregman iterative algorithms for ℓ 1 -minimization with applications to compressed sensing. SIAM Journal on Imaging Sciences , 1(1):143-168, 2008.

Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 5852-5861. PMLR, 10-15 Jul 2018. URL https://proceedings. mlr.press/v80/zhang18l.html .

Jacky Y Zhang, Rajiv Khanna, Anastasios Kyrillidis, and Oluwasanmi O Koyejo. Learning sparse distributions using iterative hard thresholding. Advances in Neural Information Processing Systems , 32, 2019.

Guangxiang Zhao, Junyang Lin, Zhiyuan Zhang, Xuancheng Ren, Qi Su, and Xu Sun. Explicit sparse transformer: Concentrated attention through explicit selection. arXiv preprint arXiv:1912.11637 , 2019.

Shuai Zhao, Qing Li, Yuer Yang, Jinming Wen, and Weiqi Luo. From softmax to nucleusmax: A novel sparse language model for chinese radiology report summarization. ACM Transactions on Asian and Low-Resource Language Information Processing , 22(6):1-21, 2023.

Junxian Zhu, Jin Zhu, Borui Tang, Xuanyu Chen, Hongmei Lin, and Xueqin Wang. Best-subset selection in generalized linear models: A fast and consistent algorithm via splicing technique. arXiv preprint arXiv:2308.00251 , 2023.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims in the abstract and introduction are supported by the theoretical results throughout the paper and the experiments in the Experiments section.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of our work in the Discussion section with an eye toward future work addressing them.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theorems are always stated with their required assumptions, and full proofs are provided in the Appendices.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All experimental details are provided in Section 5 and Appendix H, including dataset information. Code is included in the supplementary materials, and an open-source GitHub repository will be released upon publication.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The software package implementing our method and reproducing the experiments is included in the supplementary material. The GPT-2 output dataset and GSM8K dataset used are publicly available.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental settings are described in Section 5. We evaluate pretrained models without training. Section 5 specifies all datasets, decoding parameters, and metrics needed to reproduce the results.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All plots and tables report bootstrap standard errors; see section 5 and Appendix H for the corresponding figures and tables.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Section H.1 provides the details of the computational resources used in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in this paper follows the NeurIPS Code of Ethics in all aspects.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: We briefly discuss the potential societal impact of our work in our Discussion. In more detail, our work performs a foundational investigation into the theory of decoding methods, and as such, does not have immediate societal impacts; however, by laying and reexamining the foundations of decoding for language models, one hopes that it can contribute to more transparent and responsible LLM decoding methods in the future, in this way potentially leading to positive societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper does not involve the release of any pretrained models, image generators, or datasets that pose a high risk of misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All datasets and sources used in this paper are properly cited.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets. However, we provide a welldocumented software package.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core methodology in this research does not involve LLMs.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Appendix Contents

A Existence and uniqueness of dual Bregman decoding

Theorem A.1 (Uniqueness and formula for dual Bregman renormalization) . Fix a dual valid potential ϕ . Then, for any x ∈ ∆ sub ,k with ∑ i x i > 0 , the renormalization map T ∗ ϕ is uniquely defined by:

Proof. First, assume without loss of generality that 0 < ∑ i ∈ [ k ] x i < 1 . Otherwise, if ∑ i ∈ [ k ] x i = 1 then x ∈ ∆ k , so the unique unconstrained optimum, which is at x by the standard property of Bregman divergences, is also the unique optimum of our constrained projection problem.

Note that Slater's condition is satisfied for this projection problem as we are optimizing over the simplex (whose relative interior is nonempty). Therefore, in this differentiable problem, its optimal solutions can be characterized via its KKT conditions.

Introduce a Lagrange multiplier ν ∈ R for the simplex constraint, and Lagrange multipliers ( λ i ) i ∈ [ k ] for the nonnegativity constraints. Then, the Lagrangian is as follows:

Here, λ i ⩾ 0 for all i , and by complementary slackness, at optimality λ i = 0 whenever ˆ p i > 0 .

For each i ∈ [ k ] , the stationarity condition reads (except possibly when ˆ p i = 0 , where the second derivative could be infinite):

In particular, for each coordinate i for which the optimal ˆ p i ∈ (0 , 1) , the stationarity condition is:

Now, we show that ν > 0 . Indeed, observe that there must be at least one index i for which ˆ p i > x i . If that was not the case, we would get ∑ i ∈ [ k ] ˆ p i ⩽ ∑ i ∈ [ k ] x i < 1 by our assumption, contradicting that ˆ p ∈ ∆ k . In particular, then, ˆ p i > x i ⩾ 0 , and therefore we have ϕ ′′ (ˆ p i )(ˆ p i -x i ) = ν . Since ϕ ′′ (ˆ p i ) > 0 and ˆ p i -x i > 0 , we thus conclude that ν > 0 .

Having shown that ν > 0 , we now proceed to show that all ˆ p i > 0 at optimality. Note that ∂ ∂y d ϕ ( x, y ) = ϕ ′′ ( y ) ( y -x ) for y > 0 . We will now consider two cases:

If ϕ ′′ (0) is finite, ˆ p i > 0 for all i . Indeed, suppose that was not the case, and ˆ p i = 0 for some i . Then we would have: ϕ ′′ (0)(0 -x i ) = ν + λ i , or equivalently, ϕ ′′ (0) · x i + ν + λ i = 0 . Each of the three terms is nonnegative, and ν > 0 , so we arrive at a contradiction.

Next, consider the case in which lim y → 0 ϕ ′′ ( y ) = + ∞ . Then, lim y → 0 ∂ ∂y d ϕ ( x, y ) = -∞ for all x ∈ (0 , 1] . Then, since lim y → 0 ∂ ∂y d ϕ ( x, y ) = -∞ for all x ∈ (0 , 1] , for any i such that x i > 0 , setting ˆ p i = 0 would lead to ν = -∞ , hence necessarily ˆ p i > 0 . On the other hand, for any i for which x i = 0 , since lim y → 0 yϕ ′′ ( y ) = 0 , setting ˆ p i = 0 would lead to ν = 0 , which is a contradiction.

In all cases, the optimal ˆ p is in the strict interior of the simplex, so it suffices to solve (7) over this range. To show that the solution exists and is unique, we collect together the following information about Ψ from (13) with Ψ( x, y, ν ) := ϕ ′′ ( y )( y -x ) -ν for all x, y, ν . Then, for a fixed ν , (7) is equivalent to solving Ψ( x i , ˆ p i , ν ) = 0 . First, consider x > 0 . Then, we have the following:

Since the map y ↦→ d ϕ ( x, y ) is strictly convex for y ∈ [ x, 1] , it follows that ∂ ∂y d ϕ ( x, y ) = Ψ( x, y, 0) is strictly increasing for y ∈ [ x, 1] , and so is Ψ( x, y, ν ) .

We have Ψ( x, x, ν ) = -ν ⩽ 0 . Further, Ψ( x, 1 , ν ) = ϕ ′′ (1)(1 -x ) -ν ⩾ 0 , whenever ν ⩽ ϕ ′′ (1)(1 -x ) .

Hence, the map y ↦→ Ψ( x, y, ν ) has a unique zero on the interval [ x, 1] , as long as 0 < ν ⩽ ϕ ′′ (1)(1 -x ) .

Next, consider x = 0 , in which case we need to solve the equation ϕ ′′ ( y ) y = ν . Then, we have the following:

Since the map y ↦→ d ϕ (0 , y ) is strictly convex for y ∈ (0 , 1] , it follows that ∂ ∂y d ϕ (0 , y ) = Ψ(0 , y, 0) = ϕ ′′ ( y ) y is strictly increasing for y ∈ (0 , 1] , and so is Ψ(0 , y, ν ) .

By assumption, lim y → 0 + yϕ ′′ ( y ) = 0 , hence we have lim y → 0 + Ψ( x, x, ν ) = -ν ⩽ 0 . Further, Ψ(0 , 1 , ν ) = ϕ ′′ (1)(1 -x ) -ν ⩾ 0 , whenever ν ⩽ ϕ ′′ (1) .

Hence, the map y ↦→ Ψ(0 , y, ν ) has a unique zero on the interval (0 , 1] , as long as 0 < ν ⩽ ϕ ′′ (1) .

Now define M := min i ϕ ′′ (1)(1 -x i ) = ϕ ′′ (1)(1 -max i x i ) . Since by assumption ∑ i x i < 1 , it follows that M > 0 . From the above analysis, it follows that, as long as ν ∈ (0 , M ] , for each i , the equation ϕ ′′ ( y i )( y i -x i ) = ν . has a unique solution y i ( ν ) ∈ ( x i , 1] .

Furthermore, as we establish in Lemma C.2, the map ν ↦→ y i ( ν ) is strictly increasing for ν > 0 , also owing to the assumed second-argument convexity of d ϕ . In particular, define G ( ν ) = ∑ k i =1 y i ( ν ) for ν > 0 ; then G is continuous and strictly increasing, and satisfies lim ν → 0 G ( ν ) = ∑ i x i < 1 and G ( M ) ⩾ y i ∗ ( M ) = 1 , where i ∗ is any index achieving the maximum among the coordinates of x . Hence there is a unique ν ∗ ∈ (0 , M ] with G ( ν ∗ ) = 1 . Setting ˆ p i = y i ( ν ∗ ) yields a vector in ∆ k that satisfies the KKT stationarity.

Finally, note that the solution ˆ p that we just identified is unique. Indeed, we have earlier excluded boundary solutions from consideration, and then further excluded any solutions in which ˆ p i < x i for any i ∈ [ k ] ; thus, it suffices to recall that the Bregman objective is assumed to be strictly convex in the interior of the region of the simplex given by { ˆ p ∈ ∆ k : ˆ p i ⩾ x i for all i ∈ [ k ] } , thus concluding the proof.

B Proof of the primal greedy property in Theorem 3.2

We will first fix some notations. Henceforth, we will assume that the vector p has been sorted, i.e., p 1 ⩾ p 2 ⩾ . . . ⩾ p V . For any subset Q = { i 1 , . . . , i k } ⊆ [ V ] of size k , let Q c = [ V ] \ Q . Let p Q

denote the sub-probability vector with the entries of p whose indices are in Q . We define the loss L ( Q ) as

̸

Here, S Q c = ∑ j / ∈ Q d ϕ (0 , p j ) . To prove Theorem 3.2, we will show that L ( S ′ ) ⩾ L ( S ) for any S ′ ⊆ [ V ] of size k , where S = [ k ] consists of the topk indices. We will further show that strict inequality always holds if p S ′ = p S . To do this, we proceed in three steps: (1) We first simplify the form of the loss function L ( Q ) in Lemma B.1, (2) For any two subsets S, S ′ , we decompose the loss difference L ( S ′ ) -L ( S ) into three terms in Lemma B.2, (3) We individually analyze each of the terms in this decomposition and prove they are non-negative.

B.1 Decomposing the Bregman cost function on subsets

LemmaB.1. For any Q = { i 1 , i 2 , . . . i k } ⊆ [ V ] of size k , the loss function as defined in (8) simplifies to:

Proof. Observe that:

This further equals

This finishes the proof.

Let T Q ( p ) denote a minimizer of the above loss L ( Q ) , i.e.,

Note that ( a ) holds above as the term S Q c does not play any role in the location of the minimizer. However, it does contribute to the final loss L ( Q ) . Also, as the divergence is separable, once we have selected a subset Q , the ordering of its elements does not matter for the calculation of the above loss and minimizer. Thus, without loss of generality, we may assume i 1 < i 2 < . . . < i k for k ∈ [ V ] . By forming the Lagrangian and differentiating it, we obtain the primal thresholding from (4):

Here, ν Q is chosen such that ∑ k j =1 [ T Q ( p )] j = 1 .

Lemma B.2. Let S = { i 1 , . . . , i k } , S ′ = { i ′ 1 , . . . , i ′ k } ⊆ [ V ] and T S ( p ) and T S ′ ( p ) be the corresponding minimizers. Then, the following decomposition holds:

Proof. We have from Lemma B.1 that

This further equals

Now, returning to our proof, suppose S = [ k ] and S ′ = { i ′ 1 , . . . i ′ k } . We know from Lemma B.2 that

Now, consider the term II . Using (10), we can simplify this further as follows:

where ( a ) follows as ∑ k j =1 [ T S ′ ( p )] j = ∑ k j =1 [ T S ( p )] j = 1 . Also, I ⩾ 0 as D ϕ is a divergence measure.

Finally, to conclude our proof, we show that III ⩾ 0 . Since the entries of p are sorted in a nondecreasing order and as the indices in S = [ k ] and S ′ are sorted in ascending order, we have

Strict inequality holds as long as some p i ′ j is not among the topk indices of p .

C Proof of the dual greedy property in Theorem 3.3

To prove the greedy property for the two alternate conditions in Theorem 3.3, we will provide two distinct proof techniques for the two cases (A1) and (A2). The first one uses duality and the second one uses a saddle point argument. We will now recall the definition of the Legendre dual of a convex function-in this case, of the generator function ϕ -and its defining property that will help us. Below, f ([0 , 1]) denotes the image of [0 , 1] under f .

Lemma C.1 (Classical) . For a valid ϕ , let ϕ ∗ ( x ) = sup p ⩾ 0 { px -ϕ ( p ) } be the Legendre dual of ϕ , defined for all x ∈ f ([0 , 1]) . Then, we have for every x ∈ f ([0 , 1]) the identity: ϕ ( f -1 ( x )) = xf -1 ( x ) -ϕ ∗ ( x ) . Moreover ( ϕ ∗ ) ′ = f -1 , and ϕ ∗ is strictly increasing.

Proof. Since the map p ↦→ R ( p ) := px -ϕ ( p ) is continuous, it achieves a maximum on [0 , 1] . From the first order condition of the defining equation for ϕ ∗ , if the maximum is achieved in (0 , 1) , we have:

so for the maximizer p max we have f ( p max ) = x ⇒ p max = f -1 ( x ) . Now, since f is increasing and x ∈ f ([0 , 1]) , we have R ′ (0) = x -f (0) ⩾ 0 , with equality if x = f (0) . Similarly, R ′ (1) = x -f (1) ⩽ 0 , with equality if x = f (1) . Hence, it follows that the above characterization for the maximizer p max also applies on the boundaries of [0 , 1] . To conclude the proof of the identity, it suffices to observe that ϕ ∗ ( x ) = p max x -ϕ ( p max ) = xf -1 ( x ) -ϕ ( f -1 ( x )) . The expression for ( ϕ ∗ ) ′ follows by direct calculation.

C.1 Proof under Assumption (A1)

With the dual convex conjugate ϕ ∗ as per Lemma C.1, the divergence measure satisfies:

Let the loss for the dual problem be denoted as L ∗ , (the divergence measure with the arguments swapped), and let T ∗ Q be the dual renormalization map from Lemma A.1 applied to p Q , i.e.,

C.1.1 Decomposition of the loss difference

Using the form of the loss difference in Lemma (B.2) and (12), we can compute the loss difference for the dual problem as follows:

Indeed, changing the potential ϕ to ϕ ∗ , and changing all the arguments p i j , p i ′ j , T ∗ S , T ∗ S ′ to ϕ ′ ( p i j ) , ϕ ′ ( p i ′ j ) , ϕ ′ ( T ∗ S ) , ϕ ′ ( T ∗ S ′ ) respectively in Lemma (B.2) suffices. Thus, under the same setup

of the two subsets S = [ k ] and S ′ and denoting ϕ ′ = f , we obtain:

Since ( ϕ ∗ ) ′ = f -1 , this further equals

C.1.2 Analysis of terms based on the dual solution

Similar to the proof for the primal case, the term I ′ ⩾ 0 , as D ϕ ∗ is a divergence, and III ′ ⩾ 0 as ϕ ′ = f ⩾ 0 , as f (0) = 0 and f is increasing. Moreover, as f is strictly increasing, if any of the p i ′ j are not among the topk entries, then strict inequality holds.

To analyze II , we have

Since f is convex,

In the above steps, ( a ) follows as f ′ > 0 as f is strictly increasing and ( b ) follows as ∑ k j =1 [ T ∗ S ′ ( p )] j = ∑ k j =1 [ T ∗ S ( p )] j = 1 . This implies II ′ ⩾ 0 , finishing the proof.

C.2 Proof under Assumption (A2)

C.2.1 Extra notation

Define the mapping derived from solving Ψ( x, y, ν ) = 0 over y by:

It follows from the proof of Lemma A.1 that the solution ξ is well-defined. Define two auxiliary functions ψ, h that will be used in the computation of the Bregman costs below, such that for all ( x, y, ν ) ∈ D :

C.2.2 Properties of the auxiliary functions

Lemma C.2 (Derivatives ∂ξ ∂x , ∂ξ ∂ν ) . Define v : [0 , 1] × (0 , 1] → [0 , ∞ ) as v ( x, y ) = ϕ ′′ ( y ) + ϕ ′′′ ( y )( y -x ) . We have for all ( x, ν ) ∈ [0 , 1] × (0 , ∞ ) :

Proof. The proof of either identity follows by applying implicit differentiation to the function Ψ . Fix x ∈ [0 , 1] and consider

Because ϕ is C 3 on (0 , 1] , F is continuously differentiable, and

by Assumption 3.2. Hence, by the implicit function theorem, the map ν ↦→ ξ ( x, ν ) is C 1 with

For the latter identity, fix ν > 0 and define

̸

For each x 0 ∈ (0 , 1] let y 0 := ξ ( x 0 , ν ) ∈ (0 , 1] satisfy G ( x 0 , y 0 ) = 0 . We have ∂G ∂y ( x, y ) = v ( x, y ) . Assumption 3.2 gives v ( x, y ) > 0 for all 0 < y ⩽ 1 and 0 ⩽ x ⩽ y . Hence ∂G/∂y ( x 0 , y 0 ) = 0 .

̸

Since G is continuously differentiable and ∂G/∂y = 0 at ( x 0 , y 0 ) , the implicit-function theorem guarantees a C 1 map x ↦→ ξ ( x, ν ) in a neighborhood of x 0 with G ( x, ξ ( x, ν ) ) = 0 .

Differentiating G ( x, ξ ( x, ν ) ) ≡ 0 with respect to x and using ∂G/∂x = -ϕ ′′ ( y ) gives

so

When x = 0 , the same argument applies, because ∂G ∂y (0 , y ) = v (0 , y ) > 0 and ∂G/∂x | (0 ,y ) = -ϕ ′′ ( y ) is finite (the solution y = ξ (0 , ν ) is strictly positive, so ϕ ′′ ( y ) is finite even if ϕ ′′ ( y ) →∞ as y ↓ 0 ). Thus ∂ξ/∂x | (0 ,ν ) exists and the same formula holds. This completes the proof.

Lemma C.3 (Derivative ∂h ∂ν ) . Under the condition that x ↦→ u ( x ) := xϕ ′′ ( x ) /ϕ ′ ( x ) is nondecreasing from Assumption (A2), we have ∂h ∂ν ( x, ν ) ⩽ 0 for all x ∈ [0 , 1] and ν > 0 .

Proof. For the derivative with respect to ν , observe first that

Hence, by the chain rule,

Due to the defining equation ϕ ′′ ( ξ ) ( ξ -x ) = ν , this simplifies to

where the last equality uses ∂ξ ∂ν ( x, ν ) = 1 v ( x, ξ ( x, ν ) ) and ν > 0 .

Lemma C.4 (Derivative ∂h ∂x ) . Assumption (A2) implies ∂h ∂x ( x, ν ) ⩾ 0 for all x ∈ [0 , 1] and ν > 0

Proof. First recall that

Hence, with y = ξ ( x, ν ) ,

Because ξ = ξ ( x, ν ) satisfies ϕ ′′ ( ξ ) ( ξ -x ) = ν , we have

Write

Case 1: A ( ξ ) ⩾ 0 . Because ξ ⩾ x from Lemma A.1, the second term is non-negative; with ϕ ′ , ϕ ′′ ⩾ 0 the first term is also non-negative, so N ⩾ 0 .

Case 2: A ( ξ ) < 0 . Since ξ ⩾ x , we have

where u ( t ) := t ϕ ′′ ( t ) /ϕ ′ ( t ) . Indeed,

By Assumption (A2), u is non-decreasing, so u ′ ( ξ ) ⩾ 0 ; hence N ( x, ν ) ⩾ 0 in this case as well.

Because v ( x, ξ ) > 0 and N ( x, ν ) ⩾ 0 in both cases, we conclude ∂h ( x, ν ) /∂x ⩾ 0 for all x ∈ [0 , 1] and ν > 0 , thereby proving the lemma.

C.2.3 Proving the dual greedy property

Denote an arbitrary subset of the indices by: S ⊆ [ J ] . Let ν S be the corresponding Lagrange multiplier. Below, for a vector x ∈ R V and a set S ⊂ [ V ] , we denote by x [ S ] the sub-vector of x restricted to the coordinates in S . Since ϕ ′ (0) = 0 by the assumptions of Theorem 3.3, denoting Γ = ∑ J m =1 d ϕ ( p m , 0) + ϕ (0) | S | we can write for every S :

̸

Now, let us prove that the greedy property holds. Suppose S is optimal among all subsets of indices of size k but does not consist of some of the top k probability tokens. Then there exist some i = j such that i ∈ S , j ̸∈ S , and p j > p i . Denote S ′ = S \ { i } ∪ { j } .

Let ν S , ν S ′ denote the choice of ν that makes the projected probabilities sum to unity. Now since S ′ only differs from S in that it includes the larger p j > p i , we can conclude that ν S > ν S ′ .

Then, using the above formula for the value of the objective function on an arbitrary subset, we have:

.

Now, since h decreases in ν by Lemma C.3, we have that the sum is nonnegative since ν S ′ < ν S . As for the remaining term, we have:

where the first inequality is by the fact that ν S ′ < ν S and Lemma C.3, and the second inequality is by the fact that p j > p i and Lemma C.4. This concludes the proof of the dual greedy property under Assumption (A2).

D Proof of discrete convexity for primal Bregman projection

We follow the notations that were introduced in the beginning of the proof in Section B. To show that the cost function is discretely convex in k for the primal, it suffices to show that

is discretely convex in k . Indeed, the difference cost( k ) -L ([ k ]) = λk is linear in k .

To simplify notation, let us denote L ([ k ]) by L ( k ) and T [ k ] by T k . From Lemma (B.1) we know that with ˜ S V := S [ V ] -kϕ (0)

Using (10), we know that f ([ T k ( p )] j ) = f ( p j ) + ν [ k ] ∀ j ∈ [ k ] . Again, we simply denote ν [ k ] as ν k . For j ∈ [ k ] , letting x = f ( p j ) + ν k in Lemma C.1, we have:

But now, using that the nonzero entries of T k ( p ) must sum to unity, we find the following simplification:

Now, define the auxiliary function W for all j, ν for which the expression below is well defined:

where p is implicitly kept fixed. From the above calculation, we thus obtain after canceling out terms:

To prove that this is nonnegative, we leverage that W ( k, · ) is strictly concave in ν for each k , which follows as the Legendre dual mapping ϕ ∗ is strictly convex since so is ϕ . Then, observe that for every j ,

Thus,

As W ( k, · ) is strictly concave in ν , W ( k, · ) is maximized at ν k . Thus, we have: (1) W ( k +1 , ν k +1 ) ⩾ W ( k +1 , ν k ) , and (2) W ( k -1 , ν k -1 ) ⩾ W ( k -1 , ν k ) . With these in hand, we have:

Now, due to the definition of W , the last display equals

the inequality holding as p k ⩾ p k +1 , and as the mapping p ↦→ ϕ ∗ ( f ( p ) + ν k ) is increasing in p since so are ϕ ∗ and f . This concludes the proof.

E Proof of discrete convexity for dual Bregman projection

We denote θ x ( y ) = ϕ ′′ ( y )( y -x ) . As observed before, we have for all admissible x, y that

and the convexity condition for the second argument of d ϕ of Assumption 3.2 is given by:

The dual projection for any 1 ⩽ i ⩽ j ⩽ V is given (for optimal Lagrange multiplier ν j ) by:

Denote the dual Bregman objective, as a function of the selected sparsity k , as:

We now demonstrate that cost ∗ ( k ) is discretely convex in k . For this, we will directly show that the second-order differences of this function are nonnegative at every k ∈ { 2 , . . . , V -1 } . Specifically, we can write:

We now decompose this quantity into three terms corresponding to three ranges of index i ∈ [ V ] , namely i ∈ [ k -1] , i ∈ { k, k +1 } , and i ∈ { k +2 , . . . , V } . We obtain:

The last sum is identically zero, so we engage with the other two ranges of indices.

Range 1: i ∈ [ k -1] . For Range 1, recall that for any convex function ψ, it holds for any two points x, y in its domain that ψ ( x ) -ψ ( y ) ⩾ ψ ′ ( y )( x -y ) . Now, notice that for each i in Range 1, each of the two terms in figure brackets can be bounded via the convexity of d ϕ ( x, · ) in its second argument as:

and:

As a result, we may simplify the Range 1 sum as follows, using that by definition, the first j terms in the projection T ∗ j for each j ∈ { k -1 , k, k +1 } sum to unity:

Range 2: i ∈ { k, k +1 } . For Range 2, we first note that the following three types of terms cancel out: ϕ (0) , ϕ ( p k ) , ϕ ( p k +1 ) . Furthermore, terms involving ϕ ′ (0) vanish by assumption.

The remaining terms in the Range 2 sum can then be written as:

Now, we can bound

using that p k ⩾ p k +1 and the strict convexity of ϕ . We find the lower bound

By adding and subtracting the term ϕ ( p k ) twice, we have the following equivalent bound:

Returning to the main bound We can now merge the cases, resulting in the following tight lower bound of the second differential of the cost function:

Now, define the following key auxiliary function ψ k : [0 , 1] → R , such that for all x ∈ [0 , 1] :

This lets us rewrite our lower bound equivalently as:

We now establish a monotonicity property for ψ k .

Lemma E.1. For every k ∈ [ V ] the function ψ k ( x ) is increasing on x ∈ [0 , [ T ∗ k ( p )] k ] .

Proof. We consider the derivative of the function ψ k :

where we have used the connection between θ x ( y ) and ν k (see Lemma A.1).

Now, recalling that by assumption, ∂ ∂y θ x ( y ) ⩾ 0 for all y ⩾ x ⩾ 0 , and using that [ T ∗ k ( p )] k ⩾ p k by the properties of the dual projection method (see Lemma A.1), we have that:

so long as 0 ⩽ x ⩽ [ T ∗ k ( p )] k

Continuing, by the properties of the dual projection, we have:

In view of Lemma E.1, (20) implies that

This concludes the proof of dual discrete convexity of the Bregman cost function.

F Algorithmic details

F.1 Computing the dual renormalization map

Recall that when ϕ is dual valid, the renormalization map T ∗ ϕ is uniquely defined for x ∈ ∆ sub ,k with ∑ i x i > 0 by the fixed point equation (see Lemma A.1)

To compute T ∗ ϕ , recall from Section C.2.1 the function Ψ from (13) with Ψ( x, y, ν ) := ϕ ′′ ( y )( y -x ) -ν for all x, y, ν . Then, for a fixed ν , [ T ( x )] i satisfying the equation [ T ( x )] i = x i + ν/f ′ ([ T ( x )] i ) is equivalent to solving Ψ( x i , y i , ν ) = 0 for y i = [ T ( x )] i . The monotonicity properties from Lemma A.1 then suggest the following algorithm, consisting of a binary search over ν ∈ (0 , M ] , and then over each coordinate of T solving ϕ ′′ ([ T ( x )] i )([ T ( x )] i -x i ) = ν .

Algorithm 1 Dual Renormalization Map T ∗ ϕ ( x ) via Nested Binary Search

F.2 Pseudocode for algorithms

See Algorithm 3 and Algorithm 4 for pseudocode for sparse primal (resp. dual) Bregman decoding.

Algorithm 4 Regularized Sparse Dual Bregman Decoding

Require: Probability vector p ∈ ∆ V , valid convex generator ϕ , sparsity penalty λ ⩾ 0 Ensure: Sparse decoded distribution ˆ p ∈ ∆

V 1: function SPARSEDUALBREGMANDECODE( p , ϕ , λ ) 2: Sort p in descending order: p (1) ⩾ p (2) ⩾ · · · ⩾ p ( V ) 3: Define f = ϕ ′ , f ′ = ϕ ′′ 4: function COMPUTEDUALRENORMALIZATION( x ∈ R k ) 5: Solve for ν ∈ R such that: ∑ k i =1 [ T ∗ ϕ ( x )] i = 1 , where [ T ∗ ϕ ( x )] i satisfies the fixed-point equation: [ T ∗ ϕ ( x )] i = x i + ν/f ′ ([ T ∗ ϕ ( x )] i ) . 6: return ˆ p ( k ) = T ∗ ϕ ( x ) 7: end function 8: function COMPUTEDUALCOST( k ) 9: Let x = p [1: k ] 10: ˆ p ( k ) ← COMPUTEDUALRENORMALIZATION( x ) 11: Pad with zeros: ˆ p ( k ) ← (ˆ p ( k ) 1 , . . . , ˆ p ( k ) k , 0 , . . . , 0) 12: Compute D ϕ ( p, ˆ p ( k ) ) = ∑ V i =1 [ ϕ ( p i ) -ϕ (ˆ p ( k ) i ) -f (ˆ p ( k ) i )( p i -ˆ p ( k ) i ) ] 13: return cost( k ) = D ϕ ( p, ˆ p ( k ) ) + λk 14: end function 15: k ∗ ← BINARYSEARCH(ComputeDualCost, V ) 16: Recompute ˆ p ( k ∗ ) using COMPUTEDUALRENORMALIZATION( p [1: k ∗ ] ) 17: Pad with zeros to full length V 18: return ˆ p ( k ∗ ) 19: end function

G Example: α -Bregman decoding

G.1 Proof of Lemma 4.3

We first restate the lemma.

Lemma G.1. All generator functions ϕ α , α > 1 , are dual-valid and satisfy Assumption (A2).

Proof. For Assumption 3.2, we can explicitly write:

Therefore, the second derivative in y of this expression is

Now, if y ⩾ x , then using α -1 ⩾ 0 we have that the above expression is

confirming the convexity in y . Now for the condition that x ↦→ u ( x ) := xϕ ′′ ( x ) /ϕ ′ ( x ) is nondecreasing from Assumption (A2), we can observe that

Therefore, we identically have:

thus concluding the proof.

G.2 Proof of Proposition 4.2

Recall the α -renormalization map [ T α ( p )] i = ( p α -1 i + ν ) 1 α -1 , i ∈ [ k ] , where the shift parameter ν = ν ( α, p ) is chosen so that ∑ k i =1 [ T α ( p )] i = 1 . We treat each value (or limit) of α in turn.

The limit α →-∞ . Define

Because x ↦→ x 1 /β is strictly decreasing and convex on (0 , ∞ ) for β < 0 , F β is strictly decreasing and continuous on the interval ( -min i p β i , ∞ ) . Moreover, lim ν ↓-min i p β i F β ( ν ) = ∞ and lim ν ↑∞ F β ( ν ) = 0 , so a unique root ν β with F β ( ν β ) = 1 exists. Because F β (0) = S := ∑ k i =1 p i ⩽ 1 and F β is decreasing, we have ν β ⩽ 0 .

Let q ( α ) i = [ T α ( p )] i = ( p β i + ν β ) 1 /β , and i ∗ be the index where p i is largest. Using the constraint ∑ i q ( α ) i = 1 ,

̸

̸

Raising q ( α ) i ⋆ = ( p β i ⋆ + ν β ) 1 /β to the power β < 0 yields

̸

̸

For i = i ⋆ , we have ν β /p β i → 0 . Indeed, (21) implies | ν β | ⩽ p β i ⋆ ( c β -1) with c := ( p i ⋆ + δ ) /p i ⋆ > 1 . Because β →-∞ , c β → 0 , we have | ν β | = O ( p β i ⋆ ) = o ( p β i ) . Then,

̸

Summing (22) over i = i ⋆ and using ∑ i q ( α ) i = 1 gives

̸

̸

Equations (22) and (23) establish q ( α ) → T -∞ ( p ) component-wise, completing the proof.

The case α = 3 2 . Now α -1 = 1 2 , hence [ T 1 . 5 ( p )] i = ( √ p i + ν ) 2 , i ∈ [ k ] . Set s := ∑ k j =1 √ p j and A := ∑ k j =1 p j . The normalization condition becomes

Solving kν 2 +2 sν +( A -1) = 0 for the root that yields non-negative probabilities gives ν = -s + √ s 2 + k (1 -A ) k . Hence

The case α = 2 . Here α -1 = 1 , so Definition 4.1 yields [ T 2 ( p )] i = p i + ν, i ∈ [ k ] . The normalization condition gives 1 = ∑ k i =1 ( p i + ν ) = ∑ k i =1 p i + kν, hence ν = 1 -∑ k j =1 p j k . Substituting yields

The limit α → + ∞ . Write β := α -1 → + ∞ . Let ν = c β with c ∈ [0 , 1] . Then

Using 1 β log( a β + b β ) → log(max { a, b } ) as β → ∞ gives lim α →∞ [ T α ( p )] i = max { p i , c } . Choose the water level c so that ∑ k i =1 max { p i , c } = 1 . This furnishes the claimed water-filling rule.

The four cases above prove Proposition 4.2.

̸

G.3 Illustrating primal and dual renormalization

Weconsider the peaked vector v = [0 . 1 , 0 . 001 , 0 . 001 , 0 . 001 , 0 . 001] , and plot how both of its distinct constituent values get transformed by the primal and dual Bregman α -renormalization (by symmetry, all copies of 0 . 001 are guaranteed to get mapped to the same value by any of our renormalizations). The resulting plots are in Figure 4. As predicted by our theory, both renormalization families coincide at three values of the parameter, namely at α ∈ { 1 , 2 , ∞} . Furthermore, the primal family evolves more gradually than the dual family between the endpoints of the parameter interval α ∈ (1 , 2] , while the reverse behavior occurs for α ∈ (2 , ∞ ) (where both renormalizations gradually converge to the water-filling limit which, in this case, is the uniform distribution).

RenormalizationofSecondCoordinatevsα

Primal

Transformedvalue(secondcoordinate)

0.200

0.9

Dual

0.175

0.8

0.150

0.7

0.125

0.6

0.100

0.5

0.075

0.4

0.050

0.3

0.025

Primal

0.2

Dual

1

2

3

4

5

6

7

1

2

3

4

5

6

7

α

α

G.4 Illustrating general nonconvexity of dual renormalization

Figure 5 illustrates that the dual Bregman objective can in general be non-convex for large α .

PlotofDp(x,y)forx,yE[0，1]fordifferentα

α = 1.5

α = 2

α = 2.5

α =3

Dp（x,y)

Dp(x,y)

0.4

0.5

Dq(x,y)

0.4

0.3

0.25

0.30

E0

0.2

0.2

0.1

0.1

0.10

0.05

0.0

0.0

80

0.00

y

1.00.8 0.6 0.4 0.2 0.0

y

1.00.8 0.6 0.4 0.2 0.0

y

1.00.8 0.6 0.4 0.2 0.0

y

1.00.8 0.6 0.4 0.2 0.0

X

X

X

X

α = 3.5

α = 5

α =10

0.20

0.10

0.05

0.25

0.08

0.04

0.20

0.15

0.10

≤0.02

0.05

0.05

0.02

0.01

0.00

0.00

0.00

0.00

88

y

1.00.8 0.6 0.4 0.2 0.0

y

1.00.80.60.40.20.0

y

1.00.8 0.6 0.4 0.2 0.0

y

1.00.8 0.6 0.4 0.2 0.0

X

X

X

G.5 Illustrating discrete convexity

Figure 6 illustrates that the loss function cost( · ) defined in (6) is discretely convex for both the primal and dual decoding strategies. Here, we have chosen V = 80 and the regularization parameter λ as 1 / 80 . When k is close to V , the renormalization maps are all close to the true vector p , regardless of the value of α , and hence the loss primarily depends on the regularization term λk , which here equals λk = 1 for k = 80 . Thus, all curves (corresponding to different values of α ) for both the primal and dual plots, asymptote to linearity and converge to this value at k = 80 .

Dual Bregman-a

α=1.2

2.0

α=1.2

α=1.25

α=1.25

1.0

α=1.5

α=1.5

Primal Cost: cost(k)

α=2

α=3

Dual Cost: cost*(k)

1.5

α=2

α=3

1.0

0.5

0.5

0.0

0.0

0

20

40

60

80

0

20

40

60

80

k

k

G.6 The simultaneous effects of Bregman decoding and temperature scaling

α=1.5,T=0.3

α=1.5,T=1.5

2.5

α=2.0,T=0.3

α=2.0,T=1.5

!i

!

α=1.5,T=0.7

α=1,T=0.3

α=2.0,T=0.7

α=1,T=0.7

2.0

α=1.5, T=1.0

α=1, T=1.0

α=2.0, T=1.0

α=1,T=1.5

1.5

1.0

0.5

6

9

Index

Here, we provide a plot to help compare the simultaneous effects of Bregman decoding and temperature scaling. We use the same simulation setting and plotting style as in our figure from the introduction (Section 1); except we only plot the nonzero probabilities (i.e., the top k = 10 probabilities), and we plot the relative sizes of the probabilities compared to the standard topk decoding. Further, we use the same α and temperature hyperparameters used in our experiments in Table 1. The results are shown in Figure 7. Standard topk decoding corresponds to α = 1 and T = 1 . From the figure, it appears that the effect of α > 1 is to moderate/regularize the amount by which the small probabilities are pushed to zero; which could potentially be one reason why α -Bregman decoding with α > 1 can perform better at high temperatures.

H Supplementary experimental details

H.1 Compute resources

The experiments were conducted on a system running Rocky Linux 8.10, with 64 CPU cores of Intel(R) Xeon(R) Gold 6448Y processors at 2.10 GHz, 1 TB of RAM, and 8 NVIDIA L40S GPUs with 46 GB of memory each. All experiments can be done with only one GPU and multiple GPUs

5

10

15

20

25

30

35

40

45

50

k

0.84

0.86

0.88

0.90

0.92

0.94

0.96

MAUVE Score

GPT-2 MAUVE Scores

Top-k

Primal-1.5

Primal-2.0

5

10

15

20

25

30

35

40

45

50

k

0.62

0.64

0.66

0.68

0.70

MAUVE Score

Llama3.1 MAUVE Scores

Top-k

Primal-1.5

Primal-2.0

were used only to parallelize experiments. The software environment used Python 3.11.11, PyTorch 2.5.1, and CUDA 12.4.

H.2 Supplementary experimental results

In this section, we provide additional experimental results to supplement those from Section 5.

Table 2 shows results analogous to those in Table 1 for λ ∈ { 0 . 1 , 0 . 001 } .

Figure 8 presents the MAUVE scores comparing generated and human-written text under different decoding strategies. While primal decoding shows a slight advantage over topk decoding, the differences are not statistically significant. We report standard deviations estimated from 50 bootstrap resamples; a higher number of resamples was not used due to the high computational cost of MAUVE score evaluation.

H.3 Experiments for Larger models: Qwen and Phi

We repeat our experiments for Qwen2.5-14B-Instruct and Phi-3-medium-4k-instruct.

Figure 9 shows results analogous to those in Figure 3. Table 3 and 4 show the accuracy on GSM8K, analogously to Table 1 and 2. Table 5 and 6 show results for the Phi-3-medium-4k-instruct model.

10

20

30

40

50

k

3

4

5

Absolute Perplexity Diff

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

0.0

0.5

1.0

1.5

2.0

Absolute Repetition Diff (%)

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

3

4

5

6

7

Absolute Perplexity Diff

Top-k

Primal-1.5

Primal-2.0

10

20

30

40

50

k

0.0

0.5

1.0

1.5

2.0

2.5

Absolute Repetition Diff (%)

Top-k

Primal-1.5

Primal-2.0

H.4 Experiments for TriviaQA

Table 7 and 8 show accuracy on TriviaQA for LLaMA3.1-8B model. Here we choose 10% ( ≈ 1800 questions) proportion of TriviQA validation dataset for evaluation.

Table 9 and 10 show analogous accuracy results for Phi3-medium-4k-instruct on TriviaQA.

H.5 Adaptivity

In this section, we consider the adaptivity of primal decoding by presenting the mean, standard deviation and entropy of the k ∗ chosen by our method during evaluation on GSM8K and TriviaQA datasets.

In Table 11, we show the average k ∗ values (and their values rounded to the nearest integer) selected by primal Bregman decoding on GSM8K with LLaMA 3.1 8B for various temperatures, α , and λ . Table 12 shows corresponding standard deviation and entropy.

Table 13-14 show analougous adaptivity results for Qwen2.5-14B-Instruct.

Table 15-16 show analougous adaptivity results for Phi-3-medium-4k-instruct.

In Table 17, we show the average k ∗ values (and their values rounded to the nearest integer) selected by primal Bregman decoding on TriviaQA with LLaMA 3.1 8B for various temperatures, α , and λ . Table 18 shows corresponding standard deviation and entropy.

Table 19-20 show analougous adaptivity results for Phi-3-medium-4k-instruct on TriviaQA.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work establishes a theoretical framework for (generalized) top-$k$ decoding using Bregman divergence. Specifically, it introduces a general optimization objective for large language model (LLM) decoding that promotes sparsity via the $\\ell_0$ norm and investigates separable Bregman divergences. The resulting solution extends the top-$k$ decoding strategy, with a notable application being the adaptive selection of $k$ based on the target distribution $p$.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper provides detailed studies of top-k decoding through the lens of Bergman divergence.\n\n- The empirical results demonstrate good performance compared to the baselines.\n\n- The adaptivity property of the provided method is interesting.\n\nWeakness:\n- Some discussion are not well-motivated until very later sections.\n\n- Discussion on practical application is lacking. (see questions)\n\nQUESTIONS:\nSection 3.1 is overly dense and lacks intuitive explanations for the proposed setup. The authors should include a clear, high-level overview to provide context and intuition. Additionally, the role of Equation (5) in the dual renormalization mapping is unclear; a detailed explanation or example illustrating its implications would improve accessibility.\n\n In Section 3.2, Theorem 3.2 appears to say that top-$\\hat{k}$ is an optimal solution for Equation (2). If so, the introduction of $\\mathrm{cost}(k)$ in Equation (6) is confusing. Why is this cost term necessary, and why wouldn’t a naive search over all $\\hat{k}$ suffice? The authors should clarify the purpose and necessity of $\\mathrm{cost}(k)$.\n\nThe paper does not discuss scenarios where dual decoding might be preferred over primal decoding. The authors should highlight specific use cases or advantages of dual decoding, including any theoretical or practical benefits, to strengthen the contribution. \n\nAre there any experimental or theoretical results on applying this new metric to beam search? Regarding the adaptivity property, I wonder if there's any connection between this and other methods such as min-p and top-p?\n\nLIMITATIONS:\nFigure 1 should include a contour plot alongside the 3D visualization to enhance clarity and interpretation.\n\nThe final paragraph of Section 2.2 notes that when $\\lambda > 0$, a sparsity constraint is induced, which is a well-known property in the convex optimization literature. However, the unique aspect here appears to be the focus on the top-$k$ values within $p$. This distinction should be emphasized, with specific, non-trivial details (e.g., novel computational or theoretical implications) clearly articulated to highlight the contribution.\n\nThe adaptivity property should be highlighted in the introduction. \n\nAlso see Questions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors considered the well-known top-k decoding rule, and provided an optimization interpretation of the rule. It is shown that it can be viewed as the solution of a l_0-regularized Bregman divergence minimization problem. Based on this interpretation, the authors specialized it to a set of potential functions in the Bregman divergence, which induce different decoding rules other than the top-k decoding.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. It is reassuring to see the top-k decoding rule has an optimization interpretation.\n2. The theory development appears rigorous.\n\nWeakness: Though the results are reassuring, the theory development seems of limited use for decoding rules in the practice. \n\n1. The interpretation is somewhat obvious. In practice, the top-k decoding rule is mostly motivated by the simplicity of the computation, particularly in its logit form. It is then not surprising that there is an optimization problem that can yield the same solution.\n\n2. Using the Bregman divergence as the measure of divergence lends some generality to the result, but also makes it less interpretable. The main intuition seems to be already clear if we consider KL divergence or the squared l_2 norm. Though this generality leads to other decoding rules, these rules seem to be purely of mathematical interest, but of little practical advantage.\n\n3. The issue of simple computation is not considered, which is in fact the main motivation for the top-k decoding rule. The new decoding rules appear to lose this obvious advantage, and there are no clear benefits in their eventual performance. \n\n4. The numerical evaluation is relatively weak: What is the main point of having these evaluations? These decoding rules do not lead to better performance, nor are they computationally more efficient. It is not clear what the readers can learn from these numerical results.\n\nQUESTIONS:\n1. l-1 norm is often used as a surrogate for l_0 norm. In the setting considered here, will replacing l-0 with l_1 result in the same top-k decoding rule?\n\n2. If a goal of adopting the Bregman divergence is for its generality, then the following question appears important: Under the same regularized minimization setup, if the decoding rule is greedy, does it imply that the divergence must be a Bregman divergence?\n\n3. Assumption 3.1 seems redundant, since otherwise, Bregman divergence won't be well defined?\n\n4. Why is the class of permutation-equivariant maps important? It is true that the solutions given her all belong to this class, but it is not clear why the restriction of \"permutation-equivariant\" plays a role anywhere?\n\nLIMITATIONS:\nsome discussion is found in the section 7.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work discusses theoretical foundations of top-$k$ decoding through a generalized class of Bregman decoders. More concretely, they formulate “generalized top-$k$ decoding” as $l_0$ sparsity-regularized Bregman divergence minimization. The asymmetry of Bregman divergences then yields a primal and dual formulation of the problem, which are treated independently due to their differing properties. It is noted that Bregman divergences are interesting as without regularization the optimization problem yields exactly $p$. The authors proceed to derive the renormalization map of top$k$ decoding for a fixed sparsity pattern, which simplifies the problem to $\\arg\\min_{\\hat p \\in \\Delta_k} \\text{Div}(\\hat p , x)$. Whereas the authors make use of existing results which derive the primal normalization, they further derive an optimal dual renormalization. They then show that both Bregman decoding strategies are greedy, i.e., that the recovered entries of p are indeed the most probable (under some conditions on $\\phi$), and followingly show that in the greedy top-$k$ setting both Bregman decoding strategies are $k$-convex. As a practical example, the authors then discuss a class of Bregman divergence, namely, the Havrda-Charvat-Tsallis $\\alpha$-entropies, for which a few special instances exist with easy-to-compute renormalization maps. Finally, the authors compare the $\\alpha$-entropies with basic top-$k$ decoding across two models and two datasets in (1) a full decoding setting, where $k^*$ is adaptively chosen, and (2) a setting where $k$ is fixed, and the renormalization maps are evaluated.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The theoretical framing around the top-$k$ heuristic as $\\ell_0$-regularized Bregman divergence minimization is sound, novel, and exciting: it is intuitive to derive a family of optimal top-$k$ decoding strategies by minimizing divergence between $p$ and its truncated version. Further, it is fascinating to think about *families* of top-$k$ decoding strategies immediately derived from families of divergences.\n\n- The proposed method addresses two problems within the same optimization objective; adaptive top-$k$ sampling during decoding, and finding optimal normalization maps.\n\n- Even though the content is rather dense and technical (§3), the presentation is clear and well-structured. Concrete applicability is nicely shown using Havrda-Charvat-Tsallis $\\alpha$-entropies in §4 and in initial experiments in §5.\n\n---\n\n### Weaknesses\n\nAlthough the paper rightfully claims that the contributions are mainly theoretical, I think the experimental section could still benefit from some improvement to complete the paper.\n\n- Evaluates limited settings (models, datasets, regularization strengths, …) yielding uncertainty about the generalization of findings. A straightforward expansion could just be to expand it to the same settings as the referenced paper by Chen et al. (L.287).\n- Leaves out an interesting analysis about what values of $k$ are chosen by the algorithm, as a function of the regularization strength. Further, does $k^*$ significantly change during decoding? This may tell us something interesting about failure modes of fixed-$k$ decoding.\n- Only covers primal decoding, leaving open an interesting comparison between the two asymmetric formulations, although I assume this is left out as nice normalization maps only arise in the primal case?\n- I recommend a better visual grouping for both Figure 3 and Table 1 as they are a bit hard to initially navigate. \n\nFurther, it is unclear how well this approach generalizes to other generators, i.e, if the class of Havrda-Charvat-Tsallis $\\alpha$-entropies is a special outlier class or other useful such families exist. \n\n---\n\n#### Minor\n\n- L.26-29; The listing of “decoding methods” is a bit imprecise to me. As I know it, best-of-N is used for post-hoc sample selection (thus, not really a decoding method), whereas top-$k$ and top-$p$ are token level sampling strategies, “temperature scaling” just means the temperature recalibration of probabilities. A more nuanced distinction and situation of top-$k$ decoding would be useful. \n- L. 44-54: it is initially unclear (but cleared up later on) what the correspondence of $k$ and $\\alpha$ is this early-on. The passage also seems a bit out of place. I recommend reformulating this passage to avoid confusion.\n\nQUESTIONS:\n### Questions\n\n- If I understand correctly, adaptive $k$-sampling finds a $k^*$ at each decoding step, which is described in L.151 as being $O(V \\log V)$ operation for $V \\sim 10^5$. This seems quite heavy. How does this affect decoding runtime? How would you envision speedups?\n\n- Re. Fig.3: is there a theoretical/qualitative reason to expect other members of the $\\alpha$-entropies to perform better than the $\\alpha=1$ basic top-$k$ decoding? A qualitative discussion would help.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a theoretical fundation of top-k decoding as solving an ell_0-regularized separable Bregman divergence minimization problem.\n\nIt proves that (i) greedy selection of the top-k entries is always optimal and (ii) the resulting cost is discretely convex in k, enabling an exact O(V log V) binary-search algorithm to pick the optimal k. \n\nThis framework subsumes standard top-k (the KL case) and introduces a family of alpha-Bregman decoders that re-weight remaining probabilities, matching or surpassing top-k performance in open-ended text and GSM8K math reasoning tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper gives the first exact, discrete-convex formulation of top-k decoding and an O(V log V) algorithm, putting the method on solid theoretical ground.\nIt unifies top-k, sparsemax, and entmax within a single Bregman + ell_0 framework and introduces tunable alpha-Bregman decoders that broaden design options.\nGreedy optimality and discrete convexity proofs cover both primal and dual cases and are supported by indicative language-generation and math-reasoning experiments.\n\nWeaknesses:\nThe optimal k is highly sensitive to λ, yet the paper offers no principled way to set λ across temperatures, tasks, or sequence positions.\nEach time step needs full-vocabulary sorting and root-finding; despite the O(V log V) bound, latency and numerical stability may hinder use on large LLMs.\n\nQUESTIONS:\nQ1. Your guarantees rely on strong smoothness/convexity assumptions on phi (Assumptions 3.1 & 3.2) plus either (A1) phi^{prime} convex or (A2) monotone elasticity. Which common divergences or practical temperature/logit-transformations used in LLM decoding fail these conditions, and how would the framework extend to them?\n\nQ2. lambda critically governs the optimal k. Could you provide tuning rules for lambda? \n\nQ3. Equation (6) proves discrete convexity for fixed lambda. If lambda is adapted per step, does the convexity guarantee still hold, and how would this affect the binary-search algorithm?\n\nQ4. Dual Bregman decoding requires root-finding in a potentially non-convex landscape. What empirical evidence can you provide that the proposed solver always converges to the global optimum on large-scale logit vectors and mixed-precision hardware?\n\nQ5. Runtime analysis counts O(V log V) per token, dominated by a full sort. Have you benchmarked end-to-end latency on real GPUs versus standard softmax + top-k? Are approximate partial-sort methods viable without breaking optimality?\n\nQ6. For alpha-Bregman decoders, what intuition or guidance exists for choosing alpha versus lambda? Is there any interaction that can lead to degenerate cases (e.g., alpha to infty with small lambda)?\n\nQ7. Experiments report results on WebText and GSM8K, but only with LLaMA 8B.\nHave you tested on larger models or different domains (dialog, code) to confirm that the gains are not model-specific?\n\nLIMITATIONS:\nEven if the primary goal is to establish a theoretical foundation, I believe the paper would be more complete if it included experimental comparisons with other top-k decoding methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}