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
• 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

• Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata [57.27430779838529]
  We present an exact and tractable algorithm for sampling from the constrained mean-field posterior under any constraint expressible as a finite automaton.<n>The approach guarantees constraint satisfaction by construction, supports both greedy and sampling-based decoding, and is compatible with parallel and block-wise decoding.<n> Empirical evaluations on Dream-7B and LLaDA-8 show substantial accuracy gains across various tasks.
  arXiv  Detail & Related papers  (2026-07-08T05:48:57Z)
• Contextual Slate GLM Bandits with Limited Adaptivity [1.8843687952462742]
  We investigate the contextual slate bandit problem with generalized linear rewards under limited adaptivity.<n>We propose algorithms under two limited-adaptivity settings: Batched and (b) Rarely-Switching.<n>Our algorithms are computationally efficient, requiring only $textpoly(N)$ time per round despite $2(N)$ possible slates.
  arXiv  Detail & Related papers  (2026-06-30T10:24:37Z)
• Adaptation to Intrinsic Dependence in Diffusion Language Models [5.185131234265025]
  Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) approaches.<n>We introduce a distribution-agnostic unmasking schedule for DLMs that adapts to the (unknown) dependence structure of the target data distribution.<n>Our results significantly improve upon prior convergence theories and yield substantial sampling acceleration for low-complexity distributions.
  arXiv  Detail & Related papers  (2026-02-23T18:41:34Z)
• No Compute Left Behind: Rethinking Reasoning and Sampling with Masked Diffusion Models [42.158430350703505]
  Masked diffusion language models are trained to in-fill positions in randomly masked sequences.<n>We propose reasoning-as-infilling and multi-token entropy decoding.<n>Our work demonstrates that the training and compute used by MDLMs unlock many new inference and post-training methods.
  arXiv  Detail & Related papers  (2025-10-22T19:41:27Z)
• 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)
• Beyond Next Token Probabilities: Learnable, Fast Detection of Hallucinations and Data Contamination on LLM Output Distributions [60.43398881149664]
  We introduce LOS-Net, a lightweight attention-based architecture trained on an efficient encoding of the LLM Output Signature.<n>It achieves superior performance across diverse benchmarks and LLMs, while maintaining extremely low detection latency.
  arXiv  Detail & Related papers  (2025-03-18T09:04:37Z)
• 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.