Compute Where it Counts: Self Optimizing Language Models

• URL: http://arxiv.org/abs/2605.10875v1
• Date: Mon, 11 May 2026 17:27:15 GMT
• Title: Compute Where it Counts: Self Optimizing Language Models
• Abstract summary: We study dynamic budget allocation for autoregressive decoding.<n>We train the policy with group-relative policy optimization on teacher-forced episodes.<n>Our reward trades off language-model quality against soft penalties that encourage episode-average budget usage.
• Score: 10.058821474955177
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Efficient LLM inference research has largely focused on reducing the cost of each decoding step (e.g., using quantization, pruning, or sparse attention), typically applying a uniform computation budget to every generated token. In practice, token difficulty varies widely, so static compression can over-compute on easy steps and under-compute on hard ones. We study dynamic budget allocation for autoregressive decoding: learning how much computation to spend per token from within a single model. Self-Optimizing Language Models (SOL) pair a frozen LLM with a lightweight policy network that reads the LLM hidden state and selects a discrete efficiency action at each decode step. Actions can jointly control (i) token-level attention sparsity, (ii) structured activation pruning in the MLP, and (iii) activation quantization bit-width, while leaving the base model weights unchanged. We train the policy with group-relative policy optimization on teacher-forced episodes: the token sequence is fixed, while we sample multiple compute schedules (i.e., "counterfactual" schedules that vary only the efficiency actions for the same token path) and compare their likelihoods under the same supervision. Our reward trades off language-model quality against soft penalties that encourage episode-average budget usage to match a requested target. Across model variants and compute regimes, SOL improves quality at matched budget over static allocation and strong random schedule search, offering a complementary axis for inference-efficiency optimization. SOL discovers a better quality-efficiency pareto-front across all our experiments and improves MMLU accuracy by up to 7.3% over uniform budget allocation strategies.

Related papers

• Look Less, Think Faster: Joint Token-Compute Adaptation for Multimodal LLMs [16.64366862436724]
  Multimodal Large Language Models (MLLMs) have recently demonstrated strong performance across vision-language tasks.<n>Their high inference cost, arising from both the large number of input visual tokens and the heavy computation of the large language model (LLM), remains a key barrier to practical deployment.<n>We propose SmartVL, a unified adaptive inference framework that jointly controls vision token number and model compute capability.
  arXiv  Detail & Related papers  (2026-07-22T16:43:27Z)
• Cost-Aware Learning [72.31444819326795]
  We consider the problem of Cost-Aware Learning, where sampling different component functions incur different costs.<n>We propose the Cost-Aware Descent algorithm for convex functions, and derive its cost complexity to attain an error of $$.<n>We introduce Cost-Aware GRPO, an algorithm designed to reduce the cost of policy optimization while preserving performance.
  arXiv  Detail & Related papers  (2026-04-30T15:39:09Z)
• Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization [18.737087162461563]
  Test-time compute scaling has become a powerful lever for improving large language model performance.<n>Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore.<n>We formalize this as a constrained optimization problem (maximize expected accuracy subject to an average compute budget) and solve it with a two-stage Solve-then-Learn pipeline.
  arXiv  Detail & Related papers  (2026-04-16T10:39:22Z)
• IsoCompute Playbook: Optimally Scaling Sampling Compute for LLM RL [75.28528082188787]
  We study the compute-optimal allocation of sampling compute for on-policy reinforcement learning methods in large language models (LLMs)<n>We find that the compute-optimal number of parallel rollouts per problem increases predictably with compute budget and then saturates.<n>Our results recast RL scaling laws as prescriptive allocation rules and provide practical guidance for compute-efficient LLM RL post-training.
  arXiv  Detail & Related papers  (2026-03-12T16:49:21Z)
• Learning Adaptive LLM Decoding [6.643962667713069]
  We learn adaptive decoding policies that dynamically select sampling strategies at inference time, conditioned on available compute resources.<n>We introduce lightweight decoding adapters trained with reinforcement learning and verifiable terminal rewards.<n>Experiments show that the learned adapters improve the accuracy-budget tradeoff.
  arXiv  Detail & Related papers  (2026-03-10T01:15:26Z)
• $\nabla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• Informed Routing in LLMs: Smarter Token-Level Computation for Faster Inference [7.690958366125321]
  This paper introduces informed routing, a new paradigm that proactively addresses these issues.<n>We propose the Lightweight Feature Forecaster (LFF), a small predictive module that estimates a unit's output before routing decisions are made.<n>Experiments on both language modeling and reasoning tasks show that informed routing achieves state-of-the-art efficiency-performance trade-offs.
  arXiv  Detail & Related papers  (2025-10-10T09:59:36Z)
• Sparsity Forcing: Reinforcing Token Sparsity of MLLMs [40.93786579652003]
  We explicitly reinforce token sparsity in well-posed multimodal large language models (MLLMs) through a simple RL-based post-training framework named textitSparsity Forcing.<n>Our method explores the efficiency-accuracy trade-off by running multiple rollouts with different token budgets, where both efficiency (token reduction ratio) and performance (answer correctness) are formulated as joint rewards.
  arXiv  Detail & Related papers  (2025-04-23T01:45:55Z)
• 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)
• Progressive Mixed-Precision Decoding for Efficient LLM Inference [49.05448842542558]
  We introduce Progressive Mixed-Precision Decoding (PMPD) to address the memory-boundedness of decoding.<n>PMPD achieves 1.4$-$12.2$times$ speedup in matrix-vector multiplications over fp16 models.<n>Our approach delivers a throughput gain of 3.8$-$8.0$times$ over fp16 models and up to 1.54$times$ over uniform quantization approaches.
  arXiv  Detail & Related papers  (2024-10-17T11:46:33Z)
• Decoding-Time Language Model Alignment with Multiple Objectives [116.42095026960598]
  Existing methods primarily focus on optimizing LMs for a single reward function, limiting their adaptability to varied objectives. Here, we propose $textbfmulti-objective decoding (MOD)$, a decoding-time algorithm that outputs the next token from a linear combination of predictions. We show why existing approaches can be sub-optimal even in natural settings and obtain optimality guarantees for our method.
  arXiv  Detail & Related papers  (2024-06-27T02:46:30Z)

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.