Compute Where it Counts: Self Optimizing Language Models

Problem
This preprint addresses the inefficiencies in current large language model (LLM) inference strategies, which typically apply a uniform computation budget across all decoding steps. This approach fails to account for the varying difficulty of tokens, leading to over-computation on simpler tokens and under-computation on more complex ones. The authors propose a method for dynamic budget allocation during autoregressive decoding, allowing for more efficient use of computational resources.

Method
The core technical contribution is the Self-Optimizing Language Models (SOL) framework, which integrates a frozen LLM with a lightweight policy network. This policy network observes the hidden states of the LLM and selects discrete efficiency actions for each decoding step. The actions include controlling token-level attention sparsity, structured activation pruning in the multi-layer perceptron (MLP), and adjusting activation quantization bit-width. The model is trained using group-relative policy optimization on teacher-forced episodes, where multiple compute schedules are sampled for the same token sequence. The reward function balances language model quality with penalties that encourage adherence to a target budget. This approach allows SOL to optimize the quality-efficiency trade-off dynamically, improving performance across various compute regimes.

Results
SOL demonstrates significant improvements in language model quality at matched computational budgets compared to static allocation strategies and random schedule searches. Specifically, it achieves up to a 7.3% increase in accuracy on the Massive Multitask Language Understanding (MMLU) benchmark when compared to uniform budget allocation strategies. The results indicate that SOL effectively discovers a superior quality-efficiency Pareto front, enhancing the overall performance of LLMs in practical applications.

Limitations
The authors acknowledge that the performance gains are contingent on the specific architecture of the LLM used and may not generalize across all model types. Additionally, the reliance on a frozen LLM may limit the adaptability of the policy network to changes in the underlying model. The training process also requires careful tuning of the reward function to balance quality and budget adherence, which could introduce additional complexity. Furthermore, the scalability of the approach to larger models or different tasks remains to be fully explored.

Why it matters
The implications of this work are significant for the field of efficient inference in LLMs. By enabling dynamic computation allocation, SOL provides a framework that can lead to more resource-efficient deployments of language models, particularly in environments with constrained computational budgets. This approach could pave the way for more adaptive and intelligent inference strategies, enhancing the usability of LLMs in real-world applications where efficiency is critical. The findings encourage further exploration of dynamic resource allocation techniques in machine learning, potentially influencing future research directions in model optimization.

Authors: Yash Akhauri, Mohamed S. Abdelfattah
Source: arXiv:2605.10875
URL: https://arxiv.org/abs/2605.10875v1