Foundations of Top-$k$ Decoding for Language Models
Authors: Georgy Noarov
Soham Mallick
Tao Wang
Sunay Joshi
Yan Sun
Yangxinyu Xie
Mengxin Yu
Edgar Dobriban
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Models
Top-$k$ Decoding
Adaptive Decoding
Sparse Optimization
$\ell_0$-Regularization
Bregman Divergence
Primary area: theory

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 Top-k: Decoding for Language Models
Georgy Noarov Yan Sun?
Soham Mallick!* Tao Wang
Sunay Joshi'
Yangxinyu Xie
Mengxin Yu² Edgar Dobriban'
University Pennsylvania
2 New Jersey Institute of TTchnologyy' Washington University in St. Louis
Abstract
Top-k decoding is a widely used method for sampling from LLMs: at each token, only the largest next-token-probee after renormalizing 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 missing. this work, we develop theoretical framework that both explains and generalizes top-k decoding. view decoding at a fixed token as the recovery of a sparse probability distribution. We introduce Bregman decoders obtained by minimizing separable Bregman divergence (for both the primal and dual cases) with a sparsity-inducing lo ferreullarizatio:; in particular, these decoders sparsity parameter 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 can be identified in logarithmic time. We note that standard top-k decoding arises 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)
are kept, and the next token is sampled
adaptive in the sense that the
Introduction
Large language models (LLMMs) are powerful generative AI tools for producing text. When pre-trained on large text corpora and aligned according to human preferences, they can used for a wide range of tasks. On technical level, they are probability distributions over text: given any user text prompt x, an LLM samples an answer Y ~ (-11=) from a probability distribution π (-- 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 top-k 21 top-p 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 broader perspective, per-token sanplers/decoders belong to an expanding collection of post-hoc methods for improving [.LM 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-of-N or self-consistency), to variety of test-time scaling approaches; see e.g., [12 222133 32].
In this paper, we focus on decoding methods that modify each next-token-rrobabiiity distribution to induce sparsity, i.e., to keep only small number of tokens with nonzero probability. This includes the widely used top-k (1 and top-p p 32 sampling methods, among others. These methods
"Co-first authors. Correspondence to: gnoarovdaeas upenn. edu, kcillam@sharton www.u. edu, tawanOyharton. upenn. edu, dobriban@wharton upenn. edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

are motivated by the intuition that the noisy LLM probabilities need to be truncated to denoise the "unreliable tail" [2] In particular, we focus on the popular top-k decoding method, which keeps only the largest k next-token-pret probabilities pextbioken-roooabbllteee at each decoding step. These are renormalized-vve dividing by their sum-to a probability distribution from which the next token is sampled.
Despite the wide use and rich intuition behind top-k decoding, to our knowledge, a precise theoretical understanding of top-k decoding is not available (see Section 6 for discussion of related work). Therefore, in this work, we develop theoretical framework that flexibly generalizes and sheds light on the key properties of top-k decoding. For a fixed token, we view decoding as recovering a sparse probability distribution from 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 pearsit--inuccigg fn 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 top-k decoding and opens up rich field of efficient adaptive "Bregman decoding" methods with wide and tunable range of behaviors. As an example, we consider Bregman divergences generated by the a-entropies + x² /[a(a-1)] [2 [S]] and display in the adjacent figure how Bregman decoders modulate token distribution for several values of a. For a + 1, we obtain standard top-k decoding (for k = 10 here). By contrast, for a 0.25, the decoder shifts most of the mass onto the top few tokens; while for a == 2.5, the mass is spread much more uniformly across the top-k tokens. This exemplifies how our framework enables the design of novel decoders eliciting a wide range of behaviors.
... Original distribution
Probability
Top-10 decoding
0.2
Bregman-o decoding 0.25)
Bregman-o decading (x=).5)
0.0
10
20
30 Token index
40
50
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) top-k" decoding of p for an optimal k* = (p):: the utilized A* varies depending on the LLM token distribution p. and is chosen minimize the decoder's i lo-reuularized Bregman divergence from This rigorous spurse-ohjective-cmmti foundation of adaptivity 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 top-k decoding strategies as two-step: (i) select number of tokens k and (ii) renormalize the selected & tokens' entries to probability distribution (Section 1. We introduce two rich classes of decoding strategies (Section 2.2): primal Bregman decoding and dual Bregman decoding. These correspond to fo-regularized minimization of a Bregman divergence to the "raw" L.LM distribution over tokens, in its first vs. second argument
In general, lo -erreuularizttion leads to combinatorial optimization problems, for which there are no known polynomial-ting algorithms [ 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 & largest probabilities is optimal (Theorems 3.2 and 3.3 in Section 3.2 (2) k-convexity: Searching for the optimal Ao is (discretely) convex problem in & (Theorem 3.4 in Section 3.3). While simple to state and desirable, these properties are non-trivial to establish, and require range of novel structural insights into the sparse Bregman objective that could be of independent interest.
In Section we illustrate our theory by introducing a- a-Bregman decoding strategies, generated by Tsallis a-entropies x + = "[ar(a 1)]. We study how their behavior depends on a, and highlight several closed-form cases of interest. One example of the ceamiization-thhoreee decoders their convergence to water-filling as a + oo. 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 top-k decoding.
elegance of a--
2 Bregman divergences being asymmetric in general, their distinct behavior in both arguments has been widely studied in optimization and statistical learning [see e.g.. [ 10 24 56 etc].

---

2 Regularized sparse Bregman decoding
2.1 Top-k decoding preliminaries
Top-k decoding. Given a probability distribution p = (Pi...., .........py. pyy (where V stands for "vocabulary size"), and some 1 & & & V. top-k decoding first selects the indices Sy (i1,.. ig) of the largest k probabilitties, breaking ties arbitrarily. Setting all other coordinates to zero in p. one obtains the vector p|1 : k] of the & largest entries. Then, it renormalizes this vector by dividing it by its sum. Letting (P(1): P(2).- P(a)) = (Pi.... .P..) be the largest entries of p.
top-k(p) p( AA/
Σ
(PS))
One then draws a sample from the distribution top-k(p).
Decoding strategies. Next, we aim to generalize top-k decoding. We will refer to any operator Dec on probability distributions as decoding strategy; formally Dec : Dy -> Av, where Dy = {z E 10 E I; = 1} is the simplex of V-dimennional probability distributions. Observe that top-k 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 be performed in 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.
Motivated by this, we define the notion of renormalization mapping, which takes as input thresholded probability vector with k nonzero entries remaining. We consider renormalization maps that are permmutation-equivvrr output is permuted accordingly; which clearly holds for the div-iivision used in top-k. Therefore, since the sum of probabilities after selection can be less than unity, we can define them as maps from the sub-probability simplex A = {z E [0, 1] E I4 & 1} to the simplex AR. Definition 2.1 (Renormalization). For a positive integer k we call T Auub, + As a renormalization map.
Renormalization.
i.e., when their input is permuted, their
cettmuatation-qqivviiiinn map
A renormalization map can be extended to the full simplex Av, by applying it only on the nonzero coordinates We can now define generalized top-k decoding as renormalizing the top-k entries via general renormalization map.
Definition 2.2 (Generalized top-k decoding). For a fixed k. a generalized top-k decoding strategy Dec,.T Av + Av, 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 top-k elements, and renormalizes il to T(p[1 k)) E Av.
A Adaptivity. A natural extension is to choose i adaptively based on p. For this, we consider k-selector map k Av - [V] := (1,..., V}, and collection of renormalization maps TE: Asub,A -> Ae. k 1,..., V. We define an adaptive generalized top-k decoding strategy Deer : Dy - Av via p Tep (p[1: k(p))). Below, we will design specific renormalizers T and ways to choose
2.2 Regularized sparse Bregman decoding
Decoding via sparse divergence minimization. Consider divergence Div(:-): Ay x Av R between two distributions. Classical examples include the squared error Div = 1lp all2 and the KI. divergence Div g) Σj= Pý In(p:/g)) We define the decoding strategy Decpivs via rearsity divergence minimization under divergence Div, for any probability vector p as:
DecDiv (p) arg min (Div p) lollll)
repuusity decoding).
PEAv
Formally, for any p E R and S C [V]. let ps be the restriction of p to the coordinates in S. Given a vector p E Av with S [V] indexing its nonzero entries, renormalization map T'pp) can be extended to Av by embedding it into the original coordinates: [I'()]] [ [2((s))]; for j E and [T((p), = 0 otherwise. 'In our examples of interest, we will show that this optimization problem is well-defined, When there are multiple minimizers, we assume that one is selected in arbitrary measurable way.
3

---

Here, the fo-pseudonorm Pllo is the number of nonzero entries of p. and À > 0 is sparsity cost hyperparameter. As increases, the optimal solution ò = p" gets increasingly more sparse.
Separable Bregman divergences. In this work, we shall instantiate Div in Problem Sasemmwwitt separable Bregman divergences [ [0]. We will see that this class is expressive enough to induce top-k decoding and many fruitful generalizations of it. For a convex domain Dom S R and convex differentialle function é Dom + R. the one-dimensional Bregman o-divergence de is defined as: dœ 3) <(x) Φby) o (x) y), for E Dom. The separable V--dimensional Bregman d-divergenne D Dom V IR is then defined as:
Dz y) Σ do( for Da....,IV... ......... yv) DoD (E(V]
well-known property Bregman divergences is that (2 y) > 0 for all x, y, with equality y/ when strictly convex, y in fact becomes the unique minimum.
0.07
Objective Value
Objective Velue
0.10
0.06
0.05
0.08
0.04
0.06
0.03
0.04
0.02
0.02
1.0
0.8
0.6 A
0.4
0.2
0.0
0.6
0.4
0.2
0.0
1.0
0.8
0.8
0.6
0.4
P1
0.2
0.0
Figure 1: Illustration of the landscape of the sparse Bregman objective for the primal (left) and dual (right) cases. We choose V dimensional example where the target vector ps (0.1,0.01, 10.1111 show an a- a-Bregman divergence (see Section 4 with a = 10 and À = 0.01.
Primal and dual Bregman decoding. Since Bregman divergences are generally non-symmetric in their arguments, we may instantiate the sparse Bregman decoding Problem in two substantially distinct ways: by placing the estimand in the first (primal) or second (dual) argument:
Div p) == p) (primal decoding),
Div(p,p) p) = (p (dual decoding).
Both formulations possess a sound theoretical motivation. Bregman projections are commonly defined minimization in the first argument, while Bregman- baased proper scoring rules for mean elicitation ctc].
correspond to minimization in the second argument [see e.g.,
The landscapes of primal and dual decoding are illustrated in Figure [ 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 Eo penalty. While in general these decoding objectives could be combinatorial problems that may be hard solve, we will show 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 i solved at و = (and uniquely so if ó is strictly convex), with the intuition that absent sparsity requirements the best guess 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 show for both primal and dual decoding that, under mild conditions on D the optimal solution keeps exactly the top-k* coordinates of p. for an objective-choser k* = k* (p). This yields principled and broad generalization of top-k 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 i one would have to search over all (combinatorially many) size-k sparsity patterns to decide which & probabilities to keep; and one would have to try all E [V] to determine the optimal 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. optimal to select the top & tokens. (2) k-convexity: The sparse Bregman objective is (discretely) convex 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 S [V] (of some size |S| k) has been identified, the vector x: which denotes the sub-vector of p restricted indices in S needs to be projected onto the simplex 4g. Since the to regularization term becomes fixed to λk, Problem 2 becomes equivalent to: arg minşee, Div zz). This is 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 d; compared to minimal set of assumptions for a Bregman divergence to 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). 1).
then imply that for a primal valid potential é denoting f ø' (and extending its inverse fl so that f (x) = for x < f0) and (z) = for x > f(1), making it continuous and non-decreasing on all of R), the primal renormalization map To given for Asub,R by:
Existing results 33
34
+v) for all i E [k), where V E R is chosen so that -TT( -1 (4)
Since v + f-((() + v) is non-decreasing in v, the solution can be found efficiently using off-the-shee root-finding algorithms such as Brent's '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 []. To pave the road towards dual Bregman projections, we will therefore rely on additional structure in 6 and dg. expressed as the following dual validity condition.
Assumption 3.2 (Dual validity). The map ó is thrice differentiable on (0, 1] with lim zx"(x) = 0.
z-00
For (= (..00.1..ydddde is strictly convex for y E(xx,1),and y (0.y) is strictly convex for y E(0,1).
We establish in Theorem A. (see Appendix A that subject to dual validity, the dual renormalization
map Tg is uniquely defined for any Auu, with + 0% by the following implicit equations:
h [T))), = + /f for i E [k], with 1² E IR chosen so that Σ [1rg())) 1. (5)
This transformation is interpretable despite its implicit nature: For every index E [k)) Equation 5 has the effect of increasing the corresponding probability x; by a positive additive amount regulated by an auxiliary variable v²; the latter is chosen to make the increased top-k probabilities sum 1.
Assumption 3.2 short of requiring global convexity of de (2, .) on [0, 1], only enforces it for y E [x, 1]. To enable this relaxation, the proof of Theorem A. carefully excludes optimal solutions belonging to the region I I 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 a-decoding is uniquely defined for all a > 1, not just a E (1, 2]: as plots in Appendix G.4 demoosstattee, a-Bregman divergences are nonconvex for y S x for a >> 2.
'It is strictly increasing for E [=f(2)), [(=)). ) but the required may lie outside this range.
5

---

0.0200
0.0200
Original
Original Rescaled x
Ronoralized Value
0.0175
Rescaled
Renormalized Vélue
0.0175
0.0150
e 1.5
0.0150
0.5
0.0125
0.0125
0.0100
0.0100
0.0075
0.0075
0.0050
0.0050
0.0025
0.0025
0.0000
0.0000
0.000 0.002 0.004 0.006 0.008 0.010 0.012
0.000 0.002 0.004 0.006 0.008 0.010 0.012
Original
Original
Figure 2: Comparison of primal (left) and dual (right) Bregman a-renormalization maps (sce Scction on imput vector Des [1, ...... #] E Auub, with 100. We plot the renormalized values against the original coordinate values of a.
See Section F for algorithmic details on computing the dual map, as well as pseudocode for our algorithms. Figure Sattreiillusee iillustrates the primal and dual renormalization maps for o- o-Bregman divergences (introduced in Section 4). In this concrete example, Ta and T appear similar; however, for different, e.g. more "peaked", inputs Aub, they are more distinct, as we illustrate in Appendix G.3
3.2 Greedy property: Justifying top-k selection
The viewpoint that lower-probability tokens can be considered noisy [32] suggests that it would be natural and indeed desirable for decoding strategy to be renormalize over the top-k-probability tokens, for some k [V. We formalize this follows.
that is optimal to
Definition 3.1 (Greedy decoding). decoding strategy Dec Ay + Av is called greedy if for every p E Av, the set of nonzero entries of Dec(p) is set of top-k entries of p. for some k = k(p).
While many popular decoding methods are greedy 12 21
38 some not 22 36 justifications
for non non-greediness, i.c. the ability to occasionally throw out some of the top-k tokens, include that this 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 greedy for any primal valid potential áb
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 handle. Unlike Theorem 3.2 our next result requires further conditions, which we state as a menu of two options. The relationship between the extra assumptions intricate; Assumption (A2) implied by, but is strietly weaker than, log-convexity of d'.
Theorem 3.3 (Dual Bregman decoding is greedy). The dual Bregman decoding strategy from is
greedy for any dual-valid & with ø' d0)0) 0 that further satisfies either of the following conditions: (AI) d' convex;
(A2) The ma defined u(x) x (x)) for E (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.aondixC..... C. leverages the decomposition from the primal case along with the change of variables de (z, y) = dg- (v)(y) o' (x)). where o* is the convex conjugate of ó. For Condition (A2) develop saddle-point proof approach in Appendix C.2 For that, perform a sensitivity analysis of both the renormalized values [T (p)] and of the per-coordinate Bregman loss terms, relative to hypothetical changes in the dual Lagrange multiplier =² and in the entries pt of p; we carry this out via implicit differentiation of the defining equations
In the economics literature, u(æ) 26" (x)) /o' (xx) is referred the elasticity of the function 6'.

---

3.3 k-convexity: Speeding up the search for optimal adaptive k
We have seen that for fixed k, greedily selecting the top & tokens optimal. However, without further structure, we would still have to search over all & [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, 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 E Ay, and hyperparameter X. We denote the regularized cost of selecting the top-k entries of p, as a function of [V], by:
cost (k) min (Div ((é. -v)) + Mk3.
D
Recall that a function h : [V] + R is discretely convex if for all k E ....... . 13. its discrete second derivative A (k)(k) == Ah(k (k + +1) (h(k) {h(k + - {h(k) h(k 1)} > O.
Theorem 3.4 (Discrete primal and dual cost convexity). cost(-) is discretely convex in k E for:
1. L.viv(p,)) p) = Do(P.)) p). if & is primal valid; 2. Div = Dz(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 a-decoding (defined in Section Satiiooonbbllo beloww for assorted o.
Implications for efficient computation. As a corollary of Theorem 3.4, an optimal k- is provably identifiable by scarching for k for which Acost(k) & 0 and Acost + 1) >> 0, for which repeated bisection (binary search) over S k S V suffices - and thus, only (log V) tries of & 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 ky on & (c.g. ky = 50), then identifying an optimal E [ky] requires (Ru) tries. Secondly, one may requires only Oloo 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 each step.
exponential search/ instead of binary search over k: this
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 [ As its cornerstone, we use the Legendre dual mapping o 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.20 V - Σ
-i=1= (d'(p)) + +v) This
expression is concave in v. and its unique optimizer is 1/k, the optimal Lagrange multiplier for renormalizing the top k probabilities of p from (4). Using this, we then establish A c 2 0.
Dual k-convexity. The proof is in Appendix The above dualization strategy does not directly apply. Instead, we lower bound A 2cost" (k) by regrouping the loss contributions of the indices E [K- 1]. and -via intricate term rearrangement and bounding-redum reduce to proving the local concavity of special transformation (Equation 20 that turns out to hold by our dual-validity assumption.
4 Example: Bregman a-decoding
We now consider, as an illustration, single-paramesd family of Bregman decoding strategies, which arises via the generators of the Saaaaa-Chhaaaataaaeeeeaa a-entropies B
29
45 51 52]:
Φ (z) x /[a(a - 1)), E [J, for E J ;= (= (
When 0 and x = O, set _º +oo that (a) co. For (z) = x log(z), which corresponds to the Shannon entropy, arising in the limi as a - 1. Observe that @a primal valid for all of 0, as (x(x)) = -a² This yields the following primal family of renormalizations, which we will index by a rather than Φi:
one defines
Definition 4.1 (Primal Bregman a-decoding). Fix a J, [VV.. The renormalization map Ta is gor E gub, as: (p))= -I v) for fori with v E IR chosen so that Σ [T(P)] [[]p)]=]1 =1 1.
[E[A]
Firsst, identify true upper bound K" on k" by sequentially trying kn = 1,2,4, and then perform binary search in (log kū) rounds.
*One conventionally detines the entropies via (x x)) ra/(a(a - 1)). in which case the Shannon entropy is obtained in the limit as + 1. In our case, we use the definition &.. (x) = x= /[a(a- 1)] so that some technical
conditions (such & (0) 0) hold in the proofs. Both definitions lead to the same decoding strategies in 4
7

---

Note that for = 1, we have (x) = log x + 1. Hence, 44 implies e" Li=1 P; 1. and we obtain the "standard" renormalization: [Ti (p)]. P for E [k]. Therefore, primal Bregman I I-decoding is top-k decoding, showing how one recovers top-k in our framework. It turns out that some further values of a also lead to renormalization maps of special interest. For any fixed p. we let (p) lim inf T (p) and T (p) = lim inf (P), where the limits are entrywise.
P2).
Proposition 4.2 (Special primal a-renormalization maps). We have the following special instance the primal Bregman a-renormalization map, defined for all E [R] as follows:
[D)) = pi + 11-= -i= 1- ビ, Pi) assuming that arg max; p; = {
[Ti.(())] = (vol + 2 - [k) where r Σÿ= /ps and s = Σラ=
[Ta())]+ = P. + (1 (1- 1
(Tec (p)): max (p:, v). where TR is the "water level" for which I 11
[T
Along with the primal family, the dual a-decoding family can also be defined based on da. Unlike a-decoding, the dual Bregman sparse decoding Problempcaa can be non-convex, as displayed in Figure
above. Figure in Appendix further demonstrates the nonconvexity of Dø., on the unit square for some Yet, we can still show that any dual a-decoding with a > 1 is valid, greedy and k-convex:
Lemma 4.3. All generator functions &a a > 1, are dual-valid and satisfy Assumption (A2)
We give an illustration contrasting primal and dual a-decoding for various a > 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 develop the theoretical foundations of top-k decoding, our aim in this section is simply to illustrate that the performance of our novel decoding schemes can be competitive with standard top-k decoding. In particular, we do not aim 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 top-k decoding, which coincides with the case of our primal a-decoding family described in Section 4 we illustrate primal a-decoding strategies for a = 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 '- always upper bounded, we set a maximum k* & -mmx == 50. (2) In the partial evaluation approach, we instead directly evaluate-for each fixed choice of le in the grid l (5, 10,..., (0)....,.)..0 50]-our proposed renormalization strategies along with standard top-k renormalization.
Models and benchmarks. We conduct experiments using the GPT-2 Large 43 and Llama 3.1 8B 255 models. We evaluate on two benchmarks: (I) 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 Chain- of-Thought benchmark [13]. Additional experiments with larger models, Qwen2.5-144 14B--IIsttuctt and well evaluations the TriviaQA benchmark, are presented in
10--Medium-44-nmee 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 I. 32 38) etc]:
(I) Perplexity difference, which measures the perplexity (according to base model Prase) ) of human text compared to that obtained from a decoding strategy Paecoding derived from the base model, where lower is better. This equals Ex-D Bx-D|EY ((y= (Prase (Y 1/-// Ey -pssonil x)) (Posse (Y | where X ~ D is
"In particular, o (p). (i.)..i.aal Ti.((p), Ta(p) do not require solving for v in Definition 4.1 enabling fast implementation just like in the case of the canonical top-k renormalization.

---

Table 1: Accuracy on GSM8K LLaMA 8B using Bregman primal decoding E (0.01, 0.0001). a E (1.5, 2.0}) and lop-k decoding, for various temperatures. For top-k: & equals the averaged & from primal decoding with matching temperature, A, and a. Standard deviations are over 1000 bootstrap resamples.
0.01
À 0.0001
Temp
Top-k (A 0.01)
Top-k (A 0.0001)
=
1.5 a = 2.0
1.5
Cx = 2.0
0.3
85, 35.14=080
84.38
84.38:1.003..22tu..
84.69+0.99
04.69=099
83,14 84.46=1.00 83.62
0.7
13.24_102
1.. 83.78 84.69
82.03 12.
12.1 1.07
1.0
81.20=100
80.97=10s 81.20 81.20
77.41=1115 77.26=1.15
79.23+1.12 78.5
1.5
1..0
1.. 75.97
75.97
1. 64.97
1..26 58.53
prompt drawn from the dataset, Y ~ (-) denotes human- -written continuation drawn from the dataset, and Y ~ Pdeccdiing (. X) denotes a pondel-generated continuation using specific decoding strategy. Here, Y is the length of the continuation.
(2) Repetition difference: EX~D ( (+IX) (Yep(Y)) - xx (rep((Y)) where rep(Y) 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 TigureBreport 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 top-k terms of both metrics. frequency. The marked decrease in repetitiveness for E {1.5,2} relative to standard top-k (a = 1) consistent with the theoretical behavior of a-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 top-k set, thereby increasing sampling diversity.
particular, a 2.0 has the smallest gaps in perplexity and repetition
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 À E {0.01, 0.0001} and primal decoding parameters a E {1.5, 2.0}. Results for other settings are in Appendix H To ensure a matched comparison, we run top-k with k = k" for the Bregman decoding run with the same temperature, and a, rounded to the nearest integer, see Table [1 in Appendix H As seen in Table [ across all temperature settings, primal decoding with adaptive k* achieves accuracy comparable to top-k. higher temperatures (such as 1.5), the performance of top-k decoding degrades more rapidly than that of primal decoding.
Top-k
Neperttion Feuq Danva
12.
Topk
Top
2.5
Papiocity Drt
100
Primal-1.5 5.
Primal-1.5
Reroles Dift
Primal-1.5
Pri mal---1.5
Primal-2.00
Primal-2.0
Primal-2.0
P..mal2200
Pre
Depetition
5.5
1.0
30
60
50
15
30
60
10
30
10
20
30
40
Figure 3: Perplexity and repetition frequency differences between generated and human-written text for GPT2- large (left two panels) and LLaMA 3.1 8B (right two panels), various & values. We show top-k decoding and primal decoding with & E (1.5, 2.0}. Standard deviations are estimated using 1000 bootstrap resamples.
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-Perpine
(53] simplified and distilled the proof. (35) further studied the projection as a method for generating sparse probability predictions in multiclass prediction problems.
33 3 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 top-k
decoding. contrast to these works, we: (1) consider Bregman projections under to regularization, and (2) offer, to the best of our knowledge, novel analyses of dual Bregman projections.

---

l regularization. Regularization via the fo-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.. 18
2
58 61 etc]. In contrast, the algorithms we propose are exact within numerical precision for the specific class of problems we consider.
20
30
41
49
50
Bregman divergences. The properties of Bregman divergences [0] have been widely studied; see, c.g.,
39 47 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.,
[6] 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.
LLM decoding. There is vast range of work L.M sampling (or decoding), see e.g., (54) and references therein. Classical methods include greedy sampling and beam search. Sparse sampling methods such as top-k sampling [I are motivated by intuition that the "unreliable tail" of low-
probability tokens is mis-estimated [32]. In particular, [32] propose top-p sampling, and [38] propose
min-p sampling. Other sampling methods were proposed in 4 22 (6]. [2] 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] [6]] In contrast, our goal is to develop a deep theoretical understanding of top-k 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 c.g., [10, 24]. This property is highly desirable in probabilistic forecasting and prediction, incentivizing forecaster to predict the true distribution in order to minimize their loss. (2) The fo-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.. (6 23 28 etc]. Combining these two lines of thought provides the motivation for studying Eu regullaad Bregman divergence minimization.
7 Discussion
This paper develops a theoretical foundation for top-k 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 top-k. Our analysis treats per-step decoding as jointly selecting a support size k and applying renormalization via Bregman projection. A natural extension that future work may address is to develop analogous optimization-theore foundations for other popular truncation rules, such top-p and min-p, which also follow the template "truncate tail, then renormalize".
Quantifying and controlling adaptivity. A central feature of our framework is that the parameter ks (p) is chosen adaptively with p. A useful next step is to characterize how ke (p) varies with the hyperparameter À and with properties of the token distribution p (e.g.. concentration or tail decay). For a-decoding, our preliminary calculations in stylized settings suggest that there may be natural scaling laws relating k* to (o:, A); an interesting open problem is to rigorously derive such laws under distributional models that reflect L.M 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 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 a-decoders are similar yet subtly different for corresponding a. 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 top-k might be found even amongst a-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- lg-regularized separable Bregman projection onto the simplex-the nevertheless admit exact, efficient solutions via greedy support selection and discrete convexity in the support size. We conjecture that the resulting purrroe- efficient optimization primitive for sparsifying probability vectors could be applicable well beyond LLM decoding.
10

---

References
[1] Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191.
American Mathematical Soc., 2000.
[2] Sohail Bahmani, Bhiksha Raj. and Petros T Boufounos. Greedy sparsity-constrained optimization. Journal of Machine Learning Research, 14(1):807-841, 2013.
[3] Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, and Joydeep Ghosh. Clustering with Bregman divergences. Journal of Machine Learning Research, 6(0st):1755-1749 2005.
[4] 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 htttps://openreviorr
[5] Heinz H Bauschke and Patrick L Combettes. Iterating Bregman retractions. SIAM Journal on Optimization, 1554):115910 117,, 2003.
[6] Dimitris Bertsimas, Angela King, and Rahul Mazumder. Best subset selection via modern optimization lens. The Annals of Statistics, 44(2):813, 2016.
[7] Lucien Birgé and Pascal Massart. Gaussian model selection. Journal of the European Mathematical Society, 3(3):203-26) 2001.
[8] Mathieu Blondel, André F.T. Martins, and Vlad Niculae. Learning with Fenchel- losses. Journal of Machine Learning Research, (1 (635))1669. 1-69,16699699 2020. URL http::/jmlrr..... org/papers/ v21/19-021 .html 02m.mmll
[9] Thomas Blumensath and Mike E Davies. Iterative hard thresholding for compressed sensing. Applied and Computational Harmonic Analysis, 27(3):265-22 274, 2009.
[10] 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.
[11] Emmanuel Candes and Terence Tao. Decoding by linear programming, IEEE Transactions on Information Theory, 51(12):4203-215 2005.
[12] 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 htttpss//openrevieet
[13] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reüchiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
[14] 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 Natural Language Processing (EMNLP-IJCNLP)) pages 2174-2184 2019.
[15] 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. http:///m org/ httpp//r/r/.aappe//////29//
.html
[16] David L Donoho and Iain M Johnstone. Ideal spatial adaptation by wavelet shrinkage. Biometrika, (1(3):425 455, 1994.
[17] John Duchi, Shai Shalev-Shwartz Yoram Singer, and Tushar Chandra. Efficient projections onto the ball for learning in high dimensions. In Proceedings of the 25th International Conference Machine Learning, pages 272-279, 2008.

---

[18] M' M'hamed Essafri, Luca Calatroni, and Emmanuel Soubies. Exact continuous relaxations of fo-regularized criteria with non-quadinn non-quadratic data terms, 2024. UtL/tttps://arxxworr abs/
2402. 06483)
[19] M' M'hamed Essafri, Luca Calatroni, and Emmanuel Soubies. On &o Bregman-relaxations for Kullback-Leibler sparse regression. In 2024 IEFE 34th International Workshop Machine Learning for Signal Processing (MISP), pages 1-6. IEEF, 2024.
[20] Mhamed Essafri, Luca Calatroni, and Emmanuel Soubies. Box-constrained Eg Bregman- relaxations, 2025. URI.httpss//aaxiv wwrreiiii.rrreeee www./agg//////o///oo/ooo//aaeee/ abs//503... 15083
[21] 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 Linguisties, 2018.
[22] 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 (tLLhttts::// http:///ppeeeee..ee//.eee 1d=dONpC9GL1o
[23] Dean P Foster and Edward I George. The risk inflation criterion for multiple regression. The Annals of Statisties, 1974):1947 1994.
[24] Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, (02(47))335--378 2007.
[25] 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.
[26] Sebastian Gruber and Florian Buettner. Uncertainty estimates of predictions via a general bias- variance decomposition. In International Conference on Artificial Intelligence and Statisties, pages 11331-11354, PMLR, 2023.
[27] 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.
[28] Trevor Hastie, Robert Tibshirani, and Martin Wainwright. Statistical Learning with Sparsity: The Lasso and Generalizations. CRC Press, 2015.
[29] Jan Havrda and František Charvát. Quantification method of classification processes. concept of structural a-entropy. Kybernetika, 3(1):30-35, 1967.
[30] Hussein Hazimeh, Rahul Mazumder, and Tim Nonet. LOlearn: A scalable package for sparse learning using fy regularization. Journal of Machine Learning Research, 24(205): 1-8, 2023.
[31] John Hewitt, Christopher D Manning, and Percy Liang. Truncation sampling as language model desmoothing. In Findings of the Association for Computational Linguisties: EMNLP 2022, pages
2022.
[32] 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:///penrevies. http:/pprrriiiiieeeeffrrrrrrrr wtg.:/orrrmndddyyyyyytiyutt rygGQyrFvHl
[33] 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.
[34] 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-121. Cadiz, Spain, 09-11 May 2016. PMLR
12

---

[35] Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-labe classification. In International Conference on Machine Learning, pages 1614- 1623. PMLR, 2016.
[36] Clara Meister, Tiago Pimentel, Gian Wiher, and Ryan Cotterell. Locally typical sampling. Transactions of the Association for Computational Linguistics, 11:102-122 2023.
[37] Christian Michelot. A finite algorithm for finding the projection of point onto the canonical simplex of R Journal of Optimization Theory and Applications, 50:195-20M 1986.
[38] Nguyen Nhat Minh, Andrew Baker. Clement Neo, Allen G Roush, Andreas Kirsch, and Ravid Shwartz-Ziv. Turning up the heat: Min-p sampling for creative and coherent L.LM outputs. In The Thirteenth International Conference on Learning Representations, 2025. URL https://opeereviie htwp:/rwwwiie..ttooooeeeeeeeeeee mwst/fforumm PierruunutdPBB
[39] Frank Nielsen. An elementary introduction to information geometry. Entropy, 22(10) 2020.
gpt-2-ottputt-daaaase [40] OpenAL. GPT-2 output dataset. http://paar/pp-------uu---aaaeaa datasst
2019. URL https: //github. cww www.:penai/gppt.ttoo.eoeo 2--utpt
[41] Jianting Pan and Ming Yan. Efficient sparse probability measures recovery via Bregman gradient. Journal of Scientific Computing, 102(3):66, 2025.
[42] Christos H Papadimitriou and Kenneth Steiglitz. Combinatorial optimization: Algorithms and complexity. Courier Corporation, 1998.
[43] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
[44] Luca Ragazzi, Paolo Italiani, Gianluca Moro, and Mattia Panni. What are you token about? Differentiable perturbed top-k token selection for scientific document summarization. In Findings of the Association for Computational Linguistics ACL 2024, pages 9427-9440, 2024.
[45] Daniel Reem, Simeon Reich, and Alvaro De Pierro. Re-examination of Bregman functions and new properties of their divergences. Optimization, 611:27799334, 2019.
[46] Michael Eli Sander, Joan Puigcerver, Josip Djolonga, Gabriel Peyré, and Mathieu Blondel. Fast, differentiable and sparse top-A: A convex analysis perspective. In International Conference on Machine Learning, pages 29919-29936 PMLR, 2023.
[47] Leonard Savage. Elicitation of personal probabilities and expectations. Journal of the American Statistical Association, 66(336):783-80, 1971.
[48] Shai Shalev-Shward Yoram Singer, Kristin P Bennett, and Emilio Heranndez Efficient learning of label ranking by soft projections onto polyhedra. Journal of Machine Learning Research, 7(7), 2006.
[49] 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.
[50] Emmanuel Soubies, Laure Blanc-Féraud, and Gilles Aubert. A continuous exact to penalty (cel0) for least squares regularized problem. SIAM Journal on Imaging Sciences, 863):1607- 1639, 2015.
[51] Constantino Tsaliss. Possible generalization of Boltzmann-Gibbs statistics. Journal of Statistical Physics, 52:479-487, 1988.
[52] Constantino Tsallis. Introduction to nonextensive statistical mechonics: Approaching complex world. Springer, 2009.
[53] 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: 309. 1541, 2013.
13

---

[54] Sean Welleck, Amanda Bertsch, Matthew Finlayson, Hailey Schoelkopf. Alex Xie, Graham Neubig, Ili Kulikov, and Zaid Harchaoui. From decoding meta--eneratton. Inference-ting algorithms for large language models. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL https:/tt:///pppt opwp:evviepereiiieeeeoe wwwp//orrmidddeeeeMeMeee
Survey Certification.
[55] Robert C Williamson, Elodie Vernet, and Mark D Reid. Composite multiclass losses. Journal of Machine Learning Research, 17(222):1-5, 2016.
[56] Wotao Yin, Stanley Osher, Donald Goldfarb, and Jerome Darbon. Bregman iterative algorithms for f--minimization with applications to compressed sensing. SIAM Journal on Imaging Sciences, 1(1):143-168, 2008
[57] 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 htpp:tst////rrooediingg mlr. ppress/800 //////aaann1111 html
[58] Jacky Y Zhang, Rajiv Khanma, Anastasios Kyrillidis,, and Oluwasanmi Koyejo. Learning sparse distributions using iterative hard thresholding. Advances in Neural Information Processing Systems, 32, 2019.
[59] Guangxiang Zhao, Junyang Lin, Zhiyuan Zhang, Xuancheng Ren, Qi Su, and Xu Sun. Explicit sparse transformer: Concentrated attention through explicit selection. arXiv preprint arXiv: 1637, 2019.
[60] Shuai Zhao, Qing Li, Yuer Yang. Jinming Wen, and Weiqi Luo. From softmax to nucleusmax: novel sparse language model for chinese radiology report summarization, ACM Transactions on Asian and Low-Resoure Language Information Processing, 22(6):1-21, 2023.
[61] Junxian Zhu, Jin Zhu, Borui Tang, Xuanyu Chen, Hongmei Lin, and Xueqin Wang. Best-subset selection in generalized linear models: fast and consistent algorithm via splicing technique. arXiv preprint arXivv2308.00233 2023,
14

---

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 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 long 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 eye toward future work addressing them.
Guidelines:
The answer NA means that the paper has 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 (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, 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 complete (and correct) proof?
15

---

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 cross- referenced.
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
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All experimental details are provided in Section and 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, No answer 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 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 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 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 replicate the results, access a hosted model (e.g.. in the case of 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 nature the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or 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 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 possible for other researchers to have some path to reproducing verifying the results.
16

---

reproducible, they 5. Open access to data and code
Question:
Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The software package implementing 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 NeurlPS code and data submission guidelines (https://nips.co cc/
httpiic//guides//doeeeuubm..ass/ wtt.eeesuuummmiisoooooooooo..... 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 (c.g., for 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
gewiii/gguee//eeeeeoeooi. CodeeSumiiiseeeeee onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
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?
Question:
Answer: [Yes]
Justification: The experimental settings are described in Section [ We evaluate pretrained models without training. Section 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 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, 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.
17

---

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 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 eπror 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 tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section H.l 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 compute required for each of the individual experimental rums well 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 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 htttts/s//errrrs..c//oc//t/ttttt
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
18

---

Justification: We briefly discuss the potential societal impact of our work in our Discussion. In more detail, our work performs foundational investigation into the theory decoding methods, and as such, does not have immediate societal impacts; however, by laying and re- examining 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 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 unintentional) misuse of the technology.
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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not involve the release of any pretrained models, image generators, or datasets that pose high risk of misuse.
Guidelines:
The answer NA means that the paper poses 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 best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (c.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 properly cited.
19

---

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 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, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's '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, provide a well- documented 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.
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 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 crowdsoureing 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 the NeurIPS Code of Bthies, 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 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?
20

---

Answer:
[NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such 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, 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, non-standard components.
Please refer to our LLM policy hatto:///aaurpss..cc////oeeereep for what should or should not be described.
21

---

Appendix Contents
Existence and uniqueness of dual Bregman decoding
22
Proof of the primal greedy property in Theorem 3.2
23
Proof of the dual greedy property in Theorem 3.3
26
Proof of discrete convexity for primal Bregman projection
30
Proof of discrete convexity for dual Bregman projection
31
Algorithmic details
33
Example: c-Bregman decoding
36
Supplementary experimental details
39
A Existence and uniqueness of dual Bregman decoding
Theorem A.1 (Uniqueness and formula for dual Bregman renormalization). Fix dual valid potential 6. Then, for any x E Asub, with Σ, > 0, the renormalization map T is uniquely defined by:
&
[T)x) = x,,-0 /f,"= for all i E [ki, where v² E IR is chosen so that
-
Proof. First, assume without loss generality that Σ EiEla 2t Otherwise, E 2; 1 then x Ag, so the unique unconstrained optimum, which is 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 we 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 & E R for the simplex constraint, and Lagrange multipliers (A.)ke(M) for the nonnegativity constraints. Then, the Lagrangian is follows:
C(p.v)
- (bh) (':)) -(0
Ps 1)
App
Here, A 3 0 for all i, and by complementary slackness, at optimality A = whenever p. > 0. For each i E [k). the stationarity condition reads (except possibly when P. O, where the second derivative could be infinite):
dc Op.
0
- o" (p:)) x;) = = A;-
In particular, for each coordinate i for which the optimal fs E (0, 1), the stationarity condition is:
Φ" (p:) 24.)
+ o" (p.)
Dz f(()))
Now, we show that // > O. Indeed, observe that there must be at least one index for which P. > Ti. If that not the case, would get E(e P. EvE(n) 2 < 1 by our assumption, contradicting that P E Ag. In particular, then, P. x; > and therefore we have (p.) x.)) Since (p. > 0 and p. I₁ > O, we thus conclude that v > 0.
Having shown that 0, 2 ds d"(v) (y x) for y O. We will now consider two cases:
proceed to show that all p. > 0 at optimality. Note that
by
22

---

1. 6 (0) is finite;
2. limy- - o"(y) +00
If (0) is finite, 7; > for all i. Indeed, suppose that was not the case, and p;: = for some i. Then would have: (0))) x.)
Ai equivalently, (0()) 2i + X 0. Each of the three
terms is nonnegative, and V > 0, so we arrive at contradiction.
a do(2, y) -o for all Next, consider the case in which limy-oo (" = +oc. Then, limy-+0 By
(0, ll. Then, since limg-0 2 d y) -o for all x E (0, 1). for any i such that x; > 0, setting 0 would lead to i for which x; = O, since limy-0 uv²(() = 0, setting Pa = 0 would lead to v = 0, which is a contradiction.
2o, hence necessarily Pr 0. On the other hand, for any By
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 v from (3) with (z y, v) = ""(v) (v) - x) - v for all 2, y, v. Then, for a fixed v, is equivalent to solving V (2 Pr. 2) O. First, consider 0. Then, we have the following:
1. Since the map y + (z, y) is strictly convex for y E [z. 1], follows that 2 (x, v (z, y. 0) is strictly increasing for y E [z, 1], and so is (x, y. v).
2. We have (a x, p) = - & O. Further, (1x, 1) = o0(1)( x) - > 0, whenever
V < o" (1) z).
Hence, the map y + (z, v.v) has unique zero the interval [x, 1], long 0 <VV (1(1)( x).
Next, consider x = 0, in which case we need to solve the equation o² d"(v)y- = v. Then, we have the following:
1. Since the map y + (0 y) is strictly convex for y E (0, 1], follows that By d0, = (0, 0) = x" (v)' is strictly increasing for E (0, 1]. and so is (0,
2. By assumption, lim yo" 0, hence have l-m (I, = - < 0. Further, (O, 1, = =0((1(( - x) - v 2 0, whenever v & (1(1..
Hence, the map y + (0, y, v) has unique zero on the interval (0, 1]. as long as 0 < v & (1).
Now define M == min, (1 x.) = d0 max; æz). Since by assumption >. < 1, it follows that M > From the above analysis, it follows that, as long as E (0, M), for each i, the equation ø" (yx) - xx.) = v. has a unique solution (v) E (x,, 1].
Furthermore,
establish in Lemma C.2 the map + y. (v) is strictly increasing for v/ > 0, also owing to the assumed second -aruddeagmmettt convexity of d In particular, define G(v)) Et (v) for v > O; then G is continuous and strictly increasing, and satisfies lim G(v) - ((v)=>,,,
and G(M) > ys- (M) = 1, where i" is any index achieving the maximum among the coordinates of x. Hence there is a unique v* E (0, M] with G(v*) 1. Setting Pt yi(u²) yields a vector in As that satisfies the KKT stationarity.
Finally, note that the solution و that we just identified is unique. Indeed, we have earlier excluded boundary solutions from consideration, and then further excluded any solutions which P. < x; for any E [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 E A : P. > v; 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. Pi P2 > 2 pv. For any subset Q {i,...., ) S [V] of size k, let Q [V] Let po
23

---

denote the sub-probability vector with the entries of whose indices are in Q. We define the loss L(Q) as
((Q) = min ((p -x), (PQ. (Pq)))r))= PEAx
min PEAr
d do Pi,, + Sqr.
(8)
j=1
Here, Sqr = Eee dD, Pj)- To prove Theorem 3.2 we will show that L(S") 3 L(S) for any S' C [V] of size k. where S = [k] consists of the top-k indices, We will further show that strict inequality always holds if ps + ps. 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
Lemma B.I. For any to:
fii, 12,
ig}
[V] of size the loss function as defined in ( simplifies
L(Q)
Σ - ((b,)) SMM (Q)
j=1
Proof. Observe that:
k (())==(((00 (po..)) Σ
L(Q) = 0v
a((To(p)),PP.) + Se
1=1
k
Σ (o(T(()))) o(p,,) d'(p,,) Pi,)) + Sq
- Σ
--p(p s(p,))ps)))))) Scr
This further equals
Σ
Σ do,0,,;; + SQ= (Q)6(0)
JEQ
[o((z())))) - ('o, SQ + Sq 1Q1á(0)
o1c, SSv] - (Q))(0).
This finishes the proof.
Let Tp(p) denote minimizer of the above loss L(Q), i.e.
To(p) arg min (v (po, PCr (a) arg min Σ de. PEA;
PEA,
Note that (a) holds above as the term Sgr does not play any role in the location of the minimizer. However, it does contribute to the final loss I.(Q). Also, as the divergence is separable, once we have selected 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 in < ig <... .. < in for k [V]. By forming the Lagrangian and differentiating it, we obtain the primal thresholding from (4):
d' d'(TQ(p)),) ø' (p:,) + uQ V [k.]
(10)
Here, uQ is chosen such that [T))(
=
24

---

Lemma B.2. Let S = fi......i4..5 = ( M and Ts(p) and (p) be the corresponding minimizers. Then, the following decomposition holds:
- D
(TTs (Ts(P)),) (d'(( o' d'(p.,))
Is (p)ls (d'( d' (v::)
+
(II)
Proof.
We have from Lemma
B.1 that
I(S") - L(S) = Σ (TTTs (wl)) -'(rx;)(Ts - > [I((Ts(P)))) - ((oi,)
(v)ss (1TT::/)))))-(((/0000)0o00))0130 ('(R))(Ts (p))
This further equals
[o(Tss: (p)),) <((T<(p))) ((((((())))) 0 [Ts())b)]
(p)l |o (((TT(())) d'(pe)) - [T>())], oe ([(s()))))
- d'(r,)))
1=1
DD((Tss +
(pTs (1s(p)) (o' ([Ts(p)))) - o'(o,))
(T' (d((ps) - o((p;))
(d(p(,))'(pp)))
j=1
Now, returning to our proof, suppose S = [k] and S' = fi',...%% We know from Lemma B.2 that
L(S") L(S) = D Do (p) + ([Ts: - [T(()))) (d' (TT(()) d'(p,))
Ts (p)ls (d'(pi,) o' o'(p)))
Now, consider the term II. Using 10
we can simplify this further as follows:
h
*
II (p)s - [Ts(p),,) vs = vs Σ (P)( Eissp))
(0) o.
where (a) follows measure.
1Ts 1T
Also, I 0 Dg divergence
Finally, to conclude our proof, we show that III > 0. Since the entries of p are sorted in non- decreasing order and as the indices in = [k]] and S' are sorted in ascending order, we have
V [K]. j < in + j [k] p(1,) > p(())
Ts (p)l, (d((pi,)) o' d'(p;)) IIII 0.
25

---

Strict inequality holds as long as some Piç is not among the top-k 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 (AI) 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 defining property that will help us. Below, [(0,, 1J) denotes the image of [0, 1] under f.
Lemma C.1 (Classical). For valid o, let o² (x)x) suppoolp <(p)) be the Legendre dual of Ø, defined for all f([0, 1)). Then, we have for every f((0, 1]) the identity: ((.f=1(a) xf (x) - x² (z). Moreover (Ø")² = f-', and 6² is strictly increasing.
Proof. Since the map - R(p) px <p) is continuous, it achieves maximum on [0. 1]. From the first order condition of the defining equation for d", if the maximum is achieved (0, 1), we have:
aR = - x'(p) = f(p)= 0,
0p
Pmax (x). Now, since f is increasing and x E f[0,, 1]). we have (0) = x f(0) > 0. with equality if x = f(0). Similarly. R'(1) = - ((1) O, with equality if æ (1).. Hence, follows that the above characterization for the maximizer Pmax also applies on the boundaries of [0, 1]. To conclude the proof of the identity, suffices observe that o² (2x) Pmax ((mma) x ( d( (x)- The expression for (o*Y follows by direct calculation.
so for the maximizer Pmux have (Pmax)
C.1 Proof under Assumption (A1)
With the dual convex conjugate &² as per Lemma C.I the divergence measure satisfies:
dp(p. = d (Ø' (4), (p)).
(12)
Let the loss for the dual problem be denoted as I,", (the divergence measure with the arguments swapped), and let T be the dual renormalization map from Lemma A. appllid to pQ- i.e.,
L²) min P(P (p,0v pEA:
min AEAr
de (Pi, Dy) + Six. where Six = Edo(p,
jEQ
Σ dolp;, (Tt())b)) + Sie.
C.1.1 Decomposition the loss difference
Using the form of the loss difference in Lemma for the dual problem as follows:
B.2 and
12
we can compute the loss difference
( Σ
de (pe,
[T (p)l,)
Σ
(duc
de (0' ((T (w)li), o(.))
Σ d (d am1
Indeed, changing the potential & to 0², and changing all the arguments Piy Pcy T T to o' ((..)),((p))) ), o' (((T), o' d'(T3) respectively in Lemma (B.2) suffices. Thus, under the same setup
26

---

of the two subsets 5 = [k] and S' and denoting o' = f, we obtain:
('(8') - = D (F(T3
- ( (6") (69) (f(r,)))
Σ f (TT)) (10)(fo)))-
Since (0-)= f- this further equals
Div (P))
((123))))))-13300)) ((TG)) (17(6)), - P(.)
Σ (p)l)) (py-pu)
III
C.1.2 Analysis of terms based on the dual solution
Similar the proof for the primal case, the term I' 0, as Dø- divergence, and III' > 0 as d' = > 0, as ((0) = 0 and / is increasing. Moreover, as f is strictly increasing, if any of the Pi not among the top-k entries, then strict inequality holds.
To analyze II, we have
II (((T)))))) ((TG())) ((T(()), PG)
from Lemma (I(TT())))) ( ('(((())))))
vs
Since f is convex,
(f(T))()))) ((T > ff (7(10)))
(
f' (TT(p))) (C(TSS (p)) - ((T3()) > (T) - (T3())
(
(TT(() (f(T)( - ((T(())
(173 (p3()),)) = 0.
In the above steps. (a) follows as f 0 as f is strictly increasing and (b) follows as ()
-
= This implies II' > 0, finishing the proof.
C.2 Proof under Assumption (A2)
C.2.1 Extra notation
Since @
by dz v) = (0( (v - x) for y > 0, we define for (I, E D J0, 1] (0, 1] oo).
( v) = (") 2)) -v.
(13)
Define the mapping derived from solving (x, y, v) over by:
{(x, v) [o, 1] x (0, x)) + (0, 1], such that [T()]: = E(P,,)) v) for all i, and for optimal
follows from the proof of Lemma A. that the solution & is well-defined. Define two auxiliary functions db, h that will be used in the computation of the Bregman costs below, such that for all (x, y, v) E D:
(1z.)) =  d'(v)( x), and hdd,v) v(z,E(z,v)).
27

---

C.2.2 Properties of the auxiliary functions
Lemma C.2 (Derivatives
ag oc). Define u [0.1] [0 + [0. 00) as v(x,)) = o²(1)
az Bu " (y) x). We have for all (z, v) [0, 1] x (0,
0E dv (r. v) = v(z, E(x, and ae
(I. v) = )"(((x,))) v)) dr
(14)
E(x,
Proof. The proof of either identity follows by applying implicit differentiation to the function V. Fix x E [0, 1] and consider
v) = (Z d'())(v-))- for (v, E (0. 1] x (0,
Because @ is C2 on (0, 1], F continuously differentiable, and
OF
oy (v.2) = (")()) z) + o0)) = v(2,8) >
by Assumption 3.2 Hence, by the implicit function theorem, the map + ((x, is C¹ with OF
0g dy (x,v) OF/ey v(2,{(z,V)
For the latter identity, fix v > 0 and define
G(z, v) == (x. y. v) =
- 2) - (x. E [0, 1] x (0,
For each zo E (0, 1] let yo == ((). E (0. 1] satisfy G(zo. yo) = We have aG (xx y) = v(x v).
By
Assumption 3.2 gives vzz, v) > 0 for all 0 < y & 1 and 0 < x & y. Hence 0G/Oy yo) + 0.
Since G is continuously differentiable and 0G/0y * 0 at (.o, yo), the implicit-function theorem guarantees C' map + E(x, v) in a neighborhood of zo with G(I,E(z; v/) O.
Differentiating Gzz E1x, v)) = 0 with respect to and using 8G/0z- -d" (v) gives
8G
0G oe
= & ("(((x,)) uJ) + u(z E(x, v)) dc
dr
dy dr
Or.
80
0E (z,v) v(x, ((x, v))
0"({(z,1))
dx
When x = 0, the same argument applies, because a0 (0, v) = v(0, y) > 0 and 8G aG/Oxl(o.)) = By
(v) is finite (the solution y EO,, v) is strictly positive, d'(y) is finite even if d² 000 as y 4 0). Thus 0 ( exists and the same formula holds. This completes the proof.
Lemma C.3 (Derivative Du (A2) we have de (x, < for all x E [0, 1] and 0.
Under the condition that + u(x) = =0"(x)/0((r) (x) non-
decreasing from Assumption
Proof. For the derivative with respect to v, observe first that
di (x, w) x'(v) [0" (v) o'(w)] + x x =0"(y) (x v).
Ty
Hence, by the chain rule,
a dy (E, = He (s. (x, v)) 0E (x. v) = ('<((x,v) v)) (o E(x, ui]
(r. v). -
Oc
Fy
Oy
Iv
Due to the defining equation x² (5 - x) = u, this simplifies to
0h dv (x.v)
0
2 (z, (2, ((x,
0
dy (z,v) v) d(I 0,
V
dr
2 (I, (I)
where the last equality uses
u(x, E(xx( v)) and > 0.
Ov
28

---

Lemma C.4 (Derivative g Assumption (A2)
implies
0h 8x
(z. > 0 for all z E [0, and > 0.
Proof. First recall that
Oyb dy
((r = o d'(v) (y -x)
=
(z, y) = d'()).
(z.y) = (v) (x - -y).
dx
Hence, with Exx v).
0h (z. v) OU dx
(z. + Oip (x.
H (z, (I, = = (()))
0E
[r (I.
dx
Dy
dr
Or
Because { = {(2,1) v) satisfies (5) (5 - z) = v, we have
Ōh
(x, v) = [3(()) (E(
(x, (E) Fx
Ox
v(
Write
N(x, v) = (E)) + (5 x) 0' ( -() = ('()) (() + (5 - (x) A(E),
where A(t) = o' (1) - (t)²
Case 1: A(E) > O. Because { > x from Lemma A.I the second term is non-negative; with d',6" 6" > 0
the first term is also non-negative, N O.
Case 2: A(E) < 0. Since & 3 I, we have
(x, > (E) = o'(E)2
where = tt"))/) Indeed,
= d (t) (t)² o' o + + (t) - -(")()"=''''""" - + ((d(t))))((((@@tit)eee (t())'()))]]
By Assumption (A2) u is non-decreasing, so u (E) O; hence (xx > 0 in this case as well.
Because u(z, <) > 0 and N(r, v) > in both cases, we conclude 0h / > 0 for all x E [0.
and 0, thereby proving the lemma.
C.2.3 Proving the dual greedy property
Denote an arbitrary subset of the indices by: S C [J]. Let vy be the corresponding Lagrange multiplier. Below, for vector RV and set C [V] we denote by x[[S] the sub-vector of restricted to the coordinates in S. Since d'(0) 0 by the assumptions of Theorem 3.3 denoting
I= I (p 0) + (())SS) we can write for every S:
(p. Σ p(pm) ( (((((()) (pm - (T())]m) +
Σ mE[J](S
do 0)
mES
Σ
(d((T())),.) (aai))-((((() (m((())))) ([T(D) pPm) +T
(pm, [(p)) Σ hh(p, +T.
Now, let us prove that the greedy property holds. Suppose S is optimal among all subsets of indices of size & but does consist of some of the top & probability tokens. Then there exist some i + such that E S. g S, and > ps. Denote 5' = 5 (1)
Let Ps, y denote the choice of y that makes the projected probabilities sum to unity. Now since S only differs from S in that it includes the larger > Pi, we can conclude that vs > Vsr.
Then, using the above formula for the value of the objective function on an arbitrary subset, we have:
(p. ((S)) (p (pS) h(pj, vs:) h((p.
Σ mES)
(h(pm, - h(pm; vs),
29

---

Now, since h decreases in v by Lemma C33 we have that the sum is nonnegative since 1/s Vg. As for the remaining term, we have:
hp, Us*)
h(ps, vs)
> h(p,, vs):
where the first inequality by the fact that 1/gr Vg and Lemma C.3 and the second inequality by
the fact that p; > 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 the cost function is discretely convex in l for the primal, it suffices to show that
B
To show that
[((k)) min Do((á, -v-)). = D& (o((T (p). -v.k), $E
discretely convex in k. Indeed, the difference cost(k:) L((E)) = Al is linear k.
To simplify notation, let us denote L((k) by L(k) and Twx by Tx. From Lemma B. we know that
with Sy SiM1 ko)
L(k) = Σ 5=1
+ Sv.
Using 110. we know that f[]TT f(p;)
+
VM vj E [A:]. Again, we simply denote Mwn as Va-
For j E [c] letting x = f(p,) + 1/8 in Lemma
C.1
we have:
- -
+ v2) - /(D))/ -(05)) - vk)
- = (x- ( (f
+ vk).
But now, using that the nonzero entries of (p) must sum unity, we find the following simplification:
&
L(k) = Σ
[p]Tb(p)]s 6" (f((p)) + vn)} + 35
1=1
= Vv Σ
[w - -0' (f(pe) + +v) Sy = 1k - Σ ((((pp)) + mx) + +Sv (15)
-
J=1
1=1
Now, define the auxiliary function W for all j, v for which the expression below well defined:
W(K,v) = Σ /e (s(p;) + v),
(16)
1=1
where p is implicitly kept fixed. From the above calculation, we thus obtain after canceling out terms:
L(k: + 1) 2L()) + L(k 1) W(k 1, 2W (W,,, (k, vR) + W(k 1,
To prove that this nonnegative, leverage that W(K,)) is strictly concave in for each which follows as the Legendre dual mapping o' is strictly convex since so is c. Then, observe that for every j.
i W (k: (k) (0 =1 Σ
(0 - Σ (((ps) +V)
(17)
Ov
Thus,
(k. 1 = 1- Σ
(-((((s)) + =1 1- [T(()),, = O.
dv
30

---

As As (k,.) is strictly concave in (k..)) is maximized at Vs. Thus, we have: (1) W +k+++ 1,+1, Mc+1)
W(k++ + 1, 1,);,, and (2) W(k 1,1%-1) > W(k 1, 1.). With these in hand, we have:
L(k + 1) - 21 + [(k-1)
- - W(k + 1*+1 (k, + W( 1, PR-1)
(18)
7 [W 1,,) W (k. w(,, W(k - 1, 1,
Now, due to the definition of W, the last display equals
- - (f(px+1) + vx) + (f(px) + v/)) > 0,
(19)
the inequality holding pt p,+1, and the mapping p o' (f(p) vv)) is increasing in p since so are o- and f. This concludes the proof.
E Proof of discrete convexity for dual Bregman projection
We denote 0z(v) = (v) - x). As observed before, we have for all admissible x, y that
i
dz v) (v))
δy
and the convexity condition for the second argument of de of Assumption 3.2 is given by:
a
dy (v) 7 + +
+ (v) x) > for all y > x
The dual projection for any 1 و < V is given (for optimal Lagrange multiplier vz)) by:
pp. (1T) (p)l.) V o"
(p)): p.) Vy-
Denote the dual Bregman objective, as function of the selected sparsity k. as:
cost" (k) D& (p, (T2 (p), 0v + Ak.
We now demonstrate that cost" (k) is discretely convex in For this, we will directly show that the second-order differences of this function are nonnegative at every k E (2...., V.. - 1}. Specifically, we can write:
A (k)) == cost" (k- + 1) - 2cost" (k) + cost" (k- -
D& (p. (Tz+1 Ov -v-k- 2D (p, (T (p), ov) + D& (p, (T) Ov -b+1)
We now decompose this quantity into three terms corresponding to three ranges of index i E [V].
namely E [k- 1], i (kr k + 1). and i {k 2...., V). We obtain:
k-1
(T. (p))a) + (do (do (T))( ( (.) = Σ { (do(p) Th+1 d
- de (7 }
1=1
+ (0(pk) - 6(0) - (z) ((())) ((zT - - ((TT()))) - (TT()) (px-
+ (((Pk) -
(pr (7 (p)]x))
+ (0(px+1) ((0) (0(0)
Pk+ (0(p++1) - -(0) -0((0) pk+1)
+ (Ø(pH+I)= 6([T211 -((T)+1 - (po.1 (T (p)l+1)))+ }
V
Σ -) - 2d. (p. 0)}.
The last sum is identically zero, engage with the other two ranges of indices.
31

---

Range 1: i E [k - 1]. For Range 1, recall that for any convex function v. it holds for any two points z, y in its domain that (x)) (v()) > (z) y). Now, notice that for each in Range each of the two terms in figure brackets can be bounded via the convexity of do(x, :) in its second argument as:
= = (p))) (174+1(0) b (2% = va ( >= ((TT+1())) (TE())))
(pt, (T++())))) - delp,, > ((
By
and:
(P))(()) d > - ( T-TT ((T-1(P)): - (T))
( us ((TE- (p)):
=0p. (TT)())) ((TE)(P))
- (T)() (1)
As a result, we may simplify the Range 1 sum as follows, using that by definition, the first j terms in the projection 7 for each {K- - k, & 1} sum to unity:
Range Sum
Uk ([TT+ - [T(())) (TTI(())
17 ())1))
k-1 Σ 1iTR+1 - Σ [T]x)li+
R-1
k-1
Σ(TT
u/k ((1 ( (
2(1 (p))a) +1)
= " (2[TZ(0)l - (T++1 - (T+1())))++)
For Range 2, we first note that the following three types of terms cancel out: @(0). (Pk)), d(p++1). Furthermore, terms involving o'(0) vanish by assumption.
Range 2: i E k 13.
The remaining terms in the Range 2 sum can then be written as:
Range Sum >
( o' ((TT() (pk - (TT()) + (P)1))
- (px - (p)l)))))) } + - (00))
(D6+1 [T (1))e4))))))
Now, we can bound
0 ([T+41())))-)) pk+1 (((Ti-1(P)))+)) pk,
using that p& > pk+1 and the strict convexity of ø. We find the lower bound
Range 2 Sum
2 (((T) o' ((TT (ex (+))a)
((12+110)) - ((((2+1())))) (pk (p)+))
(PT))))))) &' (' (p))+1))) (px - (Te (P)++)))))))
By adding and subtracting the term ((P)) twice, we have the following equivalent bound:
> -2( 1000----(7 ((T o (TT2())))) (pe (Tx (p))))
Range 2 Sum
+ { p(pN) - o( (1((741())))) (px - [TT (p)lk))
+ (((pk) -
- o' (' (p)1a+1) (pt - The (p)k+t)) -22 (px. (r))s)) + d (px. (T (p))k) + +dg (Px, (7)((4))44))
32

---

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:
A (k) > D& (2(T) [T (p)la - (p)]k+1)++
2d (Pk: (T())) + do (pe, + (p)]&) + do (DRs (T (p)]a+1).
Now, define the following key auxiliary function the [0. 1] - R, such that for all E [0, 1]:
4 (z) = x d$ (pk.
This lets rewrite our lower bound equivalently as:
A (1)) > d (P)))) - db ((T++(P))k)) - 4 (1T)+1
(20)
We now establish monotonicity property for thk.
Lemma E.1. For every k [V] the function (/((x) is increasing [0, (p)]a).
Proof. We consider the derivative of the function thg
0
① (Pk, = Uk - Pp (x)) = Pp. (ITE)))&) Opa (x):
dr (xx(x)
dr
where we have used the connection between (v) and Ms. (see Lemma A.
Now, recalling that by assumption, By (v) for all y 2 > 0, and using that [T 3 Pt by properties of the dual projection method (see Lemma A. we have that:
a [x(x) = 8pc. (pT&)) Ppc (z) > 0,
dz
so long as 0 < x (D)).
Continuing, by the properties of the dual projection, we have:
(T)] > (T)+ > (T)+1
In view of Lemmal E.1
20) implies that
(k)) > (bb (p))) - v (T++)())))) (TTE)))) - (1T+1())++1) 20+ = 0.
This concludes the proof of dual discrete convexity of the Bregman cost function.
Algorithmic details
F.
Computing the dual renormalization map
Recall that when & is dual valid, the renormalization map 7% is uniquely defined for x E Asub, with Σ, x; by the fixed point equation (see Lemma A.I
(x),, = =,,/=" /f for all E [k], where 2² E R is chosen so that Σ [T)( =
(2) To compute Tz recall from Section C.2.1 the function v from 13) with =" x)) -v for all vr y, v. Then, for a fixed v. [ (z)), satisfying the equation [T(x)]; = -; [ is equivalent to solving (T Vi, v) = 0 for yı [T(x)]; The monotonicity properties from Lemma A.1 then suggest the following algorithm, consisting binary search over v E (0, M), and then
over each coordinate of I solving o ((T xx))
33

---

Algorithm Dual Renormalization Map T (at) via Nested Binary Search
Require: Convex generator with derivatives f = 0', f Φ"; input vector a E Auu, with
Σ 24 < 1; tolerance
Ensure: Renormalized vector p (x) A
1: function DUALRENORMALIZEG 9, E)
2:
- length of
3:
f" + 6"
4:
M M+ -1 (1- max; 24)
Upper bound on feasible
5:
Initialize Mow t- 0, high f- M
6:
while Mhigh Mow 2 do
7:
f VMow + 1/
8:
for SOLVEROOT(S,, u, ,"
do
9:
f x|i]
y
11:
end for
12:
G Σí
- vld
13:
iGG
then
14:
Vow t
15:
else
16:
Fhigh +- V
17:
end if
18: end while
19: return y
20:
end function
21: function SOLVEROOT(H; v, f
22:
a - Z₁, b+1 while a > E do
24:
m + (a + b)/2
V - ('(m)- - (m x34)
26:
if v < 0 then
28:
else
29:
b - m
30:
end if
end while
32:
return (a + b)/2
33:
end function
F.2 Pseudocode for algorithms
See Algorithm 3 and Algorithm for pseudocode for sparse primal (resp. dual) Bregman decoding.
34

---

Algorithm 2 Discrete Binary Search for Unimodal Cost Minimization
Require: Callable function COMPUTECOST, maximum support size V
Ensure: Optimal support size k" minimizing COMPUTECOST(K)
1: function BINAEYSEANCCHCOMPPEEEOO
V)
2: C1 - COMPUTECOST()
3: C2 < COMPUTECOST(2)
4: if C2 - c1 2 then
5:
return 1
6:
end
7:
ev-1 < COMPUTECOSTIV 1)
8:
cy f- COMPUTECOSTT())
9: cv -V-1 < 0 then
return
11: end if
12: Initialize L 4- 1, R V
13: while R - do
14: 17 f (L+ R2/2]
15:
Gm COMPUTECOST(M)
16:
Cm+1 + CCOMPUTECOST(n + 1)
if Cm++ Cm. then
18:
R +
else
20:
I. -- m
end
22: end while
23: return R
24:
end function
Algorithm 3 Regularized Sparse Primal Bregman Decoding
Require: Probability vector p E Áv, valid convex generator Φ sparsity penalty
Ensure: Sparse decoded distribution á E Ay
function SPARSEPRING CARRSEERRIMALBBEE BRE.MAAAABBEEEEEEDDEEEEEE é, A)
2: Sort in descending order: P(1) > P(2) ... > P(V)
3: Define f = o'
4: function COMPUTERRENORMALZZTIONGG E R$)
5: Solve for v E R. such that Σx f- v)
6: return ((k) with [((*)] = f f (f(xz) + v) for i [k]
7:
end function
8:
function COMPUTECOST(A)
9:
Let p[1:k]
10:
((*) COMPUTERENORMALIZTTINNG
Pad with zeros: ((0)) + (p) ,...,PP (k)
.,.....)
Compute D& (f(()) Ek= 5(p (b) ((p:) f(p)) - pi)
13:
return cost (k) = (p(k)) p) + Ak
end function
15:
k* <
V)
16:
Recompute pf**") using COMPUTRRENORRMAAZZTTONNIY..
17:
Pad with zeros to full length V
18:
return f(k*)
19:
end function
35

---

Algorithm 4 Regularized Sparse Dual Bregman Decoding
Require: Probability vector p E Av, valid convex generator Φ, sparsity penalty À > 0
Ensure: Sparse decoded distribution & E Ay
1: function SPARSEDUA AL BEEEEEDAABBREEANEEEEE REGMAN
Ø, X)
2: Sort p descending order: P(1) 2 P(2) 3 > P(V)
3: Define f = 6',
o"
4: function COMPUTEDDAALRENRRRALAAATOOOOO
E IR*))
5: Solve for E R such that: Σx= [T][T]] - where [T]())] satisfies the fixed- fixed-point
equation: [T) = + // (/T
6:
return (5(k) (r)
7:
end function
8: function COMPUTEDUALCOSTTE)
9:
Let p[1:k]
(()) + COMPUTEDUULRRROOMAALAATAOOOO
11:
Pad with zeros: ((0) + (p) ,h..,,.
(k) 0. .......)) 0)
Compute (p, (p.))(()) = E/=1 -1 p(p:) .(D. - f(( - -b(40)
12:
13:
return cks)(k) Do(p, ((() Ak:
14:
end function
15:
k* <- SINARYYSEERCCCCCCOpuuDDDDDCC
V)
16:
Recompute p(k*) using COMPETEDDUALRRRORRAIITTTTTT///
Pad with zeros full length V
return p(k
19: end function
G Example: a-Bregman decoding
G.1 Proof of Lemma
We first restate the lemma.
Lemma G.1. All generator functions Dos O >> are dual-valid and satisfy Assumption (A2)
Proof. For Assumption 3.2 we can explicitly write:
do(x, alla 1) a(a 1)
xaah
(x 3r
x
xa
a(a 1)
Therefore, the second derivative in of this expression is
(a -1yy (a 2)xy² = 3 Now, if x, then using 11 0 have that the above expression is
(v(a 1) - I(a - 2)) 1/8 - (y(a - 1) + x(2 -a))
- 1) + x(2 a)) -ーー3 2 0,
confirming the convexity in Now for the condition that + u(x) = xδ" (z)/2)/0'(2 decreasing from Assumption (A2) we can observe that
is non-
aa-1
22-1
d'(x)*"'(x) d"(x)² a - -aa 2) (2a
a a-1
Therefore, we identically have:
x. 2a-3
x 20-1
(((()))) (x)))
=0,
a -
0 -
thus concluding the proof.
G.2 Proof of Proposition 4.2
Recall the a-renormalization map [Ta (p)); + -ー -او
E [Rt], where the shift parameter
v(a, p) is chosen so that [T. 1. We treat each value (or limit) of a in turn.
36

---

The limit a + -00. Define
Fa(u) E +v) 1/8
Because z z1/8 is strictly decreasing and convex on (0, for Fg is strictly
decreasing and continuous on the interval (-min; pf, 00). Moreover, lim min, p. Fạ(1) = 00 and limotoo F2(v)
so unique root Va with Fø(ps) 1 exists. Because (3((0) S E Ps < and Fg is decreasing, we have Ug & O. Let qi En (0) 1,
[pa(:]]: (p + v2) 1/8 and i" be the index where pj is largest. Using the constraint
qi.
(a)
=1 Σ
= 8 P + E(p - (6000)
Pp
3 Ps- + 8.
Raising 9i- (o) = (p + v2))
1/8
to the power 3 < 0 yields
V3= = (pr. + R3) pi. Rs = E(E q: (00))
E [0,
(21)
For i + 8", we have U3 ug/p, _ + O. Indeed, implies |us| < 11 (c -11 with e: = (px- +0)/ps-
Because 3 + -C, 00 + O, we have |ua|
(p..)
= =dp,) Then,
pi (1
Mg
1/3
+ Pi,
if..
(22)
Pf
Summing (2 over i + i" and using Σ, (a)
= 1 gives
% (0)
Σ P: P Ps+ + infii
(23)
Equations 22 and 23 establish (a) + T T (p) component-wise, comppppetttiise.ccomptnennnng the proof.
The case a R. Now a 1 2. hence (Ti.(pp)]; (Vri+ +)
[k]. Set
E Pj and A = j== Pj. The normalization condition becomes
1 Σ +V) A 2sv + kv?
Solving kv² 2sy (A (A-1) 1) 0 for the root that yields nen-negative probabilities gives
/2 (1- A) Hence
2 (vo v + k - A)
[Tis(P)]e
i [k).
k
The case a = 2. Here a - 1, so Definition 4.1
yields [TT(p]] = ps + v,
i i [k] The 1- Σ Ps
normalization condition gives 1 Σk (pi + +V) Substituting yields
Σt Ps + kv, hence v
1 E-E [T(()]]: pt
i [k).
k
The limit +00. Write 3 a: + +00. Let = e with E [0. 1]. Then
[T)p)]] = (v. + o = =exp log(p 3 + cd) }²)} Using 1/2 log(a + b + log(max(a, (a, b)) as B + oc gives lim, (P)lp max(p,, c).
Choose the water level so that Ek= max(pr.c) = 1. This furnishes the claimed water-filling rule.
The four cases above prove Proposition 4.2
37

---

G.3 Illustrating primal and dual renormalization
We consider the peaked vector = [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 a-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 Figure + As predicted by our theory, both renormalization families coincide at three values of the parameter, namely at E (1, oo). Furthermore, the primal family evolves more gradually than the dual family between the endpoints of the parameter interval a E (1, 2]. while the reverse behavior occurs for E (2, oo) (where both renormalizations gradually converge to the
water-filling limit which, this case, is the uniform distribution).
Renormalization Top Coordinate a
Renormalization of Second Coordinate vs.
Thanformed value (vecond coontinat))
0.200
Transfomed value (tep coodmmate
0.9
0.175
0.8
0.150
0.7
0.6
0.5
0.4
0.050
0.025
Dual
Figure 4: Comparison of primal and dual renormalization maps: The transformation of the larger
value (0.1, left) and of the smaller value (0.001, right).
G.4 Illustrating general nonconvexity of dual renormalization
Figure iillsstrate that the dual Bregman objective can in general be non-convex for large a.
Plot of Dø(x. for [0. 1] for different
a 1.5
2.5
15
5.0
1.0
022
5.03
0.4
035
02
64
22
04 02
34 02
3.5
0 10
0 20
0.3
2.25
8
000
504
0.00
S082
0.2
8.0
100
84
24
Figure 5: Nonconvexity of the Bregman dual landscape on the square (x, E [0, 13
G.5 Illustrating discrete convexity
F iillustrates 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 X as
1/80. When à 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 Ak, which here equals Al: for k: 80. Thus, all curves (corresponding to different values of a) for both the primal and dual plots, asymptote to lincarity and converge to this value at & 80.
38

---

Primal Bregman-a
Dual Bregman-o
12 am
g=1.2
u=1 1.255
1.5
1=1.5
Primal Cost: cost(k)
o-2
Dual Cost: costr(k)
QF2
1.0
20
50
Figure 6: Discrete convexity of the function ++ cost(k) for primal and dual Bregman a-decoding.
G.6
The simultaneous effects of Bregman decoding and temperature scaling
Normalized Decoded Probabilities
a=1.5, T-0.3
a=1.5, T=1.5
a=2.0,
T=1.5
2.5
a=1.5, T=0.7
a=1, T=0.3
a=2.0, T=0.7
a=1, T=0.7
2.0
a=1.5, T=1.0
a=1, T=1.0
a=2.0, T=1.0
a=1, T=1.5
1.5
1.0
0.5
10
Index
Figure 7: Comparison with changing the temperature.
Here, provide plot to help compare the simultancous effects of Bregman decoding and temperature scaling. We use the same simulation setting and plotting style as in our figure from the introduction (Section [: except we only plot the nonzero probabilities (i.e., the top k 10 probabilities), and we plot the relative sizes the probabilities compared to the standard top-k decoding. Further,
the same and temperature hyperparameters used in experiments in Table The results are shown in Figure 7 Standard top-k decoding corresponds to a = 1 and T 1. From the figure, it appears that the effect of the small probabilities are pushed to zero; which could potentially be one reason why a-Bregman decoding with can perform better at high temperatures.
is moderate/regullriize the amount by which
H
Supplementary experimental details
H.1 Compute resources
The experiments were conducted on system running Rocky Linux 8.10, with 64 CPU cores of Intel(R) Xeon(R) Gold 6448Y processors GHz, TB of RAM, and NVIDIA L40S GPUs with 46 GB of memory each. All experiments can be done with only one GPU and multiple GPUs
39

---

GPT-2 MAUVE Scores
Llama3.1 1 MAUVE Scores
0.96
0.70
MAUVE Score
MAUVE Score
0.92
0.90
0.88
0.64
Top-k
Top-k
0.86
Primal-1.5
0.62
Primal-1.5
Primal-2.0
Primal-2.0
0.84
5
10
20
25
35
40
50
5
10
15
20
25
30
35
40
45
50
k
Figure 8: MAUVE scores results between generated and human human-written text for GPT2-large (left panel) and LLaMA 3.1 8B (right panel), for various values. We show top-k decoding and primal decoding with (1.5, 2.0]. Standard deviations estimated using 50 bootstrap resamples
were used only to parallelize experiments. The software environment used Python 3.11.11, Py PyTorhh 1.5.1, and CUDA 12.4.
H.2 Supplementary experimental results
In this section, we provide additional experimental results to supplement those from Section
Table shows results analogous to those Table for {0.1, 0.0013.
Table 2: Accuracy on GSM8K for LLaMA 3.1 8B using Bregman primal decoding (A E {0.1, 0.001.
(1.5, 2.0}) and top-k decoding, across different temperature settings. For top-k, k equals the averaged optimal ke from the corresponding primal decoding run (matching temperature, A, and o). Standard deviations estimated using 1000 bootstrap resamples.
À >= 0.1
À = 0.001
Temp
Top-k 0.1)
Top-k (A 0.001)
1.5 a 2.0
1.5
a = 2.0
0.3
13
84,4621.0
84,69
84.69 2099
83 93.99=1101 85.29
83.62+1.00 83.6221102
0.7
83.47
85
84,69 84.69 2099
8 a..
83.78+1.z 83.78
1.0
84.46+100
1.
84,69
84.69
1..
a0s
78.5441.13 81.20+1.0
1.5
13.78 14.38 84.69
84.69 20.9
69 22
73.92 64.67 75.97
Figure ( presents the MAUVE scores comparing generated and human-writte text under different decoding strategies. While primal decoding shows slight advantage over top-k decoding, the differences are not statistically significant. We report standard deviations estimated from bootstrap resamples; 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 Owwen.55114B-Inniii and Phi-3-mnedinn-4k-insrnuutt
Figure 9 shows results analogous to those in Figure Table and show the accuracy on GSMSK, analogously to Table and Table 5 and show results for the 1ew.---meiumm----niiieeee model.
40

---

Tapk
2.0
bop Pimal-15
Aoselute Rerplestty Dir
Top-
3
Primal-1.5
Aosolute Depalll Dir
Primal-1.5
Aovelate Depllsity
Prmall2.0
Repelition
1.5
Primazo0
Primal-2-0
1.0
0.5
+
0.0
20
50
10
40
50
30
40
Figure 9: Perplexity and repetition frequency differences between generated and human-wrilten text Phi-3- medium-4k-instree (left two panels) and Qwen2.5-114B-Isstte (right two panels), for various & values. We show top-k decoding and primal decoding with E (1.5, 2.0}. Standard deviations are estimated using 1000 bootstrap resamples.
Table Accuracy on GSM8K for Qwen2.5- 148--nsttuct using Bregman primal decoding (A E {0.1, 0.01}, E (1.5, 2.0}) and top-k decoding, for various temperatures. For top-k, & equals averaged k from primal decoding with matching temperature, x, and a. Standard deviations are over 1000 bootstrap resamples.
X 0.1
0.01
Temp
Top-k (A 0.1)
Top-k (A 0.01)
a = 1.5
a=
2.0
α 1.5 a 2.0
0.3
82.7 1:04 12.26:1105
81.42 81 1mm
12.6411.04 12.18
81.43 81,43
0.7
81.73=1.0
81.43+107 81.05 81.43
1. 80.21
80.21 81.43+1.97
1.0
80.59
81.
81.43+1107 81.43
78.85+1112
80.29
79.30 81.43
1.5
80.89 1.. 81.43
1.0 81.43
17.18 18.99 17.48 81.431157
Table
Accuracy on GSMSK Qwen2.5-114B-Isstrii using Bregman primal decoding (A [0.001, 0.0001]. J. a E (1.5, 2.0}) and top-k decoding, for various temperatures. For top-k, k equals the averaged &"' from primal decoding with matching temperature, X, and a. Standard deviations are over 1000 bootstrap resamples.
X = 0.001
À = 0.0001
Temp
Top-k 0.001)
Top-k 0.0001)
1.5
2.0
1.5 a a= 2.0
0.3
82.11 1.06
1..
8..4 82. 80.06 80.21
81.88 82.26
82.03 82.41
80.21
79.7611111
79. 79.76
79.98 80.06
1.0
78.92
1.14
79.30 +91
78.47+113
79.30
78177 +
1.5
76.72 1.16 78.01+1.14
75.89 77.48
+1.9 74.91
71.19 75.89
Table 5: Accuracy on GSM8K for Phi-3-mediim--4kinneee
using Bregman primal decoding (A E {0.1, 0.01), E (1.5, 2.0}) and top-k decoding, for various temperatures. For top-k, & equals the averaged k from primal decoding with matching temperature, μ4, and a. Standard deviations over 1000 bootstrap resamples.
0.1 Temp = 1.5 a = 2.0
0.01
Top-k (A 0.1)
Top-k (A 0.01)
= 1.5 = 2.0
0.3
86.81 #093 87.87 85,97 =096 85,97 0.96
87 17,41
=0.11
87.04
87 07.2261992 87 27.22449.92
0.7
86.964095 88, 0. 85.97
-5.97
85.67
86.88
88. 10.4.9 88.10+a.8
1.0
9695 87. 2092
85.97 :096 85.97
0.96
84.99 20.98
13.93
85.44
85.44+1097
1.5
86.58 85.97 10.6 85.97 9.96 82.94 13.70
80. 1. 80. 10.14
Table 6: Accuracy on GSM8K for 1if-3-meddium---instnee using Bregman primal decoding (A
{0.001, 0.0001}, E (1.5, 2.0}) and top-k decoding, for various temperatures. For top-k, k equals the averaged &* from primal decoding with matching temperature, μ., and Standard deviations over 1000 bootstrap resamples.
A = 0.001
À = 0.0001
Temp
Top-k (A = 0.001)
Top-A (A 0.0001)
1.5
2.0
1.5
2.0
0.3
88.11
4092
86,50 86
36.81#099
87.49
87.49
86.20 86.50+a95
0.7
86.81 209 86.50
85. 85.67
84.99
84.91
85.60.1097 85.29 98
1.0
83.62+112
82.34 10s
82,71 82,79
82.71 1.04
+
81.35
82.71 1:00
1.5
1.19
78.92
69.75
23
22.25
16.04-1.18
62.62 65.814155 41.1
41

---

H.4 Experiments for TriviaQA
Table 7 and X show accuracy on TriviaQA for LLaMA3.1-8B model. Here we choose 10% (As 1800 questions) proportion of TriviQA validation dataset for evaluation.
Table 7: Accuracy on TriviaQA for LLaMA 3.1 8B using Bregman primal decoding (A E {0.1, ).01),
E (1.5, 2.0 and top-k decoding, for various temperatures. For top-k, k cquals the averaged A from primal decoding with matching temperature, X, and a. Standard deviations over 1000 bootstrap resamples.
À = 0.1
Top-k (A 0.1)
= 0.01
Top-k (A 0.01)
Temp
a
=
1.5 a a= 2.0
0 a=
a = 2.0
a 1.5 0 2.0
a a= a = 2.0
0.3
67.80 11
67.58+1.11 66.57+1.11 61
66.74.1.11 66.69 66.74.1.1 61.1 65.01
0.7
65.68
67.58 66.35
64.23 63.8
1.0
67.58
67.58
61.06=1115 61.17
62.67 62.67
1.5
64.85 66.96.1.11 1.15 61.1
59.78
60184
11.5 60.84
Table 8: Accuracy on TriviaQA for LLaMA 3.1 8B using Bregman primal decoding (A E {0.001, 0.0001 }. E (1.5, 2.0}) and top-k decoding, for various lemperatures. For top-k, & equals the averaged k from primal decoding with matching temperature, A, and a. Standard deviations over 1000 bootstrap resamllss
A = 0.001
Top-k (A = 0.001)
À = 0.0001
Top-k (A = 0.0001)
Temp
1.5
2.0
1.5
2.0
= 1.5 a 2.0
== 1.5 = 2.0
0.3
66.85
67.58
67.1 17.1
66.69.1.11 1.10
67.19 67.58
0.7
63.40=1.14
63.18+1.11
64.68 64.79
62.73+1.11 62.73
63.79+1113 63.68+111
1.0
59.00 59.00
11
62.23
57.99
59.11
58.55
60. 1116
55.04 55.71 11.18 56.38
19.19 52.59 50.19 51.31
Table and 10 show analogous accuracy results for Phi3-medium-4--nnstruccc on TriviaQA.
Table 9: Accuracy on TriviaQA for Phi-3-medium -olluum-----iee insinutruct using Bregman primal decoding (A E {0.1, 0.01). (1.5,200} and top-k decoding, for various temperatures. For top-k. & equals the averaged K' from primal decoding with matching temperature, X, and a. Standard deviations are over 1000 bootstrap resamples.
0.1
Top-k (A = 0.1)
A 0.01
Top-k (A 0.01)
Temp
Q
1.5 -= 2.0 a 1.5 a 2.0
1.5
0 2.0
1.5 a 2.0
0.3
58.44
59.67
59.05 1 60.50
59.33
59.22+111
19.11 59.39 1.6
0.7
57.44=1.17
58.22+1.16 56.77 60.50=111
55.21 55.88
55.54 56.77
1.0
56.60 56.94.1.17
54.54 60.50
52.09
51.75
50.31 52.37
1.5
57.16-1.17 58.22:111,,00114: 60.50
49.47+11. 50.19 43.55 45.29
Table 10:
Accuracy TriviaQA for Phii3--nedium----inntuue using Bregman primal decoding (A 0.0001,, }, E {1.5, 2.0}) and top-k decoding, for various temperatures. For top-k, le equals the averaged &*' from primal decoding with matching temperature, X, and a. Standard deviations are over 1000 bootstrap resamples.
X 0.001
Top-k 0.001)
X 0.0001
Top-k (A 0.0001)
Temp
a = 1.5 a = 2.0
a = 1.5 a = 2.0
== a= 2.0
a = 1.5 a = 2.0
0.3
58.61+1.16 59.44+1.16 54.82=1.17
59.72=1.16
59.22=1.16 (9.e:l..69993311111
59.83+1.16 59.39
59.4441.16 59.44+1116
0.7
14.04 53.70 54.60
54.54
54.43
56.21 54.71
1.0
48.13=1.18
19 49.58 50.64
48.58
48.6441.18 48.6411118
1.5
42.51 14.18
39.55 42.67
39.94
1..0 37.72
H.S
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 GSM8K and TriviaQA datasets.
In Table 11 we show the average A values (and their values rounded to the nearest integer) selected by primal Bregman decoding on GSM8K with LLaMA 3.1 8B for various temperatures, a, and Table 12 shows corresponding standard deviation and entropy.
42

---

Table 11: Mean (and rounded) average k- values on GSM8K with LLaMA 3.1 8B for various temperatures, a, and X.
Temp
À
=
0.01
A 0.001
A= 0.0001
a 1.5
a 2.0
0
1.5 a= 2.0
a 1.5 a = 2.0
a 1.5
a = 2.0
0.3
1.2231(1) 1.1537
1.620 (2) 1.4453 (1) (2) 1.7964 (2) 2.8578 (3) 2.2112 (2)
0.7
1.2295 (1) 1.1554 (1) 1.6689 (2) 1.4794 (1) 2.3193 (2) 1.9048 (2) 3.2554 (3) 2.4974 (2)
1.0
1.2287
(1) 1.1594 (1)
1.7519 (2) 1,5048 (2) 2.7231 (3) 2.0234 (2) 4.6926 (5) 3.0924 (3)
1.5
1.2331 (1) 1.1566 (1) 1.8106 (2) 1,5189 (2) 4. 4.18 (4) 2.4067 (2)
14.2539 (14) 5.6002 (6)
Table 12: Standard deviation (and entropy) of average k* values on GSM8K with LLaMA 3.1 8B for various temperatures, Ot, and X.
0.1
A 0.01
À 01001
À 0.0001
Temp
a
1.5
2.0
0
a
a= 1.5
2.0
1.5
0
0.3
0.46 (0.82) (0.62))))
1.07 (1.55) 0.77 (1.28)
1.89 (2.08) 1.31 (1.77)
3.11 (2.58)
2.00 (2.16)
0.7
0.47 (0.84) 0.36 (0.62)
1.12 (1.62) 0.80 (1.34)
2.21 (2.24)
1.47 (1.89)
3.98 (2.78)
2.53 (2.37)
1.0
0.47 (0.84) 0.37 (0.63)
1.23 (1.72) 0.83 (1.38)
3.03 (2.49)
1.65 (2.00)
7.31 (3.21) 3.69 (2.69)
1.5
0.47 (0.85) 0..66))33)
30 (.779) 0.84 (1.40) 5.37 (3.13) 2.19 (2.32)
18.01 (4.04)
7.77 (3.51)
Table 13 14 show analougous adaptivity results for Qwen2.5- 1.B.-114B-Instruec
Table 13: Mean (and rounded) average k' values on GSM8K with Qwen2.5-1453 In-trrcctt for various temperatures, a, and A.
X = 0.1
A= 0.01
A x= 0.001
A = 0.0001
Temp a = a = 2.0
a= 1.5 a 2.0
1.5 a 2.0
a 0= 1.5 a= 2.0
0.3
1.0993((1 1.0660(()
1.4899(1) 1.3425(1) 2.7614(3) (.7425)))..76144)3))
1.9317(2) 5.4537(5) 3.1726(3)
0.7
1.1010(1) 1.0672(()
1.5043(2) (.7783)) 1.9522(2) 5.5047(6) 3.1911(3) 1.5171(2)
1.0
1,1000(1) 1.0666(1)
1.3591(1) 2.7985(3) 1.9723(2) 6.5603( 3.2493(3)
1.5
1.1008(1) 1.0662(1) 1.5211(2) 1.3628(1) 2.8761(3) 2.0028(2) 5.7831(6) 3.3285(3)
Table 14: Standard deviation (and entropy) of average k-* values on GSM8K with Qwen2.5-14B- Instruct under À= 0.0001 and varying temperatures.
Temp
1.5
2.0
0.3
10.75 (2.81)
4.88 (2.26)
0.7
10.71 (2.86)
4.85 (2.29)
1.0
10.70 (2.90)
4.88 (2.34)
1.5
10.75 (3.03)
4.90 (2.42)
Table 15 show analougous adaptivity results for Phi-3-medium-44kinntteee
Table 15: Mean (and rounded) average hst values on GSM8K with Phi-3-medium-44k-nntteee for various temperatures, a, and
X 0.1
== 0.01
X 0.001
À 0.0001
Temp
a = 1.5 a 2.0 a == 1.5 a= 2.0
a= 1.5 a = 2.0 a == 1.5 a: = 2.0
0.3
1,4048(1) 1.2609(1) 2.4123(2) 1.9287( 4.7116(5) 3.1299(3) 8.6473(9) 5.2889(5)
0.7
1.4074(1)
1.2601(1) 2.4337(2) 1.9409(2) 4.6706(5) 3.1307(3) 8.6958(9) 5.3697(5)
1.0
1.4073(1) 1.2603(1) 2.4541(2) 1.9364(2) 4.7772(5) 3.1792(3) 8.8501(9) 5.4394(5)
1.5
1.4098(1) 1.2575(1) 2.4667(2) 1.9498(2) (.9289(5) 3.2335(3) 9.4782(9) 5.6113(6)
Table 16:
Standard deviation (and entropy) of average Ae values GSM8K with Phi-3-medium-4e instruct under À= = 0.0001 for varying temperatures and o.
Temp
1.5 Q 2.0
0.3
12.09 (3.83)
6.77 (3.32)
0.7
12.01 (3.89)
7.23 (3.61)
1.0
11.98 (3.98)
6.74 (3.45)
1.5
11.79 (4.24)
7.29 (3.79)
43

---

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, a, and Table 18 shows corresponding standard deviation and entropy.
Table 17: Mean (and rounded) average ks values on TriviaQA with LLaMA 3.1 8B for various temperatures, a, and X.
X 0.1 1.5
== 0.01
= 0.001
À 0.0001
Temp
2.0
a= 1.5 a 2.0
= 1.5 a= 2.0
== 1.5
2.0
0.3
1.1536(1)
1.1452(1) 1.9135(2) 1.5291(2) 2.5753(3) (155753())69440)) 6.9406(7) 4.5149(5)
0.7
1.2265(1)
1.1275(1) 2.0109(2) 1.6265(2) 3,8877(4) 2.7593(3) 8.8845(9) 5.1892(5)
1.0
1.2138(1) 1.1324(1) 2.0273(2) 1.6818(2) 3.9715(4) 2.9759(3) 8. 8.4552(8) 5.7381(6)
1.5
1.2013(1) 1.1384(1) 2.0289(2) 1.7032(2) (.1749(4) 2,9398(3) 8.4399(8) 5.5166(6)
Table 18: Standard deviation (and entropy) of average k* values on TriviaQA with LLaMA 3.1 8B for various temperatures, a, and A.
À
0.1
À 0.01
== a= 1.5 a = 2.0
À 0.0001
Temp
a
1.5
a= 2.0
&
1.5 a = 2.0
a= 1.5 a = 2.0
0.3
0..4 (0.65) 0.35 (0.60) 1.37 (1.90) 0.86 (1.42) (2.85) 2.09 (2.42)
10.36 (3.63) 5.35 (3.28)
0.7
0.48 (0.83) 0.33 (0.55) 1.44 (2.00) 0.93 (1.56) 4.24 (3.09) 2.20 (2.53)
12.18 (4.10) 5.98 (3.56)
1.0
0.47 (0.81) 0.34 (0.56) 1.42 (2.01) 0.98 (1.63) 4.42 (3.03) 2.43 (2.68)
12.07 (3.77) 6.54 (3.68)
1.5
0.46 (0.78) 0.35 (0.58) 1.42 (2.01) 1.00 (1.66) 5.07 (3.02) 2.46 (2.62)
12.90 (3.34) 7.18 (3.35)
Table 20 show analougous adaptivity results for
20
on TriviaQA.
Table 19: Mean (and rounded) average 1** values on TriviaQA with hew.--medium--k--innsse for various temperatures, a, and X.
À 0.01
= 0.001
0.0001
Temp
a = 1.5 a 2.0
a
1.5 a= 2.0
== & = 2.0
a == 1.5
& = 2.0
0.3
1.7393(2)
3.6184(4) (.414(()).6614(()) 2.81843)
1.4142(1)
9.2976(9) 5.2226(5)
18.7026(19)
10.4901(10)
0.7
1.7148(2)
1.4288(1)
3.6134(4) 2.6381(3) 8.4512(8) (58 16.8627(17) 9.3718(9)
1.0
1.7348(2)
1.4216(1)
3.6840(4) 2.6050(3) 8.3500(8) (.8924(5)
16. 16.7567(17)
9.6411(10)
1.5
1.6687(2) 1.4378(1) 3.6081(4) 2.6601 8.6007(9) 5.1906(5) 18.2735(18) 9.7162(10)
Table 20: Standard deviation (and entropy) of average Ae values TriviaQA with Phi-3 thi-3-medium-44 instruct for various temperatures, a, and A.
À =
0.1
À 0.01
A=
0.001
A= 0.0001
Temp
1.5
2.0
a
1.5
a 20
1.5 a= 2.0
1.5 a 2.0
0.3
0.87 (1.43) 0.49 (0.98) 2.84 (2.65) 1.76 (2.26) 8.19 (4.18) 3.99 (3.32)
16.54 (5.06) 8.88 (4.33)
0.7
0.87 (1.41) 0.49 (0.99) 2.69 (2.76)
1.70 (2.22) 7.50 (4.16) 3.78 (3.28)
15.26 (5.16) 8.38 (4.31)
1.0
0.87
(1.43) 01.49 (0.98) 2.68 (2.82) 1.65 (2.24) 7.04 (4.22) 3.62 (3.37)
13.81 (5.27) 7.75 (4.46)
1.5
0.84 (1.40) 0.50 (0.99) 2.58 (2.84)
1.63 (2.27) 6.60 (4.30) 3.58 (3.45)
13.94 (5.37) 7,51 (4.51)
44


Reviews:
[{"review_id": "YDgQ3iYe1N", "reviewer": "Reviewer_HJHQ", "summary": "The 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.", "questions": "Q1. 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?", "limitations": "Even 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The authors addressed my concerns.\nI would like to keep my high score.", "rebuttal": "Thanks so much for your careful reading and interesting questions! We now address your Weaknesses comments and then your questions.\n\n**Q:** Optimal k is sensitive to $\\lambda$, how to set?\n\n**A:** In our empirical evaluations, the optimal $k$ is not highly sensitive but is rather reasonably robust to the choice of $\\lambda$. As can be seen e.g. from Table 2 (in the Appendix), even for values of $\\lambda$ on different orders of magnitude, the average optimal $k$ does not fluctuate too much (in particular, it's never larger than $\\approx 50$ in our experiments). That being said, there is a reasonable ballpark way to set $\\lambda$ to get reasonable optimal $k$: please see our answers below to your Q2, and to Reviewer 2, about the relationship between the optimal $k$ and the hyperparameter $\\lambda$.\n\n**Q:** New decoders need full sorting and root-finding, which can hinder use on large LLMs.\n\n**A:** Actually, we would argue that the computational and numerical aspects of our decoders are quite a bit more optimistic than this assessment, in a few ways, and that they shouldn't cause big issues when used on large LLMs (certainly they don't cause issues in our empirics).\n\nAs for root finding/numerics: First, some important special cases like primal $\\alpha \\in \\{1.5, 2\\}$ decoding (and the original top-$k$ decoding) have closed-form renormalization maps hence don't require root finding. In the other cases when root finding is required, the functions whose root we need to find are scalar, increasing, and well-behaved, precluding numerical issues that could occur for nonmonotonic or nonsmooth functions. \n\nAs for full-vocabulary sorting: This is also not quite required in practical settings. To see that, note that the optimal $k^\\*$ for practical purposes is usually not greater than (e.g.) $V' \\in [50, 200]$, which is $\\ll V$. We can certainly enforce that we limit to the top $200$ tokens if need be. Even without enforcing this, our methods can in fact run faster in such settings.\nSpecifically, they can run in $O(V + V' \\log V')$ time where $V'$ is a true upper bound on the optimal $k^\\*$.\nThis can be significantly more efficient than the worst-case bound $O(V \\log V)$. \n\nTo see this, we note: \n(1) identifying the top $V'$ probabilities doesn't require full vocab sorting and can be done in $O(V)$ time (e.g. using a heap);\n \n(2) instead of binary searching for $k \\in [1, V]$, one can search for $k \\in [1, V']$. If $V'$ is unknown, then one can use *exponential search* (gallop by trying $V' = 1, 2, 4, \\ldots$, and then perform binary search once a true upper bound $V'$ is found---i.e., when the discrete difference is positive according to the convexity property) to find the correct $V'$ and binary search within $[1, V']$. \n\nThis will take $O(V' \\log V')$ time (see also our answer to the similar question from Reviewer 2).  While this paper was written mainly from the theoretical standpoint, which was the reason we gave the worst-case bound $O(V \\log V)$, we will be happy to clarify this in a remark in the revision.  \n\n---\n\n**Q1:** On our assumptions on Bregmans.\n\n**A:** Assumption 3.1 suffices for all our primal decoding results, and is very mild (it only requires a well-defined Bregman divergence with first-order smoothness) hence is satisfied by KL (the main Bregman score used in practice) and in fact all our primal rules (Havrda-Charvat-Tsallis entropies) for $\\alpha \\neq 0$ (Section 4). The dual conditions look complex, bur are also relatively mild and are satisfied for all Havrda-Charvat-Tsallis entropies with $\\alpha > 1$. Intuitively, all our assumptions will hold for Bregmans whose generators are \"not too singular\" at the boundary ($x \\to \\{0, 1\\}$). \n\n**Q2:** On tuning $\\lambda$ to control $k$.\n\n**A:** While clearly $\\lambda$ increases, $k^\\*$ decreases and vice versa, deriving the asymptotics of the optimal $k^\\*$ as a function of $\\lambda$ and $\\alpha$ (with or without assumptions on $p$) turns out to be an intricate task that does not to our knowledge immediately fall out of known techniques. \nFor a concrete observation, one can derive via a brief calculation that $k^\\*$ is the floor or the ceiling of $1/(\\lambda \\alpha)^{1/\\alpha}$ when $p = (1/V, \\ldots, 1/V)$.\nBut in fact, our preliminary experiments suggest that the scaling $k^\\* \\sim 1/(\\lambda \\alpha)^{1/\\alpha}$ holds much more broadly. It also faithfully approximates the optimal $k^\\*$ in expectation over drawing uniformly random (or other well-behaved Dirichlet) random probability vectors $p$ from the $V$-simplex (we can't include the plot by new NeurIPS rules). Thus, $k \\sim 1/(\\lambda \\alpha)^{1/\\alpha}$ should serve as a good guideline for tuning $\\lambda$ to ensure the desired ballpark of $k^\\*$. (We had not included this information to keep this paper focused and fit in the page limit.)\n\n**Q3:** On discrete convexity for adaptive $\\lambda$.\n\n**A:** The convexity of the cost function in $k$ indeed holds regardless of $\\lambda \\geq 0$ (as $\\lambda k$ is linear in $k$). That being said, the shape and the location of the argmin $k^\\*$ of the cost function will change, and binary search may need to be redone. Thus, we recommend selecting a fixed $\\lambda$ for each invocation of the decoder. \n\n**Q4:** Non-convexity in dual Bregman decoding?\n\n**A:** One of the novel structural insights into the dual case the optimization of $\\hat{p}$ reduces to the case where $\\hat{p}_i \\geq p_i$ for $i \\in [k]$, and in this regime *the problem becomes convex*! So in fact, we show that the optimum must lie within a convex subdomain of the generally nonconvex problem, and we are able to solve the problem within that region.\nThus, both theoretically and practically the problem becomes convex.\n\n**Q5:** Runtime of decoding methods, improvements to sorting algorithm?\n\n**A:** For new Bregman decoders ($\\alpha \\in \\{1.5, 2\\}$) that we evaluated, the runtime is actually (nearly) the same as for standard top-k ($\\alpha=1$). It is important to note that the algorithmic primitive here is binary search (for $k \\in [1, V]$) rather than sorting. Practical speedups thus can be obtained by searching for best $k$ in a smaller range (e.g. $k \\in [1, V']$ for $V' < V$, e.g. $V' = 50$); and retrieving the top $V'$ tokens for some constant $V' \\ll V$ takes only $O(V)$, not $O(V \\log V)$, time (e.g. via heaps).\nFurther, in practice, approximation algorithms could be tried, and should perform well if relatively accurate.\n\n**Q6:** Interactions between $\\alpha$ and $\\lambda$ in the limiting cases? Choosing them jointly?\n\n**A:** Any pair of $\\alpha$, $\\lambda > 0$ can be used; e.g. if some $k^*$ is targeted, could use 3-way formula btw $k, \\alpha, \\lambda$ in Q2. \n\nThe limiting behavior for both $\\alpha$ and $\\lambda$ is well-defined and well-interpreted, and their limiting interaction is nicely disentangled: \nIntuitively, $\\lambda$ controls how many tokens $k$ end up selected while $\\alpha$ controls the shape of renormalization on these $k$ tokens. When $\\alpha \\to \\infty$: regardless of $\\lambda$, renormalization converges to ``water-filling'' (i.e., it adds mass to the smallest selected token probabilities). When $\\lambda \\to 0$: regardless of $\\alpha$, *all tokens* ($k=V$) will be selected in the limit; conversely, when $\\lambda \\to \\infty$, only one token ($k=1$) will be selected.\n\n**Q7:** Experiments: larger models, different domains (dialog, code).\n\n**A:** In addition to the experiments presented in the main paper, we have conducted further evaluations to strengthen our findings:\n\n1. Evaluated larger models \n— **Phi-3-medium-4k-instruct (14B)** and **Qwen2.5-14B-Instruct** \n— on both open-ended text generation and GSM8K.\n\n2.Added **TriviaQA** evaluations, using LLaMA-3.1-8B-Instruct, \n**Phi-3-medium-4k-instruct (14B)** and \n**Qwen2.5-14B-Instruct**.\n\nDue to the character limit in this response, we can only include a subset of our results. \nWe would be happy to provide the full results upon request.\n\n**Selected Evaluation Results:**\n> Open-ended text generation: \n\nPrimal decoding matches top-k in PPL and repetition diff. \n\nPrimal-2.0 consistently yields lowest PPL gap.\n\n### Phi-3-medium-4k-instruct\n\n| k| PPL Diff|||Rep Diff %|| |\n| --|--------|-----|-------|----------|-----|------- |\n|| Top-k| P-1.5|P-2.0|Top-k| P-1.5|P-2.0 |\n| 10|4.8| 4.2|**4.2**|1.1| 0.8|**0.6** |\n| 20|4.1| 3.7|**3.4**|0.5| 0.5|**0.2** |\n| 30|3.7| 3.3|**2.9**|0.3| 0.6|**0.3** |\n| 40|3.4| 3.0|**2.6**|0.3| 0.3|0.4 |\n| 50|3.2| 2.8|**2.4**|**0.1**| 0.4|**0.1** |\n\n---\n\n### Qwen2.5-14B-Instruct\n\n| k| PPL Diff|||Rep Diff %|| |\n| --|--------|-----|-------|----------|-----|------- |\n|| Top-k| P-1.5|P-2.0|Top-k| P-1.5|P-2.0 |\n| 10|6.3| 5.9|**5.5**|1.2| 0.8|1.2 |\n| 20|5.3| 4.8|**4.3**|1.4| 1.3|**1.2** |\n| 30|4.8| 4.1|**3.6**|1.0| 1.6|**0.8** |\n| 40|4.3| 3.7|**3.1**|1.0| 1.2|1.1 |\n| 50|4.0| 3.3|**2.7**|1.2| 1.2|1.4 |\n\n---\n\n> GSM8K: We only report \\lambda=0.0001. \nPrimal decoding matches top-k. \nTop-k uses average k\\* from primal. \nAt high temp (e.g. 1.5), top-k degrades faster.\n\n### Phi-3-medium-4k-instruct\n\n| T|P-1.5|P-2.0|Top-k-1.5|Top-k-2.0 |\n| ---|-----|-----|---------|--------- |\n| 0.3|87.5| 87.5| 86.2| 86.5|\n| 0.7|85.0| 84.9| 85.6| 85.3|\n| 1.0|82.7| 82.1| 81.4| 82.7|\n| 1.5|72.2| 76.0| 62.6| 65.8|\n\n---\n\n### Qwen2.5-14B-Instruct\n\n| T|P-1.5|P-2.0|Top-k-1.5|Top-k-2.0 |\n| ---|-----|-----|---------|--------- |\n| 0.3|81.9| 82.3| 82.0| 82.4|\n| 0.7|79.6| 79.8| 80.0| 80.1|\n| 1.0|78.5| 79.3| 78.8| 79.4|\n| 1.5|74.9| 74.9| 71.2| 75.9|\n\n---\n\n> TriviaQA: We only report \\lambda=0.0001\n\n### Llama3.1-8B-Instruct\n\n| T|P-1.5|P-2.0|Top-k-1.5|Top-k-2.0 |\n| ---|-----|-----|---------|--------- |\n| 0.3|66.7| 67.1| 67.2| 67.6|\n| 0.7|62.7| 62.7| 63.8| 63.7|\n| 1.0|58.0| 59.1| 58.6| 60.1|\n| 1.5|49.2| 52.6| 50.2| 51.3|\n\n---\n\n### Phi-3-medium-4k-instruct\n\n| T|P-1.5|P-2.0|Top-k-1.5|Top-k-2.0 |\n| ---|-----|-----|---------|--------- |\n| 0.3|59.8| 59.4| 59.4| 59.4|\n| 0.7|54.5| 54.4| 56.2| 54.7|\n| 1.0|48.7| 48.6| 48.6| 48.6|\n| 1.5|38.2| 39.9| 36.0| 37.7|"}, {"review_id": "LlQaJFrDro", "reviewer": "Reviewer_wNkD", "summary": "This 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.", "questions": "### 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.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "The paper is theoretically sound, and provides a novel, and exciting perspective on top-k decoding as $\\ell_0$-regularized Bregman divergence minimization. As shown in the discussions, the experimental section is not fully satisfying in terms of the evaluated settings. However, after engaging with the authors, they provide some additional experiments and vastly extending the empirical section may be beyond the scope of the current study.", "rebuttal": "Thank you so much for your careful reading, favorable evaluation, and insightful question and comments! We now respond to your points in the Weaknesses and Questions sections.\n\n-----------------------\n\n**Weaknesses:**\n\nFirst, we agree with your presentational comments about Ll. 26-29 and 44-54, these will be improved!\n\n**Q:** Evaluates limited settings (models, datasets, regularization strengths, …). Expand it to settings in Chen et al.?\n\n**A:** We have added experiments with larger models and new datasets, see response to reviewer 1 Q7. Our open-ended generation setup directly follows Chen et al.: same dataset (WebText test split) and evaluation metrics. Since we evaluate under fixed truncation size k (partial evaluation), other decoding methods (e.g., top-p, min-p, typical decoding) effectively reduce to top-k. Thus, our comparisons isolate the effect of normalization under various pre-defined k values.\n\n**Q:** Leaves out interesting analysis about values of $k$ chosen by algorithm as a function of $\\lambda$. Does $k$ significantly change during decoding?\n\n**A:** Thank you for these insightful questions! For the latter question we will make sure to include various plots of adaptivity in $k^\\*$ to the manuscript: as one would anticipate from their adaptivity properties, the Bregman decoders do choose diverse values of $k^\\*$ for different token distributions throughout our experiments. Plotting the histograms of the distributions of $k^\\*$ per-task (we unfortunately are not allowed to include images anymore by the new conference rules), one can observe distributions of $k^\\*$ over $k \\geq 1$ that are peaked at $k = 1$ and then near-monotonically decay until modest values of $k \\in [20, 50]$. In other words, the decoding methods choose a diverse set of small $k$s. (One can stretch these distributions further out by making the regularization parameter smaller, $\\lambda \\to 0$; and conversely, can concentrate $k^\\*$ on $1$ by increasing $\\lambda$.) \n \nTo illustrate this numerically without being able to show histograms, below we provide the standard deviations and Shannon entropy for the distributions of $k^\\*$ on a subset of the Llama and Phi evaluations, on GSM8K and TriviaQA respectively, for $\\lambda = 0.0001$ and $\\lambda = 0.001$:\n\n**Table: Std dev (and entropy) of average $k^\\*$ on GSM8K using LLama3.1-8B-instruct with $\\lambda = 0.0001$.**\n\n| Temp | $\\alpha=1.5$ | $\\alpha=2.0$ |\n|------|-----------------|-----------------|\n| 0.3 | 3.1 (2.6) | 2.0 (2.2) |\n| 0.7 | 4.0 (2.8) | 2.5 (2.4) |\n| 1.0 | 7.3 (3.2) | 3.7 (2.7) |\n| 1.5 | 18.0 (4.0) | 7.8 (3.5) |\n\n**Table: Std dev (and entropy) of average $k^\\*$ on TriviaQA using Phi-3-medium-4k-instruct with $\\lambda = 0.001$.**\n\n| Temp | $\\alpha=1.5$ | $\\alpha=2.0$ |\n|------|-----------------|-----------------|\n| 0.3 | 8.2 (4.2) | 4.0 (3.3) |\n| 0.7 | 7.5 (4.2) | 3.8 (3.3) |\n| 1.0 | 7.0 (4.2) | 3.6 (3.4) |\n| 1.5 | 6.6 (4.3) | 3.6 (3.5) |\n\nNow, while it is clear that as $\\lambda$ increases, $k^\\*$ decreases and vice versa, deriving the asymptotics of the optimal $k^\\*$ as a function of $\\lambda$ and $\\alpha$ (with or without assumptions on $p$) turns out to be an intricate task that does not to our knowledge immediately fall out of known techniques. \n\nFor a concrete observation on that front, one can derive via a brief calculation that $k^\\*$ is the floor or the ceiling of $1/(\\lambda \\alpha)^{1/\\alpha}$ when $p = (1/V, \\ldots, 1/V)$.\nBut in fact, our preliminary experiments suggest that the scaling $k^\\* \\sim 1/(\\lambda \\alpha)^{1/\\alpha}$ holds much more broadly. It also faithfully approximates the optimal $k^\\*$ in expectation over drawing uniformly random (or other well-behaved Dirichlet) random probability vectors $p$ from the $V$-simplex (unfortunately we cannot include the plot confirming so by the new conference rules). Thus, this relationship $k \\sim 1/(\\lambda \\alpha)^{1/\\alpha}$ should serve as a good guideline for tuning $\\lambda$ to ensure the desired ballpark of $k^\\*$. Once again, we had not included this information to keep this paper focused and fit in the page limit.\n\n**Q:** Only covers primal decoding, leaving open comparison between primal and dual?\n\n**A:** We agree that a more detailed empirical study of the dual family and comparisons to the primal family are very interesting directions! For this manuscript, we mainly aimed for a minimalistic evaluation with the goal to show that the new decoding rules are worth a more substantial empirical investigation in future work, and in this sense aimed to pick out just a couple natural primal decoding strategies, that we also gave as an example in Section 4 right before that, which would demonstrate some of the ``first order effects'' of using them instead of vanilla top-$k$. \n\nThere indeed appear to be some subtle connections and correspondences between the primal and dual families, which we thought was best to leave to follow-up work to keep the presentation focused and to do proper justice to this question. For instance, primal and dual decodings not only coincide at $\\alpha=2$ by symmetry, but also have similar but subtly different renormalizations for other values of $\\alpha$. As another observation, the primal and dual families have different sensitivities to changing the parameter $\\alpha$ for various values of $\\alpha$; so hypothetically, a practitioner with prior knowledge about the approximate range of the optimal $\\alpha$ may want to correspondingly choose primal or dual decoding. \n\n**Q:** Better visual grouping for Figure 3, Table 1.\n\n**A:** We would be very happy to improve the visual grouping of the empirical results' displays (Figure 3 and Table 1). As simple improvements, we could bring them closer together in the document (e.g., to appear on the same page one under the other), and we could arrange the sub-figures in Figure 3 in a $2 \\times 2$ grid for better interpretability (given that extra space would be available later). As for Table 1, arranging it in a non-confusing way was admittedly tricky for us, given that more than $2$ dimensions had to be presented; if you have any desired presentational tweaks in mind on that front, please let us know.\n\n**Q:** Generalization to generators beyond Havrda-Charvat-Tsallis-entropies.\n\n**A:** To us, going beyond $\\alpha$-decoders is an exciting direction for future work. Indeed, $\\alpha$-decoding is a rich, concrete and interpretable class of decoders, but we don't necessarily believe it ``spans'' the whole space of decoding behaviors one can elicit from Bregman divergences --- and we expect and hope that other decoding-appropriate Bregman scores will be designed/discovered in follup-ups!\n\nFor our part, to help kickstart such future work, we have put effort into refining assumptions under which our structural results provably hold --- and as a result, our theory (i.e., the greedy and k-convexity properties) does apply broadly to Bregman scores beyond Havrda-Charvat-Tsallis-entropies.\n\n-----------------------\n\n**Questions:**\n\n**Q:** Adaptive-sampling $O(V\\log V)$ complexity for $V\\sim 10^5$ seems heavy, how does this affect runtime? Envisioned speedups?\n\n**A:** There are indeed some natural speedup options. First and foremost, in practice one would usually restrict attention to $k \\leq V'$ for some $V' < V$ (e.g. $V' = 200$ or $V' = 50$ or even smaller) --- thus the complexity becomes $O(V' \\log V')$, which is very manageable. (In fact, in our evaluations we selected $V'=50$, hence getting barely any runtime overhead.)\n\nSecondly, if one doesn't want to commit to a particular $V'$, one can still take advantage of the prior knowledge that $k^\\*$ is likely small, while keeping the rigorous guarantees: Instead of binary search, just perform *exponential search* (gallop by trying $V' = 1, 2, 4, \\ldots$, and then perform binary search once a true upper bound $V'$ is found) --- this will give the time complexity $O(k^\\* \\log k^\\*)$. \n\n**Q:** Re. Fig.3: qualitative discussion about why other Bregmans can outperform vanilla top-$k$?\n\n**A:** This is a great question, and we will be happy to add explicit discussion of this to the empirical section. Theoretically and qualitatively (and as schematically shown in the Figure in lines 44-54), if we keep $k$ fixed and gradually increase $\\alpha$, the $\\alpha$-Bregman renormalization/projection will go from boosting higher-probability tokens (for small $\\alpha$), to boosting lowest-probability tokens among the top-$k$ (for large $\\alpha$ and in the limit as $\\alpha \\to \\infty$). (By ``boosting'' we mean how the renormalization redistributes the residual mass of the non-top-$k$ tokens between the top-$k$ ones). For this reason, we believe that e.g.\\ the observed marked decrease in repetitiveness that the $\\alpha \\in \\{1.5, 2\\}$ decoding exhibits over standard top-$k$ (i.e., $\\alpha=1$) is likely due to higher $\\alpha$ leading to higher diversity of decoding."}, {"review_id": "TKWWfw9KOM", "reviewer": "Reviewer_4MPc", "summary": "The 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.", "questions": "1. 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?", "limitations": "some discussion is found in the section 7.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "In general, I feel the paper provides a reassuring result and has some value, but I didn't find it to be exciting. Moreover, I'm not convinced that the generality via Bregman divergence can bring any practical benefits. I have taken into account the rebuttal/discussion and decided to keep the score.", "rebuttal": "Thank you for your effort in evaluating this manuscript. Before answering your concrete questions, we’d like to comment on this paper's goals: It aims to provide a theoretical foundation for top-$k$ decoding, to achieve the following: \n(1) Through a rigorous optimization/loss-theoretic framework, explain the critical design choices inherent in top-$k$ decoding, including: How to adaptively choose the optimal $k$? How to re-normalize the chosen probabilities of the top-$k$ tokens? etc.; \n(2) Illustrate the framework on the rich class of Bregman losses (as elaborated below, this let us elicit many novel decoding intuitions/insights); \n(3) Identify, and solve, the mathematical challenges that arise when rigorously establishing basic properties of a decoding method: namely, what we call its ``greedy property'', and the convexity of the optimization problem for $k$;\n(4) Confirm, through a few realistic simulations, that decoding rules generated by this framework are sensible (and indeed, we would argue, promising) for further study.\n\nNow, we'll address your points in the weaknesses and questions sections:\n\n---\n\n**Weaknesses:**\n\n**Q:** Interpretation is obvious. Top-k: easy computation. Optimization problem not surprising.\n\n**A:** We'd like to begin by remarking that we do not wish to claim that our general sparse optimization formulation is extremely \"surprising\", but that it's: \n\n(1) Very natural, as it formalizes and expands on the major motivation behind top-$k$: Sparsity. \n\n(2) Novel and fills an important knowledge gap: There are currently very few theoretical investigations into the foundations of LLM decoding in general, with the field being predominantly empirical and not resting on rigorous foundations.\n\nTo the point of our framework being \"somewhat obvious\": While the general L0 formulation is natural as explained above, though we still wouldn't call it obvious, the tractable structures that we uncover when instantiating this optimization problem (with the expansive class of Bregman divergences) are not at all obvious. \n\nIndeed, we start with a regularized problem where the regularization is the nonsmooth + nonconvex L0 pseudonorm, and in which the divergence objective itself is not even necessarily convex (e.g. in the dual Bregman case). In these circumstances, our finding about the presence of a doubly favorable optimization landscape, where: \n(1) One only has to check linearly many sparsity patterns out of exponentially many (the \"greedy\" property), and\n(2) The optimization of sparsity norm $k$ is discretely convex in $k$ --- is arguably both surprising and deep. Proving greediness and, even more so, $k$-convexity, required us to develop nontrivial and multifaceted proof techniques (which we feel will be useful even beyond this setting).\n\n**Q:** KL and $l_2$ enough, Bregman less interpretable.\n\n**A:** From our perspective, the generalization of vanilla top-$k$ decoding (KL) to other Bregman divergences actually makes top-$k$ *more*, not less, interpretable. \n\nFor instance, the *general interpretation of renormalization as a simplex projection* is an insightful way to think about this ingredient of generalized top-$k$ decoding, which we believe adds rather than subtracts interpretability. \n\nRevisiting your assessment that the main intuition is already clear from KL and L2 cases, we would like to argue the opposite. Looking at these two cases, it may not be obvious at all if these decoders are just two distinct heuristic choices --- where you either normalize top-$k$ by dividing by the sum, or you add back the remaining mass in equal chunks to all top-$k$ tokens --- or if there is a sense in which one can smoothly interpolate between them, creating \"intermediate decoders\". \n\nIn fact, our primal and dual families provide two distinct ways to put KL and L2 decoding on a continuum --- with the parameter $\\alpha$ (with $\\alpha=1$ giving KL and $\\alpha=2$ giving L2) being nicely interpretable: as $\\alpha$ increases, our family of Bregman decodings ranges, for any fixed $k$, from \"less diverse\" to \"more diverse\" (see also more on this below in this answer). This e.g. lets us easily read off from this continuous family that L2-decoding ($\\alpha=2$) is more diverse than vanilla top-$k$ ($\\alpha=1$), and even that there is an \"intermediately diverse\" closed-form decoding for $\\alpha=1.5$. \n\nThe limiting cases are also insightful. When $\\alpha \\to \\infty$, our insight (see L. 259) says that not only do decoders not diverge in this case as one might suppose, but they in fact converge to the beautiful *water-filling* limit in which the remaining mass is interpreted as \"water\" that gets \"poured\" on the lowest-probability tokens out of top-$k$. On the opposite end, if one takes $\\alpha \\to -\\infty$, one gets (see L. 256) another clean, \"least diverse\", limit in which the remaining mass gets placed \"on top of\" the top token. \n\nIn the context of decoding, the most diverse option is commonly thought of as a uniform distribution over the $k$ tokens, and the least diverse option as simply putting all mass on the top token.\nBy contrast, the limits just described are substantially different.\n\n**Q:**  New rules only of theoretical, not practical interest.\n\n**A:**  We do believe that our empirical evaluations show some practical promise for future work within our framework --- please see the response to Question 5 for a discussion of our empirical evaluations.\n\n**Q:**  Simple computation isn't considered but motivates top-k; new rules not simple, no clear performance gains.\n\n**A:**  We would like to clarify that we did in fact carefully consider the issue of simple computation, from both theoretical and empirical angles, making sure that the new decoding rules do not lose the scalability that's inherent to vanilla top-$k$ decoding. \n\nWe have established the discrete $k$-convexity of the optimization problem; please see Section 3.3. \nThese theorems demonstrate that rather than having to naively grid search over all $V$ (vocab size) possible values of $k$, it provably suffices to check (via binary search) only $O(\\log V)$ of them. This is a drastic improvement from the overhead on computing the best $k$ that one would expect a-priori. \n\nWhat's more, in fact even $O(\\log V)$ tries are unnecessary in typical practice, since via exponential search, we would only need $O(\\log k^*)$ tried values of $k$ (see responses to Reviewers 1 and 2), which is barely any overhead at all over the vanilla top-$k$ decoding, given that in practice, optimal $k$ is far smaller (on the order of at most 50 or so) than the vocab size.\n\nOn the empirical side, for important special cases of Bregman decoders, such as the choices $\\alpha \\in \\{1.5, 2\\}$ that we tested in the empirical evaluations, we actually match the speed of vanilla top-$k$ decoding, benefiting from vectorized computations of projections. This is due to the closed forms for renormalization in those cases, as well due to improved search for the optimal $k^*$.\n\nMoreover, we reiterate that we do see some promise in the empirical results, as explained in the next answer.\n\n**Q:**  Numerical evaluation is weak?\n\n**A:** We position this paper as studying the theoretical foundations of top-k decoding (with the corresponding title, exposition, and NeurIPS category). Thus, we view \nthe purpose of our simulations as simply to confirm that our theory leads to decoders that are both meaningful and reasonable, and that further hold promise for future study. We believe our evaluations clear that bar; and we believe our approach matches the standards of other theoretical papers published in NeurIPS. \n\nRegarding the performance of our methods, we would like to clarify that there are in fact a few promising aspects: (1) $\\alpha=2$ broadly achieves better perplexity in our experiments, and repetitiveness reduction is observed.\nMoreover, (2) on GSM-8k the performance of our decoders appears more robust across temperatures (i.e., accuracy doesn't degrade as badly as the standard top-$k$).\n\nWe have already mentioned these claims in our paper (lines 303 and 310), but didn't emphasize them too much, due to our focus on the theoretical foundations. However, we will be happy to add some additional emphasis to the claims.\n\nIn the manuscript we made sure to appropriately caution the reader to not over-interpret these preliminary results but instead to view them as a starting point towards a fuller evaluation, on more models, tasks and metrics.  \n\n---\n\n**Questions:**\n\n**Q:** Replace $l_0$ with $l_1$?\n\n**A:** The $l_1$ norm of a probability vector (which $\\hat{p}$ must be) is always 1 --- so this regularization would not actually lead to a top-$k$-sparse optimization problem (e.g. for Bregman losses it would just output the original probability vector back).\n\n**Q:** If decoding rule is greedy, must its divergence be Bregman?\n    \n**A:** The greedy property is not limited to this class---e.g. from some additional investigation (performed after submitting the paper), we can also obtain greedy decoders by minimizing certain $f$-divergences. \n\nIn general, the greedy property appears to be a natural one from intuitive considerations, and we expect it to hold for a variety of reasonable loss functions. \n\n**Q:** Assumption 3.1 redundant?\n\n**A:**  Assumption 3.1 adds the extra requirement of first-order smoothness (precisely, continuous differentiability) of the generator --- which is not necessary in the weakest definitions of Bregman divergences. One can operate them even in the nondifferentiable setting by using subgradients. We will clarify this in the revision.\n\n**Q:** Why permutation-equivariance?\n\n**A:** We talk about permutation-equivariance mainly for expository purposes, to underscore that in this paper we focus on decoders which view different tokens symmetrically. You are right that we don't have to enforce it separately since it automatically holds for our setup. If you suggest, we would be happy to remove it."}, {"review_id": "NBt9bV8qFR", "reviewer": "Reviewer_rK4x", "summary": "This 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$.", "questions": "Section 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?", "limitations": "Figure 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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)", "was_revised": true, "final_justification": "The paper provides good theoretical contribution to top-k decoding, which is interesting for the community and LLM development.", "rebuttal": "Thank you so much for your careful reading and interesting questions. As a high-level comment, we'll make sure to incorporate, your suggestions for improving the readability of the manuscript (including: de-congesting parts that are dense (Section 3.1, 3.3); highlighting adaptivity in the intro; highlighting the specific optimization challenges we solve; discussing dual vs primal decoding; adding contour plot lines; etc). While we've put a lot of effort into making the presentation accessible and intuitive, the page limits have forced us to exclude some bits of intuition --- the future extra page will let us to bring them back into the manuscript. \n\n**Q:** Section 3.1 overly dense, lacks intuitive explanations. Include clear, high-level overview with context/intuition. Role of Equation (5) in dual renormalization mapping unclear; detailed explanation/example of its implications?\n\n**A:** We are committed to refining and making more accessible the presentation of our results in this subsection. As briefly described in lines 156-160 at the top of this subsection, the focus of this subsection is on projecting, i.e., turning the selected $k$ tokens' probabilities into a valid probability distribution. This is necessary as the sum of the probabilities of the selected tokens is generally less than $1$, and so the remaining probability mass (coming from the non-selected tokens) needs to be somehow \"added back in\". The rest of the subsection defines conditions/assumptions under which these projection maps $T$ and $T^*$ (for the primal and dual Bregman rules) are unique and well-defined, and gives formulas for computing these projections --- an explicit one in the primal case (Equation 4) and an implicit one in the dual case (Equation 5). Both these formulas are defined in terms of the monotonic function $f$ which specifies the Bregman divergence used to generate the decoder. We will make sure to include this additional description in the updated version.\n\nFor Equation (5) specifically, its role is to define, for a *dual Bregman* decoder, the transformation that the probability $x_i$ of each individual selected token $i$ will undergo when renormalizing the top-$k$ tokens' probabilities to a valid distribution (in parallel to how Equation (4) defines the same transformation for a *primal* Bregman decoder). This equation for each $i$ is implicit, as both of its sides feature [T^\\*_\\phi(x)] _i.\n\nHowever, as we rigorously show under our dual validity assumptions, it still uniquely specifies [T^\\*_\\phi(x)] _i, whose value can be found via a simple root-finding procedure. \n\nFor additional intuition about this set of equations, we can make sure to explicitly mention the following structural properties in the updated version: \n(1) This dual transformation redistributes the probability mass of the non-top-$k$ tokens among the top-$k$ tokens in a way that increases the probability of each top-$k$ token: that is, [T^\\*_\\phi(x)]_i $\\geq x_i$ for all $i \\in [k]$; \n(2) The parameter $\\nu^\\*$ governs how much is added to each $x_i$ by the transformation (the larger $\\nu^\\*$, the larger each transformed $x_i$); \n(3) Thus, between $\\nu = 0$ (when no mass gets added to any $x_i$) and $\\nu \\to \\infty$ (when very large mass is added to each selected token's probability), there exists a (unique) intermediate value $\\nu^\\*$ at which the transformed probabilities sum up to $1$, thus forming the desired valid distribution over top-$k$ tokens. This intermediate value $\\nu^\\*$ can itself be identified via simple root-finding.\n\n**Q:** Theorem 3.2 says top-$k$\n is optimal solution for Equation (2). Why is Equation (6) cost term necessary, why wouldn’t naive search over all k \n suffice? Purpose and necessity of cost($k$)?\n\n**A:** Indeed, as Theorem 3.2 establishes, the optimal solution to Equation (2) is to select *some number* $k$ of top tokens (and then renormalize), so the naive search over all $k$ would suffice to identify the optimal number of top tokens. However, the challenge that we address in Section 3.3 is that an exhaustive search over all possible $k \\in [1, V]$ is computationally prohibitive, as $V$ is large and each tried value of $k$ requires projecting the top-$k$ tokens (and computing the objective value). Practically, one could settle for grid search over some subset of values of $k$, but this will not give provable theoretical guarantees. \n\nSo instead we exhibit a surprising extra structure in the optimization problem: Namely, that *it is convex with respect to k*, and thus one only needs to try $O(\\log V)$, rather than all $V$, candidate values of $k$ to find the optimal one! The notation $cost(k)$, thus, is to help cleanly formalize this convexity --- it makes explicit the dependence of the objective value on the choice of $k$ (while keeping implicit the token distribution $p$).\n\n**Q:** Scenarios where dual might be preferred over primal? Specific use cases/advantages of dual decoding, theoretical or practical benefits?\n\n**A:** The question of when dual decoding can be preferred over primal one is indeed interesting and subtle. While we currently already provide some visual comparisons between the renormalizations induced in the primal and dual case  --- such as in the dedicated subsections G in the Appendix (and e.g. on figures 4 and 6) --- we will be happy to discuss their implications more prominently in the main part. (That said, we believe the subtle nature of the comparison merits its own detailed investigation beyond the scope of this paper.)\nBelow are some initial considerations for when dual decoding may be preferred, which we are happy to include in the updated version and hope will guide future work: \n\n(A) *Inducing larger range of optimal $k$:* Dual decoders can induce a wider range of k* than primal ones: e.g. on the plot of the objective as a function of k in Figure 6, one can notice that the minimum is located in a wider span of locations for the dual (k* $\\in [1, 40]$) than for the primal (k* $\\in [1, 20]$) methods. Thus, when a broader adaptivity range (i.e., how many tokens get selected) is desired, dual decoding may prove favorable.\n\n(B) *Less effort in tuning k:* Also in the context of choosing optimal $k^*$, choosing dual decoding may be preferable when one *doesn't* want to carefully tune the optimal $k$: As can be seen in Figure 6, the minima of the objective are flatter and misspecifying $k$ is of less impact than it would be for primal decoding. \n\n(C) *Tuning $\\alpha$:* The $\\alpha$-parameterization of the primal and dual families exhibits different degrees of sensitivity to changes in $\\alpha \\in (1, \\infty)$ --- e.g. see Figure 4 for how the top-token probability changes with $\\alpha$. Thus, an experimenter who knows an approximate range/location (but not the exact value) of $\\alpha$, should prefer the dual decoding whenever it has lower sensitivity to $\\alpha$ in that range, to cut down on hyperparameter tuning and to ensure robustness to misspecification of $\\alpha$.\n\nWe point out that the above considerations are not distribution-free (i.e. they depend on the underlying distribution $p$, unlike the main framework results that we develop in this paper in a distribution-free manner) --- and hence will require task-specific quantification and careful empirical studies beyond the scope of the present foundations paper.\n\n**Q:** Experimental or theoretical results on applying this metric to beam search? Regarding adaptivity, any connection to methods like min-p and top-p?\n\n**A:** We agree that improving the efficiency of beam search by performing decoding with our new methods is an interesting question for future work. For instance, empirically, the structural property of the Bregman parameter $\\alpha$ --- that it controls the diversity of the decoding --- may result in *controllably diverse* sequences in the beam, helping better explore options within the beam. Mathematically, in this paper we focus on per-token decoding (which already presents interesting challenges to resolve), and building this up to a theory that would provide insights into beam search decoding will require one to develop models of how local, per-token decoding properties translate to global, sequence-level decoding properties.\n\nRegarding the adaptivity property, this is indeed an insightful question --- indeed there is a (non-trivial) connection/extension that one can build to make our framework cover methods such as top-$p$ and min-$p$. The structural reason for which this is possible is that just like top-$k$, these decoding methods also follow the template of selecting some number $k'$ of top tokens, *which depends on the underlying distribution over tokens and hence is adaptive*, and then renormalizing them back to a proper distribution. The difference is that the parameter $k$ in these other methods is reparameterized in terms of the total desired probability mass (for top-$p$) or in terms of the desired minimum included probability (for min-$p$).\n\nWe have not explicitly detailed this extension in this paper (as we wanted to make sure the presentation stays focused for the reader's benefit), but our theory does address both fundamental questions: how to *adaptively* pick the number of tokens to keep, and how to renormalize them back to a proper distribution once they are picked --- which are the essential considerations for top-$k$ and top-$p$ and min-$p$ likewise. \nWe will mention the development of the full theory as a direction for future work in our paper."}]

Author rebuttal:

Meta-review: This paper develops a theoretical framework for top-$k$ decoding by formulating it as an $\ell_{0}$-regularized Bregman divergence minimization problem, proving that greedy selection is always optimal and that the cost function is discretely convex, enabling efficient binary search for the optimal $k$. This unifies existing decoding rules and introduces new Bregman-based decoders, some of which show improved robustness and perplexity in preliminary experiments. 

The proposed theoretical foundations of top-$k$ decoding are valuable in placing a widely used heuristic on solid theoretical ground and offering a principled generalization, though concerns remain regarding limited empirical validation, sensitivity to hyperparameters, and unclear practical benefits relative to the simplicity of standard top-$k$ (which should be included in the updated version). 

During rebuttal, the authors addressed these concerns with additional experiments on larger models and datasets, clarified runtime considerations, and expanded the discussion of primal vs.\ dual decoding and adaptivity. Overall, the strong theoretical contributions and clarified empirical evidence make this a solid and timely paper. 

Thus, I recommend acceptance.
Decision: Accept (poster)