Tokenized Bandit for LLM Decoding and Alignment

• URL: http://arxiv.org/abs/2506.07276v1
• Date: Sun, 08 Jun 2025 20:32:08 GMT
• Title: Tokenized Bandit for LLM Decoding and Alignment
• Authors: Suho Shin, Chenghao Yang, Haifeng Xu, Mohammad T. Hajiaghayi,
• Abstract summary: We introduce the tokenized linear bandit (TLB) and multi-armed bandit (TMAB)<n>We propose algorithms with regret $tildeO(LsqrtT)$ and $tildeO(LsqrtT2/3)$ for TLB and TMAB, respectively.
• Score: 32.22367277030496
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Foundations of Top-$k$ Decoding For Language Models [19.73575905188064]
  We develop a theoretical framework that both explains and generalizes top-$k$ decoding.<n>We show how to optimize it efficiently for a large class of divergences.
  arXiv  Detail & Related papers  (2025-05-25T23:46:34Z)
• Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling [90.86991492288487]
  evaluating constraint on every token can be prohibitively expensive.<n> LCD can distort the global distribution over strings, sampling tokens based only on local information.<n>We show that our approach is superior to state-of-the-art baselines.
  arXiv  Detail & Related papers  (2025-04-07T18:30:18Z)
• Let the Code LLM Edit Itself When You Edit the Code [50.46536185784169]
  underlinetextbfPositional textbfIntegrity textbfEncoding (PIE)<n>PIE reduces computational overhead by over 85% compared to the standard full recomputation approach.<n>Results demonstrate that PIE reduces computational overhead by over 85% compared to the standard full recomputation approach.
  arXiv  Detail & Related papers  (2024-07-03T14:34:03Z)
• Order-Independence Without Fine Tuning [18.020492646988746]
  We present Set-Based Prompting, a technique that guarantees the output of an LLM will not have order dependence on a specified set of sub-sequences.<n>Despite our inputs being out of distribution, the impact on expected accuracy is small, where the expectation is over the order of uniformly chosen shuffling of the candidate responses.
  arXiv  Detail & Related papers  (2024-06-04T16:09:13Z)
• Can Large Language Models Play Games? A Case Study of A Self-Play Approach [61.15761840203145]
  Large Language Models (LLMs) harness extensive data from the Internet, storing a broad spectrum of prior knowledge. Monte-Carlo Tree Search (MCTS) is a search algorithm that provides reliable decision-making solutions. This work introduces an innovative approach that bolsters LLMs with MCTS self-play to efficiently resolve turn-based zero-sum games.
  arXiv  Detail & Related papers  (2024-03-08T19:16:29Z)
• Recursive Speculative Decoding: Accelerating LLM Inference via Sampling Without Replacement [11.91629418177851]
  Speculative decoding is an inference-accel method for large language models. Recent works have advanced this method by establishing a draft-token tree. We present Recursive Speculative Decoding (RSD), a novel tree-based method that samples draft tokens without replacement.
  arXiv  Detail & Related papers  (2024-02-21T22:57:49Z)
• Provably Efficient High-Dimensional Bandit Learning with Batched Feedbacks [93.00280593719513]
  We study high-dimensional multi-armed contextual bandits with batched feedback where the $T$ steps of online interactions are divided into $L$ batches. In specific, each batch collects data according to a policy that depends on previous batches and the rewards are revealed only at the end of the batch. Our algorithm achieves regret bounds comparable to those in fully sequential setting with only $mathcalO( log T)$ batches.
  arXiv  Detail & Related papers  (2023-11-22T06:06:54Z)
• Linear Bandit Algorithms with Sublinear Time Complexity [67.21046514005029]
  We propose to accelerate existing linear bandit algorithms to achieve per-step time complexity sublinear in the number of arms $K$. We show that our proposed algorithms can achieve $O(K1-alpha(T))$ per-step complexity for some $alpha(T) > 0$ and $widetilde O(stT)$ regret, where $T$ is the time horizon.
  arXiv  Detail & Related papers  (2021-03-03T22:42:15Z)
• Stochastic Bandits with Linear Constraints [69.757694218456]
  We study a constrained contextual linear bandit setting, where the goal of the agent is to produce a sequence of policies. We propose an upper-confidence bound algorithm for this problem, called optimistic pessimistic linear bandit (OPLB)
  arXiv  Detail & Related papers  (2020-06-17T22:32:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.