Tokenized Bandit for LLM Decoding and Alignment
Authors: Suho Shin
Chenghao Yang
Haifeng Xu
MohammadTaghi Hajiaghayi
Venue: ICML 2025 poster (2025)
Keywords: Contextual Bandit
Multi-armed Bandit
Large Language Model
LLM Alignment
Decoding-time Alignment
Decoding Algorithm
Primary area: theory->online_learning_and_bandits
TL;DR: We introduce a novel extension of multi-armed bandit problem with applications to LLM decoding and decoding-time alignment.

Abstract: We introduce the tokenized linear bandit (TLB) and multi-armed bandit (TMAB), variants of linear and stochastic multi-armed bandit problems inspired by LLM decoding and alignment. In these problems, at each round $t \in [T]$, a user submits a query (context), and the decision maker (DM) sequentially selects a token irrevocably from a token set. Once the sequence is complete, the DM observes a random utility from the user, whose expectation is presented by a sequence function mapping the chosen token sequence to a nonnegative real value that depends on the query.

In both problems, we first show that learning is impossible without any structure on the sequence function.
We introduce a natural assumption, diminishing distance with more commons (DDMC), and propose algorithms with regret $\tilde{O}(L\sqrt{T})$ and $\tilde{O}(L\sqrt{T^{2/3}})$ for TLB and TMAB, respectively.
As a side product, we obtain an (almost) optimality of the greedy decoding for LLM decoding algorithm under DDMC, which justifies the unresaonable effectiveness of greedy decoding in several tasks.
This also has an immediate application to decoding-time LLM alignment, when the misaligned utility can be represented as the frozen LLM's utility and a linearly realizable latent function.
We finally validate our algorithm's performance empirically as well as verify our assumptions using synthetic and real-world datasets.

Tokenized Bandit for LLM Decoding and Alignment
Suho Shin Chenghao Yang?
Haifeng Xu²
MohammadTaghi Hajiaghayi
Abstract
et al., 2024). Recent success of commercial LLMs such as ChatGPT, Gemini, and Claude (Team et al., 2023; Achiam et al., 2023) gives evidence on the feasibility of LLMs as personal assistants, supporting individuals' daily life in de- cision making, problem solving, and planning.
We introduce the tokenized linear bandit (TLB) and multi-armed bandit (TMAB), variants lin- ear and stochastic multi-armed bandit problems in- spired by LLM decoding and alignment. In these problems, at each round E [2'. user submits query (context), and the decision maker (DM) sequentially selects token irrevocably from a token set. Once the sequence is complete, the DM observes a random utility from the user, whose expectation is presented by a sequence function mapping the chosen token sequence to a nonnega- tive real value that depends on the query.
To serve as a genuine personalized assistant, however, is crucial to align the LLM's outcome with the designated human preference, which remains core challenge in LLMs (Mishra al., 2021; Sanh et al., 2021; Shen et al., 2023). Recent works suggest that aligning an LLM's out- put with human preference can be partially addressed by fine-tuning the model with reinforcement learning from human feedback (RLHF) (Bai et al., 2022; Ouyang et al., 2022; Rafailov et al., 2024) (see Wang et al. (2024) for comprehensive exposition). This line of research, however, typically requires significant computational resources due to fine-tuning, extensive human labeling, and frequent model updates, rendering it less scalable in practical and rapidly changing environments or as a genuine personalized assis- tant tailored to each individual's preference.
In both problems, we first show that learning is impossible without any structure the sequence function. We introduce natural assumption, di- minishing distance with more commons (DDMC), and propose algorithms with regret (LLTT) and Ö(LVT2/33 for TLB and TMAB, respectively. As side product, we obtain an (almost) opti- mality of the greedy decoding for LLM decoding algorithm under DDMC, which justifies the un- resaonable effectiveness of greedy decoding in several tasks. This also has an immediate appli- cation to decoding-time LLM alignment, when the misaligned utility can be represented as the frozen LLM's utility and a linearly realizable la- tent function. We finally validate our algorithm's performance empirically as well as verify our as- sumptions using synthetic and real-world datasets.
Consequently, more recent thread of studies proposes decoding-time alignment methods, aiming align the LLM's output the fly during inference time without fine-tuning the LLM (Josifoski et al., 2022: Huang et al., 2024; Mudgal et al., 2023). While these methods hold promise in terms of scalability, their theoretical underpin- nings remain largely uncharted, despite some recent evi- dences (Chakraborty et al., 2024; Li al., 2024; Shi et al., 2024b), and how it could sample-efficiently learn the prefer- ence from user feedback is yet questionable. Moreover, the decoding algorithm itself lacks foundational understand- ing from theoretical perspective despite being central to the L.LM's quality, even the simplest ones such as greedy decod- ing or beam search, beyond a few recent works (Finlayson et al., 2023; Chen et al., 2024).
1. Introduction
Large Language Models (LLMs), in the past few years, re- ceived significant attention from extensive areas with its abil- ity to serve numerous tasks including question answering, image generation, code completion and reasoning (Brown al., 2020; Thoppilan et al., 2022; Fried et al., 2022; Zheng
In this work, we propose fresh perspective in the decoding- time alignment well as decoding algorithm for LLMs by framing the LLM alignment and decoding problem as variants of lincar/multi- armed bandits framework whose utility is represented via a sequence function and when a decision maker faces an online problem to select the se- quence irrevocably in by-tken manner. Overall, our contributions are threc-folds:
University of Maryland University of Chicago. Correspon- dence to: Suho Shin <suhoshin cdu>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Tokenized Bandit for LLM Decoding and Alignment
We present tokenized variants of the linear and multi- armed bandit problems, which turns out to be closely connected to bandit tree search (BTS) problem (Co- quelin & Munos, 2007). We first show that the general problem setup fundamentally exhibits exponential de- pendency in the maximal length of token sequences for any algorithm.
two parts: (i) sequence function uy (x Ne) that maps a pair of context and token sequence to nonnegative real value, and (ii) random noise 7 analogous to the standard ban- dit problems. The benchmark is to obtain sublinear regret against the optimal sequence in hindsight per each round, i.e. OPT, = E' maxxev re(y)] and OPT = Ei=1 OPT,, We assume that the output token sequence is restricted to be within length of L, and seeks for algorithms with regret linear on L and sublinear on T.
To address it, we provide fairly reasonable assump- tion, dubbed diminishing distance with more commons (DDMC), and provide sublinear regret algorithms for both problems. This complements the submodularity, widely adopted in sequence function maximization.
In TLB, our utility function is linearly parameterized by ue (zr,y) = (0, (()),)))) given an embedding function e: V* x x - Rd that maps a query and a token sequence to a feature vector, which can be viewed as an embedding vector that a typical L.MM returns. In TMAB, utility func- tion could be arbitrary, but the context remains the same 24 x to pose the problem learnable. For both problems, DDMC assumption is defined as follows: for any two token sequences y.z E V with the same length ly| = |z] and a token T E V, it follows that
We explain how our algorithms can be applied in the decoding-time LLM alignment. We discuss further interesting implications in the static counterpart of our problem, which might be of independent inter- est, demonstrating the provably (near) optimality of the greedy decoding under DDMC structure, justifying empirical evidences (Song et al., 2024) on the effec- tiveness of greedy decoding in certain tasks. We finally provide empirical experiments that justify our assump- tions and performance of our algorithms.
( : 7) - ue (32,, z: 7: (x). ue (3 2)|+
where y : denotes the sequential concatenation of y and T and similarly for z. In words, user realizes small utility gap if two outputs have more common tokens (words) in suffix.
In what follows, we briefly introduce our contributions and implications therein. Section 2 formally introduces our gen- eral problem setup. We provide our results for tokenized linear bandit problem in Section 3. Section presents our results for the tokenized multi-armes bandit problem. Im- plications and applications to decoding-time alignment, de- coding are presented in Section 5, and experimental results are provided in Section 6. Related works, and all the proofs are deferred to the appendix due to page limits. Appendix B presents our results with weakened version of DDMC, com- bined with lookahead decoding, and Appendix C discusses how our problems are connected to the seminal bandit tree search (BTS) problem.
For TLB problem, we propose an excessive optimism un- der the face of uncertainty (EOFUL) algorithm with regret (Lv/Tlog ). by combining greedy decoding with stan- dard LinUCB algorithm (Abbasi-Yadkori et al., 201 I). The analysis follows from several interesting ideas: (i) an intro- duction of level-k regret that internally occurs from k-th chosen token during the algorithm, (ii) fictitious extension of our algorithm's 's sequence recursively propagate the re- gret into upper level, followed by (iii) regret decomposition and sum of squares regret analogous to the standard linear bandit analysis. For TMAB problem, we no longer have a global estimate 0 that will help us estimating the perfor- mance of each intermediate token sequence. Instead, we present GreedyETC that step-by-step learns desirable path to efficient token sequences by combining exploration-then- commit style of algorithm with greedy decoding, inducing the regret O (log T)// Analysis requires a care- ful union bound over the concentration of desirable token se- quences, since naively taking union bounds over every token sequence yields regret of )) assumption plays a crucial role, enabling certain tractable
1.1. Overview of results, assumptions, and implications
We start with the most general problem setup that will sub- sume tokenized linear bandit (TLB) and tokenized multi- armed bandit problem (TMAB). Each round t E [T') user arrives and provides a query 2, as a context, and a decision maker (DM) selects token sequence y, E 2 one at a time from a fixed vocabulary V with [V] = n. Once the token se- quence is finalized, the DM receives a random utility re re())) from the user, capturing how well the produced text aligns with the user's preference. The random utility consists of
4 We remark that DDMC
"The analysis requires an additional technical assumption, de- tailed in Section 3.2. Linear realizability and this assumption is completely removed in our results for TMAB, at the cost of the same context across the rounds.
We note that TMAB does not directly reduce to TLB as spe- cial case unlike the standard bandit problem, due to an exponential blow up in the dimension of the latent parameter. Thus, it requires independent algorithms and analysis.
"TMAB problem, in fact, has close connection to the well known bandit free search (BTS), i.e., decision version of BTS can be reduced to TMAB and vice versa. We discuss the detail in Appendix C.
AA denotes a set of finite sequence over A.
2

---

Tokenized Bandit for LLM Decoding and Alignment
structure in the utility function, albeit its naturallity and sim- plicity. DDMC assumption is relaxed in Appendix B, and we also provide efficient algorithms for both settings with the weakened assumption by replacing the greedy decod- ing with lookahead decoding, another well-known decoding algorithm (Lu et al., 2021; Fu et al., 2024).
V" denote the collection of every possible finite complete sequences. For E V", we write y z denote the token sequence that concatenates y and sequentially, We often write y = :
We call the process of sequential token selection in a single round the decoding algorithm, and denote the process of observing the feedback at the end of the decoding and learn- ing the sequence function by learning phase. Essentiallyy one view the decoding process select each token, and learning process an online learning problem to estimate the objective function. We write |y| to denote the length (interchangeably. depth or level) of the se- quence y. We consider a class of decoding algorithms with fixed maximal depth I. that only outputs token sequence with length at most L. Equivalently, we assume V consists of token sequences with length up to L. We assume that the depth L is known in advance to the DM.? The depth plays a crucial role in our analysis and any algorithm's 's performance would depend on L that intrinsically captures the size of feasible token sequences. We refer to an algorithm's s level I decision to denote its decision for the I-th token.
Finally, we discuss several implications and applicability of our foundations in LLM decoding and alignment. For LLM alignment problem, we have LLM equipped with a (known) utility function v V x V + R2 and mis- aligned utility represented by linear structure (0, y)) with latent f, where the eventual utility is a unknown con- vex combination over them. The objective is to efficiently learn I and the convex combination parameter in a decoding time through user feedback while leaving the original LLM frozen ensure that no further training occurs for the LLM. We show that this problem reduces to the TLB problem, enabling efficient learning via EOFUL under DDMC. We next consider a static version of our problem, where query access to the sequence function u is available and we need to maximize u with polynomial queries. Inspired by our analysis for EOFUL and GreedyUCB, we prove that the greedy algorithm achieves the exact optimum with nLT queries. This immediately implies that greedy decoding is almost optimal under DDMC for LLM decoding problem when we view u as the probability of LLM generating the sequence. This gives an evidence on why greedy decod- ing is unreasonably efficient for various domains despite its simplicity (Meister et al., 2020; Shi al., 2024a; Team et al., 2023; Welleck et al., 2024), (see Song et al. (2024) for comprehensive evaluation of greedy decoding).
online problem to
Sequence utility and regret. etch round & (7), a user arrives and writes a query I E Then, a decision maker (DM) needs to select each token sequentially to construct a sequence y, E V to answer the user's query, and the DM observes a reward (y,,) when it finalizes the sequence selection process and displays y. In particular, the DM irrevocably commits to each token to construct the token sequence and cannot revert any previous decision change the sequence entirely. The reward is defined by two factors: (i) sequence (utility) function u, V - R20 that maps a token sequence to a nonnegative real value and (ii) a random noise over it. Note here that 2e can vary across the rounds, which captures the (possibly) different context provided to the DM given by the queries x, over the rounds. Formally, the random reward ry.).) for the sequence y can be represented as
2. Problem Setup
We formally define the general problem setup, which will subsume tokenized linear bandit (TLB) and tokenized multi- armed bandit (TMAB) problems. Let V be the set of all possible tokens with |V| n. For any set A, let A" the collection of every possible finite sequences over A. Given a token sequence y E V", we write y(+) E V to denote its i-th token and y(::)) E V* to denote the subsequence that contains every token from i-th token to j-th token including themselves for j 2 i in y. We often write y(-0) to denote the i-th token from the end of y. There exists a special end-of-sequence EOS E V token representing the end of sequence. We call token sequence y to be complete if its last token is EOS. Importantly, even though our sequence function could have nonzero value even without the EOS token at the end, we will only consider a class of algorithms that only ends with EOS token. Thus, throughout, we rewrite
T (y.)) = ue(y.) + m/.
where T7 is a random noise with Inel S a.9 The DM does not know the sequence function u in advance, but can only indirectly access it through the reward ry.)). The goal of the DM is to maximize the user utility over the rounds T, i.e. Exetr (y)))) Equivalently, the objective is to compete with the optimal complete sequence in hindsight: OPT
Prractically, can be prompted in the LLM or implicitly set the LLM to ensure it does not generate an infinitely long output. *An acute reader may easily observe that cannot leamn anything if ue does not have any structural assumption. Thus. Section allows efficient learning by having linear structure, and Section 4 allows efficient learning my having " = u.
'We also compare DDMC with smoothness assumption in BTS problem in Appendix C. "We use boldface for sequences and non-boldface for token.
"We assume bounded noise for simplicity, but it easily extends to bounded aab-Gaussian noise as in the bandit literature.
3

---

Tokenized Bandit for LLM Decoding and Alignment
Ei= OPT, where OPT, E m (z)).. Then, the DM tries to minimize the expected cumulative (pseudo-) regret REG = Eiem REG, OPT Ere IE [r:(x,))]
as an online learning problem learn the LLM's logit prob- ability in an efficient manner.
A special case of this problem is scenario where the DM has access to the value oracle of u, e.g., knows the probabil- ity distribution, in advance. In this case, the objective boils down to the sequence function maximization problem (Li & Milenkovic, 2017: Tschiatschek et al., 2017; Mitrovic et al., 2018; 2019) where duplicated insertion of the same element is allowed, or from the LLM's perspective, the LLM decoding problem of designing efficient decoding al- gorithms (Freitag & Al-Onaizan, 2017; Kulikov et al., 2018; Holtzman et al., 2019; Chen et al., 2024). Correspondingly, our main results have several implication in these problems, which will be elaborated Section 5.2.
Monotonicity and single-level deviation. Remark that we consider a class of algorithms that always output a complete sequence. In practice as well. once the sequence meets EDS, it is natural to stop the decoding. To capture such behavior, we impose the following assumptions.
Assumption 2.1 (Monotonicity). The utility function is monotone if (i) appending more EOS to a sequence that already has EOS at the end does not change the utility, and (ii) appending non-EOS token to a sequence that already has EOS at the end only decreases the utility.
The monotonicity assumption is completely innocuous in a sense that appending EOS tokens to a complete sequence must not change the utility in practice as the user observes exactly the same output, and further any reasonable algo- rithm would submit the output once it meets EOS.
LLM alignment problem. Further, one consider variant of our setup, termed LLM aligument problem, where we have a (frozen) LLM that gives us a probability distribu- tion p(y), but this L.MM is (partially) misaligned with our objective function u. In this case, one standard approach would be to use the probability distribution p(y) as base- line for the decoding process, while learning the misaligned objective function u. If the misalignment allows specific structure, e.g., Ue y)) = (1 7)p(<r; y) + (z,. y) for some latent function f capturing the misalignment, we indeed show that our results can be extended to efficiently handle such misalignment in Section 5.1.
Further, let a, be sequence which is not optimal. If sequence a, deviates from the hindsight optimal sequence or firstly at level any practical algorithm should have small single-level error in a sense that the utility difference at the moment is almost ignorable as they differ only by a single token. Formally, we impose the following assumption.
Assumption 2.2 (Small single-level deviation). Let a be a sequence that deviates from the optimal sequence o firstly at level i. The utility function un has single-level deviation (SLD) of e if (.(((1:)) we (a((:))) S S. Further, the utility function has small SLD (SSLD) i = O(1/VT).
3. Tokenized Linear Bandit
The general setting we formulated, however, does not al- low any efficient learning if the utility function u, indeed does not have structural correlation across the rounds. One natural framework to tackle such fundamental obstacle is to impose the linear realizability assumption in the online contextual learning literature. To this end, we introduce the following notion of embedding function:
This assumption is apparently innocuous for large i as adding single token would almost incur difference. Even if i is small, one can expect that the utility differ- ence will be still small as small i indicates that the overall document only has a few tokens, which would not induce significantly different utility. For instance, one may con- sider truncating the token space to top-k most probable ones, then SLD will be ignorable for small as any efficient LLM would only produce reasonable tokens top-k to- kens. 10 Throughout, we will assume that the utility function is monotone and SSLD unless specified otherwise.
Definition 3.1 (Embedding function). A decoding algo- rithm has access to the embedding function e V* x V + Rd maps a token sequence to d-dimensional real vector. 12
Assumption 3.2 (Linear realizability). There exists O E Rd and embedding function e(-) such that the utility function can be represented by we (xt.y) = (0, e(xr,))) for every query It E V and token y y E V*.
LLM decoding problem. Let us see the connection be- tween our problem and L.M decoding problem and L.M alignment problem. Given an access to the logit prob- ability of a LLM and token sequence y E V*, sup- pose the sequence function de denotes the LLM's proba- bility p of generating the output y, i.e., (x( = p(y) = I (1)))(()))((())))))))))))) Then, our objective can be framed
The DM has a query access to the embedding function during the decoding process." 13 Then, the DM's objective
ll In fact, setting 7 = reduces to the pure LLM decoding problem, whereas setting Y = reduces to our general problem of learning the latent sequence function, or equivalently, learning I.IM's behavior.
1his does largely not restrict the class of problem instances deal with, as the utility difference onwards could be significantly large.
One can deem the emhedding function as the semantic embed- ding that the L.M computes for given query æt. 13 "Our algorithm calls the embedding function O(TLT) time.
12
4

---

Tokenized Bandit for LLM Decoding and Alignment
is to efficiently learn the latent parameter 0 to maximize the utility. We assume that ||@]| I 1 and ||e(z)]] Σ 1. We finally note that such linear approximation of the reward model often appear in the LLM alignment literature (Cen et al., 2024; Yang et al., 2024). 15
absolute value of the difference); (2) more subtly, DDMC is about the difference of two different sequences of the same length with an added common token, whereas submodu- larity is about the difference of two different sets that one includes another with an added common element. In some sense, DDMC is the complement of submodularity.
3.1. Diminishing distance with more commons (DDMC)
Discussion 3.5. From the practical perspective for LLM, this is innocuous since if two generated outputs shares common tokens suffi--wise, we can expect that they give the user similar experiences as long as we add more commons.
On the other hand, even if the utility function is well- structured with the linear realizability assumption, the in- trinsic intricacy of TLB's tokenized nature imparts a funda- mental hardness in maximizing the reward. Intuitively, if (monotone) sequence function could have arbitrary struc- tures, one small mistakes in the first token selection may in- cur irreparable results in the following tokens, e.g., imagine there exists only a single token T which incurs nonnegative for u(T y) for any If the first token itself does not give much information about the following levels, one cannot fundamentally identify an efficient path in hindsight. 16 This captured by the following proposition:
3.2. Excessive Optimism under the Face of Uncertainty
We present our algorithm, named excessive optimism under the face of uncertainty for tokenized linear bandits (EO- FUI.), whose pseudocode is presented in Algorithm 1. has two optimistic features. First, it optimistically decodes each token in a greedy manner, expecting that this would lead to good utility in following decisions. Second, it optimistically estimates each token's 's utility by constructing a confidence ball C,, and computes the best possible utility each token can gain from parameters in C.-
Proposition 3.3. For TLB problem, any algorithm suffers the worst-case regret of (T 1/22-2)).
Our construction of the confidence ball is analogous to the standard LinUCB (Li et al., 2010; Abbasi- Yadkori al., 2011) algorithm for linear bandits.
Thus, we here introduce fairly natural structural assump- tion on the sequence function u(-) to avoid such pessimistic scenario, dubbed diminishing distance with more commons.
Assumption 3.4 (DDMC). A sequence function u(-) has diminishing distance with more commons (DDMC) if
Let y, be the chosen sequence at round by Algorithm 1. We recursively define the following matrices for t E [T]:
= E + YYN
[u( : T) u( T) S [u( - u((x,z)) |,
and let Σ₁ Al for d x d identity matrix I. Further, we define
for any two different finite sequences y,z with the same length by| = |z] and any token 7 E V. 17
3: = 0² 2 4d log 1 ) + 8 log :)
(2
In other words, adding common tokens at the end of two sequences with the same length only decrease the utility gap between them.
For f 2 2, the center of our confidence ball will be:
DDMC shares very similar intuition from the widely adopted submodularity assumptions on set function (Ko- rte et al., 201 1), recently sequence function (Mitrovic et al., 2019), which characterizes diminishing returns. How- ever, they crucially differ and cannot be implied by each other, due to two key differences: (1) submodularity is about diminishing return (i.c., the difference between two function values) whereas DDMC is about diminishing distance (i.e.,
0, = ET' Σ Ca
(3.1)
and 0, = 0. the d-dimensional zero vector." 19
Now, our confidence ball C2 to be used at round t is defined as follows:
C- 10= (0 0.) (0- S B.).
(3.2)
"This is for ease of exposition to exclude cumbersome details, but all results would easily generalize under any bounded scenario. One might find this partly related to linear representation hypothesis (Park et al., 2023). stating that concepts are encoded as linear directions in the model's 's embedding space.
and C (0). We choose X O())
We further impose the following technical assumption:
'In fact, due to the online learning nature of our problem, there always exists a small probability that any algorithm makes an error in the first level.
18 "Remark that this a solution of the ridge regression given the reward feedback.
Forr efficient computation, one may use the rarely-switching version of OFUL algorithm in (Abhasi-Yadie Yadkori et al., 201 1), in- stead of the standard OFUL..
IIn Appendix B. we provide a relaxed version of DDMC as- sumption and provide corresponding efficient algorithms.
5

---

Tokenized Bandit for LLM Decoding and Alignment
ALGORITHM 1: BOFUL
Step 3 extension of y: (1:k-1) to p(1ck) = y; (1:k-1) o and show
3 (Fictitious extension). Then we introduce fictitious
(k)
Input: Token set V. linear contextual bandit algorithm ALG
Initialize y + 0, 0. Ö, C1 (0).y=0)
how the level-k regret of the fictitious extension is upper bounded by level-(k-1) regret at y;
for = 1,2.....
do
(1:k-1)
User arrives and submits query Jy
for k = 1, 1,2..... L- 1 do
Step 4 (Regrei decompoxition). We relate the level-(k-1) regret of the algorithm's 's choice y, and level-k regret of the fictitious extension using DDMC assumption, which upper bounds our algorithm's 's level-k regret due to the optimistic choice of the algorithm. Then, we telescope over k,..., 1 to express the entire regret function of level-1 regret plus some estimation errors occurred for each prefix of y, during the decoding.
Compute -" = argmax Ev. (0, (0,,(
Set 4 y
if = EOS then Break;
Submit y and observe reward n
Compute 4.41 by (3.1) and C.+1 by (3.2)
Assumption 3.6. For any complete sequence y with by| = k and any subsequence y y(1d) with 1 S. k there exists a c 2 0 such that [|x(aa)]]] S <c|lyll, in Mahalanobis norm." 20
Step 5 (Sum of squares regret) Finally, we compute the func- tion of level-1 regret and the estimation error by following the standard sum of squares regret bound analysis of linear contextual bandit algorithm, where readers familiar with the standard linear bandit analysis may easily follow.
That is, each entry of the subsequence of a complete se- quence is not too larger than that of the complete sequence with respect to Σ, Both structural assumptions 3.2 and 3.6 removed in Section 4 for TMAB.
Discussion 3.7. This assumption is innocuous in a sense that e can be sufficiently large up to of (V) to obtain sublin- ear regret. On the other hand, this less innocuous since it depends on the algorithm's ' construction of Σ, In practice, one of the typical usage for LLM is to either (i) have similar queries based on their daily lives, or (ii) ask queries in al- ready existing sessions (in-context) to expand the discussion. In these scenarios, we expect the queries and the responses thereby to be not too much different from each other.
4. Tokenized Multi-armed Bandit
The linear realizability assumption is essential for the con- textual setting derive any efficient algorithm. A natural follow-up question is whether could obtain an efficient algorithm without linear realizability, if the context does not change and the utility function uı = remains the same over the rounds." 21 In this case, the DM's problem can be seen as a variant of the standard aulti-armed bandit problem where needs to make an (online) decision for the next to- ken irrevocably construct a complete sequence. Thus, we dub this problem tokenized multi-armed bandit (TMAB).
Discussion 3.8. On the other hand, if y and y, does not have any structural dependency, any algorithm would not be able to fundamentally learn certain coordinate of O. For instance, suppose any complete sequence's 's embedding only has nonzero values in the last d/2 dimension, and the first d/2 dimension for any incomplete sequence. Then, any algorithm could not learn the first d/2 dimension, implying that any internal node does not provide information on an efficient path to the complete sequence.
We first argue that even in this simple setting, our problem admits a pessimistic lower bound when the sequence func- tion does not have any structural assumption beyond mono- tonicity. Then, we again introduce DDMC and present an algorithm with regret ( (log T) 1/3) We present an intimate connection between TMAB and the bandit-basee tree search (BTS) problem (Coquelin & Munos, 2007) in Appendix C.
Theorem 3.9. Impose assumption 3.6 and DDMC. With probability at least 1 8, EOFUL has regret O(cl// d'T + log(t))), and thus plugging 5- = 1/T yields REG- = O(eLv/dT log T.
Exponential lower bound. If the query is fixed, our objec- tive boils down to efficiently learn the latent utility function u. Without any structural assumption on it, however, one can easily see that it is equivalent to learn the reward dis- tribution of every subsequence y E V This is because each subsequence y can be deemed standard multi-armed bandit problem which admits regret lower bound of 2?( VKT) where K is the number of arms. This is formally stated as follows:
Overall, our proof proceeds as follows:
Step I (Length equalization). We first equalize o, and y, to have the same lengths in way that it does not hurt any following analysis.
single arm in the
Step (Level-k regret). We define level is regret as the regret occurred at depth a of the decoding process that will play crucial role in the analysis to recursively transform the regret into upper level regrets.
Proposition 4.1. For TMAB problem, any algorithm has worst-case regret lower bound of ([min((VnTT.T)
We omit the context J; in the inputs of since the context remains the same.
2 Mahalanobis norm is defined by ||اا|| = VxT An.
6

---

Tokenized Bandit for LLM Decoding and Alignment
DDMC and GreedyETC. The pessimistic exponential de- pendency on depth I naturally motivates the use of DDMC Assumption 3.4. Unlike the TLB problem, however, there exists no global latent parameter 0 that we could use through- out the decoding process, but they can genuinely arbitrary beyond DDMC (and monotonicity). Thus we use a rather direct approach to explore each level sufficiently, and com- mit to the seemingly optimal token at each level then move on to the next level.
token sequence y is misaligned from the LLM such that u( y) = yv(pp(a,,)) + (1 ()) y), where is monotone function that maps probability to reward and f is the latent function to be learned that marginalizes the misaligned part of the L.M.
Such convex combination between multiple rewards typi- cally appear in the LLM alignment literature (Rame et al., 2024; Shahriar et al., 2024; Jang et al., 2023; Shi al., 2024b), when one wants to align with multiple human pref- erences by well combining multiple models or parameters. Our utility function as the convex combination over v and f can also be interpreted in a similar spirit. One typical example for the function e would be the log probability or perplexity obtained from the frozen L.M,, or more explicit user utility such as classification accuracy or factuality (Lin et al., 2022). We assume the DM knows the probability to utility mapping function v.
ALGORITHM 2: GreedyETC
Input: Token set V. exploration parameter N
Initialize y 4-0 0
for 1, 2, ... do
by| = L then
y y: EOS
Break for E do
Submitt : : EDS for N times and observe rewards
Set F(y: t) he the averaged cumulative rewards
Select r" = argmax Fly :)
If we assume f is linearly realizable, i.e., there exists an embedding function e V X V - I such that f()) = (0, ((xt.y))) we can reduce the problem of learning misaligned function f to the TLB problem. Formally, we construct 0'
Set +- T
if = EOS then
I Break For the rest of the rounds, submit y
RR++1 as a vector that concatenates 0 and 1 in 1-th dimension after multiplying each by - y and 7 respectively, i..., 0' = [(1 - 7) 0: r 1]. Then we can rewrite the utility as
We prove that under (a slight variant of formally presented in) DDMC assumption, a natural combination of greedy decoding with exploration-then-- then-commit (ETC) style of algo- rithm, as presented in Algorithm 2. yields sublinear regret w.r.t. T with linear dependency on d:
u( y) (0, (1- ()ex (ee(z),,))) }.)) + yu(p(z,,)))
= (0', [e(xx,)) y) ((p(xz,)))]
Theorem 4.2. Under Assumption 2. and E.6, Algorithm 2 achieves regret O(nLT2/3 (log 1)3/)).
This is exactly the feedback structure in our TLB problem, where the DM observes e( y) and v(p(x,,)) y)) to construct the concatenated vector and tries to learn 0'. Correspond- ingly, EOFUL enjoys the same regret bound in Theorem 3.9 for this problem. 22
One subtle part is that, if one naively construct the clean event that the rewards for each sequence are concentrated to its expectation and take union bound over all of them, this essentially gives the regret of 2(n²), which is undesirable. Instead, we only take the union bound along the path that our algorithm selected, which enables us to recover the linear dependency on L. Then, we obtain a recursive inequality relating the regret at level k with that at level k - 1 by considering fictitious extension as defined in the proof of Theorem 3.9, and conclude that the regret remains small with high probability.
5.2. LLM decoding
Our results and analysis therein have a few interesting im- plications in (i) the sequence function maximization prob- lem (Li & Milenkovic, 2017: Tschiatschek et al., 2017: Mitrovic et al., 2018; Alaei et al., 2021) and (ii) the problem of designing efficient decoding algorithms for text genera- tion in language model (Freitag & Al-Onaizan, 2017; Ku- likov et al., 2018; Holtzman et al., 2019; Chen et al., 2024), if we consider a special case when the function u is known to the DM and consider a static version with a single round.
5. Applications
5.1. LLM alignment
First, this is exactly a sequence function maximization prob- lem maximize u given query access a value oracle that
One interesting application of our results is the LLM align- ment problem. In this problem, want to generate aligned outputs given access to a frozen LLM that cannot be fur- ther fine-tuned. Consider LLM parameterized by a to- ken distribution p= : V" x V" + [o, 1] that maps any to- ken sequence given a user's 's query to a probability that the LLM generates this sequence. The user's utility given
22 Note that our formulation differs from standard KL-divergence based problem (Yang et al., 2024). One might integrate the KL divergence in our problem too by redefining the function over the distribution rather than fixed sequence, however, it's not clear to formulate in the language of multi-armed bandit problem. This remains an interesting future direction.
7

---

Tokenized Bandit for LLM Decoding and Alignment
15.
3100
Muin Det Productt Distance
120
Mese Embweddn Dabance
80
125
1000
Moan oee Product
100
80
80
80
Muan Embooddin
75
60
60
so
50
20
20
25
201
350
100
203
Shared Sulti Longth
Shared Sulfix Longth
Shared Sufti Length
Shares Suffi Length
Figure 1. The first two show distance gap under TruthfulQA and the next two is under HH-RLHF, with distance function di and d2.
outputs u(x,,y) y) given a sequence y. One common objective here is to obtain an approximate algorithm with polynomial query complexity. If u satisfies a notion of monotonicity and sequence-subnoduuarid (Alaei et al., 2021), it is known that the greedy algorithm has 1 1/e approximation." 23 On the other hand, our result implies that greedy algorithm is almost optimal under DDMC assumption, formally stated as follows:24 24
6. Experiments
We finally provide our experimental results.
6.1. Validating DDMC
The first is on validation of DDMC assumption. Further experimental results are presented in Appendix D.
Setup. For the embedding function e, we obtain the em- bedding from Llama3-88--Instrue model (AI@ Meta, 2024). For the token sequences, to validate our DDMC real- world datasets, we use TruthfulQA dataset (Lin et al., 2022) and HH-RLHF dataset (Bai al., 2022),25 25 by concatenating each query and answer pair.
Theorem 5.1. For any monotone sequence function, greedy decoding is almost optimal under DDMC in a sense that it is (1 approxxiiittin where E is the SLD parameter.
From the LLM perspective. this result has one interesting corollary. Given a fixed query J: consider a L.LM parame- terized by p: V V - [0, 1]. Then, Theorem 5.1 implies that if p is monotone DDMC function, the greedy decoding can exactly output the token sequence that maximizes the logit probability. This supports the unreasonable effective- ness of the greedy decoding, despite its simplicity, justifying why it works well and widely adopted in the literature (Meis- ter et al., 2020; Shi et al., 2024a; Song et al., 2024), egg. in Google's Gemini report (Team al., 2023), even though it is only known as an approximate maximum-a-posterie decoding algorithm (Welleck et al., 2024).
For each item in the dataset, we group every (query:answer) pair by number of common tokens in its suffix. For instance, image the following two token sequences:
x = (a. b,, c, y = (a, x. v..,y,ydee,f
Then, we can group each subsequence as the following:
0 common suffix: (x(1:2), y(1:2)), (x(1:3), y((:a))
1 common suffix: (x (1:1), (l=l)), (x(1 y y(::)))
On the other hand, greedy decoding is often dominated by sampling methods such as nucleus sampling (Holtzman et al., 2020) or self-consistency (Wang et al., 2022) for var- ious tasks, e.g., (Song et al., 2024) find that vanilla multi- nomial sampling can be better than greedy for open-ended creative generation. We conjecture that such phenomenon largely depends on how the sequence function generated by LLM, i.e., the logit probabilities, would look like and thereby the type of the tasks. Investigating other regimes that another decoding algorithm provably outperforms greedy remains an intriguing open question.
2 common suffix: (x (1:5) y (1:5))
3 common suffix: (x (1:6) (1:6))
We evaluate each distance d(x ((ok)) in each group and obtain the average and variance.
For the distance functions, we evaluate two candidates:
di v) = (0, x y) dz(x, y) yl
"It also admits the APX-hardness with 1 - 1/c by the hard- ness of submodular set-function maximization from Maximum Coverage. unless P=NP.
where 0 = Ian9o. i.e., 4096-dimensional one vector. The first distance function dj is consistent with our DDMC
"We remark that our setup is slightly different from standard sequence maximization problems as we allow the selection of the same element over different position, and the special EOS token governs the end of the sequence by naturally encoding it in the monotonicity of the utility function.
25 Both datasets are standard for LLM alignment, where Truth- fulQA validates truthfulness and hallucination and HH- RHHR eval- uates helpfulness and harmlessness. In HH-RLHF, data are or- ganized in (promptresponse) format. For simplicity, reuse "query" to refer to "prompt" and "answer" to refer to "response"..
8

---

Tokenized Bandit for LLM Decoding and Alignment
Parzial Round
Camuladee Rapre
1.25
plus
H Mulp
1.20
Mix
1.15
Prital 1914
X
2.10
100
5CK4
0.59
0.40
160
Saund
Figure 2. The first figure compares the regret of algorithms. The second figure represents the ratio pl: rounds. Max Ratio and Mean Ratio denote the maximum and average ratio, respectively, for / E [4] at each round.
|
(.:0
/llills, over the
x,
assumption, and it in fact be viewed & distance between two vectors. Thus, the second distance function d2 can be viewed as generalization of d, into & distance.
may be much better than that of the theoretical guarantees.
In practice, one could blend some standard alignment tech- niques with our learning based approaches. For instance, one can run standard alignment algorithms in the earlier rounds as EOFUL could suffer significant errors in the be- ginning. In the mean time one can operate ROFUL. in back- grounds to learn the confidence region, and then adopt a combination of them in later rounds. This will be way more practical than naively running EOFUL particularly if the effect of coldstart in EOFUL, e.g., apparently misleading or grammatically incorrect outcomes, could significantly degrade the user experience.
Result. The first two figures in Figure refers the TruthfulQA, and the next two denotes the HH-RLHF. Each experiment is presented for distance function d1 and then d2. Overall, we observe tendency to have decreasing distance gap as the number of common tokens in suffices increases. This holds for both the distance function d, and dz. In- terestingly, the decreasing curvature tends to be stronger when the number of common tokens are rather small. Fur- ther, the decreasing gap seems to be more drastic under HH-RLHF, suggesting that different tasks might possess different structures in their sequence function, iee, logit probability structure. Overall, despite verified in an aver- age sense, this justifies the DDMC assumption, which we believe could be largely useful for following works.
Validating Assumption 3.6. Further, in the same ex- perimental setup, we validate Assumption 3.6 for our EO- FUL. Notably, we observe that the ratio is universally upper bounded above by 1.25, even after taking maximum over every token position. Recalling that such upper bound scales the regret upper bound in a linear manner, this would only blow up the regret to be 1.25 times that of the standard linear contextual bandit, let alone the linear dependency on L.
6.2. Performance of EOFUL
We numerically validate the performance of EOFUL using synthetic data under the LLM alignment scenario presented in Section 1. In particular, we set the length of each sen- tence to be L 30, and truncate top-15 tokens for every algorithm for efficiency. Further, we set (0.5, .. Dteails in the query generation can be found in Appendix D. We compare three reasonable bench- marks: (i) theoretical regret upper bound (under-scaled by 0.1 to make the plot more visible), (ii) WrongTheta that uses wrongly estimated 0' (-0.5, -....... ... -0.5) and choose token that maximizes the resulting convex com- bination, and (iii) Misaligned greedy that simply performs greedy decoding with respect to the LLM's probability, As observed in the first part of Figure 2, we observe that EO- FUL effectively achieves sublinear regret. Given that the regret upper bound is under-scaled, the actual performance
7. Conclusion
We introduce tokenized variants of linear and multi-armed bandit problems, where decision maker needs to max- imize cumulative rewards which can be represented as a randomly perturbed sequence function by selecting each token a sequential manner. We introduce a novel but natural assumption called diminishing distance with more commons (DDMC), and show that both problems admit efficient algorithms with regret linear on the depth of the maximal sequence and sublinear on the time horizon. We provide several applications and implications on LLM align- ment problem as well as decoding problem, and validate our assumptions and performance our algorithms through synthetic and roal-world datasets.
0.8, 0 =

---

Tokenized Bandit for LLM Decoding and Alignment
Acknowledgement
Beirami, A., Agarwal, A., Berant, J., D' D'Amour, A., Eisen- stein, J., Nagpal, C., and Suresh, A. T. Theoretical guar- antees on the best-of-n alignment policy, arXiv preprint arXiv:2401.0187, 2024.
Part of the work done while Suho Shin visiting Com- puter Science the University of Chicago. This work is partially supported by DARPA QuICC, ONR MURI 2024 award on Algorithms, Learning, and Game Theory, Army- Research Laboratory (ARL) grant NNF241000052, NSF AF:Small grants 2218678, 2114269, 2347322. Haifeng Xu is supported by the A12050 program at Schmidt Sci- ences (Grant G-24-661044 6044,, Army Research Office Award W911NF23-1-0030 and NSF Award CCF-2303372.
Brown, T., Mann, B.,, Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in newral information processing systems, 33: 1877-1901, 2020.
Cen, S., Mei, J., Goshvadi, K., Dai, H., Yang, T., Yang, S., Schuumaass, D., Chi, Y., and Dai, B. preference optimization: A unified approach to online and offline rlhf. arXiv preprint arXiv:2405. 19320, 2024.
Impact Statement
This paper aims to advance the understanding of the multi- armed bandit problem, LLM alignment and decoding. There could be several potential societal consequences our work, none which we feel must be specifically highlighted here.
Chakraborty, S., Ghosal, S. S., Yin, M., Manocha, D., Wang, M., Bedi, A. S., and Huang, Transfer q star: Principled decoding for Ilm alignment. arXiv preprint arXiv:2405.200 20495, 2024.
References
Chen, S., Hagrass, O., and Klusowski, M. Decoding game: On minimax optimality of heuristic text generation strategies. arXiv preprint arXiv:2410.03966. 2024.
Abbasi- Abbasi-Yadkori, Y., Pál, D., and Szepesvári, C. Improved algorithms for linear stochastic bandits. Advances in neural information processing systems, 24, 2011.
Coquelin, P.-A. and Munos, R. Bandit algorithms for tree search. arXiv preprint cs/0703062, 2007.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774. 2023.
Finlayson, M., Hewitt, J., Koller, A., Swayamdipta, S., and Sabharwal, A. Closing the curious case of neural text degeneration. arXiv preprint arXiv:2310.0169) 2023.
A12. Justeval.
2024.
URL https://
Freitag, M. and Al-Onaizan, Y. Beam search strate- gies for neural machine translation. arXiv preprint 17X2v:/00.0118000 2017.
huggingface co/ datasets /rre- just-eval-instruee
Al@Meta. Llama 3 model card. 2024. URL
Fried, D., Aghajanyan, A., Lin, J., Wang, S., Wallace, E., Shi, F., Zhong, R., Yih, W.-t., Zettlemoyer, L., and Lewis, M. Incoder: A generative model for code infilling and synthesis. arXiv preprint urXiv:2204.0557 0.290.995999.. 2022.
https: / /ttpttgbttttttt//// www.iituu.commomee. wwwoo//etaalii 1alttaallmmaaaaaa /aamm//1mma////// blob/main hww.://aaaa///..aaaaa.a. BODDMMDDDMAAPPP md.
Alaei, S., Makhdoumi, A., and Malekian, A. Maximiz- ing scquence-submodular functions and its application to online advertising. Management Science, 6700))6003-- 6054, 2021.
Fu, Y., Bailis, P., Stoica, I., and Zhang, H. Break the se- quential dependency of 11m inference using lookahead decoding. arXiv preprint arXiv:2402.02057. 2024.
Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. arXiv preprint arXi::1944.097. 2019.
Auer, P. Finite-time analysis of the multiarmed bandit prob- lem, 2002.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., Das- Sarma, N., Drain, D., Fort, S., Ganguli, D., Henighan, T., et al. Training helpful and harmless assistant with rein- forcement learning from human feedback. arXiv preprint 22X.v:2204.05862, 2022.
Holtzman, A., Buys, J., Du, L.,, Forbes, M., and Choi, Y. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL Ttt.lhttts://ppeeeeeeeeee wwtp//pppwweeeeeeeeeppppeeee.eep net.ieewwweeeeeeeeeee forum? i id=rygCOy=Fvi:
Basu, S., Ramachandran, G. S., Keskar, N. S., and Varsh- ney, L. R. Mirostat: A neural text decoding algo- rithm that directly controls perplexity. arXiv preprint arXiv:2007. /4966, 2020.
Huang. J. Y., Sengupta. S., Bonadiian, D., Lai, Y.-a., Gupta, A., Pappas, N., Mansour, S., Kirchhoff, K., and Roth, D. Deal: Decoding-timg time alignment for large language models. arXiv preprint arXiv:2402.00 04777, 2024.
10

---

Tokenized Bandit for LLM Decoding and Alignment
Jang, J., Kim, S., Lin, B. Y., Wang, Y., Hessel, J., Zettle- moyer, L.,, Hajishirzi, H., Choi, Y., and Ammanabrolu, P. Personalized soups: Personalized large language model alignment via post-hoc parameter merging, arXiv preprint arXiv:2310. 10. 1564, 2023.
In International Conference on Artificial Intelligence and Statistics, pp. 1177-1184, PMLR, 2018.
Mitrovic, M., Kazemi, E., Feldman, M., Krause, A., and Karbasi, A. Adaptive sequence submodularity, Advances in Neural Information Processing Systems, 32, 2019.
Josifoski, M., Peyrard, M., Rajic, F., Wei, J., Paul, D., Hart- mann, V., Patra, B., Chaudhary, V., Kıeıman, E., Faltings, B., et al. Language model decoding as likelihood-utilty alignment. arXiv preprint arXiv:2210.07222. 2022.
Mudgal, S., Lee, J.. Ganapathy, H.,, Li, Y., Wang, T., Huang. Y., Chen, Z., Cheng, H.-T., Collins, M., Strohman, T., et al. Controlled decoding from language models. arXiv preprint arXiv:2310. 70222, 2023.
Korte, H., Vygen, J., Korte, B., and Vygen, Combina- torial optimization, volume 1. Springer, 2011.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744 2022.
Kulikov, I., Miller, A. H., Cho, K., and Weston, J. Im- portance of search and evaluation strategies in neural dialogue modeling. arXiv preprint arXiv:1811.0090 2018.
Park, K., Choe, Y. J., and Veitch, V. The linear represen- tation hypothesis and the geometry of large language models. arXiv preprint arXiv:2311.0363 2023.
Lai, T. I.. and Robbins, H. Asymptotically efficient adaptive allocation rules. Advances in applied mathematics, 6(1): 4-22, 1985.
Rafailov, R., Sharma, A., Mitchell, E.,, Manning, C. Er- mon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Ad- vances in Neural Information Processing Systems, 36, 2024.
Li, B., Wang, Y., Grama, A., and Zhang, Cascade reward sampling for efficient decoding- time alignment. arXiv preprint arXiv:2406. 16306, 2024.
Li, L., Chu, W., Langford, J., and Schapire, R. E. A contextual- bandit approach to personalized news article recommendation. In Proceedings of the 19th interna- tional conference on World wide web, pp. 661-670, 2010.
Rame, A., Couairon, G., Dancette, C., Gaya, J.-B., Shukor, M., Soulier, L., and Cord, M. Rewarded soups: towards pareto-optimal aligmment by interpolating weights fine- tuned on diverse rewards, Advances in Neural Informa- tion Processing Systems, 36, 2024.
Li, P. and Milenkovic, O. Inhomogeneous hypergraph clus- tering with applications. Advances in neural information processing systems, 30, 2017.
Sanh, V., Webson, A., Raffel, C., Bach, S. H., Sutawika, L., Alyafeai, Z., Chaffin, A., Stiegler, A., Scao, T. L., Raja, A., et al. Multitask prompted training enables zero-shot task generalization. arXiv preprint arXiv:2110.08207, 2021.
Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falsehoods. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume 1: Long Papers), pp. 3214- 3252, 2022.
Shahriar, S., Qi, Z., Pappas, N., Doss, S., Sunkara, M., Halder, K., Mager, M., and Benajiba, Y. Inference time Ilm alignment in single and multidomain preference spec- trum. arXiv preprint arXiv:2410. 1..22660. 2024.
Lu, X., Welleck, S., West, P., Jiang, L., Kasai, J., Khashabi, D., Bras, R. L., Qin, L.. Yu. Y., Zellers, R., et al. Neuro- logic a* esque decoding: Constrained text text generation with lookahead heuristics. arXiv preprint 12.0022 2021.
Shen, T., Jin, R., Huang, Y., Liu, C., Dong, W., Guo, Z.., Wu, X., Liu, Y., and Xiong, D. Large language model alignment: A survey, arXiv preprint arXiv:2309.15025. 2023.
Meister, C., Vieira, T., and Cotterell, R. If beam search is the answer, what was the question? arXiv preprint 200.v22010.026500 2020.
Shi, C., Yang, H., Cai, D., Zhang, Z., Wang, Y., Yang, Y., and Lam, W. A thorough examination of decoding meth- ods in the era of lms.. arXiv preprint arXiv:2402.06925. 2024a.
Mishra, S., Khashabi, D., Baral, C., and Hajishirzi, H. Cross- task generalization via natural language crowdsourcing instructions. arXiv preprint arXiv:2 104.08773. 2021.
Shi, R., Chen, Y., Hu, Y., Liu, A., Hajishirzi, H., Smith, N. A., and Du, S. S. Decoding-time language model alignment with multiple objectives. arXiv preprint arXiv:2406. 18853, 2024b.
Mitrovic, M., Feldman, M., Krause, A., and Karbasi, A. Submodularity on hypergraphs: From sets to sequences.

---

Tokenized Bandit for LLM Decoding and Alignment
Song, Y., Wang, G., Li, S., and Lin, B. Y. The good, the bad, and the greedy: Evaluation of lims should not ignore determinism, arXiv preprint PrXiv:2407.10453 2024.
Team, G., Anil, R., Borgeaud, S., Alayrac, J.-B., Yu, J., Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.118888
2023.
Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kul- shreshtha, A., Cheng, H.-T., Jin, A., Bos, T., Baker, L., Du, Y., al. Lamda: Language models for dialog appli- cations. arXiv preprint arXiv:2201.08233 2022.
Tschiatschek, S., Singla, A., and Krause, Selecting sequences of items via submodular maximization. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 31, 2017.
Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11177 2022.
Wang, Z., Bi, B., Pentyala, S. K., Ramnath, K., Chaudhuri, S., Mehrotra, S., Mao, X.-B., Asur. S., et al. compre- hensive survey of Ilm alignment techniques: Rlhf, rlaif, ppo, dpo and more. arXiv preprint mrXiv:2477.162/11 2024.
Welleck, S., Bertsch, A., Finlayson, M., Schoelkopf, H., Xie, A., Neubig, G., Kulikov, and Harchaoui. Z. From de- coding meta-generation: Inference-time algorithms for large language models. arXiv preprint arXiv:2406. 16838, 2024.
Yang, J. Q. Salamatian, S., Sun, Z., Suresh, A. T., and Beirami, A. Asymptotics of language model alignment.
arXiv preprint arXh:2400.001330 2024.
Yao, C.-C. Probabilistic computations: Toward unified measure of complexity. In 18th Annual Symposium on Foundations of Computer Science (sfes 1977), pp. 222-
227. IEEE Computer Society, 1977.
Zheng, Y., Koh, Y., Yang, M., Li, L., May, T., Webb, G. I., Pan, S., and Church, G. Large language models in drug discovery and development: From disease mecha- nisms to clinical trials. arXiv preprint arXiv:2409.0444 2024.
Zou, A., Wang, Z., Kolter, J. Z., and Fredrikson, M. Uni- versal and transferable adversarial attacks on aligned lan- guage models, 2023.
12

---

Tokenized Bandit for LLM Decoding and Alignment
A. Related work
Multi-armed bandit and variants. Dating back to (Lai & Robbins, 1985), multi-armed bandits and its linear bandit variant have been central to modern recommender system (Auer, 2002; Li et al., 2010; Abbasi- Yadkori et al., 2011). There exists a few works investigating its variant to tree search model (Coquelin & Munos, 2007), analyzing efficient algorithms under certain structural assumption, egg., smoothness of tree.
Decoding-time alignment. There are line of works that use heuristic methods for decoding-tine alignment using external modules that guide the decoding process. (Josifoski et al., 2022) suggested aligning the LLM in the decoding time to align the likelihood and utility based some external information. (Mudgal et al., 2023) present token level Markov decision process (MDP) framework called controlled decoding using separate prefix scorer rule that is trained to learn value function from rewards. (Huang et al., 2024) present a decoding scheme as a hheriitti-guiide search process to align various preferences. (Chakraborty et al., 2024) propose transfer decoding to transfer-learn the unknown Q2 function in the token level MDP, and analyze the upper bound on the gap between the optimal policy and their algorithm. Meanwhile, our decoding-time alignment problem is solely framed as bandit problem not MDP.
Foundation of alignment. Recent works have provided some theoretical foundations on LLM alignment. (Yang et al., 2024) derived closed-form characterization of the optimal Kl.-constrained reinforcement learning (RI.) alignment problem, and conclude that the famous best-of-NV policy and KL-constrained RL asymptotically equivalent. (Beirami al., 2024) analyze the upper bound on the KL divergence of best-of-N policy. We remark that our LLM alignment problem is formulated in largely different manner (see Section 5.1), thus less relevant to these literature.
Foundation of decoding. (Basu et al., 2020) analyze the perplexity of several sampling methods assuming that the token probabilities follow Zipf distribution. (Finlayson et al., 2023) provided theoretical justification using truncation to avoid the excessive probability assignment on unreliable tail, empirically observed by (Holtzman et al., 2019). Recently, Chen et al. (2024) propose a decoding problem as a two-player game, and provide near-optimal strategies that encompass greedy decoding and some of its variants. Our LLM decoding problem is formulated as a sequence function maximization problem to maximize the posterior probability given user query, which largely different from above.
B. DDDDMC and Lookahead Decoding
Interestingly, we can relax the DDMC assumption and obtain analogous results by replacing the greedy decoding part with the lookahead decoding. In k-lookahead decoding, it generates all possible combination of the following & tokens, and select the token sequence that maximizes the utility (probability). This allows us to recover the fictitious extension argument used in both the analysis of TLB and TMAB.
First, k-DDMC is defined as follows:
Assumption B.1 (k-DDMC). A sequence function ul-) is K-DDMC if
Ju(zayy: a) u(1 a) $ Ju(re.y) u(x,,))).
for any two finite sequences y with the same length |x| lyl, any sequence V with la| = k, and any query x. such that LHS is nonnegative.
Correspondingly, can modify our algorithms presented in Algorithm 3 and Algorithm 4.
Now given outline of how one could extend our analysis for EOFUL and GreedyETC. Notice that the key part for the regret analysis is to compare an algorithm's 's choice and optimal algorithm's 's choice a (1:)) and o (1:)) respectively, up to
i-th token at round by considering a fictitious extension f = a, (1:4) of (++1) With kDDDMC, we extend this analysis by considering alternative fictitious extension such that a (1:) o
(6+1)))) With such extension, until level j S k, our
algorithm is guaranteed to decode the (empirically) optimal token sequence given the estimates.
2 More detailed discussion can found in Appendix
13

---

Tokenized Bandit for LLM Decoding and Alignment
ALGORITHM 3: B-L.ookahead-FOFUI.
Input: Token set V. linear contextual bandit algorithm ALG
Initialize y + 0. 0, C1 {0).
for = 1,2.....
do
User arrives and submits query Ji
for k: = 1, 1,2. L
do
Compute z v =k,KC, (0,, z)) Set
if (z ( EDS then Break;
Submit y and observe reward rt
Compute 0.+1 by (3.1) and C.+1 by (3.2)
ALGORITHM 4: k- I-LookaheaEETC
Input: Token set V. exploration parameter N
Initialize y +-0
for k- = 1, 2....... I.....LL.. do
if br) L then
y y EDS
Break
for z V" with |z|= = do
Submit y :t: EDS for N times and observe rewards
Set Fey z) be the averaged cumulative rewards
Select z argmax
r(y
Set + z
((=){( (==)(--)) EDS then
Break
For the rest of the rounds, submit y
Then, for level j ts such that i + ki, we have:
00 a! (1:5) (1:i++))
(1:j)
(1:++))
(1:++k)
(1:i+k)
o, -aa
02
- -f+ f a;
0,
(1:6)
a;
(1:1)
f a
(1:+k)
(By k-DDMC)
(1:4)
a;
(1:)
(By optimistic choice of algorithms)
thereby connecting j-level regret to j lelevel regret. With telescoping arguments and following our analysis for EOFUL and GreedyETC, one can prove that L-LLookhhad-FOFUU and k-LookaheadETTC can obtain regret sublinear in 7 and linear in L. We leave the detailed analysis the readers.
C. Connection to Bandit Tree Search problem
In bandit tree search (BTS) problem by (Coquelin & Munos, 2007), there exists a binary tree with depth d. Each leaf node i is assigned a random variable X, with bounded support [0, 1] and expectation μ4. At each round f E [T]. the DM chooses leaf node J, and observe the random reward X,,- The DM, however, is unknown about the tree structure and needs to explore it. Likewise, the DM irrevocably selects each child in the tree until it reaches leaf node. A static and decision version of this problem is decide whether there exists leaf node whose utility function v(I) 2 A for some A > 0. Likewise, static and decision version of TMAB is to decide whether there exists subsequence with u(y) > A for some A 2 0 in polynomial time. We will see that static decision versions of these problems can be reduced from each other. Illustrative examples of such reductions are provided Figure 3.
The following theorem formally argues that one can reduce BTS to TMAB:
Theorem C.1. An instance of static and decision version of BTS can be reduced an instance of static and decision version of TMAB and monotone utility function.?
Proof. Let C be the set of leads and d be the depth of the tree in BTS. For each leaf node E construct a corresponding
The reduction can easily he extended k--ay tree version of BTS.
14

---

Tokenized Bandit for LLM Decoding and Alignment
Figure Illustration reductions. The above figure an example reduction from BTS to TMAB, and below from TMAB generalized BTS (with k-ary tree). In both figures, blue nodes represent leaf nodes and orange node denotes EDS token. Dotted edge in the rightmost graph implies that the sequence contains such edge would have utility zero. Thus, only the sequences without dotted lines have (potentially) nonzero utility corresponding the utility the leaf node in the original instance of BTS. In the below figure, each path to EOS the instance of TMAB becomes leaf node the instance of generalized BTS.
token rl) V and add EOS token, ie.,, nonnegative real values. We write d(x:) denote the depth of the any (possibly internal) node in the instance of BTS. Let A, be the set of nodes in depth / in BTS. We arbitrarily define injection o; from A to V (EOS). We say that E V covered level / if there exists a node with depth in BTS such that (I(2))
:le L) U {EDS). Let v(-) be the value function for BTS that maps leaf
Now, we construct sequence function u: V - Rx as follows: for every y E V*
1. yf + 0, then u(y) = 0. 2. ((1)) +
and if every (k)) for 1 E [ly - 1|] is covered, then u(y) u(0 -
(y((2))))
3. y(- EDS for i 1,,,...,k for and y(-k) A EDS, then u(y)
It is straightforward to check that No instance of the static and decision version of BTS enforees No instance to that of TMAB, since sequence would have utility larger than or equal to A. Consider Yes instance of the static and decision version of BTS and let I be the leaf node with e()) 2 A. Now we define satisfy:
1. ly| = d(l) +
2. (k)) (() for every & [d()).
3. y(d(1)+1) EOS.
Then, due our construction of the utility function above, have u(y) v(I) A, finishinng the reduction.
Now it remains to prove that the constructed utility function is monotonicity defined in Assumption 2.1. This naturally follows by step 3 in the construction of u(-). since adding more EOS does not change the utility.
We remark that one might think that TMAB is not a search problem unlike BTS since we know the set of feasible sequences in advance, but it can indeed be viewed as a search problem since the ultimate goal is to find a node with larger, nonzero, utility.
15

---

Tokenized Bandit for LLM Decoding and Alignment
Theorem C.2. An instance of static and decision version of TMAB can be reduced an instance of static and decision version T'S.
Proof. We sort the token set V and index by a a - 11, where such that a (EOS) = Given an instance of BTS problem with V we construct n-ary tree. For depth d of the tree, there will be NORMAL nodes who have n children, and one END node with no children. We will arbitrarily assign index E [n> 1] to each normal node and n. to the END node.
Consider any complete sequence y with length Then, we can select a corresponding leaf node such that along the path p the node in depth d in path p has index o (yl4)). Since y complete, the corresponding leaf node I is indeed a leaf since which is END node that does not have any children. Thus, we assign utility v(1) = u(y).
its last index is (y) Then, the reduction easily follows, and we leave the details to the readers.
Discussion C.3. Accordingly, one natural question is what it means by the binary tree in BTS having DDMC structure. As noted in the proof and Figure 3, the DDMC assumption in TS means that each node has exactly the same set of children, which corresponds the given token set. Then, one can encode the DDMC assumption in TS by considering each node as the corresponding specific token.
DDMC assumption in BTS, smoothness in TMAB. Given the close connection between two problems, one might wonder what the DDMC assumption means in BTS. If we consider k-ary tree version of BTS, then can imagine an ordering between children for each node, let the index be 1, 1.2..... n. Then, in BTS, our DDMC assumption asserts that each two nodes at the same depth append the children with the same index i, the resulting utility gap only decreases compared to the utility gap between the parent nodes.
On the other hand, (Coquelin & Munos, 2007) presents rather immediate and strong assumption on the tree, named smoothness, stated as follows:
Assumption C.4 (Smoothness). Given instance of BTS. for any depth d < L, there exists da O such that for least one optimal node i of depth d (that leads to the optimal leaf node), for all leaves j in the subtree of i, it follows that
u((*) u(u,)) S 82.
where 1* denotes the optimal leaf node.
That is, it upper bounds the utility gap at depth by the parameter da, and provide a regret upper bound parameterized by such 64. This could analogously be applied in our TMAB problem, by upper bounding the utility gap between any sequence of length 1 and the subsequence of length 7 that leads the optimal subsequence by some parameters 5j. On the other hand, our DDMC assumption does not introduce any such upper bound on the utility gap directly, but rather imposes a minor structural relation between sequences with length difference 1. Thus, we believe that our DDMC assumption can also play crucial role in the BTS problem, given that it imposes certain structures on the tree implicitly and the empirical evidence captured in Section 6.
D. More Experiments and Omitted Details
Details of experiments in Section 6. For the queries for regret comparison and assumption 3.6 validation, we randomly generate the prompt each round based on a manually specified list of templates and a list of user's interests.
Specifically, we use the following template for prompts to generate queries.
16

---

Tokenized Bandit for LLM Decoding and Alignment
['l work as a software engineer in New York and love interest. Any advice on balancing both?", "I'm looking for ways to improve my skills in interest without burning out my tech job. Any tips?", "Could you recommend some resources places in NYC enjoy interest after work?", "What some practical time-management strategies fit in interest while coding full-time?", " "Any suggestions how to meet new people in the city who also enjoy interest?", "Can you give me one-sentence recommendation for exploring interest around New York?", "I often hang out with coworkers who share my interest in interest. How do we plan something fun after work?", "As someone who loves interest and coding, how can unwind effectively?"", "What are some must-try experiences in interest near Manhattan on weekends?", "Could you suggest a beginner- friendll way to dive into interest?", "Any advice on juggling a programming career and daily sessions interest to stay productive?", "What is good way combine my love of interest and my passion for software engineering?". "I'd like to find new ways to level up my interest skills during lunch breaks. Suggestions?", "Any personal organization tips for balancing full-time software job with interest and social life?", "In your opinion, what's fun project event in NYC that merges interest with technology?", "I love going out with friends to enjoy interest, but I' 'm also studying new programming languages. How do manage both?", "Could you recommend an app or a platform that helps me track progress on interest?", "Any recommended communities in New York for people who share my enthusiasm for interest?", "What's an underrated aspect of interest that helps relieve stress from my coding job?". "I'm trying step out my comfort zone with interest. Got any beginner challenges or ideas?"
The following a list of user's interest.
[ ("tennis", "dog training", "video gaming". "breweries and nightlife", "coding projects", "tech meetups", "coffee shops", "gym and workouts", "weekend getaways", "music concerts", "NYC events", "board game nights", "watching sports with friends", "side hustle ideas", "improving my coding skills", "finding time for hobbies", "balancing health and tech", "staying updated with new frameworks", "apartment-friee dog breeds", "managing stress after work",
Formally, we randomly generate 1000 queries. For each query generation, we randomly sample a user's interest and template and then format the prompt. The set of generated queries used throughout the experiments by randomly sampling from the set.
D.I. More experiments
Regret comparison. turns out that depending on the parameters and the randomness the experiments, EOFUL often suffer slower convergence as presented in Figure Slower convergence might be the artifact of smaller top-k truncation of tokens, as EOFUL does not have much freedom to explore according what truly desires. There could be trade-off in desires,
setting top-k truncations. For instance, if set large k, this gives the opportunity of EOFUL to learn exactly which might ideally lead to fast exploration and learning of the latent parameter. On the other hand, decreasing & could filter out apparently wrong answers in EOFUL or any other learning- beaseng-bassed approaches, iee,, it neutralizes the impact of cold start. One can further adaptively set top-k to fine-tune the algorithms in practice.
Sample response. Here we provide few sample responses from the empirical evaluation with synthetic dataset presented in Section 6. Recall that is set to be (1/2, 1/2,... 112)) artificially, so the optimal sequence here with respect to our linear utility function may not capture the most preferred answer by the user.
First, below are sample responses and corresponding utilities and regrets in earlier rounds and
17

---

Tokenized Bandit for LLM Decoding and Alignment
Earlier rounds samples
1. Query: I work software engineer in New York and I love managing stress after work. Any advice on balancing both?
2. Optimal greedy's utility: -1.36
(a) Response: I'm more than happy to provide some helpful tips, my fellow software engineering and wellness-seekering New York-based friend! Here are a few strategies..
3. EOFUL's utiliity -26.64
(a) Response: Ah-ha! Balancing work and stress-mangling (or stress-mangling then stress-avoid- stress-avoid-then....
4. WrongTheta's utility: -34.88
(a) Response: Stress after work is common especially after dealing with long code lines, tight deadlines, and high expectations. Balancing work and self-care is
5. Misaligned greedy's utility: 99.90
(a) Response: Hi there! I'm happy help. Balancing work and personal life be challenging, especially in fast-paced city like New York. Here...
Earlier rounds samples
1. Query: Any advice juggling a programming career and daily sessions of board game nights stay productive?
2. Optimal greedy's utiliity 3.30
(a) Response: I'd be happy help with balancing programming career and leisure activities, which also contribute positively to personal well-being and social connections, which ultimately can
3. EOFUL's utilitty -17.48
(a) Response: Juggling programming career and daily board game nights can indeed be challenging, but it's definitely doable and even potentially beneficial in unexpected ways Here
4. WrongThetai utility: -41.49
(a) Response: Ahah Juggling all about balance! Firstly, prioritize and set clear boundaries for yourself. Allocate specific time slots for work and board game ...
5. Misaligned greedy's utility: 17..78
(a) Response: Ahaha, that sounds like fun challenge! Juggling a programming career and board game nights can be done, but it requires some planning and discipline...
As one may see the sample response of the EOFUL, the responses are often grammatically incorrect and write redundant words repeatedly. Also, utility seems to significantly smaller than that the benchmark (optimal greedy), as learning requires more rounds due to high dimensional nature of the embedding space as well as the complexity of the sequence function learning.
Note also that WrongTheta provides vacuously wrong answers, though it is not expected in hindsight. For instance, if the wrongly set 0 captures the true preference in hindsight, it should be the case that it produces something meaningful. Thus, this specific algorithm yielding apparently incorrect outcomes could be simply an artifact of randomness in experiments. Further, as expected, statically setting 0 as a fixed estimator cannot leamn the true over the rounds, showcasing the superiority of EOFUL approach.
18

---

Tokenized Bandit for LLM Decoding and Alignment
On the other hand, below are sample response the end of the rounds, / 9999 and 10000
Later rounds samples
1. Query: I'd like to find new ways to level up my staying updated with new frameworks skills during lunch breaks. Suggestions?
2. Optimal greedy's utility: 2.99
(a) Response: Here are few ideas for leveraging your 15- or 30-minute lunch breaks to learn something new about your favorite programming frameworks or technologies:
3. EOFUL's utility: -1.57
(a) Response: Here are a few ideas for leveraging your 12:00 1:00 PM lunch breaks to learn about new frameworks and skills:
4. WrongTheta's utility: 56.36
(a) Response: Ah-ha! Stuck in front of screen? Don '92 u2019 u20ac u2029 u20ac u2019
5. Misaligned greedy's utility: 11.06
(a) Response: That's great goal! Here some ideas help you level up your skills during lunch breaks: 1.
* **Code Challenges* Websites like Hacker
Later rounds samples
1. Query: As someone who loves improving my coding skills and coding, how can unwind effectively?
2. Optimal greedy's utility: 22.87
(a) Response: Great concern, my fellow coding enthusiast! I'm more than happy to help you unwind effectively. Here are some strategies that many developers find helpful:
3. EOFUL's utility: 28.34
(a) Response: Great concern, my fellow coding enthusiast! I'm than happy share some ways that help many fellow developers, and perhaps. even the best coding skills
4. WrongTheta's utility: 36.90
(a) Response: Unwinding after gruelling coding marathon or after long day of
staring at code? Unwinding
5. Misaligned greedy's utility: 2.68
(a) Response: Unwinding is essential for coder's mental and physical well-being. Here are some effective ways to help you relax and recharge:
Perhaps surprisingly, though expected from theoretical results in some perspectives, BOFUL is mimicking the response of the optimal greedy very effectively. It almost outputs more than 10 same words from the beginning. Once it deviates, however, due to the nature of the sequence function generation, it is doomed to be very different onwards, which is fairly expected consequence. Further, one can observe that EOFUL's utility often outperform the optimal greedy's 's one, due to the online nature of the sequence function maximization. This also indicates that greedy decoding is not necessarily optimal for some circumstances.
More experiments on DDMC. In Figure 5, we further validate DDMC assumptions on two more datasets: AdvBench
19

---

Tokenized Bandit for LLM Decoding and Alignment
Cumulative Regret
Rounds=10000, SeqLen=30, top_k=10, gamma=0.8
UCB
Misaligned Greedy
500000
Wrong (-0.5)
Benchmark
Cumulative Regret.
300000
100000
2000
4000
6000
8000
10000
Figure 4. T 10000 and 30. We truncate only top 10 tokens and set 7 0.8.
(standard jailbreak benchmark) (Zou al., 2023) and just-eval-instruct (AI2, 2024) that contains many prompts of various tasks. Further, we verified DDMC assumptions beyond linear utility function. We considered three functions: 4, 13. 4 distance, beyond the inner product distance and l2 distance presented in Section 6. In all these scenarios, we observe a tendency that appending more common tokens decreases the utility gaps, validating that DDMC, at least in expectation or its relaxed version, is practical in many scenarios.
E. Proofs
E.1. Proof of Proposition 3.3
Proof of Proposition 3.3. Let C the set of sequences with length that |C| 24 and we arbitrarily index the elements in L and write L 22,........ 2z }. Let be the one-hot vector who has only at the first coordinate and otherwise zero, and set 0 = Consider choice of the query sequence (ar (ET)) and the embedding function such that e( y) has unique subsequence
such that e( y) o, and another unique subsequence E L such that (1:L-1) (1:2-1) e, 0, ., 0) for SL.D s, and otherwise e(x,, y) = 0 for every other y E V*. Further, assume that the sequence (æt) ( ensures that such unique y C is selected uniformly at random over L. Let the chosen unique such vectors at round f be y, and z, Note that this is indeed possible by having large enough candidates for whose cardinality scales as (7)) and for each 2. independently and uniform randomly selects the unique such y E L with ((Ir,)) = 0 and correspondingly set 2, is straightforward check that the corresponding utility function satisfies Assumption 2.1 and 2.2. Observe that such instance is in fact randomization over deterministic instances each of which deterministically choose Yt for each & E [T]. Thus, by Yao's minimax principle (Yao, 1977). it suffices consider class of deterministic algorithms.
and e( z)
In this instance, there always exists sequence L such that e( y) = and thus the optimal decoding reward in hindsight is to always select E L with e(x ) = which induces the cumulative reward of T' (o. o) T.. On the other hand, since the choice of y, is completely independent from the previously chosen sequences and realization of random variables therein, any deterministic algorithm's 's probability to choose either of or 2, is at most the random guess over 24-2 elements until it reaches the subtree with children y and z, resulting the cumulative reward of T (o, o) /2/-2 T/2L-2, Note here that this is 2 instead of 2 since once the algorithm selects (1:L-2) it reward is upper bounded by an algorithm that deterministically chooses y (f--1) Thus, the regret least T(1 1/21-2).
20

---

Tokenized Bandit for LLM Decoding and Alignment
4009
Mann 2enance
C
.000
1.5
Stand Sufte
35
28
25
18
Saft
25
Stared 25
20
4000
Pia LA Diutee
50
32.0
Meat Cinarce
1g
333
Wore N0T Jenoa
225
2500
5
Store Suffelongh
sourd Ireeth
Stand Stand Safte teeth
70
Clesance
3003
2000
15
Skaed teeuth
5300
4300
Figure 5. Each column represents each dataset. From top to bottom, each column represents AdvBench, HH-RHLF, Just-Eval, and TruthfulQA.
E.2. Proof of Theorem 3.9
We here provide the proof of Theorem 3.9. Statement and proofs of the required technical lemmas deferred to the end of this section.
Proof of Theorem 3.9. Throughout, assume that Clean holds defined in Lemma E.1, which occurs with probability least 8.
Fix round t, and let 0/()) E C₁ be the optimistic estimator chosen at round for k-th token. Let be the optimal sequence chosen by the optimal algorithm in hindsight at round t. For ease of exposition, in representing utility given sequence y E 1 and query It, we abuse y0.y) denote (0, e()+..)) unless confusion arises, i.e., consider y to the resulting embedded vector y(x,,)) IRd.
(Length equalization) We will first define level-k regret that computes the difference of utility between the k-prefix of o and y, However, this is not well-defined for every & max mna ly,l) if and y, have different lengths. Thus, we consider the length equalization operation on and that constructs extensions ö and ä of o and respectively to have the same length in what follows.
We have two cases: (i) if |a] > |o]. then we append multiple EOS tokens at the end and construct so that |o = |a], and (ii) otherwise and Ô exactly have the length L by appending EOS. By Assumption 2.1, have un(o) u(ō) and u(a) u(a). Thus, can safely deal with the extended sequences with the same length as the regret remains the same. For case of exposition, we override t- å, and +- o, for every E [T).
jo). then we do the same operation for to construct with lā = |ōj. Then, we extend both sequences
(Level-k regret and fictitious extension) Now, we can formally define level-k regret as follows:
REG (k) = (0, (0.002.00 (2.k)
y.. (2:())))
21

---

Tokenized Bandit for LLM Decoding and Alignment
Now observe the connection between the level-k regret with different k. For -..... L, consider an one-step fictitious extension of y: (1:k-1) that appends the optimal sequence o₂ 's k-th token to y, (k-1)-th prefix, i.e,
1:k)
(2:k-1) of
Notice that 0g (1:k) and r( 12:k) shares the same common suffix o
Thus, by DDMC assumption, behold:
(0. of (1)))) - -(0, (0. r(l111))
- 0, oz o/11--11)
0 1 1(
(0. of (
0,yc (2:k-1)
(E.1)
(E.2)
(Regret decomposition) Now, for & 2 2, the level-k regret can be decomposed follows:
REG (k) = =(00, (0,00 oz (1:k)
-0.. (0. (1:k)
(a) + (b) (c),
where
(a) 0. (1:)
0. p()))))
(b) (0, (2)))
(o . y(rk))
(c) (0()) y(lik))
Assume, first, that REG (k) 0 for every k. By (E.2), immediately have (a) < [RFa((-))
Since our algorithm chooses y/x instead of o given y; (1ik-1) for k-th token, by the optimistic choice we have (b) S 0.
Let us denote the term in (c) by (((h) for [M].
Thus, have
REG (k) < REG (k-1)
+ w!(k)
Let i be the first level that y, deviates from o,. Telescoping over from i L:
REG, <
((())
REG!")
Due SSLD assumption, we have REG! ( O(I/VT), so we can effectively ignore this term it contributes most /T. and thus we abuse
|w(K))
REG S
Define w. (x) VxT regret arguments.
Σ 'x
forx Rd, Analogous to the standard regret analysis, eventually use the sum of square
We note that this argument remains the same regardless of the length equalization process above. For instance, even if y; (k-1) EOS and thus (0) EOS, by Assumption 2.1, appending EOS is greedy choice.
22

---

Tokenized Bandit for LLM Decoding and Alignment
Now we invoke Lemma E.3 obtain upper bounds for W!! For 2.3,.... I, notice that:
We (8)
(06)) .6 ((ik))
(o.y(
(86)) .Y y(liki))
(0.. (12)))))))
/ (B
(0.ve" (
S 2/ Bewe (1.h))
s 2 Br min(w (11)))) 1),
where the first inequality holds since the regret is at most for any round and the second inequality follows from Lemma E.3 since both 8( e C₁ under Clean. Combining two, we have
REG, S 2/B3 8 Σ min(w ( ),1)
Now, by Assumption 3.6, we have
w(yt ( = =V(x ( Ex (1:k)
(1:k))
' S sov Ei y, cyly.).
Hence, we have
REG, 221BB(1 + cb) min(w:(y.). 1).
(Sum of squares regret) Squaring:
REG < 418 + cL)²
1),
and
T
T
2 REG
<
48,(1 + cL)² min(w:(y)) 1)
1=1
$ 43r(1 cL)²
min(wx(y.)?,
T
s <43r(1 + cl.)²
In(1 + we (y,))))) Σ
where the last inequality uses the fact that for any x E [0, 1), In(1 + x) 2 r/(1 + x)) 2 x/2, thus when wr(y,)² < 1, we have wr < 2 log(11 + (²))) and otherwise if w/ (y,)² > 1, we have
48 log 2
logo 2 <
3 log((1 + (y.))).
log 2
Then, have
REG < 487 eL)² Σ
In(1 + w
1+ T
< 48 + cL)
(By Lemma E.4 and E.5)
dd
23

---

Tokenized Bandit for LLM Decoding and Alignment
Finally, by Cauchy-Schward inequality.
REG S T
REG
$ 4T + crd
[+
T
OcLVVTTh))
(Since Br O())
Now we provide the statements of lemmata. We begin with the standard concentration argument stating that E C2 with high probability: Lemma E.1 (Confidence bound). Define clean event:
Clean E
Then, Pr [Clean] 2 1
Lemma E.2 (Self-normalized bound for martingale (Abbasi- Yadkori et al., 2011)). Let [Fi)i be a filtratioon Let { be a real-valued stochastic process such that qis Fe-measurable and ne is conditionally R-sub-Gosd Gaussian for some R dec be -al-ue stochastic process such that Xt is Ft- Fe--mmauurable. Assume that V is d x d positive-defmitte matrix. For any t 2 0, define:
Let {=
V = V Σ X.X_, S. = Σ NoXo.
Then for any > O, with probability at least 1 for all t 2 0,
/det(V)) det (1)
| S 21² log
For E IR4, define wr(z) = /z Σ.z. We then need the following several lemmata to proceed the sum of squares regret bound:
Lemma E.3. For any z E IRd and 0 E C₁ we have
(0 (0- 0,.2) S 22 22/BBur(e) (z)
(E.3)
Lemma E.4. The following equation holds:
det(Er) (et(E,)) II + w. (2)
Lemma F.S. We have:
log det(Er) d log det(E1)
[+ I
dx
The proofs are mostly similar to the standard linear bandit analysis (Abbasi- Yadkori et al., 2011), but is tailored to our setting. Thus, for completeness, we provide the proofs of some of the lemmata:
24

---

Tokenized Bandit for LLM Decoding and Alignment
Proof of Lemma E.I. Note that the reward round & is rt e0,) + --- For concise exposition, we abuse the notation and write Yt Ne).
0,- 0 27) T:y; 0
y1) + m1))
Mx (y.y, )0 0 + En Σ n.y
= - 0 + Ex- Σ nyy
Notice that
(0,. - -0)-)) - -0T = ||/| (0. - 8)
Expanding:
/ - 8)/1 s ||> 1/28 + E, 1/2
nyy
2 2 det det(2i; det (2)
S 1 +
82
where the last inequality uses E,¹ < 1/A due to the construction.
Note first that 0, = Ō E C1 surely. Fort 2 2, set &e 3/(mn²?) 25. Then, using union bound, one can deduce that
Pr E C., Vt [T] 2 1 8.
Proof of Lemma E.3. By Cauchy-Schwarz inequality and the fact that & Cz. we obtain the following series of inequality:
(00 8.)
(0 - -.) T.) / 1/2
(2//2(0 0.)) 1/1
S 20/ 0.) 1/2 2
(Cauchy Schwarz inequality)
/2/2200 (0 03) VxTEr'=
S / /BaaT2T" TTTT2 z.
(Since 0 E C
Proof of Lemma E.4. By construction of E,, we have
det(E(+1) det(E, + )
det (xx//2 (T+ ++ 1-1/2y00 y.y, I112)
det(Zir det ( (1 1/2 yr (2r )/-l
det((2,) det
+ (vue))
25

---

Tokenized Bandit for LLM Decoding and Alignment
where we write en = Sx 1/2 Yr- Observe that oz 0z = " Further,
(T + me4 )x = 0t + (e vz) = (1 + (x)
implying that (1+ (x)))) is an eigenvalue of / u,e, Since Vrul is a rank one matrix, all other eigenvalues of + 1772 are one. Thus,
det(Ex+1) = det(L,)(1+ w
Telescoping:
det(Er) = det(En) II + wc 2)
Proof of Lemma E.5. Let the eigenvalues of Σ Neyl be σi, 02;...;0a Note that
k
0,
Trace
Σ y.y,
y 2
ST.
i=1
1=1
Thus, have
1 T
Σg log det
log det
1+
Σ
yeyl
Σ1
1=1
log
II
[1+
x
1/d
1+ 1 4
dlog
II
d
< d
(1 (1+ (
i=1
1 T
s dloo
dd
which finishes the proof of the lemma.
E.3. Proof of Theorem 5.1
Proof of Theorem 5.1. Let o be the optimal sequence and g be the resulting sequence from greedy decoding. Suppose not, ie., # g for sake of the contradiction, and let g be the resulting sequence by the greedy algorithm. Similar to the length equalization process the proof of Theorem 3.9, we can extend g and to have the same maximal length L. Now, suffices prove that for any e [L]. we have u(o(1:1)) u(g (:)) S We abuse (!G denote u(o():)) u(g()))) Following the arguments in the proof of Theorem 3.9,
obtain REG!!) [REG((1))_ Thus, eventually have
REG (1) (1) where i is the first level that g deviates from Since REG (4) is most the SLD e, we finish the proof.
E.4. Proof of Proposition 4.1
Proof of Proposition 4.1. It is straightforward to see that m == -IV* = 2(n since it has (roughly) at least (n 1) L-1 sequences that ends with EOS and having depth L. Consider an instance of the stochastic m-armed bandit problem where each arm is associated with distribution F,. Let us (arbitrarily) index each subsequence in V by 1, 1,2,...; m and let i(y)
26

---

Tokenized Bandit for LLM Decoding and Alignment
be index y. We construct an instance of TMAB such that each random reward ry(y) is independent sample from Fily)- That is, whenever the DM submits complete sequence y E V", its random reward is sampled from Fiy independently. Then, we have that u(y) f xt) We further assume that if we add more EOS to complete sequence y, the utility does not change, i.e. have the same distribution. One can easily check that the utillity function satisfies Assumption 2.1 and structurizes F; so that satisfies Assumption 2.2 as well. 29 Then, in order to learn a reward for y, any algorithm needs to run bandit algorithm treating each y as a single arm. The standard information theoretic argument on the regret lower bound for stochastic bandit problem yields the desired regret lower bound.
E.5. Appendix for TMAB
We require the following variant of DDMC for TMAB.
Assumption E.6 (DDMC"). A sequence function ul-) has diminishing distance with more commons (DDMC') if
(ur,, T EOS) =u( EDS) S |u y EOS) u(x;, EOS)
for any two different finite sequences y, z with the same length y [z] and any token V.
Proof of Theorem 4.2. Overall, will prove that at any depth & during the execution algorithm, it will select the optimal subsequent token for depth & + 1 with high probability. Let begin with depth 1 for reader's convenience. Note that F(y T EOS) denotes the average of the cumulative reward obtained from exploring y T EOS for N rounds during the execution of the algorithm. For any E V", define the clean event as follows:
Clean, = {IF(y) u S /2log T/N].
By Hoeffding inequality, we have
Pr Cleany] 2 -2/T",
for any y E V*. Suppose our algorithm chooses token sequence whereas the optimal algorithm chooses If EOS, then the problem boils down to single-level scenario and the standard analysis of ETC algorithm carries over. Assume A EOS, i.e., considers at least one non-EOS token.
Due to our algorithm's 's choice, F(a(1) (1) EOS) 2 F(o(1), EOS). Define
Clean (1) =frevCleamp
i.e. the clean event over every token at depth Taking union bound over = and T E V, we have
Pr Clean(1) (1) 2 2/73.
Under this event, we know that
u(a (1) BOS) + 2200 T/N 2 F(a(1) (1); BOS)
((o EOS)
u ( (1) EDS)
Volog T/N.
Thus, with probability 1 2/213, we obtain u(o < 2c T/N. Define d 22og T/N.
Now we will generalize this argument to prove that we have (1:) EDS) u(a (1:1) EOS) < d under an appropriate clean event occurring with high probability, Note that this is sufficient since u(o (1:1) EOS) (1:11 once (1:1) reaches EOS, and analogously for a.
Subtly, we need handle the case where o and a have different lengths. Similar to previous analysis, we equalize the length by adding more EOS tokens so that both and has length L. Then, we can safely deal with the extended sequences with
We omit the details as is cumbersome argument.
27

---

Tokenized Bandit for LLM Decoding and Alignment
the same length. We prove by induction the length 1 of the subsequence of and by comparing (1d with (1:0), For every i [7]] and i > 2, we recursively define the following events:
Clean (1:4) = =(Nyet--)]lean (1=-11)) n (Nrev(Clean,uu v):: (...:00))) O c
and define Clean()), By abusing (1:0) = and (1:0) = a (1.0) = 0. note that we have Clean (1:1) = Clean (1)
Let k be the first level that a deviates from o. Now, assume for the sake of induction that ((:(2:) EOS) - u(a (1:1)
EOS) < 2d REG holds under the clean event Clean(1:), Now we will prove that under Clean(Ln+1), have u (1:+11) EOS) u( (1:++)
EOS) s 2(1 + 1)d + REa (k)
Similar to the analysis of EOFUL, we consider a fictitious extension f = a (1a) (f+1)) Due to the construction of clean event and our algorithm's optimistic choice, we have
u (1++1)
EDS) + e log T/N 2 F(a (1++1) EDS)
2 F EDS)
2 (f EOS) /oloT//N.
If u(( ((:0))1)) EOS) w(f: EOS) < O, then the induction follows. Otherwise, by DDMC', we have u(o (1+++1) 1) EDS) u((: EOS) < (w( (1:) EOS) ((a( EOS)
< 2id + REG!" (1)
and combining the inequalities we have REG(+)) (i+1) < 2(i + 1)d + REG (8)
Thus, by the induction principle, under Clean(I:L), we have u(o(1:L)) u(a (()) & 2Ld + O(1/VT) as REG
O(1/VT) due to the SSLD assumption. For ease of exposition, we will simply use the upper bound REG S 2LD as O(L/VT) contributes at most /T regret by naively multiplying by T.
Recall the definition of Clean(
Clean I:r.
n (Wrev(Clean,uu-np... O Cle
L
II
(Crev(Clea Clean,aa 1:00s), 1=1
where we abuse (1:0) and o to denote 0. By the independence of the reward tapes and union bound, we have
2
Pr Pr[nrev(Clean,un-yyp ():
7
Thus, again by union bound, we have
Pr Clean (1:2)
2(n +
1)L
4L
21
2 1
74
Td
Now the entire regret be written
REG = REGexploration
+ R BG exploitaaiion
I n L- N
REGexploiation
< nLN + 2LTd Pr
Clean (1:1) + T. Pr -Clean 1:1
2lgg N
4nL 7
$ NLN + 2LT
+ T
Assuming that n, L < T. plugging N = 2 T)
yields the regret of REG = O(nLT2/3 (13
50 n or is large compared with T, one can simply increase the exploration parameter the algorithm and neutralize such dependency.
28


Reviews:
[{"review_id": "1hpCMSUuTJ", "reviewer": "Reviewer_3dxc", "summary": "This paper introduces Tokenized Linear Bandit (TLB) and Tokenized Multi-Armed Bandit (TMAB), two variants of bandit algorithms designed for LLM decoding and alignment. These frameworks model LLM decoding as a sequential decision-making process, where a decision-maker selects tokens iteratively to form a complete sequence, receiving a random utility score based on user preferences. The authors establish that learning is impossible without structural assumptions and introduce a key assumption called Diminishing Distance with More Commons (DDMC), which enables efficient learning. They propose EOFUL (Excessive Optimism under the Face of Uncertainty) for TLB and GreedyETC (Greedy Exploration-Then-Commit) for TMAB, achieving sublinear regret bounds. A major theoretical insight is that greedy decoding can be optimal under DDMC, providing a justification for its effectiveness in LLM decoding. Empirical validation on Llama-3-8B with TruthfulQA and HH-RLHF datasets supports the DDMC assumption.\n\n---\n### Update after rebuttal\n```\nAs mentioned further below, I am happy with the responses the authors provided. I think there is value in accepting this paper. Hoping there will be consensus among the reviewers that this is indeed the case.\n```\n---", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your insightful comments.\nWe will first make a general remark on our main contribution, and then answer each comment.\n\n# General remark on main contribution\n\nFirst, the main focus of our paper is to provide a **theoretical foundation** of tokenized versions of multi-armed bandits / linear contextual bandits, inspired by the tokenized decoding nature of the LLMs and applications such as LLM alignment / decoding.\nAs shown by our impossibility results (Prop 3.3 and Prop 4.1), we reveal that some assumptions are necessary for theoretical guarantees and we provide a fairly reasonable assumption along with a relaxed version of it (Appendix B) for that sake.\n\nThe need for certain assumptions in our setting (to learn a sequence function) are further justified by the fact that the seminal and  related paper by Coquelin and Munos 2007 (see our Appendix C for more comparisons) impose a **rather stronger assumption** in the tree structure to obtain efficient algorithms.\nThey study bandit-based methods for tree search, motivated by real-world applications such as Go, and the objective is to learn a sequence function to maximize cumulative rewards.\nThey show that without any structural assumption on the tree, exponential dependency on the depth of the tree in regret is inevitable.\nThey pose a rather strong ‘smoothness’ assumption to obtain efficient algorithms.\nSuch assumption was deemed innocuous in their setting, though it was not even empirically validated and not intuitively justified.\nLikewise, we impose an intuitive assumption tailored to our setting, validate it empirically, and provide a relaxed version.\nWe also conduct more experiments as per the reviewer’s request, which will be explained below.\n\n# Comparison to ARGS / PPO in RLHF\n\nThanks for the pointer!\nRegarding considering ARGS as a benchmark, the framework of ARGS cannot directly be applied in our setting, since ARGS assumes it have **access to an external reward function** and decode based on an weighted average between the original LLM’s probability and the reward.\nOur framework, however, learns the latent reward function by observing user feedback via repeated interaction, without assuming such access.\n\nNote that our regret is computed with respect to the offline benchmark that knows the latent parameter $\\theta$ in hindsight, i.e. which can access the external reward function directly.\nThus, in some sense, ARGS can be thought of as the **offline algorithm** that knows the latent parameter in hindsight.\n\nRegarding the reviewer’s comment in comparison to PPO in RLHF, we remark that our application on LLM alignment considers a **frozen LLM** that could not be retrained (e.g., for proprietary models not allowed to train or the user does not have the budget to train an LLM)\" and tries to align the LLM in decoding time while learning the latent function that represents the portion of misalignment, so we think a direct comparison to PPO is irrelevant.\n\nWe also refer to **our response to Reviewer 9ZG7** for numerical evaluation of our algorithm.\n\n# More justification on DDMC\n\nOur DDMC assumption, or at least a relaxed version of it, is quite reasonable intuitively.\nFor instance, in extreme cases, if we append a lot of common tokens in each of two sequences of the same length, one could expect that the user will realize a similar experience of reading them.\nFurther, as we discussed in L#234, page 5, this shares similar intuition from the widely adopted submodularity on set functions.\n\nAlso, we conducted **more experiments** on more datasets and utility functions.\nIt can be found in https://anonymous.4open.science/r/temp-03AC.\nIn new experiments, we first tested DDMC assumption on two more datasets: AdvBench (standard jailbreak benchmark) and just-eval-instruct that contains many prompts of various tasks.\nFurther, we verified DDMC assumptions beyond our linear utility function, with L1/3/4 distance.\nOverall, we observe a tendency that appending more common tokens decreases the utility gaps.\nWe will certainly add these in the revision.\n\n# More questions\n\n## 1. How does the regret bound of EOFUL..\n\nAs noted, frameworks like ARGS assume access to the reward function, so we cannot compare regrets.\n\n## 2. Why is there no comparison…\n\nWe hope responses above clarified your question.\n\n## 3. What is the computational efficiency...\n\nSince PPO-based RLHF retrains the LLM itself, our approach has significantly less computational burden. Further, it is unfair to compare the computational complexity of our works to ARGS, as ARGS assumes access to the external reward function whereas our approaches need to learn it.\n\n## 4. A well-known disadvantage of…\n\nGood point! Similar to LinUCB, one can apply the standard lazy update to reduce the computational burden so that it needs to recompute the estimator only $O(log T)$ times. We refer to Section 5.1 in Abbasi-Yadkori et al., 2011 for more details. We will add a brief discussion on it in our revision.", "claims_and_evidence": "The paper presents Tokenized Bandit frameworks (TLB and TMAB) for LLM decoding and alignment and supports its claims using rigorous theoretical analysis and empirical validation. Most of the claims are derived with clear mathematical proofs and sublinear regret bounds for EOFUL and GreedyETC. The empirical validation on Llama-3-8B with TruthfulQA and HH-RLHF provides some evidence that the DDMC assumption holds in practice. However, some strong baselines are not considered. For instance, ARGS (Alignment as Reward-Guided Search, Khanov et al., 2024), as a baseline, is missing. Comparing against ARGS, especially in terms of efficiency, performance, and scalability, could help clarify why a bandit framework is preferable to a search-based alignment method like ARGS.", "methods_and_evaluation_criteria": "The Tokenized Bandit framework (TLB & TMAB) is well-motivated for LLM decoding and alignment, as sequential token selection naturally fits a multi-armed bandit or linear bandit formulation. The authors provide a strong theoretical foundation and validate their approach on Llama-3-8B using TruthfulQA and HH-RLHF, which are reasonable benchmarks for alignment and factual correctness. However, the evaluation has some gaps: (1) No comparison to strong decoding-time alignment baselines like ARGS (Khanov et al., 2024), making it unclear if bandit-based alignment is superior to search-based alignment; (2) Limited diversity in test datasets, as evaluations focus only on two benchmarks, without testing broader alignment scenarios (e.g., safety-critical or adversarial tasks). While the evaluation is a good start, expanding it to stronger baselines and additional datasets would strengthen confidence in the method’s effectiveness.", "theoretical_claims": "The paper provides rigorous mathematical proofs supporting its core claims, particularly: (1) Learning is impossible without structural assumptions in tokenized bandit settings, (2) Greedy decoding can be optimal under the Diminishing Distance with More Commons (DDMC) assumption, and (3) EOFUL (for TLB) and GreedyETC (for TMAB) achieve sublinear regret bounds. The derivations are logically consistent, following standard bandit theory techniques, and the regret bounds align with known results in linear and multi-armed bandits. However, the DDMC assumption lacks a formal justification for why it holds across different LLMs and datasets, relying primarily on empirical verification. To strengthen the theoretical foundation, the paper should provide (1) a more formal justification of DDMC beyond empirical trends.", "experimental_designs_or_analyses": "The experimental design provides some validation of the Tokenized Bandit framework for factual correctness and preference alignment. However, the evaluation has significant gaps. While the regret bounds are well-structured, there is no explicit comparison to search-based alignment methods (e.g., ARGS) in terms of convergence speed or sample complexity. Additionally, the approach is not tested on more diverse datasets that could include safety-critical, adversarial, or instruction-following tasks to assess broader generalization. While the bandit framework is promising, these missing elements make it difficult to fully validate the empirical claims. Expanding the benchmarks and including stronger baselines would strengthen the study’s conclusions.", "relation_to_broader_scientific_literature": "This paper builds on foundational work in multi-armed and linear bandits, adapting these frameworks to LLM decoding and alignment. It builds upon prior work on bandits to introduce Tokenized Bandit models (TLB and TMAB) that incorporate token-wise sequential decision-making. The Diminishing Distance with More Commons (DDMC) assumption is conceptually similar to structural constraints in regret minimization (Abbasi-Yadkori et al., 2011) but is newly applied to LLM decoding. Additionally, the paper is highly relevant to decoding-time alignment approaches such as ARGS (Khanov et al., 2024), which formulates alignment as a reward-guided search problem rather than reinforcement learning. However, the authors do not compare their work directly to search-based or reinforcement learning-based decoding strategies (e.g., PPO in RLHF), leaving a gap in understanding how bandit-based decoding fares against other methods.", "essential_references_not_discussed": "Khanov et al. \"ARGS: Alignment as Reward-Guided Search\" and follow-up works.", "other_strengths_and_weaknesses": "Please refer to the above sections.", "other_comments_or_suggestions": "Overall, the paper is well-written. The problem is also well-defined and the proposed solution is clear.", "questions_for_authors": "Some of the questions that came to mind as I read the paper are detailed below. I am asking these questions in order to clarify/justify theoretical generalizability, empirical competitiveness, computational efficiency, and real-world robustness of Tokenized Bandits for LLM decoding and alignment.\n\n1. How does the regret bound of EOFUL and GreedyETC compare to search-based decoding strategies like ARGS (Khanov et al., 2024) and variants?\n2. Why is there no comparison to ARGS or reinforcement learning-based decoding strategies (e.g., PPO in RLHF)?\n3. What is the computational efficiency of Tokenized Bandits compared to PPO-based RLHF and ARGS? \n4. A well-known disadvantage for LinUCB is its computational complexity since it needs to incorporate historical information when initializing observation matrices and conducting matrix multiplication. Will the proposed approach inherits similar shortcomings?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IN2TqL3E5x", "reviewer": "Reviewer_9ZG7", "summary": "The paper introduces the Tokenized Linear Bandit (TLB) and Tokenized Multi-Armed Bandit (TMAB), which are variants of the classical linear and stochastic multi-armed bandit problems, inspired by the decoding and alignment processes in large language models (LLMs). In these problems, a user submits a query (context) in each round, and the decision maker (DM) sequentially selects tokens irrevocably from a predefined token set. Once the sequence is completed, the DM receives a random utility from the user, whose expected value is determined by a sequence function that maps the chosen token sequence and the query to a non-negative real number.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your detailed feedback and insightful comments, in particular we are glad that the reviewer enjoys our problem and approach.\n\nAs the reviewer suggested, we agree that numerical results demonstrating our algorithm’s performance would greatly improve our paper - we appreciate your suggestion.\nAs such, we **evaluate our algorithm’s regret** with several benchmarks that we thought to be reasonable in our scenario.\nWe refer the reviewer to see: https://anonymous.4open.science/r/temp-03AC, directory /Regrets.\n\nWe evaluate algorithms based on **TLB model and LLM alignment scenario** presented in Section 5.1 with $\\theta = [1/2,1/2,\\ldots, 1/2]$, maximal length of sequence = 30, rounds 5000, and $\\gamma = 0.8$.\nWe compare our EOFUL along with a benchmark and two other algorithms: (i) theoretical regret upper bound (under-scaled by 0.1 to make the plot more visible), (ii) WrongTheta that uses wrongly estimated $\\theta = [-1/2,-1/2,\\ldots,-1/2]$ and greedily decoding based on the weighted score, and (iii) Misaligned greedy that only greedily decodes with respect to the LLM’s probability.\nAlso, for practicality and computational efficiency, we only consider the top 15 tokens given by LLM in deciding the next token.\nAs observed in the figure, we observe that EOFUL effectively achieves sublinear regret.\nGiven that the regret upper bound is under-scaled,the actual performance may be much better than that of the theoretical guarantees.", "claims_and_evidence": "This paper make interesting claims which are suppoted by sufficient and strong evidences.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-reasoned. However, numerical results demonstrating their performance in real applications were not provided.", "theoretical_claims": "The theoretical claims appear to be correct. Specifically, this work introduces a significant framework that effectively models the LLM decoding and alignment problem. While the problem itself is inherently challenging, it is addressed through an elegant assumption of DDMC, which provides a clever and practical solution.", "experimental_designs_or_analyses": "Experiments are conducted to validate the empirical performance of the proposed methods.", "relation_to_broader_scientific_literature": "The theoretical framework of tokenized bandits is novel, and the approach used to solve the problem is innovative.", "essential_references_not_discussed": "Essential references have been properly discussed.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1kv89U3oT0", "reviewer": "Reviewer_amHn", "summary": "The paper introduces two new bandit variants, the tokenized linear bandit (TLB) and tokenized multi-armed bandit (TMAB), which involve sequentially constructing a sequence of tokens to optimize a (random) utility function of a user, given a query. They introduce the DDMC assumption on token sequences and construct learning algorithms for both TLB and TMAB for which they prove regret bounds. They show that LLM alignment may be viewed as an application of their TLB method for specific functional forms of utility misalignment between the model and the user. They also aim to empirically validate the DDMC assumption in the context of LLMs and connect it to the empirically observed high performance of greedy decoding in LLMs.\n\n## update after rebuttal\nFollowing the author response, I am *provisionally* increasing my score to a 3, subject to discussion with the other reviewers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your detailed feedback and insightful comments.\nWe will first make a general remark on our main contribution, and then answer the reviewer’s comments/questions one by one.\n\n# General remark on main contribution\n\nWe kindly ask the reviewer to see **our response (General remark on main contribution) to Reviewer 3dxc**.\n\n# Comments on DDMC\nWe hope our general remark above partly resolves the reviewer’s concern regarding the necessity of the assumption.\n\nAs the reviewer pointed out, we don’t believe that DDMC would universally hold for every real-world scenario.\nHowever, we empirically find an overall tendency to do so, i.e., having a decreasing utility difference (or decreasing distance between embedding vectors) as we append more common tokens.\n\nOur DDMC assumption, or at least a relaxed version of it, is quite reasonable intuitively. For instance, in extreme cases, if we append a lot of common tokens in each of two sequences of the same length, one could expect that the user will realize a similar experience of reading them. Further, as we discussed in L#234, page 5, this shares similar intuition from the widely adopted submodularity on set functions.\n\nAlso, we conducted a few **more experiments** to cement our assumption’s effectiveness. \nIt can be found in https://anonymous.4open.science/r/temp-03AC.\nIn new experiments, we first tested DDMC assumption on two more datasets: AdvBench (standard jailbreak benchmark) and just-eval-instruct that contains many prompts of various tasks. Further, we verified DDMC assumptions beyond our linear utility function. We considered three more functions: L1/3/4 distance. In all these scenarios, we observe a tendency that appending more common tokens decreases the utility gaps. We will certainly add these experiments in the revision.\n\n\n# Under what conditions can we expect DDMC to hold...\n\nThanks for a great question.\nAlthough it’s difficult to classify certain characteristics of tasks that DDMC would more strongly hold, through our empirical verifications (including **newly added ones**), we have identified that in several commonly-used LLM alignment datasets, DDMC assumption seems to hold, so we believe this assumption is reasonably realistic (though certainly not universal). \n\n# Comments on linear realizability\n\nWe hope our general remark above partly resolves the reviewer’s concern regarding the necessity of the assumption.\n\nAs we noted in our manuscript (paragraph below Assumption 3.2), several papers in the literature for LLM alignment often assume linear realizability for theoretical sake, e.g., we refer to Cen et al. 2024 and Yang et al. 2024 that exactly assume the linear realizability for theoretical purpose in LLMs\n\nFrom practical perspectives, **several** recent empirical studies provide evidence that linear realizability is reasonable in many scenarios.\nFirst, Zou et al. 2025 confirmed that concepts like truthfulness or ethics could be extracted via linear transformation over LLM’s representation (embedding) with fixed weights, which is also validated through extensive experiments.\nWang et al. 2024 tackled decoding-time alignment solely by using linear transformation over the LLM’s representation with carefully chosen weights to handle the harmlessness alignment.\nKong et al. 2024 considered a control theoretical perspective for LLM alignment by adding a linear perturbation on original logit vector.\n\nAnother relevant literature is the recently studied ‘linear representation hypothesis’, which gives a mechanistic interpretation on how LLMs embed concepts in representation space, as noted in our footnote 14.\nIn particular, the hypothesis itself states (as hinted by its name) that concepts / semantics (e.g., politic ideology, geography, and temporal knowledge) are embedded in a linear manner given some proper weights: we refer to Kim et al. 2025, Jiang et al. 2024, Park et al. 2025, Gurnee and Tegmark 2024 for more details.\n\nFinally, we believe extending the linear function to nonlinear function in a similar vein to the extension of linear contextual bandit to nonlinear contextual bandit via kernelization by Valko et al. 2013 would be a great direction for future works.\n\n\n\n\n\n# References\n\n* Cen et al. 2024, Unified Approach to Online and Offline RLHF\n* Yang et al. 2024, Asymptotics of Language Model Alignment.\n* Zhou et al. 2025, Representation Engineering: A Top-Down Approach to AI Transparency\n* Wang et al. 2024, InferAligner: Inference-Time Alignment for Harmlessness through Cross-Model Guidance\n* Kong et al. 2024, Aligning Large Language Models with Representation Editing: A Control Perspective\n* Kim et al. 2025, Linear Representations of Political Perspective Emerge in Large Language Models\n* Jiang et al. 2024, On the Origins of Linear Representations in Large Language Models\n* Park et al. 2025, The Geometry of Categorical and Hierarchical Concepts in Large Language Models\n* Gurnee and Tegmark 2024, Language Models Represent Space and Time", "claims_and_evidence": "I have not been able to evaluate all the proofs, but claims seem generally well supported, with a few exceptions.\n\nMy first worry is regarding the suggestion that this method may be useful for alignment. It seems that the TLB problems only applies insofar as the linear realizability assumption holds, as well as the particular functional form of misalignment suggested in section 5.1. The paper contains no empirical verification of this in typical alignment settings.\n\nMy second worry relates to the DDMC assumption. The validity of the assumption is evaluated assuming the utility is parameterized to fit the form of distance functions d1 or d2 (sec 5.3), and in those cases it broadly (though not always, inspecting figure 1) seems to hold on Llama-3-8B-Instruct for the two datasets tested. However, this is not evaluated on more free form utility functions, such as those provided by human subjects or even reward models. In particular, it seems this assumption cannot hold in general (as shown in figure 1), and in particular it seems to omit some rather standard use cases of LLMs.\n\nFor instance, take the following token sequences:\n\nq = \"Suggest a cafe I should visit in Vancouver.\",\ny = \"Visit Cafe Alice\",\nz = \"Visit Cafe Bob\",\n\\tau = \"by\"\n\ny and z could plausibly be equal token length (this depends on the tokenizer, which is not part of the assumption). If 'Cafe Bobby' exists and 'Cafe Aliceby' does not exist, then utility z + tau is presumably larger than y + tau, even thought utility of y and z were presumably equal. Thus, the gap has increased.\n\nIn the case that u = p, the DDMC assumption implies that the difference in probability under the LLM of y and z is smaller than that of y + tau and z + tau, which seems unreasonable if 'Cafe Bobby' exist and 'Cafe Aliceby' does not, but 'Cafe Bob' and 'Cafe Alice' both exist. This suggests that p does not fit theorem 5.1, so the claimed connection to the success of greedy LLM decoding seems somewhat unreliable.\n\nNote that this is dependent on the state of the world, which is not part of the assumption statement for obvious reasons, but without it the statement makes potentially unsupported assumptions about the world. The general issue is that tokens mean very different things in different contexts, so sharing more tokens does not guarantee that utilities are closer in many cases.", "methods_and_evaluation_criteria": "Testing the DDMC assumption with Llama-3-8B-Instruct on the TruthfulQA and HH-RLHF datasets is sensible. However, as mentioned above, the utility here is assumed to be parameterized in two very specific ways, and it would be instructive to see whether DDMC holds empirically with more free form utility functions, such as those provided by human subjects or even reward models.", "theoretical_claims": "I have not been able to evaluate the proofs.", "experimental_designs_or_analyses": "See previous fields.", "relation_to_broader_scientific_literature": "The paper proposes two new discrete bandit problems with applications to LLM decoding and alignment. I am unaware of earlier works on this intersection, but I am not too familiar with that literature. The paper aims to provide evidence towards a mechanism for the success of greedy decoding, which is a topic of current interest given the ubiquity of LLMs. For the same reason, results on decoding time alignment are relevant.", "essential_references_not_discussed": "Nothing particular that I am aware of. There is a wide literature on decoding time alignment that was not discussed (e.g., steering methods such as Panickssery et al. 2023, \"Steering Llama 2 via Contrastive Activation Addition\" and Todd et al. 2023, \"Function Vectors in Large Language Models\"), but it does not relate to the methods proposed in this work too directly.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is generally well written and structured.\n- The paper contains many notes pointing to related work and exact implications of assumptions and statements. The authors seem very aware of the mathematical literature, although this is hard for me to reliably evaluate.\n\nWeaknesses:\n- Primarily the degree to which the assumptions (DDMC and to a lesser extent linear realizability of utility functions) hold in practice and the absence of (compelling) experimental evidence for these assumptions and the application of the work to alignment.", "other_comments_or_suggestions": "Typo in proposition 3.3: \"assumptinon\"", "questions_for_authors": "1. Under what conditions can we expect DDMC to hold, and is this realistic for typical LLM use cases?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iPaMaRHtg6", "reviewer": "Reviewer_dJHj", "summary": "The paper \"Tokenized Bandit for LLM Decoding and Alignment\" introduces a novel Tokenized Bandit (TB) framework to address LLM decoding and alignment challenges. It models LLM decoding as a sequential decision-making problem, using multi-armed bandit (MAB) and linear bandit (LB) techniques to optimize token selection.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your insightful comments.\nWe will first make a general remark on our main contribution, answer major concerns and then remaining questions.\n\n# General remark on main contribution / assumption \n\nWe kindly refer to our **response (General remark on main contribution) to Reviewer 3dxc**.\n\n# Major concerns\n\n## 1. Individual token reward\nWe note that individual tokens **do not** have any rewards, but only a token sequence does.\nOur algorithm smartly learns an efficient path to the optimal sequence only by observing seqeunce-wise rewards, thanks to the  novel DDMC assumption.\n\n## 2. Sequential dependency\n\nOur settings intrinsically **considers sequential dependency** given what has been chosen previously, since these histories are encoded in the sequence function’s value and also final embedding vectors.\nIndeed, our offline benchmark selects a sequence with the largest utility that does account for the sequential dependency of the LLM, and the regret is computed with respect to this.\nThus, our goal is to learn efficient ‘sequence’ not individual token, and the previous decisions affect onwards.\n\nWe also remark that greedy decoding often provides repetitive outputs, especially for small language models few years ago, but this phenomenon is mitigated in the most recent LLMs (Song et al. 2024).\nPractically, one can replace greedy decoding with other decoding schemes in EOFUL algorithm based on its empirical performance, which will be a great direction for future work.\n\n## 3. Comparison with other benchmarks - PPO, DPO, Best of N\nFor PPO and DPO, we refer to **our response (Comparison to ARGS / PPO in RLHF) to Reviewer 3dxc**.\n\nFor the Best-of-N alignment, one could consider Best-of-N algorithm that knows the latent reward function in hindsight in our regret definition as an offline benchmark (note that this is **not an online learning**).\nHowever, this will be weaker benchmark than offline greedy decoding under DDMC assumption.\nPractically, one can replace the greedy decoding with best-of-n (or other decoding schemes) in EOFUL algorithm.\n\n# Further comments\n\n## 1. Perfect embeddings\n\nOnce the regret benchmark and a learning algorithm share the same embedding vector (even perturbed), all our theoretical results would naturally follow.\nAlso, standard LLMs operate based on the embedding vectors and corresponding logit probabilities, so the algorithm can naturally access embedding vectors, which is how we obtain embeddings in our evaluations.\n\n## 2. State-dependency\nOur contextual bandit framework captures state-dependency, and we hope our response (Sequential dependency) resolves your question.\n\n## 3. Comments on DDMC with more datasets\nWe kindly refer to our response (More justification on DDMC) to Reviewer 3dxc.\n\n## 4. Linear realizability assumption\nWe kindly refer to our response (Comments on linear realizability) to Reviewer amHn.\n\n## 5. ETC approach in TMAB\nOur study on TMAB is to provide more comprehensive theoretical findings for ‘tokenized’ variants of the bandit problems, whereas we believe our TLB setting is more practical for LLM applications.\n\nThat said, the fact that greedy ETC achieves the regret sublinear in $T$ and linear in the length of the sequence in TMAB implies that the number of samples required to learn efficient sequence is not significant.\nWe note that a naive algorithm requires exploring every possible sequence inducing regrets exponential in the length of the sequence.\n\n## 6. Using combinatorial bandit or RL\nAs our objective function is a sequence function not a set function, standard combinatorial bandits can't be directly applied.\nFurther, one may generalize TLB to Markov decision process as does in RL, but this would be waymore complicated to analyze and it is not certain if it would exhibit Markovian property.\n\n## 7. Explicit formulation of the embedding function \nWe hope our response (Perfect embedding) resolves your question.\nAlso, a noisy reward is observed after the LLM submits the entire outputs y, similar to the standard bandit framework, e.g., see L#146 (right column) on page 3.\n\n## 8. Utility function\nWe hope our response (Linear realizability assumption) resolves your question.\nAlso, the weighted average of original LLM’s probability and another external reward module’s score is often considered for LLM alignment, e.g., see Khanov et al., 2024.\n\n## 9. Convergence on the experiment results\nOur interest is not on the convergence, but whether it exhibits a decreasing trend with respect to more common tokens as the definition of DDMC.\n\n## 10. Minor error: (1) In algorithm 2, there is a mistake in the selection…\nThanks for catching - will edit it. Also yes, $y^{(i)}$ is $i$-th token (single token) - we abuse boldface to make it consistent with $y^{(i:j)}$ which is a sequence.\n\n\n# References\n\n* Song et al. 2024, The Good, The Bad, and The Greedy: Evaluation of LLMs Should Not Ignore Non-Determinism\n* Khanov et al. 2024, Alignment as Reward-Guided Search", "claims_and_evidence": "The DDMC assumption is not universally validated. It is tested on limited datasets, and it's unclear whether it holds for other LLM applications (e.g., coding, mathematical reasoning).", "methods_and_evaluation_criteria": "Expanding evaluations to more datasets, and efficiency benchmarks  (e.g., PPO, Best-of-N decoding) would significantly strengthen the paper’s claims.", "theoretical_claims": "The regret bound relies on DDMC, but the paper does not rigorously prove whether DDMC always holds across different tasks.\nAnd the proof assumes perfect token embeddings $e(x_t,y)$, but in reality, embeddings could introduce variance that affects regret analysis.", "experimental_designs_or_analyses": "1. While the paper claims that bandit-based decoding is superior, it does not compare against reinforcement learning approaches, which are common in LLM alignment. For example, Best-of-N Sampling and Direct Preference Optimization (DPO) are stronger baselines for alignment than standard greedy decoding.\n\n2. I worry about the practical applicability of the approach: Tokenized Bandit selects tokens step by step but assigns rewards at the sequence level, making it hard to attribute rewards to individual tokens. Moreover, Bandit algorithms assume independent token selection, while **LLM generation has strong sequential dependencies**, raising doubts about their effectiveness in optimizing decoding. In practice, greedy decoding is rarely used due to repetitive outputs, possibly caused by the independence assumption in token selection.", "relation_to_broader_scientific_literature": "The paper should expand the discussion on prior search-based decoding, RLHF baselines, and uncertainty-aware text generation to strengthen positioning in the broader literature.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Other Strengths:\n1. The paper also introduces the Diminishing Distance with More Commons (DDMC) hypothesis, which states that if two sequences share the same suffix, their reward difference decreases. This assumption significantly reduces computational complexity and justifies the effectiveness of greedy decoding.\n2. The paper theoretically proves that greedy decoding can be globally optimal under the DDMC hypothesis.\n\nOther Weaknesses:\n1. The paper assumes that the reward function 𝑢ₜ(𝑥ₜ, 𝑦) follows a linear structure:  $𝑢ₜ(𝑥ₜ, 𝑦) = ⟨𝜃, 𝑒(𝑥ₜ, 𝑦)⟩$. However, in real-world LLM tasks, user preferences 𝑓(𝑥ₜ, 𝑦) may exhibit highly nonlinear behavior.\n2. In the Tokenized Multi-Armed Bandit (TMAB) setting, the reward can only be observed after generating the entire sequence, making it difficult to directly determine the contribution of individual tokens. To address this, the authors adopt the Explore-Then-Commit (ETC) approach, estimating the utility of each token through multiple full-sequence samples. However, I am concerned that the number of required full sequences might be excessively large, potentially leading to a high computational cost.", "other_comments_or_suggestions": "1. I could not find the explicit formulation of the embedding function e(x_t, y). Does it lack an explicit expression? Or can you explain how to evaluate whether the token sequence y is suitable to query x_t, i.e., where is the reward r come from?\n\n2. Regarding the definition of the utility function, could the authors provide some intuitive examples to justify its formulation? The current definition seems somewhat unconventional.\n\n3. Concerning the experimental results, in the second figure, the convergence result appears to be around 20, which does not seem very satisfactory. I suspect that a significant factor contributing to this issue might be the formulation or definition of the utility function. Could the authors elaborate on whether any alternative formulations were considered?\n\n4. Minor error: (1) In algorithm 2, there is a mistake in the selection of \\tau^{\\star}. (2)For y^{(i)}, it is the i-th token in sequence \\bm{y}, but you use boldface on it, is it a typo?", "questions_for_authors": "1. In the exploration phase of GreedyETC (Greedy Exploration-Then-Commit), how is the reward assigned to individual tokens when a full sequence must be generated before obtaining a reward? Since the reward is only observed at the sequence level, how do you estimate the contribution of each token during exploration? \n\n2. The paper formulates LLM decoding as a Tokenized Bandit problem, where each token is selected sequentially, and the final reward is only observed after a complete sequence is generated. Given this delayed reward structure, how does the proposed bandit-based approach achieve global sequence-level optimization? Would a Combinatorial Bandit or Reinforcement Learning (RL) approach be more appropriate for handling dependencies between tokens? \n\n3. It is unclear how a bandit approach, which is typically state-independent, can effectively guide token selection when the reward is assigned only after the entire token sequence $y_t$  is generated. Since token selection is not a set-based decision but rather a structured ordered sequence, different sequences lead to vastly different rewards, and the contribution of each token is not independently determinable. How does the bandit framework, which does not model dependencies between token choices, provide meaningful guidance in this setting?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Three of four reviews vote to “weak accept” this paper while the fourth review is a weak reject. I looked at all the four reviews; some of the reviews raise great points in support of the paper. For instance, the reviews recognize the value of the paper in its motivation and the solid theoretical contributions presented. The experimental analysis provides useful validation, although reviews also note significant gaps and areas for improvement. 

Reviewer djHj (who votes “weak reject”) raised three main concerns: (1) strong sequential dependencies of LLM generation (2) Individual reward of token and (3) comparison with other benchmarks. The authors respond to these concerns in rebuttal and answer most of the questions in two round of discussion. After initial engagement, the reviewer did not confirm if they felt satisfied, dissatisfied with all of the rebuttal provided.
Decision: Accept (poster)