Compute Where it Counts: Self Optimizing Language Models

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

• The paper introduces a learnable controller for dynamic compute allocation that adjusts token-level attention, MLP pruning, and quantization during LLM inference.
• It employs teacher-forced counterfactual schedules with GRPO optimization to achieve lower perplexity and higher accuracy while meeting efficiency targets.
• Empirical results demonstrate improved quality-efficiency trade-offs and reduced KV-pollution, showing SOL's practicality for real-world LLM deployment.

Self-Optimizing LLMs: Dynamic Compute Allocation for Efficient LLM Inference

Motivation and Background

The static per-token computation allocation in standard LLM inference regimes leads to suboptimal efficiency and quality trade-offs. Most compression techniques (quantization, pruning, sparse attention) either overcompute for trivial tokens or undercompute for tokens requiring complex reasoning or long-range context, causing both immediate and delayed degradation in output quality. The phenomenon of KV-pollution, wherein compressed token representations persist in the KV cache and harm future generations, is particularly problematic and exacerbated by over-aggressive uniform compression schemes.

A broad prior spectrum of efficiency research either focuses on compression at a hardware or token level (e.g., [Tang et al., 2024] for attention sparsity, [Yao et al., 2022] for quantization) or on adaptive “early-exit” mechanisms, which tune depth rather than the intra-layer compute modalities. The paper "Compute Where it Counts: Self Optimizing LLMs" (2605.10875) introduces an orthogonal approach: a learnable controller policy that dynamically allocates the compute budget per token during autoregressive decoding, operating within a frozen LLM via discrete efficiency actions without adjusting model parameters.

SOL Architecture and Methodology

The Self-Optimizing LLMs (SOL) framework augments a frozen LLM with a lightweight policy model—a single-layer autoregressive transformer controller—that, at each decoding step, receives the LLM's hidden state, token embeddings, and explicit progress/budget features. The controller selects discrete actions governing three primary axes:

• Token-level attention sparsity: Adjusts the fraction of attended tokens via Quest masking.
• Structured MLP activation pruning: Selects active channels in MLP via TEAL-style gating.
• Activation quantization bit-width: Applies per-token symmetric quantization to MLP outputs.

Actions are tuples controlling these compute modalities, broadcast across all layers per decode step. The episode-based structure (control horizon, default $T=16$) constrains compounding error, with optional KV-refreshes after each episode to mitigate pollution.

Training leverages teacher-forced counterfactual schedules: for each input, multiple compute schedules are sampled, all sharing the same token trajectory, isolating the effect of efficiency actions. The policy is trained with Group-Relative Policy Optimization (GRPO), optimizing log-likelihood rewards penalized by episode-average deviations from user-specified (and sampled) budget targets—a squared hinge loss within tolerance bands incentivizes adherence.

Empirical Evaluation and Results

SOL is assessed across varying action space sizes (8 to 1560 joint actions), horizons (4–64 steps), and LLM sizes (1B–8B parameters). Strong empirical results are documented:

• Strictly improved quality-efficiency trade-off: SOL consistently achieves lower perplexity than static schedules at matched net keep-rate, tracking the lower envelope of random search and outperforming strong best-of-500 random schedule baselines (see normalized perplexity gap, Table 1 Appendix).
• Robustness to action space granularity: Gains persist even as the discrete action space grows substantially, indicating non-trivial learned allocation strategies beyond random or hand-crafted heuristics.
• Adherence to target budgets: The controller reliably matches user-specified efficiency targets across axes; deviations are rare and generally occur only under extreme compression where quality otherwise degrades sharply.
• Improved downstream task accuracy: On downstream benchmarks (arc_easy, piqa, winogrande, MMLU, GSM8K), SOL increases average accuracy by up to 7.3% over uniform compute allocation, even under the same budget envelope.

KV-pollution experiments show that sparse attention inflicts the highest tail pollution, supporting the necessity for policy-driven modulation rather than static compression.

Practical and Theoretical Implications

SOL demonstrates practical deployment-time controllability: the policy can be conditioned on requested efficiency regimes, enabling fine-tuned trade-offs in real-world serving environments. The inference-time overhead is negligible compared to the LLM forward pass, and controller training cost is modest (20 GPU-hours for 8B models). By exposing monotonic proxies (keep-rates, bit-widths), SOL aligns with realistic hardware/resource constraints but could be further optimized for latency or power with enhanced reward signals.

Theoretically, dynamic compute allocation across multiple compression axes opens a new paradigm in LLM inference—moving from static, agnostic budgets toward contextually aware, adaptive inference. This approach is distinct from depth-adaptive early-exit schemes: SOL fixes model depth and instead modulates context, precision, and sparsity per decode step, handling delayed effects and long-range dependencies endemic to autoregressive generation.

Future directions may involve integrating per-layer or per-head allocations (complexifying the action space), optimizing for real latency, memory footprint, or power, and leveraging real-time deployment statistics to further adapt policies. The framework also interfaces well with “offload” strategies for routing difficult tokens to larger models, potentially synergizing with model collaboration paradigms.

Conclusion

"Compute Where it Counts: Self Optimizing LLMs" (2605.10875) establishes a robust methodology for dynamic compute allocation in LLM inference via a learnable controller, yielding superior quality-efficiency Pareto frontiers relative to static allocation and uninformed random schedules, with practical and theoretical benefits. SOL's foundations avoid the pitfalls of over/under-computation, mitigate delayed KV-pollution effects, and provide a scalable, controllable axis for serving environment adaptation. The paradigm is transferable to larger models and downstream tasks, suggesting broad applicability and critical future exploration along adaptive inference lines.