Foundations of Top-$k$ Decoding for Language Models

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.

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?

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.

For 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.

However, 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.

For additional intuition about this set of equations, we can make sure to explicitly mention the following structural properties in the updated version: (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]$; (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$); (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.

Q: Theorem 3.2 says top-$k$ is optimal solution for Equation (2). Why is Equation (6) cost term necessary, why wouldn’t naive search over all k suffice? Purpose and necessity of cost($k$)?

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.

So 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$).

Q: Scenarios where dual might be preferred over primal? Specific use cases/advantages of dual decoding, theoretical or practical benefits?

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.) Below 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:

(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.

(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.

(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$.

We 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.

Q: Experimental or theoretical results on applying this metric to beam search? Regarding adaptivity, any connection to methods like min-p and top-p?

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.

Regarding 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$).

We 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. We will mention the development of the full theory as a direction for future work in our paper.

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.

Weaknesses:

First, we agree with your presentational comments about Ll. 26-29 and 44-54, these will be improved!

Q: Evaluates limited settings (models, datasets, regularization strengths, …). Expand it to settings in Chen et al.?

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.

Q: Leaves out interesting analysis about values of $k$ chosen by algorithm as a function of $\lambda$. Does $k$ significantly change during decoding?

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$.)

To 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$:

Table: Std dev (and entropy) of average $k^\*$ on GSM8K using LLama3.1-8B-instruct with $\lambda = 0.0001$.

Table: Std dev (and entropy) of average $k^\*$ on TriviaQA using Phi-3-medium-4k-instruct with $\lambda = 0.001$.

Now, 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.

For 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)$. But 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.

Q: Only covers primal decoding, leaving open comparison between primal and dual?

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$.

There 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.

Q: Better visual grouping for Figure 3, Table 1.

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.

Q: Generalization to generators beyond Havrda-Charvat-Tsallis-entropies.

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!

For 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.

Questions:

Q: Adaptive-sampling $O(V\log V)$ complexity for $V\sim 10^5$ seems heavy, how does this affect runtime? Envisioned speedups?

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.)

Secondly, 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^\*)$.

Q: Re. Fig.3: qualitative discussion about why other Bregmans can outperform vanilla top-$k$?

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.

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: (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.; (2) Illustrate the framework on the rich class of Bregman losses (as elaborated below, this let us elicit many novel decoding intuitions/insights); (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$; (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.

Now, we'll address your points in the weaknesses and questions sections:

Weaknesses:

Q: Interpretation is obvious. Top-k: easy computation. Optimization problem not surprising.

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:

(1) Very natural, as it formalizes and expands on the major motivation behind top-$k$: Sparsity.

(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.

To 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.

Indeed, 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: (1) One only has to check linearly many sparsity patterns out of exponentially many (the "greedy" property), and (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).

Q: KL and $l_2$ enough, Bregman less interpretable.

A: From our perspective, the generalization of vanilla top-$k$ decoding (KL) to other Bregman divergences actually makes top-$k$ more, not less, interpretable.

For 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.

Revisiting 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".

In 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$.

The 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.

In 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. By contrast, the limits just described are substantially different.

Q: New rules only of theoretical, not practical interest.

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.

Q: Simple computation isn't considered but motivates top-k; new rules not simple, no clear performance gains.

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.

We have established the discrete $k$-convexity of the optimization problem; please see Section 3.3. These 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.

What'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.

On 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^*$.

Moreover, we reiterate that we do see some promise in the empirical results, as explained in the next answer.

Q: Numerical evaluation is weak?

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 the 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.

Regarding 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. Moreover, (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$).

We 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.

In 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.

Questions:

Q: Replace $l_0$ with $l_1$?

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).

Q: If decoding rule is greedy, must its divergence be Bregman?

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.

In 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.

Q: Assumption 3.1 redundant?

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.

Q: Why permutation-equivariance?

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.

Thanks so much for your comments!

We agree with your comments regarding simplicity and computation; and sorry for giving the impression of misunderstanding---we think we are on the same page here. Even though simplicity is the main motivation for top-k, there are also other simple and good methods such as top-p and min-p. So what is special about top-k? Our investigation reveals some answers, which were certainly non-obvious a priori. So in this sense, we think our work adds some insight. We will add additional discussion about this to the paper.

Additionally, the Bregman framework is not only general, but we would also point out that it provides a new way for future researchers to generate novel task-appropriate top-k methods as they please; all that they will need is decide on the shape of the Bregman loss function that their application needs. Indeed, even within the alpha-family that we consider as an example, there is no a-priori reason that the hyper parameter setting alpha=1 (vanilla top-k), of all alphas, should uniformly dominate all other choices in practice. So we view our example Bregman families not just as a generalization for its own sake, but as an actionable generalization.

Regarding the comment that our new method "does not bring any new practical benefit", we would like to reiterate our point that (1) alpha-decoding with alpha=2 is more robust than top-k at higher temperature, and (2) achieves better/lower perplexity than top-k decoding. The level of comprehensiveness of these experiments (especially with our new models---Qwen, Phi-3, and new dataset---TriviaQA) is close in comprehensiveness to that of even the most thorough recent empirical papers in the field, for instance the min-p paper (Turning Up The Heat, ICLR 2025 Oral). We will add more emphasis to these results in the revision.

Regarding the Bregman framework, to us it is quite exciting that we can develop a unified theoretical framework at this level of generality. Everyhing goes through, and we get interesting examples that show some promise as well (alpha-Bregman). From a theoretical standpoint, we believe that this gives ample evidence that this is a good framework; though we understand that opinions can differ, especially from more applied perspectives. We will explain our motivation in more depth in the revision.

Again, we are very grateful for engaging with our paper, and for giving us the opportunity to clarify our points. While we certainly understand where you stand, we do hope that, if indeed all your concerns are addressed and "there is nothing wrong with this paper", then you can reconsider the evaluation.

Sincerely, The Authors

Thanks so much for your careful reading and interesting questions! We now address your Weaknesses comments and then your questions.

Q: Optimal k is sensitive to $\lambda$, how to set?

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$.

Q: New decoders need full sorting and root-finding, which can hinder use on large LLMs.

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).

As 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.

As 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. Specifically, they can run in $O(V + V' \log V')$ time where $V'$ is a true upper bound on the optimal $k^\*$. This can be significantly more efficient than the worst-case bound $O(V \log V)$.

To see this, we note: (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);

(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']$.

This 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.

Q1: On our assumptions on Bregmans.

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\}$).

Q2: On tuning $\lambda$ to control $k$.

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. For 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)$. But 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.)

Q3: On discrete convexity for adaptive $\lambda$.

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.

Q4: Non-convexity in dual Bregman decoding?

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. Thus, both theoretically and practically the problem becomes convex.

Q5: Runtime of decoding methods, improvements to sorting algorithm?

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). Further, in practice, approximation algorithms could be tried, and should perform well if relatively accurate.

Q6: Interactions between $\alpha$ and $\lambda$ in the limiting cases? Choosing them jointly?

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.

The limiting behavior for both $\alpha$ and $\lambda$ is well-defined and well-interpreted, and their limiting interaction is nicely disentangled: Intuitively, $\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.

Q7: Experiments: larger models, different domains (dialog, code).

A: In addition to the experiments presented in the main paper, we have conducted further evaluations to strengthen our findings:

1. Evaluated larger models — Phi-3-medium-4k-instruct (14B) and Qwen2.5-14B-Instruct — on both open-ended text generation and GSM8K.

2.Added TriviaQA evaluations, using LLaMA-3.1-8B-Instruct, Phi-3-medium-4k-instruct (14B) and Qwen2.5-14B-Instruct.

Due to the character limit in this response, we can only include a subset of our results. We would be happy to provide the full results upon request.

Selected Evaluation Results:

Open-ended text generation:

Primal decoding matches top-k in PPL and repetition diff.

Primal-2.0 consistently yields lowest PPL gap.

Phi-3-medium-4k-instruct

Qwen2.5-14B-Instruct

GSM8K: We only report \lambda=0.0001. Primal decoding matches top-k. Top-k uses average k* from primal. At high temp (e.g. 1.5), top-k degrades faster.

Phi-3-medium-4k-instruct

Qwen2.5-14B-Instruct

TriviaQA: We only report \lambda=0.0001

Llama3.1-8B-Instruct

Phi-3-medium-4k-instruct