Compute Where it Counts: Self Optimizing Language Models
Yash Akhauri 1 Mohamed S. Abdelfattah 1
arXiv:2605.10875v1 [cs.LG] 11 May 2026
Abstract
1. Introduction Deploying LLMs at scale has driven efficiency research to make each decoding step cheaper via quantization (Chen et al., 2025; Hooper et al., 2024; Liu et al., 2024b; Zhao et al., 2024), pruning (Liu et al., 2023b; Akhauri et al., 2024; Child et al., 2019), low-rank compression (Chang et al., 2025a;b; Zhang et al., 2024), sparse attention patterns (Choromanski et al., 2020; Zaheer et al., 2020; Akhauri et al., 2025a; Zhang et al., 2023; Li et al., 2024), speculative decoding (Leviathan et al., 2023), and related techniques. Most of these methods, however, apply essentially the same compute budget to every generated token. In practice, token difficulty varies widely: some steps are locally predictable, while others depend on long-range context or precise intermediate computations. A fixed per-token budget therefore tends to over-compute on easy steps and under-compute on harder steps.
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 undercompute 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 paretofront across all our experiments and improves MMLU accuracy by up to 7.3% over uniform budget allocation strategies. [Code] 1 Cornell University. <[email protected]>.
A fixed per-token budget drives the model to often either over, or under-compute. For example, with token sparsity; some tokens require longer-range dependencies whereas others do not, providing an opportunity to dynamically determine computation budget through the number of tokens attended to with each generated token. This also applies to other LLM compression methods such as quantization, weight sparsity, and low-rank factorization. Controlling the compression rate based on token generation difficulty has the potential to unleash a new axis of efficiency in LLMs. One main challenge is that per-token efficiency decisions interact across time. Each generated token can be attended to by future tokens, therefore, aggressively compressing or dropping information that seems safe at step t can cause degradation in quality several steps later. This motivates treating compute allocation as a sequential decision problem rather than a per-step or step-agnostic heuristic. To address this, we propose Self-Optimizing Language Models (SOL). The base LLM is kept frozen and we add a lightweight policy model that reads the LLM hidden state and simple progress/budget features, then selects a discrete efficiency action at each decoding step. Each action instantiates a compute regime inside the same model by controlling (i) token-level attention sparsity, (ii) structured MLP activation pruning, and/or (iii) activation quantization bit-width—without changing the base weights.
Correspondence to: Yash Akhauri
Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
A key design goal is deployment-time controllability. Rather 1
Compute Where it Counts: Self Optimizing Language Models Sparsit y Pruning Quantization
10 % 10 % 16- bit
10 % 10 % 16- bit
50 % 60 % 4 - bit
30 % 30 % 8- bit
30 % 30 % 8- bit
Inference Strategy Context
Context LLM
Policy
Policy
Policy
Policy
Policy
LLM
Optimize Policy 40%
Policy
LLM Decides
Eff iciency
Per
Token!
Where Policy 10 %
LLM
Policy
PPL (? )
Eff iciency (? )
8.32
62%
8.4 0
38%
8.58
40%
9.25
4 6%
It Policy 10 %
LLM
Counts Policy 90 %
Figure 1. (Top Left) The policy is a small autoregressive transformer that uses the LLM’s own hidden state to decide how much compute to allocate per token. (Bottom Left) At training time, the policy samples multiple counterfactual compute schedules (trajectories) and learns to optimize quality under a requested budget target. (Right) At inference, the controller incurs a minimal, constant overhead (proportional to the control horizon, e.g., 16 decode steps), enabling fine-grained control over per-step efficiency.
2. Related Work and Motivation Prior research on methods to optimize LLM inference focus on (i) reducing the memory and computational footprint of the LLM, by quantizing the weights, activations (Lee et al., 2022; Huang et al., 2021; Dotzel et al., 2023), or pruning linear projections and feed-forward network (FFN) layers (Liu et al., 2023b; Akhauri et al., 2024; Feng et al., 2024), (ii) context compression by sparsifying the attention itself, using importance-based techniques to retain only important tokens in the KV-Cache (Zhang et al., 2023; Li et al., 2024; Xiao et al., 2023; Liu et al., 2023a; Xiao et al., 2024), or exploiting the low-rank nature of the KV-Cache to reduce its memory footprint (Chang et al., 2025a;b). (iii) reducing memory bandwidth by employing query-aware sparsity (Akhauri et al., 2025a; Tang et al., 2024; Wu et al., 2024). These methods typically apply a uniform budget to every token, meaning, the compression level is the same for every generated token.
• We introduce learnable per-token budget allocation for LLM decoding: a lightweight policy allocates compute at each step while keeping the base LLM frozen. • Our policy acts on unified action space that controls multiple efficiency mechanisms: token-level attention sparsity, structured MLP activation pruning, and activation quantization. SOL tracks a better qualityefficiency pareto-front when compared to fixed budget allocation across all our experiments.
A second line of work focuses on how much compute to spend per token, but skipping layers of computation. These “early-exit” methods dynamically decide how many layers to execute for each token, stopping computation once a confidence criterion is met; this predates transformers and has been applied to convolutional networks as well. In sequence-to-sequence and autoregressive LMs, this “depthadaptive” execution is orthogonal to our focus: we keep the depth fixed and instead adapt how much context, precision, and sparsity to use for each generated token. (Schuster
• We present a practical training recipe based on grouprelative policy optimization with teacher-forced counterfactual trajectories, conditioning the policy on requested budget targets and using soft budget-matching penalties to achieve controllable efficiency and model quality trade-offs. 2
Policy Rew ard: qualit y - ? (overspend)
than hard-coding a single operating point, we condition the policy on a requested compute budget. During training, we sample per-sequence budget targets and provide them as part of the policy observation, so a single policy can learn to operate across a range of compute budgets. We train the policy using group-relative policy optimization (GRPO) with teacher-forced episodes: for each input we sample multiple counterfactual compute schedules that share the same token path, and optimize a reward that balances language-model quality against soft penalties that encourage the episodeaverage budget usage to match the requested targets. Our contributions are:
Compute Where it Counts: Self Optimizing Language Models
et al., 2022b; Elbayad et al., 2019; Elhoushi et al., 2024) introduce a token-wise early exit strategies. (Sukhbaatar et al., 2019) demonstrates that later layers actually demonstrate a higher attention span than earlier layers, and over long-decode tasks, early-exit of several tokens can harm the global-attention pattern as tokens that exit early do not explicitly commit to KV-Cache of subsequent layers. Beyond early exit of tokens, recent research has also proven that splitting token generation across models (Akhauri et al., 2025b; Yu et al., 2025; Fu et al., 2025; Li & Goyal, 2025) by learning to offload difficult tokens or spans to larger models can significantly improving reasoning performance.
of compression can cause delayed perplexity spikes after returning to dense decoding, with token sparsity exhibiting the strongest tail effects (Appendix A, Fig. 7). Episodes and control horizon. The impact of an efficiency decision can materialize several steps later (e.g., through degraded KV entries that future tokens attend to), so purely myopic, single-step feedback at time t can underestimate the true cost of aggressive compression. We therefore cast decoding-time control as short episodes of length T (default T =16): after a dense prefill initializes a clean cache and state s0 , the policy selects a sequence of discrete actions at ∈ {1, . . . , A} for t=1, . . . , T . Each action instantiates that efficiency method inside the frozen LLM via per-step knobs (κ(at ), ρ(at ), η(at )) (e.g., attention keep fraction, MLP keep fraction, and normalized precision), giving the next state st+1 = f (st , at ) and an episode trajectory τ = (s0 , a1 , . . . , aT ) whose objective aggregates reward over the horizon, e.g.,
Our contribution is a controller that operates inside a single frozen LLM, allocating per-token inference compute by selecting discrete actions over multiple compression methods (token-level attention sparsity, structured MLP activation pruning, and activation quantization). Unlike early-exit / layer-skipping approaches that adapt depth, SOL modulates how each decode step is executed through these orthogonal efficiency knobs. We train the controller to reliably hit a requested budget even in large, combinatorial action spaces by conditioning on budget targets and optimizing over teacherforced, counterfactual compute schedules—motivated by the fact that the effectiveness (and downstream KV impact) of sparsity/quantization/pruning is token- and contextdependent.
R(τ ) =
T X
r(st , at ),
t=1
enabling credit assignment for delayed effects within the episode. Continuing indefinitely under aggressive actions can compound KV-pollution, so we bound its persistence with periodic KV-cache refresh: after every T decode steps, we optionally run a fast, fully-dense pass over the alreadyknown last T tokens to rebuild their KV states to their dense equivalents before starting the next episode; this refresh recomputes cache entries but does not resample or regenerate tokens. In practice, prefill is substantially higherthroughput than token-by-token decode, so this refresh has limited overhead. Unless mentioned otherwise, we use T =16 to align with page-based inference settings and to balance two goals: (i) limiting the lifetime of polluted cache entries, and (ii) providing a long enough horizon for the policy to learn non-myopic trade-offs; T can be increased (or refresh removed) if optimizing for different deployment constraints such as memory footprint. This is consistent with our KV-pollution measurements showing delayed degradation even after switching back to dense compute (Appendix A, Fig. 7). The refresh period is a deployment knob rather than a requirement of SOL: our horizon study evaluates T ∈ {4, 16, 64} and shows that the policy advantage over fixed allocation persists, and in fact becomes stronger, as the horizon increases.
3. Background Decoding is a stateful process. A transformer decoder maintains a key–value (KV) cache that accumulates per-layer representations for previously processed tokens. At step t, the model produces the next-token distribution from the cur(j) (j) rent token and by attending to {(Kℓ , Vℓ )}j<t, ℓ∈[1..L] in the cache, where L is the number of decoder layers. The choice of computation at step t (full, sparse attention, pruned heads, low-precision arithmetic) does not only determine the next token produced, but also impacts the committed KVCache entry for that token, which future steps will attend to. If a step is decoded with less information, its committed KV states deviate from an identical, dense model. We refer to this deviation as KV-pollution, where the tokens cached states are approximated, harming future predictions when those states become important for the current query. This phenomenon is not specific to attention sparsity, it also arises when we skip layers, prune heads (missing KVCache for that token on layers or heads) or prune/quantize neurons (lower quality of the hidden-state for that token). Further, this impact is delayed by nature: the next-token loss may appear unaffected because the model relied on its entire past context, however, the effect of this approximation may appear several steps later when future tokens attend to the approximated token. Empirically, even short bursts
4. Method We use a language model with parameters fθ and learn a lightweight policy πϕ that allocates inference time compute per decode step. The base model parameters θ are never 3
Compute Where it Counts: Self Optimizing Language Models
• κ(a) ∈ [0, 1] is the token-attention keep fraction (token sparsity),
updated; only the parameters of the policy ϕ are trained. Notation. We optimize over episodes of T decode steps (the control horizon). A compute schedule is an action sequence a1:T , where each action selects a compute regime inside the frozen LLM. For each input/prefix, we sample K counterfactual compute schedules under teacher forcing (the GRPO group size) and update the policy by comparing these schedules.
• ρ(a) ∈ [0, 1] is the structured activation keep fraction (channel pruning in the MLP), • q(a) ∈ {4, 5, . . . , 16} is the activation quantization bitwidth (we use the normalized ratio η(a) = q(a)/16 ∈ [0, 1]). Unless stated otherwise, the same action tuple (κ, ρ, q) is broadcast uniformly across all transformer layers at a given decode step. Thus, this paper studies temporal compute allocation across decode steps rather than per-layer or perhead allocation. A per-layer formulation is compatible with SOL, but would substantially enlarge the action space and predictor cost, so we leave it to future work.
Policy inputs and architecture. We decode episodes of length T (we use T =16). At each decode step t ∈ {1, . . . , T }, the policy is fed: (i) the final-layer hidden state from the previous step ht−1 , i.e., the output of the last transformer block before the LM head for token xt−1 , and (ii) the embedding of the current input token e(xt ). In addition, the policy receives a 8-dimensional scalar feature vector st ∈ R8 that makes the budget control explicit: t/T , 1eff,t , Cκ , C ρ , C η , | {z } | {z } |{z} progress effective-step flag requested targets st = . κ̄<t − Cκ , ρ̄<t − Cρ , η̄<t − Cη | {z }
The action space may include any subset of these axes; if an axis is not enabled (only one possible value), it simply becomes constant and incurs no budget penalty. We use “keep fraction” here for clarity and consistency across different compression methods. We train under teacher forcing: within each episode the token sequence is fixed and only the efficiency actions vary across counterfactual compute schedules. Let pat t (·) be the next-token distribution produced by the frozen LLM at step t when executing action at , and let yt denote the ground-truth next token. We define the per-step task reward as the token log-likelihood: rttask = log pat t (yt ) = −CE pat t , yt . (1)
running deviations from target
Here 1eff,t ∈ {0, 1} indicates whether step t counts toward the budget. In our sparse attention implementation, we always keep Ts sink tokens and the most recent Tw window tokens dense (we use Ts =4 and Tw =2 in our main experiments; Appendix B); when the controllable region beyond these always-dense tokens is empty (e.g., for very short prefixes), the chosen action has no effect and we set 1eff,t = 0. In our main experiments with a 1024-token dense prefill, 1eff,t = 1 for all t. The quantities κ̄<t , ρ̄<t , η̄<t are running averages over previous effective steps in the episode (defined below; if no effective step has occurred yet, we set these running averages to the corresponding targets).
To encourage the policy to meet a requested compute target, we penalize deviations of the episode-average realized compute from the requested budgets. Let 1eff,t ∈ {0, 1} indicate whether step t counts toward the budget (i.e., whether the controllable region beyond the always-dense sink/window tokens is non-empty). For an episode of length T , define the realized averages
The policy is a small autoregressive transformer that maintains a KV cache across steps within the episode; it also conditions on the previous action via a learned action embedding. At each step it outputs logits over the discrete action set. During training, we sample actions at ∼ πϕ (· | ht−1 , e(xt ), st ) with a temperature of 1.3. Unless stated otherwise, we select actions greedily at evaluation time: at = arg maxa πϕ (a | ht−1 , e(xt ), st ). The policy state (KV cache of the policy itself) is reset at the start of every episode. In our largest setting (Llama-3.1-8B-Instruct), the controller has 8.56M parameters, making it lightweight relative to the frozen base model.
PT κ̄ =
t=1 1eff,t κ(at ) , PT t=1 1eff,t PT
η̄ =
PT ρ̄ =
1 η(at ) t=1 PT eff,t t=1 1eff,t
t=1 1eff,t ρ(at ) , PT t=1 1eff,t
.
where κ(at ) is the token-attention keep fraction, ρ(at ) is the structured activation keep fraction (e.g., MLP channel keep-rate), and η(at ) is the normalized activation precision (e.g., η = q/16 for q-bit quantization). Let (Cκ , Cρ , Cη ) be the requested targets (provided to the policy as inputs). We use a tolerance band τ and a squared hinge penalty outside the band:
Actions and controlled compute. Each discrete action a ∈ {1, . . . , A} corresponds to a tuple of compute knobs in the available optimization actions a 7→ (κ(a), ρ(a), q(a)), where:
ψ(∆; τ ) = 4
2 max{0, |∆| − τ } .
Compute Where it Counts: Self Optimizing Language Models
The episode-level compute penalty is C = ακ ψ(κ̄−Cκ ; τ ) + αρ ψ(ρ̄−Cρ ; τ ) + αη ψ(η̄−Cη ; τ ), (2) with nonnegative trade-off weights (ακ , αρ , αη ). This objective encourages the policy to match the requested budgets (within tolerance), rather than simply minimizing compute. Because efficiency decisions can have delayed effects within an episode, we use a discounted return-to-go over the task rewards to assign credit: Gt =
T X
γ u−t rutask ,
(3)
across all (t, k) in the batch. Using these advantages, we update the policy with a clipped policy-gradient objective (PPO-style) and an entropy bonus for exploration.
5. Experiments Model families and naming. We evaluate SelfOptimizing Language (SOL) controllers that allocate compute dynamically at each decode step, keeping the base LLM frozen. We first describe our naming convention, a model name encodes (i) which efficiency axes are controlled, (ii) the granularity of the action space and (iii) the control horizon (episode length):
u=t
SOL-
where γ ∈ (0, 1] is a discount factor (we use γ = 0.85). We then combine task return and compute penalty into the per-step signal used for GRPO: rt = Gt − C.
X |{z}
efficiency axes
− |{z} G − |{z} Tk . granularity
horizon
We have three key parts here: Axis tag X. Context (C) controls token-level attention sparsity (keep-rate κ). Quant (Q) controls activation quantization (bit-width ratio η). Prune (P) controls structured MLP activation pruning (keep-rate ρ). Joint (J) controls all efficiency axes simultaneously. Granularity tag G. 2L and 3L denote 2 or 3 discrete levels per enabled axis (e.g., SOL-J-2L-T16 has 23 =8 joint actions; SOL-J-3L-T16 has 33 =27). Fine (FL) denotes a multi-level action space with many choices per axis (full definition in Appendix B). Episode-length tag -Tk. Models with a suffix -Tk vary the episode length (control horizon) T =k while using a fixed mid-size joint action space (336 actions; Appendix B). For example, SOL-J-FL-T4, SOL-J-FL-T16, and SOL-J-FL-T64 share the same action space but differ in how long the policy acts before a KV refresh, and is the training episode length as well.
(4)
C is computed once per trajectory (compute schedule) from the episode-average gaps and broadcast across steps. In all experiments, the token-sparsity knob κ is instantiated with Quest (Tang et al., 2024) by selecting a budgeted subset of KV pages and applying an additive −∞ mask to dropped keys. The structured pruning knob ρ is implemented as TEAL-style (Liu et al., 2024a) activation pruning: per token, we keep the top-⌈ρ dmodel ⌉ MLP input channels by activation magnitude and zero the rest (weights remain unchanged). The quantization knob q is implemented as ZeroQuantstyle (Yao et al., 2022) activation quantization: we apply symmetric per-token fake-quantization at q bits to the MLP output (dynamic range per token), with q=16 recovering the dense path. Full implementation details for all three axes are provided in Appendix E.
Training procedure. All controllers are trained with teacher-forced counterfactual compute schedules as described in Section 4. For each training input we run a dense prefill over a 1024-token prefix to fill the LLM KV cache, then decode an episode of T decode steps (default T =16) in which the policy selects an efficiency action at each step. Within an episode, the token path is fixed (teacher forcing) and only the compute schedule varies across sampled schedules, allowing us to attribute differences in language-model loss to efficiency action allocation decisions instead of sampling noise of the LLM itself. We update the controller by sampling multiple schedules per input (GRPO group size K) and computing group-relative advantages. During training we sample requested budget targets per sequence and provide them to the policy; the reward combines token log-likelihood with appropriate penalties that encourage the episode-average budget usage to match the requested targets (Eqns. 2–4). The policy state is reset at episode boundaries, so inference-time overhead is constant per decoded token
GRPO with per-sequence budget sampling. For each training input, we run a dense prefill on the context prefix to initialize the KV cache. We then create K counterfactual compute schedules of length T (the GRPO group size) that share the same teacher-forced token path; only the efficiency actions differ across schedules. This isolates the effect of compute allocation while keeping supervision fixed. To train a single policy that operates across a range of efficiency regimes, we randomly sample requested budget targets per sequence during training. Concretely, for each input we sample (Cκ , Cρ , Cη ) from user-specified ranges (or discrete lists) and provide them to the policy as part of its observation. The same sampled targets are used for all K schedules of that input and define the episode penalty C. We use process-level GRPO (Shao et al., 2024): at each time step t, we compute group-relative advantages by comparing the K schedules for the same input (mean-centering across schedules at fixed t), and then apply a global whitening 5
Compute Where it Counts: Self Optimizing Language Models
12 11 0.4
0.6 0.8 Net Keep Rates
11.5
SOL-J-3L-T16 Fixed Policy
11.0 10.5 10.0
SOL-J-FL-T16
10.4 Perplexity
13
12.0 Perplexity
Perplexity
SOL-J-2L-T16
10.2 10.0 9.8
0.5 0.6 Net Keep Rates
0.6 0.7 0.8 Net Keep Rates
SOL-J-FL-T4 10.4 10.2
SOL-J-FL-T16 Fixed Policy
10.0 9.8
0.6 0.8 Net Keep Rates
0.6 0.8 Net Keep Rates
Perplexity
8.50 8.25 8.00 7.75 7.50
Perplexity
Perplexity
Figure 2. Quality–efficiency trade-offs for jointly-controlled SOL policies with increasingly fine-grained action sets: (Left) SOL-J-2L-T16 (8 joint actions), (Middle) SOL-J-3L-T16 (27 joint actions), and (Right) SOL-J-FL-T16 (1560 joint actions).
SOL-J-FL-T64
9.50 9.25 9.00 8.75 8.50
0.6 0.8 Net Keep Rates
Figure 3. Effect of episode length T (KV-refresh period) on the quality–efficiency frontier for a joint controller with a shared action space: (Left) SOL-J-FL-T4, (Middle) SOL-J-FL-T16, and (Right) SOL-J-FL-T64.
ized token-attention keep-rate, ρ̄ the realized MLP-channel keep-rate, and η̄ = q̄/16 the realized normalized activation precision. When all three axes are enabled, we define
and negligible relative to the base LLM forward pass. Hyperparameters are reported in Appendix B. We evaluate SOL by measuring language-model perplexity under the efficiency actions selected by the controller. For each evaluation input, we run a dense prefill on the first 1024 tokens and then teacher-force the next T tokens while executing the controller’s actions inside the frozen LLM. We report perplexity on this episode segment by averaging the negative log-likelihood over the T evaluated positions (excluding the dense prefill), and exponentiating (giving us the perplexity). Unless stated otherwise, we select actions greedily (argmax) at evaluation time. Throughout, we treat keep-rates and normalized activation bit-width as monotonic proxies for compute and report quality versus the requested/realized budget targets rather than hardwarespecific latency. Unless stated otherwise, evaluation uses T =16 and the same always-dense sink/window token conventions as in training; additional evaluation details and all action-space definitions appear in the appendix.
NetKeep =
1 (κ̄ + ρ̄ + η̄) . 3
If only a subset of axes is enabled, the average is taken over that subset. Lower net keep-rate means more aggressive compression. We use this metric as an architecture-agnostic proxy for retained compute capacity, not as a universal hardware metric: the exact latency, memory, and throughput gains depend on the kernels and hardware implementation. Action spaces. Here, we ask whether a learned controller remains useful as the discrete search space grows. To test robustness to action-space complexity, we train three jointly-controlled policies with the same episode length (T =16) and the same three efficiency axes (token sparsity, MLP pruning, and activation quantization), but with increasingly fine-grained action sets: SOL-J-2L-T16, SOL-J-3L-T16, and SOL-J-FL-T16. The 2L and 3L variants expose 23 =8 and 33 =27 joint actions respectively, while the fine-grained (FL) variant uses the Cartesian product of the per-axis choices (e.g., |K|×|R|×|Q| =
Efficiency metric: net keep-rate. Most experiments report quality against a scalar net keep-rate, which summarizes the realized episode-average resource usage across the enabled efficiency axes. For an episode, let κ̄ be the real6
Token Sparsity
0.8 0.6
SOL-J-2L-T16 0.4 SOL-J-3L-T16 SOL-J-FL-T16 0.2 0.5 1.0 Target Token Keep-Rate
Pruning
1.0 0.8 0.6 0.4
0.50 0.75 1.00 Target Prune Keep-Rate
Realized Quantization Ratio
1.0
Realized Prune Keep-Rate
Realized Token Keep-Rate
Compute Where it Counts: Self Optimizing Language Models
Quantization 0.8 0.6 0.4 0.50 0.75 Target Quantization Ratio
Figure 4. Budget adherence of SOL when provided a range of requested targets. The plots compare (Left) requested vs. realized token-attention keep-rate κ, (Middle) requested vs. realized structured MLP keep-rate ρ, and (Right) requested vs. realized normalized quantization ratio η.
10×13×12 = 1560 joint actions in our default setting; see Appendix B).
SOL-J-FL-T64. In each case, the policy acts for T decode steps and is then reset; the evaluation protocol matches training, using teacher-forced trajectories and the same always-dense sink/window convention. We evaluate all horizons on the same sweep of requested budgets as in the action-space study and report perplexity versus net keeprate.
At evaluation time, we sweep a grid of requested budgets and measure the resulting quality–efficiency trade-off under teacher forcing. Specifically, we request token keep targets from 0.15–0.95, prune keep targets from 0.40–1.00, and quantization targets from 5–13 bits. To display these results, we collapse the three realized budgets into a single scalar net keep-rate proxy by averaging the realized token keep, prune keep, and normalized quantization ratio. We compare SOL against a fixed baseline that uses a static action schedule (constant over decode steps) and, when a requested budget falls between discrete action levels, mixes adjacent actions across the batch to match the requested average budget. Figure 2 shows that SOL consistently achieves lower perplexity than the fixed baseline at matched net keep-rate, and that these gains persist as the action space grows from 8 to 1560 options. Across operating points, SOL tracks the lower envelope of the random-schedule landscape and remains competitive with best-of-500 random search at matched net keep-rate (Appendix A, Figs. 8–9). Appendix C reports paired statistics over the full sweeps, including standard deviations, paired differences, significance tests, and policy win rates.
As shown in Figure 3, SOL improves the quality–efficiency trade-off over the fixed baseline across all horizons considered. While longer horizons have more risk of compounding error, the controller still learns policies that are on a better Pareto frontier, demonstrating that SOL can operate effectively under different KV-refresh schedules.
Steering policy across efficiency targets A key advantage of SOL is control over budgets at deployment: since the policy is conditioned on the requested budgets, we should be able to specify a target efficiency regime at inference time and have the policy hit the target without additional tuning. We evaluate this by measuring budget adherence: for a range of requested target triples (Cκ , Cρ , Cη ), we run the policy under teacher forcing and compute the realized episode-average budgets (κ̄, ρ̄, η̄) over effective steps. Figure 4 plots requested versus realized budgets for operating points drawn from the policy’s Pareto frontier (Figure 2). Overall, the controller tracks the requested targets closely: most points fall near the identity line, consistent with training under a budget-matching penalty with a τ =0.02 tolerance band. Deviations arise because the budget penalty has finite weight, so under extremely aggressive settings the policy trades budget violations for large improvements in language-model likelihood. We observe this most clearly for very low prune-keep requests (e.g., ρ ≲ 0.5), where perplexity degrades sharply; in this regime the controller tends to choose less aggressive pruning than requested.
Horizon The length of the control horizon determines how long efficiency decisions can compound before the next KV refresh and policy-state reset. Longer horizons may be useful in deployment (e.g., to align with larger paging schemes or to reduce refresh overhead), but may also amplify effects such as KV-pollution and increase the difficulty of credit assignment. We therefore study whether SOL remains effective as the episode length T varies. We train three joint policies that share the same joint action space (336 actions; Appendix B) but differ in horizon: SOL-J-FL-T4, SOL-J-FL-T16, and 7
Perplexity
Fixed Quest@50% ZeroQuant-5Bit TEAL@60% Policy Quest@30% ZeroQuant-8Bit TEAL@80% Llama-3.2-1B Llama-3.2-3B Llama-3.1-8B-Instruct 12 11 9 10 11 9 8 10 8 7 0.4 0.6 0.8 0.4 0.6 0.8 0.4 0.6 0.8 Net Keep Rates Net Keep Rates Net Keep Rates Perplexity
Perplexity
Compute Where it Counts: Self Optimizing Language Models
Avg. Accuracy (%)
Figure 5. Training policies for a range of model sizes. Each subplot has policy configured to SOL-J-FL-T16, with varying LLM sizes.
Llama-3.2-3B 70 65 0.7 0.8 0.9 Net Keep-Rates
Dense Policy Llama-3.1-8B-Instruct
Fixed DeepSeek-R1-Distill-Llama-8B
46
75
44
70
42
65
0.7 0.8 0.9 Net Keep-Rates
0.6 0.8 Token Keep-Rate
Figure 6. We evaluate policies against a fixed (static) compute schedule at several operating points (net keep-rate ≈ 0.4–0.7). Left: Average accuracy over arc easy, piqa, and winogrande. Middle: Average accuracy over three MMLU subjects in the continuation format. Right: Accuracy on GSM8K-cot-llama with 8-shot prompting with Token-Sparsity only (Quest page-size 16).
Scaling to larger model sizes. We keep the same policy architecture and train it on three base LLMs: Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B-Instruct, using an identical action space of 1560 discrete choices per decode step. For the 3B and 8B models we train with a prefill context length of 512 due to VRAM constraints. Figure 5 summarizes a large sweep over the efficiency search space; for readability we plot only the Pareto-optimal frontier. We show two frontiers: a Fixed baseline that uses a static, layer-agnostic efficiency setting throughout decoding, and our learned Policy that adaptively selects actions per step. To contextualize well-known static efficiency methods, we additionally annotate their corresponding operating points in the same plot: Quest token sparsity (Tang et al., 2024) at fixed keep rates (e.g., Quest@50% and Quest@30%), ZeroQuant-style activation quantization (Yao et al., 2022) at fixed bitwidths (ZeroQuant-8Bit and ZeroQuant-5Bit), and TEAL-style activation pruning (Liu et al., 2024a) at fixed pruning levels (TEAL@80% and TEAL@60%). Concretely, each annotated point applies the named method at the indicated setting while holding the other efficiency knobs to their dense configuration, providing a direct reference for how each static method trades perplexity
against net keep rate on each model. Across all three model sizes, the learned policy consistently discovers a strictly better quality–efficiency Pareto frontier than fixed strategies, including those anchored at these standard static baselines. Downstream Evaluation. We test whether SOL’s per-token compute allocation improves downstream accuracy at matched budget. Concretely, we evaluate SOL-J-FL-T16 controllers trained for Llama-3.2-3B and Llama-3.1-8B-Instruct: for Llama-3.2-3B, we report average accuracy over arc easy, piqa, and winogrande (Clark et al., 2018; Bisk et al., 2020; Sakaguchi et al., 2021), and for Llama-3.1-8B-Instruct, we report average accuracy over MMLU conceptual physics, high school chemistry, and international law (Hendrycks et al., 2021) in the continuation setting (Gao et al., 2024), where each candidate answer is scored as a textual continuation (via log-likelihood) and the prediction is the highest-likelihood completion, reflecting a true completion-style evaluation rather than decoding letters for multi-choice questions. We sweep requested operating points with net keep-rate between 0.6 and 0.9 and compare against a fixed baseline that uses a static (per-step constant) efficiency action 8
Compute Where it Counts: Self Optimizing Language Models
schedule matched to the same average budget. We also train a policy for the DeepSeek-R1-Distill-Llama-8B model for token-sparsity, with aggressive pruning using quest page-size 16, with horizon of 64 tokens before a KVrefresh. We evaluate the model on gsm8k-cot-llama with 8-shot prompting, where the model has to decode a reasoning as well as an answer, stressing long-form text generation. Figure 6 shows that SOL consistently improves average accuracy over the fixed baseline at matched net keep-rate for all models, suggesting that the learned per-token allocation transfers beyond perplexity and improves end-task correctness under the same overall compute budget.
portance is predictable. arXiv preprint arXiv:2503.07518, 2025a. Akhauri, Y., Fei, A., Chang, C.-C., AbouElhamayed, A. F., Li, Y., and Abdelfattah, M. S. Splitreason: Learning to offload reasoning. arXiv preprint arXiv:2504.16379, 2025b. Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432–7439, 2020. Chang, C.-C., Lin, C.-Y., Akhauri, Y., Lin, W.-C., Wu, K.-C., Ceze, L., and Abdelfattah, M. S. xkv: Cross-layer svd for kv-cache compression. arXiv preprint arXiv:2503.18893, 2025a.
6. Discussion and Conclusion Most inference-efficiency methods focus on how to compress an LLMs weights or activations to meet a target budget, choosing sparsity / quantization / pruning methods (block sparsity, channel pruning etc.) to deliver a single operating point. SOL instead learns how much compute to spend per token. A lightweight controller reads the models activations and selects discrete efficiency actions, balancing model quality with our budget-matching objective to meet efficiency requirements. While we use monotonic proxies for compute (keep-rate for tokens/channels/bits), our approach could leverage the surplus of efficiency data (real latency, power) that such policies can optimize for. SOL opens an orthogonal axis of efficiency optimization: learning a policy that adapts compute allocation to the difficulty of the generation process and constraints of the serving environment in which the LLM is deployed.
Chang, C.-C., Lin, W.-C., Lin, C.-Y., Chen, C.-Y., Hu, Y.F., Wang, P.-S., Huang, N.-C., Ceze, L., Abdelfattah, M. S., and Wu, K.-C. Palu: Kv-cache compression with low-rank projection. In The Thirteenth International Conference on Learning Representations, 2025b. Chen, Y., Dai, X., Chang, C.-c., Akhauri, Y., and Abdelfattah, M. S. The power of negative zero: Datatype customization for quantized large language models. arXiv preprint arXiv:2501.04052, 2025. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019. Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
7. Acknowledgments We would like to thank Xingyou Song for providing feedback on the initial drafts of this paper.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:1803.05457v1, 2018.
Impact Statement This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dotzel, J., Wu, G., Li, A., Umar, M., Ni, Y., Abdelfattah, M. S., Zhang, Z., Cheng, L., Dixon, M. G., Jouppi, N. P., et al. Fliqs: One-shot mixed-precision floatingpoint and integer quantization search. arXiv preprint arXiv:2308.03290, 2023.
References
Elbayad, M., Gu, J., Grave, E., and Auli, M. Depth-adaptive transformer. arXiv preprint arXiv:1910.10073, 2019.
Akhauri, Y., AbouElhamayed, A. F., Dotzel, J., Zhang, Z., Rush, A. M., Huda, S., and Abdelfattah, M. S. Shadowllm: Predictor-based contextual sparsity for large language models. arXiv preprint arXiv:2406.16635, 2024.
Elhoushi, M., Shrivastava, A., Liskovich, D., Hosmer, B., Wasti, B., Lai, L., Mahmoud, A., Acun, B., Agarwal, S., Roman, A., et al. Layerskip: Enabling early exit inference and self-speculative decoding. arXiv preprint arXiv:2404.16710, 2024.
Akhauri, Y., AbouElhamayed, A. F., Gao, Y., Chang, C.-C., Jain, N., and Abdelfattah, M. S. Tokenbutler: Token im9
Compute Where it Counts: Self Optimizing Language Models
Feng, Y., Lv, J., Cao, Y., Xie, X., and Zhou, S. K. Adakv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference. arXiv preprint arXiv:2407.11550, 2024.
Liu, Z., Desai, A., Liao, F., Wang, W., Xie, V., Xu, Z., Kyrillidis, A., and Shrivastava, A. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. Advances in Neural Information Processing Systems, 36:52342–52364, 2023a.
Fu, T., Ge, Y., You, Y., Liu, E., Yuan, Z., Dai, G., Yan, S., Yang, H., and Wang, Y. R2r: Efficiently navigating divergent reasoning paths with small-large model token routing. arXiv preprint arXiv:2505.21600, 2025.
Liu, Z., Wang, J., Dao, T., Zhou, T., Yuan, B., Song, Z., Shrivastava, A., Zhang, C., Tian, Y., Re, C., et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, pp. 22137–22176. PMLR, 2023b.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac’h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. The language model evaluation harness, 07 2024. URL https://zenodo.org/records/12608602.
Liu, Z., Yuan, J., Jin, H., Zhong, S., Xu, Z., Braverman, V., Chen, B., and Hu, X. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750, 2024b. Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: an adversarial winograd schema challenge at scale. Commun. ACM, 64(9):99–106, August 2021. ISSN 0001-0782. doi: 10.1145/3474381. URL https: //doi.org/10.1145/3474381.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021. Hooper, C., Kim, S., Mohammadzadeh, H., Mahoney, M. W., Shao, Y. S., Keutzer, K., and Gholami, A. Kvquant: Towards 10 million context length llm inference with kv cache quantization. Advances in Neural Information Processing Systems, 37:1270–1303, 2024.
Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay, Y., and Metzler, D. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456–17472, 2022a. Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay, Y., and Metzler, D. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456–17472, 2022b.
Huang, C.-W., Chen, T.-W., and Huang, J.-D. All-youcan-fit 8-bit flexible floating-point format for accurate and memory-efficient inference of deep neural networks. arXiv preprint arXiv:2104.07329, 2021.
Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
Lee, S., Park, J., and Jeon, D. Toward efficient low-precision training: Data format optimization and hysteresis quantization. In International Conference on Learning Representations, 2022.
Sukhbaatar, S., Grave, E., Bojanowski, P., and Joulin, A. Adaptive attention span in transformers. arXiv preprint arXiv:1905.07799, 2019.
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274– 19286. PMLR, 2023.
Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikci, B., and Han, S. Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774, 2024.
Li, A. O. and Goyal, T. Off-trajectory reasoning: Can llms collaborate on reasoning trajectory? arXiv preprint arXiv:2510.06410, 2025.
Wu, W., Pan, Z., Wang, C., Chen, L., Bai, Y., Wang, T., Fu, K., Wang, Z., and Xiong, H. Tokenselect: Efficient longcontext inference and length extrapolation for llms via dynamic token-level kv cache selection. arXiv preprint arXiv:2411.02886, 2024.
Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, D. Snapkv: Llm knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947– 22970, 2024.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023.
Liu, J., Ponnusamy, P., Cai, T., Guo, H., Kim, Y., and Athiwaratkun, B. Training-free activation sparsity in large language models. arXiv preprint arXiv:2408.14690, 2024a.
Xiao, G., Tang, J., Zuo, J., Guo, J., Yang, S., Tang, H., Fu, Y., and Han, S. Duoattention: Efficient long-context 10
Compute Where it Counts: Self Optimizing Language Models
llm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10819, 2024. Xin, J., Tang, R., Lee, J., Yu, Y., and Lin, J. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004.12993, 2020. Yao, Z., Yazdani Aminabadi, R., Zhang, M., Wu, X., Li, C., and He, Y. Zeroquant: Efficient and affordable posttraining quantization for large-scale transformers. Advances in neural information processing systems, 35: 27168–27183, 2022. Yu, B., Yuan, H., Li, H., Xu, X., Wei, Y., Wang, B., Qi, W., and Chen, K. Long-short chain-of-thought mixture supervised fine-tuning eliciting efficient reasoning in large language models. arXiv preprint arXiv:2505.03469, 2025. Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020. Zhang, R., Wang, K., Liu, L., Wang, S., Cheng, H., Zhang, C., and Shen, Y. Lorc: Low-rank compression for llms kv cache with a progressive compression strategy. arXiv preprint arXiv:2410.03111, 2024. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710, 2023. Zhao, Y., Lin, C.-Y., Zhu, K., Ye, Z., Chen, L., Zheng, S., Ceze, L., Krishnamurthy, A., Chen, T., and Kasikci, B. Atom: Low-bit quantization for efficient and accurate llm serving. Proceedings of Machine Learning and Systems, 6:196–209, 2024.
11
Compute Where it Counts: Self Optimizing Language Models
A. Appendix
perplexity (%)
Corrupt Region Token sparsity
Mean Tail Pollution Pruning
P90 Tail Pollution Quantization
102 101 100 0 0.2 0.4 0.6 0.8 0.4 0.6 0.8 5.0 7.5 10.0 12.5 Level (Token Sparsity) Level (Act. Pruning) Level (Quantization Bits)
Figure 7. Axis-dependent KV-pollution: mean direct ∆ perplexity on the corrupted steps vs. p90 peak ∆ perplexity in the subsequent dense tail.
KV Pollution. We quantify KV-pollution using teacher forced wikitext windows (Figure 7). For each window, we run a dense prefill on 512 token prefix, and decode the next tcorrupt = 4 tokens using one efficiency mechanism: token sparisty (keep fraction κ), structured MLP activation pruning (keep fraction ρ) or activation quantization (bit-width q), and then switch back to fully dense decoding (κ=1, ρ=1, q=16) for the remainder of the trajectory. We compare this run to an all-dense baseline under the same token path and report ∆ perplexity as percent change, ∆ppl(%) = (exp(∆NLL)−1)·100. Figure 7 shows three plots, where me measure the ∆ perplexity as a percentage change in the corrupt region itself (where the optimization was active), and then the impact of the corrupt region on the subsequent tokens after returning to dense compute. This tail pollution isolates delayed errors caused by the polluted KV states. While the mean tail pollution is low, the P90 tail pollution remains higher even than the perplexity change in the corrupt region. Here, we can also see that token sparsity exhibits the highest KV-pollution, which is expected, as missing tokens directly impact the information that is written to the KV-cache.
12 10 0.4
0.6 0.8 Net Keep Rates
11.5 11.0 10.5 10.0
Fixed SOL-J-3L-T16
Policy Perplexity
SOL-J-2L-T16
Perplexity
Perplexity
14
Random (KDE)
0.5 0.6 0.7 Net Keep Rates
SOL-J-FL-T16
10.50 10.25 10.00 9.75
0.6 0.8 Net Keep Rates
Figure 8. For each model variant (SOL-J-2L-T16, SOL-J-3L-T16, SOL-J-FL-T16), we sample random counterfactual compute schedules at Pareto operating points and visualize the resulting distribution as a 2D KDE over (net keep-rate, perplexity). We overlay the fixed (static) allocation baseline and the learned policy.
Random-search landscape at matched budgets. We contextualize the learned controller against the distribution demonstrated by uninformed compute schedules. For each model in Figure 2, we take the Pareto-frontier configurations produced by the policy (each corresponding to a particular budget regime) and, for each configuration, sample 30 random counterfactual compute schedules by uniformly drawing actions from the same discrete action set (8, 27, and 1560 actions for 2L, 3L, and 12
Compute Where it Counts: Self Optimizing Language Models
0.5 0.0
0.5 0.6 0.7 Net Keep Rates
Best-of-500 Random SOL-J-3L-T16
0.4
0.6 Net Keep Rates
Fixed Normalized Perplexity
1.0
Random SOL-J-2L-T16
Normalized Perplexity
Normalized Perplexity
FL). We evaluate each schedule under teacher forcing and record its realized net keep-rate (average of token, channel, and quantization keep-rates) and resulting perplexity. Across model variants, the policy consistently lies on the lower envelope of the random density cloud and improves over the fixed baseline, indicating that SOL is not simply matching a typical random schedule, but selecting compute allocations that achieve lower perplexity at comparable net keep-rates. Pooling these random schedules across all Pareto points yields a dense set of samples, which we visualize as a 2D KDE over (net keep-rate, perplexity) in Figure 8. We overlay the fixed baseline (constant action across decode steps) and the policy.
Policy SOL-J-FL-T16
0.6 0.7 Net Keep Rates
Figure 9. Comparison against 500-sample random schedule search. For selected Pareto operating points in SOL-J-2L-T16, SOL-J-3L-T16, and SOL-J-FL-T16, violins show the distribution of normalized perplexity across 500 randomly sampled compute schedules at similar net keep-rates; markers denote the fixed baseline, the best-of-500 random schedule, and the learned policy. Lower is better.
Policy vs. random search. While the KDE view summarizes the global landscape, it does not directly answer how the policy compares to strong random search at a specific operating point. We therefore run a targeted stress test: for each of SOL-J-2L-T16, SOL-J-3L-T16, and SOL-J-FL-T16, we sample four Pareto-frontier configurations from Figure 2 and, for each, evaluate 500 random counterfactual schedules (again sampling actions uniformly from the same action set), alongside the fixed baseline and the policy. Figure 9 reports the resulting per-point random distributions (violins) together with the fixed and policy perplexities (with row-wise normalization for comparability across operating points). Across the 12 test points, the policy matches or exceeds the best-of-500 random schedule in 8 cases. In the remaining 4 cases, the policy remains in the extreme tail of the random distribution (top 5.6% at worst; i.e., at most 28 out of 500 random schedules outperform it), and its regret relative to the best random schedule is small (maximum gap 0.085 perplexity points). Together with Figure 8, these results show that SOL learns a non-trivial per-token allocation strategy: it tracks the lower envelope of the random-search landscape globally, and remains competitive even against substantial pointwise random search over the same efficiency mechanisms. Designing static (non-learned) strategies for budget allocation. Beyond fixed and random schedules, we evaluate two simple hand-crafted criteria that allocate per-token compute using a scalar signal from the previous decoding step, while still steering the episode-average budgets to match the requested targets. Entropy-Matched Criterion (EMC). We compute an uncertainty score ut−1 ∈ [0, 1] from the previous step’s logits using normalized entropy Ĥt−1 (Xin et al., 2020); higher uncertainty biases the next action toward less aggressive compression (larger attention keep-rate κ, higher MLP keep-rate ρ, and/or higher activation precision η). Drift-Aware Criterion (DAC). We compute a representation-change score dt−1 ∈ [0, 1] as the cosine drift between consecutive last-layer hidden states, dt−1 = 12 (1 − cos(ht−1 , ht−2 )) (using embedding drift for the first decoded token) (Schuster et al., 2022a), and allocate more compute on high-drift steps. For both baselines, actions are chosen from the same discrete action set as SOL using a greedy budget-steering rule: at each step and for each enabled axis, we select between the two nearest discrete levels around the remaining required average, then a feasible action so the target budget remain attainable over the remaining steps. From Figure 10, we find that these hand-crafted static strategies do not outperform the policy. Jointly trading off different optimization methods allow us to hit significantly more aggressive efficiency targets than hand-designed strategies for individual optimization methods. 13
Compute Where it Counts: Self Optimizing Language Models
Normalized Perplexity
Random
Best Random (min)
SOL-J-2L-T16
1.0
Fixed
EMC
Drift-Aware (DAC)
SOL-J-3L-T16
Policy
SOL-J-FL-T16
0.5 0.0
0.4
0.6 Net Keep Rate
0.4
0.6 Net Keep Rate
0.4
0.6 0.8 Net Keep Rate
Figure 10. Hand-crafted budget-allocation heuristics (entropy- and drift-based) do not consistently improve over the fixed baseline, while the learned SOL policy achieves the best quality–efficiency trade-off.
Optimizing individual efficiency methods. To isolate whether SOL’s gains come from joint multi-axis control or from learning a per-token controller in the simplest setting, we train three single-axis controllers for each efficiency method: SOL-Q-2L-T16 (quantization only), SOL-P-2L-T16 (MLP pruning only), and SOL-C-2L-T16 (token sparsity only). In all cases, the episode length is T =16 and the action space contains only two levels for the active axis; the other axes are held fixed at their dense settings. We evaluate three binary action sets per axis (x-axis labels in Figure 11), corresponding to different choices of the two available levels (e.g., two candidate bit-widths for quantization).
Normalized Perplexity
For each binary setting, we compare the learned controller against (i) a fixed baseline that uses a static schedule over the same two levels and (ii) random compute schedules that sample one of the two actions at each decode step. We report normalized perplexity for each setting (normalized within each setting for comparability across axes). As shown in Figure 11, SOL consistently selects schedules that lie on the lower envelope of the random-search landscape, and it matches the best-of-30 random schedule in 7 out of the 9 binary settings tested. This indicates that even with a minimal two-level action space, SOL learns a non-trivial per-token allocation strategy beyond what is obtained by static or uninformed schedules.
1.0
Random Quantization Policies
Best-of-30 Random Pruning Policies
Fixed
Policy Token-Sparsity Policies
0.5 0.0 10 16 16 Q5|Q Q5|Q Q6|Q
100 0|S100 0|S100 S | 0 S3 S4 S5
1.0 20|1.0 40|1.0 0. 0.
| 0.10
Figure 11. Per-axis optimization with two-level action sets. Each subplot trains three single-axis controllers (SOL-Q/P/C-2L-T16). For each two-level action set (x-axis), the violin shows the distribution of normalized perplexity over 30 randomly sampled schedules. Lower is better.
14
Compute Where it Counts: Self Optimizing Language Models
B. Model Configurations We primarily train and evaluate six main models, whose configurations are described below. B.1. Search Space Size Ablation Models Shared configuration (used unless overridden) Base LLM: meta-llama/Llama-3.2-1B Token Sparsity Implementation: Quest (page size = 4) Budget Ranges: Token budget [0.1, 1.0]; pruning budget [0.4, 1.0]; quantization ratio budget [0.3125, 1.0] Budget Penalty Weights: ακ = 100; αρ = 100; αη = 200 GRPO group size / horizon: T =16 decode steps/episode; K=16 schedules/input; entropy coef = 0.05 Optimization: Batch size = 8; grad accumulation = 8; lr = 10−4 ; max grad norm = 2.0; epochs = 1 Context Length: 1024 Policy Network (Controller): Transformer: dmodel = 512; nheads = 4; nlayers = 1; MLP ratio = 4.0; action dim = 32; dropout = 0.0 PPO Settings: PPO clip = 0.2; PPO epochs = 1; minibatch = 512; target batch = 2048 Sink and Window Tokens: Sink tokens Ts = 4; window tokens Tw = 2 Dataset: allenai/dolma (v1 6-sample); text field: text; dataset pct: 100%; seed: 1234 Binary: SOL-J-2L-T16 (8 choices per decode step) Shared settings: See shared config in B.1. Action Space: Keep-Rate {0.1, 1.0}; Prune {s60, s100}; Quant {q5, q16} Override: Budget Ranges: Pruning budget [0.6, 1.0] (token/quant as shared) Ternary: SOL-J-3L-T16 (27 choices per decode step) Shared settings: See shared config in B.1. Action Space: Keep-Rate {0.1, 0.6, 1.0}; Prune {s40, s80, s100}; Quant {q5, q7, q16} Multi: SOL-J-FL-T16 (1560 choices per decode step) Shared settings: See shared config in B.1. Override: Token Sparsity Implementation: Quest (page size = 8) Override: Group size: K=32 schedules/input (T /entropy as shared) Action Space: Keep-Rate {0.1,0.2,...,1.0}; Prune {s40,s45,...,s100}; Quant { q5,q6,...,q16} B.2. Horizon (Episode length) ablation Shared configuration (used unless overridden) Base LLM: meta-llama/Llama-3.2-1B Token Sparsity Implementation: Quest (page size = 8) Action Space: Keep-Rate {0.1, 0.4, 0.7, 1.0}; Prune {s45, s50, s60, s70, s80, s90, s100}; Quant {q5, q6, ..., q16} Budget Ranges: Token budget [0.1, 1.0]; pruning budget [0.45, 1.0]; quantization ratio budget [0.3125, 1.0] 15
Compute Where it Counts: Self Optimizing Language Models
Budget Penalty Weights: ακ = 100; αρ = 100; αη = 200 GRPO group size / horizon (default): T =16 decode steps/episode; K=32 schedules/input; entropy coef = 0.05 Optimization: Batch size = 8; grad accumulation = 8; lr = 10−4 ; max grad norm = 2.0; epochs = 1 Context Length: 1024 Policy Network (Controller): Transformer: dmodel = 512; nheads = 4; nlayers = 1; MLP ratio = 4.0; action dim = 32; dropout = 0.0 PPO Settings: PPO clip = 0.2; PPO epochs = 1; minibatch = 512; target batch = 2048 Sink and Window Tokens: Sink tokens Ts = 4; window tokens Tw = 2 Dataset: allenai/dolma (v1 6-sample); text field: text; dataset pct: 100%; seed: 1234 Horizon 4: SOL-J-FL-T4 Shared settings: See shared config in B.2. Override: Horizon: T =4 decode steps/episode (K/entropy as shared) Horizon 16: SOL-J-FL-T16 Shared settings: See shared config in B.2. Horizon 64: SOL-J-FL-T64 Shared settings: See shared config in B.2. Override: Horizon: T =64 decode steps/episode (K/entropy as shared) B.3. Model Scaling Experiments meta-llama/Llama-3.2-3B Shared Settings: See shared config in B.1. Base LLM: meta-llama/Llama-3.2-3B Override: Context Length: 512 Override: Optimization: Batch size = 2; grad accumulation = 32 (lr/max grad norm/epochs as shared) Override: GRPO group size / horizon: T =16 decode steps/episode; K=16 schedules/input; entropy coef = 0.05 Override: Policy Network (Controller): Same as shared, except max length = 512 meta-llama/Llama-3.1-8B-Instruct Shared Settings: See shared config in B.3. Base LLM: meta-llama/Llama-3.1-8B-Instruct Override: Context Length: 512 Override: Optimization: Batch size = 2; grad accumulation = 32 (lr/max grad norm/epochs as shared) Override: GRPO group size / horizon: T =16 decode steps/episode; K=16 schedules/input; entropy coef = 0.05 Override: Policy Network (Controller): Same as shared, except max length = 512 16
Compute Where it Counts: Self Optimizing Language Models
deepseek-ai/DeepSeek-R1-Distill-Llama-8B Shared Settings: See shared config in B.1. Base LLM: deepseek-ai/DeepSeek-R1-Distill-Llama-8B Override: Token Sparsity Implementation: Quest (page size = 16) {0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0};
Action Space: Keep-Rate Prune {s100}; Quant {q16}
Override: Budget Ranges: Token budget [0.1, 1.0]; pruning budget fixed at 1.0; quantization ratio budget fixed at 1.0 Override: Budget Penalty Weights: ακ = 100; αρ = 0; αη = 0 Override: GRPO group size / horizon: T =64 decode steps/episode; K=16 schedules/input; entropy coef = 0.05 Override: Optimization: Batch size = 2; grad accumulation = 32; lr = 10−4 ; max grad norm = 2.0; epochs = 1 Override: Sink and Window Tokens: Sink tokens Ts = 16; window tokens Tw = 16 Override: Policy Network (Controller): Same as shared, except max length = 512
C. Statistical significance of policy gains For each evaluated configuration, we compare SOL and the fixed baseline at matched requested budget targets. Each row in Table 1 aggregates over the corresponding budget sweep. We report mean perplexity, the paired difference ∆ = PPLpolicy − PPLfixed , a one-sided paired t-test, and the fraction of target configurations for which the policy achieves lower perplexity. Negative ∆ indicates that SOL is better. Table 1. Paired comparison between SOL and fixed allocation over matched budget sweeps. SOL achieves significantly lower perplexity across all configurations.
Config
Policy PPL
Fixed PPL
∆
p-value
Win rate
SOL-J-FL-T4 SOL-J-FL-T16 SOL-J-FL-T64 SOL-J-2L-T16 SOL-J-3L-T16 Llama-3.2-3B Llama-3.1-8B
9.53 ± 1.87 11.26 ± 1.29 9.12 ± 0.64 11.51 ± 0.60 10.63 ± 0.77 9.14 ± 1.19 7.93 ± 0.62
9.66 ± 2.20 11.42 ± 1.60 9.34 ± 0.79 11.89 ± 0.72 10.94 ± 1.08 9.46 ± 1.28 8.20 ± 0.77
−0.127 ± 0.748 −0.160 ± 0.696 −0.212 ± 0.259 −0.383 ± 0.241 −0.309 ± 0.346 −0.327 ± 0.243 −0.270 ± 0.240
3.2×10−5 3.0×10−8 < 10−10 < 10−10 < 10−10 < 10−10 < 10−10
58.7% 70.4% 84.7% 95.3% 84.4% 95.5% 88.9%
D. Controller training cost Table 2 reports approximate training cost for the SOL controller. The base LLM is frozen in all cases; only the lightweight policy network is trained. Table 2. Approximate controller training cost on H100 GPUs.
Base model
Training cost
Llama-3.2-1B Llama-3.2-3B Llama-3.1-8B-Instruct
4 GPU-hours 7 GPU-hours 20 GPU-hours
E. Implementation details of efficiency actions SOL actions select a discrete tuple (κ, ρ, q) at each decode step, where κ controls token sparsity in attention (Quest), ρ controls structured activation pruning in the MLP, and q controls activation quantization bit-width. We implement all three 17
Compute Where it Counts: Self Optimizing Language Models
knobs as inference-time controls in a frozen HuggingFace LLaMA model by monkey-patching the attention and MLP forward paths; the base model weights are unchanged. E.1. Token sparsity via Quest (context keep-rate κ) Budgeted token selection. At decode step t, let Kt denote the number of keys available in the KV cache for the current query (i.e., the current KV length). Given a requested keep-rate κt ∈ [0, 1], we convert it to a per-sequence token budget bt = ⌈κt · Kt ⌉ , clamped to [0, Kt ]. This budget is set per sequence in the batch and broadcast across heads. Quest masking in LLaMA attention. We implement token sparsity using Quest (Tang et al., 2024) by modifying the LLaMA eager attention forward path to add an additive mask (bias) before softmax. Concretely, we group the Kt keys into contiguous pages of size S (our quest page size), pad to a multiple of S, and compute a page-level upper bound score for each head and query. Let q ∈ Rd denote the query vector and let a page contain keys {k (j) }Sj=1 . Define per-dimension page extrema: m+ = max k (j) , m− = − min k (j) . j
j
We decompose the query into positive and negative magnitudes, qpos = |q| ⊙ 1[q ≥ 0] and qneg = |q| − qpos , and score pages by the Quest bound: ⊤ ⊤ score(page) = qpos m+ + qneg m− . We then keep the top ⌈bt /S⌉ pages according to this score, restricted to keys that are allowed by the model’s attention mask (causal/padding). Expanding the selected pages yields a boolean token keep-mask of shape [B, H, Q, Kt ], which we convert to an additive bias: ( 0 if key k is kept ∆Mt (i, h, q, k) = −∞ otherwise. This bias is added to the existing attention mask and the original attention computation is left unchanged. If the requested budget keeps all allowed tokens, we bypass masking and execute dense attention. Granularity. Although κt is specified per sequence, selection is performed per head (and per query position if Q > 1). In our autoregressive decoding setting, Q = 1 and the policy changes κt once per token. E.2. Structured MLP activation pruning (keep-rate ρ) What is pruned. We implement structured pruning as input-channel gating on the residual stream entering each MLP block (HF LlamaMLP). This is a structured (channel-wise) activation mask applied at inference time; model weights remain unchanged. Per-token channel selection. Let xt ∈ Rdmodel be the MLP input activation for the token at step t (teacher-forced decoding uses a single-token forward pass, so sequence length within the call is 1). Given a keep-rate ρt ∈ [0, 1], we keep mt = ⌈ρt · dmodel ⌉ channels per example. We score channels by magnitude (max absolute activation within the current call): s(j) =
max |x(j)|,
token in call
and keep the top-mt channels. The resulting binary mask is shared across tokens inside the call (and is therefore per-token in decode), and we apply it multiplicatively: x̃t = xt ⊙ mt . We apply pruning only inside the MLP block (i.e., we do not prune the attention input for the same layer), so attention projections remain dense for that layer while the MLP update is compressed. 18
Compute Where it Counts: Self Optimizing Language Models
E.3. Activation quantization (bit-width q) What is quantized. We implement activation quantization as a fake-quantization operator applied to the output of each MLP block (the residual-stream update). The MLP matmuls are executed in full precision and the output is quantized/dequantized before being added back through the residual path. Quantizer. Given MLP output activation z ∈ Rdmodel and a selected bit-width q ∈ {5, . . . , 16}, we use symmetric uniform quantization with per-token dynamic range. Let qmax = 2q−1 − 1 and let a = maxj |z(j)| (computed per token). The scale is s = a/qmax and the quantized output is m j , −q , q · s. ẑ(j) = clip z(j) max max s For q ≥ 16 we return z unchanged. Mixed-bit batches are supported by applying this operator per example conditioned on the selected q. Normalized precision ratio. In the main text we use the normalized ratio η = q/16 ∈ [0, 1] as the policy-controlled precision knob and as the target-matching quantity in the budget penalty.
19