End-to-End Dynamic Sparsity for Resource-Adaptive LLM Inference

Abstract: LLMs inference is typically deployed under a static resource assumption, where models execute a fixed computational graph regardless of the runtime environment. However, real-world cloud infrastructure is inherently dynamic, characterized by fluctuating availability (e.g., spot instance preemption) and tiered Quality-of-Service requirements. In such volatile settings, static models are inflexible: they either crash under resource constraints or waste compute on redundant operations. To bridge this gap, we propose Learning to Allocate (L2A), an end-to-end framework for resource-adaptive inference. Unlike prior methods that condition only on input difficulty, we formulate inference as a constrained allocation problem conditioned on both the input and the runtime resource budget itself. We introduce lightweight, budget-conditioned and input-aware gating networks integrated into the LLM. These gates are trained via a unified objective that jointly optimizes task performance, logical consistency, and resource costs along three axes matching how real-world dynamics manifest: layer skipping for memory and depth pressure, head pruning for throughput contention, and reasoning-token reduction for latency tightening. This lets the model learn a budget-aware policy beyond input difficulty alone: it adaptively configures its computational footprint with respect to real-time resource dynamics, maximizing reasoning depth when resources permit while enforcing strict frugality when budgets tighten. A single L2A model traces the entire compute-accuracy Pareto frontier on Llama-3-8B and Qwen-3-4B: at up to 34% realized layer sparsity, it stays within 0.6% of the dense baseline on GSM8K, with the same gap holding zero-shot on out-of-distribution tasks, while every static or heuristic baseline requires a separately tuned model and still drops by 5-10% at comparable inference time.

• The paper introduces the L2A framework which uses budget-conditioned gates for layer skipping, attention head pruning, and reasoning truncation.
• It employs KL-distillation and resource budget penalties to maintain near-baseline accuracy, even at high sparsity levels.
• Experimental results demonstrate that L2A optimizes compute-accuracy trade-offs and adapts efficiently across diverse tasks.

End-to-End Dynamic Sparsity for Resource-Adaptive LLM Inference

Motivation and Problem Statement

Deployed LLMs must address fluctuating resource constraints inherent to multi-tenant cloud environments. Conventional LLM inference conducts computation on a fixed architecture, activating all layers and attention heads regardless of input or runtime budget, leading to inefficiencies and fragility during resource volatility. Static compression (e.g., pruning, distillation) yields inflexible models that are suboptimal for dynamic or heterogeneous workloads, while heuristic-based dynamic methods are brittle and often fail to generalize to unseen inputs or tasks. The central challenge addressed is: how to dynamically allocate computation across input samples and budgets during inference, maximizing utility under real-world constraints.

L2A: Learning to Allocate Framework

The proposed Learning to Allocate (L2A) framework comprises lightweight gating networks injected into a pre-trained LLM to control three fundamental axes of computational cost: layer skipping, attention head pruning, and reasoning length reduction. L2A's gates are both input- and budget-conditioned, receiving, in addition to the hidden state, an explicit normalized scalar budget signal $b \in [0, 1]$ that is calibrated to real-time deployment signals (e.g., remaining time, memory headroom, queue length).

L2A pursues a unified, end-to-end optimization objective that balances:

• Task loss (cross-entropy on language modeling or QA targets)
• KL-distillation from the dense frozen model to preserve logical consistency and minimize accuracy drift under sparse execution
• Resource budget matching, penalizing compute when expected FLOPs/latency exceed the set budget
• Reasoning length regularization, minimizing the number of generated reasoning tokens under limited budget

The unified architecture is shown immediately below.

Figure 1: The L2A framework integrates budget-conditioned layer-skip, head-prune, and reasoning-stop gates into a frozen LLM, enabling dynamic compute adaptation at inference time.

During training, the model is exposed to a curriculum of budgets and input types, learning to specialize its routing policy dynamically. At inference, these gates are binarized, enforcing hard computational skips with respect to the current resource regime.

Dynamic Mechanisms

Layer Skipping

Each transformer layer is assigned a gate producing a scalar $\alpha_l \in (0, 1)$ determined via a shallow MLP conditioned on the previous hidden state and budget embedding. At inference, layers are skipped entirely if $\alpha_l < 0.5$, resulting in variable-depth computation per input.

Attention Head Pruning

Similarly, gates $\beta_{l, h}$ prune specific heads in self-attention, zeroing out those deemed extraneous given the input and current budget. This affords per-head, per-layer width adaptation, complementing depth control.

Reasoning Segment Control

Reasoning-heavy tasks often require long traces before output segmentation (e.g., chain-of-thought). A dedicated transition gate produces a per-token $\tau_t$; when exceeding threshold, the model forces an <answer> boundary, truncating further computation in the reasoning segment.

Parameter-Efficient Adaptation

Rather than full model retraining, L2A leverages LoRA for parameter-efficient adaptation, injecting low-rank trainable updates atop the frozen model backbone, substantially reducing adaptation cost and memory footprint, and allowing scalability to large backbones.

Experimental Results

Overall Performance and Pareto Frontier

Empirical results on Llama-3-8B and Qwen-3-4B demonstrate that L2A achieves near-baseline accuracy on reasoning benchmarks such as GSM8K even at high sparsity ($>34\%$ layers skipped, $>24\%$ heads pruned), while static pruning or heuristic methods incur 5–10% accuracy degradation at similar compute constraints.

Notably, L2A dynamically configures its sparsity per input, yielding a superior compute-vs-accuracy Pareto frontier compared to static or one-size-fits-all adaptive baselines. Under tightening resource budgets, the mean realized allocation is modulated on-the-fly, as illustrated in the following dynamic allocation diagnostic.

Figure 2: Alignment of the time-varying budget schedule and realized compute allocation during training evidences the coupling between external budgets and internal routing decisions.

The training loss shows monotonic convergence, with no gate collapse, as the allocation policy properly tracks the budget schedule.

Figure 3: Training task loss under a dynamic budget schedule confirms stable optimization and effective trade-off learning.

Cross-Task Generalization

Cross-domain evaluation reveals that policies trained solely on simple data (e.g., OpenWebText) over-prune and collapse on reasoning or code-gen tasks, while those trained only on hard tasks are wasteful on simple data. L2A trained across domains can generalize its difficulty perception, applying high sparsity to simple instructions and allocating extra depth to logically demanding samples, confirming robust transfer and adaptive allocation beyond memorization.

Robustness and Ablation Study

Disabling the KL-distillation objective leads to catastrophic reasoning degradation under sparsity, highlighting the essential role of teacher-student consistency. Likewise, disabling budget penalties converges to dense, non-adaptive policies, and using only layer skipping (without head pruning) forces excessive depth reduction, underperforming mixed (layer + head) strategies at equivalent latency.

Practical Implications and Future Directions

L2A enables single LLM deployments that can elastically adapt to changing resource budgets, reducing operational costs and maximizing hardware utilization while maintaining user-tailored QoS. The decoupling of adaptation across depth, width, and reasoning length directly addresses scenarios encountered in real-world cloud inference, such as spot instance preemption, user-tiered service, and fluctuating batch loads.

Practically, the dynamic gating modules incur negligible inference overhead and require no specialized runtime support. Given the constraints of current hardware, the primary realized latency gains are from layer skipping and reasoning truncation; the actual benefit from head pruning depends on future kernel-level optimizations for attention computation.

Theoretical and Future Directions

The results suggest that LLMs can be trained to internalize a notion of input difficulty aligned with resource trade-offs, opening avenues for more granular and context-aware routing. Possible directions include:

• Extending the gating approach to mixture-of-expert and multimodal architectures
• Jointly learning routing and token pruning
• Integrating online system signals (e.g., queue backpressure, memory utilization) for closed-loop resource pricing
• Automated curriculum learning for better OOD robustness of difficulty estimation

Conclusion

L2A provides a unified framework for inference-time dynamic sparsity in LLMs, enabling robust, budget-adaptive deployment with minimal quality loss. Through end-to-end, jointly optimized gating and parameter-efficient adaptation, L2A efficiently navigates the compute-accuracy spectrum, outperforming static and heuristic baselines both in main metrics and robustness under practical deployment scenarios. This work establishes a new foundation for resource-adaptive, elastic AI systems capable of cost-effective real-world operation.