SOL — Self-Optimizing Language Models via Token-Level Efficiency Policies

Paper. Yash Akhauri, Mohamed S. Abdelfattah. Compute Where it Counts: Self Optimizing Language Models. ICML 2026. Cornell University. [arXiv]

0. The Picture in One Paragraph

LLMs allocate compute uniformly: every token gets the same attention computation, MLP depth, and numerical precision. This is efficient in hardware but not in reasoning — some tokens encode high-value information (“however”, entity names, key numerical quantities) while others are low-value filler. Self-Optimizing Language Models (SOL) train a lightweight policy network to read the frozen LLM’s hidden state at each token and assign a discrete efficiency action (which attention heads to sparsify, which MLP activations to prune, what quantization bit-width to apply). The policy is trained via GRPO on counterfactual trajectories: the teacher LLM generates high-quality outputs, and the policy is rewarded for matching that quality under a reduced compute budget. Three axes of optimization — attention sparsity C, MLP activation pruning P, and quantization Q — can be combined into a joint policy SOL-J. The result: +7.3% MMLU accuracy over uniform-budget baselines, with a consistently better accuracy-compute Pareto front, across Llama-3.2-3B, Llama-3.1-8B-Instruct, and DeepSeek-R1-Distill-Llama-8B.

1. The Problem — Uniform Compute Is Wasteful

Every token in a transformer forward pass receives the same computational resources:

Full attention: O(L²) computation over all context
Full MLP depth: all hidden dimensions activated
Full precision: fp16 or bf16 for all operations

This uniformity is a deliberate hardware efficiency choice, not a semantic one. From the model’s information-processing perspective, however, tokens are not created equal:

Token type	Informational role	Optimal compute
High-value (rare words, key quantities, logical pivots)	Updates reasoning state significantly	High
Low-value (common function words, repetition, filler)	Contributes little to next prediction	Low

Standard adaptive compute research addresses this at the model level (layer-skipping, early exit). SOL addresses it at the token level within a layer — a finer-grained decomposition — using a trained policy rather than a hand-designed rule.

2. The SOL Framework — Three Axes, One Policy

2.1 Three efficiency dimensions

SOL defines three independent axes of per-token efficiency:

C — Attention sparsity (κ) At each token, only κ ∈ {0%, 25%, 50%, 75%} of attention heads are computed. High-value tokens use full attention (κ = 0%); low-value tokens drop up to 75% of heads.

P — MLP activation pruning (ρ) At each token, ρ ∈ {0%, 25%, 50%, 75%} of MLP activations are zeroed before the nonlinearity. Equivalent to a form of dynamic width reduction.

Q — Quantization bit-width (η) At each token, weights and activations are quantized to η ∈ {16, 8, 4, 2} bits. Critical tokens keep full precision; less important tokens are aggressively quantized.

2.2 Action granularity — FL, 3L, 2L

Actions can be applied at different layer granularities:

Granularity	Resolution	Action space size
2L (2-layer groups)	Coarse	Lower
3L (3-layer groups)	Medium	Medium
FL (full per-layer)	Fine	Up to 1560 actions

FL is the most expressive but has the largest policy output space. In practice, 3L provides a good accuracy-complexity trade-off.

2.3 Policy network

A lightweight network (much smaller than the LLM) reads the hidden state at each token position and outputs a discrete action over the chosen granularity × dimension combination. The policy is architecture-agnostic: it can be applied to any transformer-based LLM.

2.4 Variants

Variant	Axes used	Description
SOL-C	C only	Attention-sparsity policy
SOL-P	P only	MLP-pruning policy
SOL-Q	Q only	Quantization policy
SOL-J	C + P + Q	Joint policy over all three axes

SOL-J achieves the highest gains but also the largest action space.

3. Training — GRPO on Counterfactual Trajectories

3.1 The training setup

Training the policy requires a reward signal: given a token, which efficiency action should be taken? This is not directly labeled in any dataset.

SOL uses GRPO (Group Relative Policy Optimization) on teacher-forced counterfactual trajectories:

Generate reference trajectories: the frozen LLM produces high-quality outputs on training tasks.
Counterfactual rollouts: the policy proposes efficiency actions for each token; the model forward pass is run with those actions applied.
Reward: the reward for a trajectory is a combination of (a) log-likelihood match to the reference output and (b) a penalty for total compute used (measured by net keep-rate).
GRPO update: the policy is updated to prefer actions that achieve high log-likelihood match with low compute.

The net keep-rate (proportion of activations/attention heads retained on average) serves as the compute proxy. A policy that maintains high match quality with a low net keep-rate is rewarded.

3.2 Why GRPO over standard RL?

GRPO uses group-relative rewards: within a batch of trajectories for the same input, each trajectory’s reward is normalized by the group’s mean and variance. This stabilizes training when absolute reward magnitudes vary across tasks.

4. Results

4.1 MMLU accuracy vs. uniform baseline

Method	MMLU accuracy gain vs. uniform budget
SOL-C (FL)	+5.1%
SOL-P (FL)	+4.8%
SOL-Q (FL)	+4.2%
SOL-J (3L)	+7.3%

The joint policy dominates all single-axis variants. The gain is measured at the same net keep-rate (equal compute).

4.2 Pareto frontier improvement

Across all tested compute levels (net keep-rates from 40% to 100%), SOL-J consistently sits above the uniform baseline on the accuracy vs. compute curve. There is no tested compute regime where uniform allocation beats SOL.

4.3 Models evaluated

Llama-3.2-3B (instruction tuning)
Llama-3.1-8B-Instruct
DeepSeek-R1-Distill-Llama-8B (reasoning model)

Gains are consistent across all three models, including the reasoning-specialized DeepSeek-R1 distillation. This is notable: reasoning models with long CoT chains also benefit from per-token efficiency allocation.

4.4 Granularity ablation

Granularity	Accuracy	Policy complexity
2L	Lower	Smallest
3L	Best trade-off	Medium
FL	Marginal gains over 3L	Highest

3L provides nearly all of the accuracy benefit of FL at substantially lower policy complexity.

5. Relationship to Adaptive Compute Literature

SOL sits at the intersection of two research directions:

Direction	Examples	SOL’s relationship
Layer-level adaptive depth	Early exit, layer skipping	SOL is token-level within layers, complementary
Token-level importance	Attention pruning, token merging	SOL learns to prune via RL, not hand-designed rules
Quantization	GPTQ, AWQ	SOL does per-token quantization, not uniform
RL for efficiency	AdapLeR, DynaBert	SOL uses GRPO with counterfactual trajectories

The key novelty: prior methods pick one axis and optimize statically (uniform pruning, uniform quantization). SOL jointly optimizes three axes dynamically per token using a learned policy.

6. Why It Matters

Three reasons:

Token-level compute allocation is the right granularity. Layer-level early exit discards all computation for a token; token-level SOL keeps useful computation while reducing wasteful computation. The policy learns which aspects of which tokens matter most — a fundamentally richer representation than any static rule.
+7.3% on MMLU at equal compute is substantial. MMLU is a broad, stable benchmark. A gain of this magnitude from a compute reallocation alone — without changing the frozen LLM’s weights — demonstrates that significant latent efficiency is being left on the table by uniform allocation.
The policy is backbone-agnostic and lightweight. Because the policy reads hidden states but does not modify the LLM, it can be trained once and applied to any compatible backbone. As better base models become available, the policy can be retrained on the new backbone without architectural changes.

7. Limitations Worth Knowing

Net keep-rate ≠ wall-clock speedup. The paper reports compute savings in terms of net keep-rate (proportion of activations retained). Actual wall-clock speedup depends on hardware implementation: sparse attention and quantization gains require specific hardware support (sparse tensor cores, mixed-precision units). No hardware benchmarks are reported.
Training cost. GRPO training on counterfactual trajectories requires running many forward passes with modified efficiency actions. The training overhead relative to the model size is not reported.
Reasoning token heterogeneity. For long CoT chains, the policy must distinguish reasoning tokens (high value) from filler tokens across a long trajectory. Whether the learned policy generalizes well across reasoning styles is not evaluated beyond DeepSeek-R1.
Discrete action space. The current action space (4 levels per axis) is coarse. Finer-grained or continuous efficiency actions are not evaluated.

8. The Takeaway for a First Reader

If you remember three things:

SOL trains a lightweight policy to assign per-token efficiency actions (attention sparsity κ, MLP pruning ρ, quantization bit-width η) to a frozen LLM via GRPO on counterfactual trajectories. High-value tokens get more compute; low-value tokens get less — learned, not hand-designed.
Three axes can be combined into SOL-J, a joint policy. Granularity 3L (3-layer groups) provides the best accuracy-complexity trade-off, with an action space of up to 1560 discrete actions at FL granularity.
+7.3% MMLU accuracy over uniform-budget baselines at the same compute, with a consistently better Pareto front across Llama-3.2-3B, Llama-3.1-8B-Instruct, and DeepSeek-R1-Distill-Llama-8B.

References

Akhauri, Y., & Abdelfattah, M.S. (2026). Compute Where it Counts: Self Optimizing Language Models. ICML 2026. arXiv:2605.10875.
Related on this site: Reasoning Cache — iterative summarize-and-restart for budget extrapolation, a complementary approach to efficient long-horizon reasoning; Context-Folding — adaptive context management as a learnable agent skill, paralleling SOL’s theme of learned efficiency allocation.