PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Reasoning on the Edge

Yelysei Bondarenko, Thomas Hehn, Rob Hesselink, Romain Lepert, Fabio Valerio Massoli, Evgeny Mironov, Leyla Mirvakhabova, Tribhuvanesh Orekondy, Spyridon Stasis, Andrey Kuzmin, Anna Kuzina, Markus Nagel, Ankita Nayak, Corrado Rainone, Ork de Rooij, Paul N Whatmough, Arash Behboodi, Babak Ehteshami Bejnordi

arXiv:2603.16867v2cs.LGcs.CL

TL;DR

Edge reasoning is difficult because verbose generation, KV-cache growth, quantization challenges, and model-size constraints conflict with mobile resource budgets. The paper combines LoRA adaptation, budget-forced reinforcement learning, dynamic routing, KV-cache reuse, and hardware-aware inference, reporting shorter completions with minimal accuracy degradation and practical reasoning on compact edge models.

  • Problem

    Verbose reasoning traces, large KV caches, quantization difficulties, and limited edge resources make capable reasoning models difficult to deploy on mobile devices.

  • Method

    The framework combines modular LoRA reasoning adapters, budget-forced reinforcement learning, dynamic routing, KV-cache reuse, quantization-aware adaptation, and parallel test-time scaling.

  • Results

    2.4× average completion-length reduction, with up to 8× maximum compression and minimal task-accuracy degradation.

  • Takeaways & Limitations

    A compact 4-bit weight-quantized 7B model can achieve reasoning performance comparable to substantially larger models under edge deployment constraints.

  • Takeaways & Limitations

    The current Switcher uses a supervised classifier head; learning routing through reinforcement learning remains future work.

Abstract

from arXiv · show

Large language models (LLMs) with chain-of-thought reasoning achieve state-of-the-art performance across complex problem-solving tasks, but their verbose reasoning traces and large context requirements make them impractical for edge deployment. These challenges include high token generation costs, large KV-cache footprints, and inefficiencies when distilling reasoning capabilities into smaller models for mobile devices. Existing approaches often rely on distilling reasoning traces from larger models into smaller models, which are verbose and stylistically redundant, undesirable for on-device inference. In this work, we propose a lightweight approach to enable reasoning in small LLMs using LoRA adapters combined with supervised fine-tuning. We further introduce budget forcing via reinforcement learning on these adapters, significantly reducing response length with minimal accuracy loss. To address memory-bound decoding, we exploit parallel test-time scaling, improving accuracy at minor latency increase. Finally, we present a dynamic adapter-switching mechanism that activates reasoning only when needed and a KV-cache sharing strategy during prompt encoding, reducing time-to-first-token for on-device inference. Experiments on Qwen2.5-7B demonstrate that our method achieves efficient, accurate reasoning under strict resource constraints, making LLM reasoning practical for mobile scenarios. Videos demonstrating our solution running on mobile devices are available on our project page.

1 Introduction

Reasoning-capable LLMs could improve mobile assistants, but edge deployment is constrained by memory, generation cost, and model-size limits. The proposed framework addresses these constraints with modular reasoning adapters, dynamic routing, KV-cache sharing, budget-forced training, and parallel decoding.

  • Mobile reasoning models can support multi-step assistants that plan tasks, respond contextually, and operate across apps.
  • Edge deployment is limited by DRAM capacity, token-generation costs, long reasoning traces, large KV caches, and restricted model sizes.
  • LoRA adapters enable one frozen backbone to operate in standard chat or reasoning mode, while a switcher activates reasoning when likely useful.
  • Masked LoRA training supports KV-cache sharing between base and reasoning modes during prompt encoding without significant accuracy impact.
  • The training recipe combines supervised fine-tuning with reinforcement learning that uses budget forcing to discourage excessively long reasoning traces.
  • Parallel test-time scaling targets memory-bound decoding by running multiple decoding paths and using available on-device compute units.

2 Reasoning on Edge: System Design

The system specializes a compact base LLM with LoRA reasoning adapters, controls verbosity through reinforcement learning, and routes queries according to whether multi-step reasoning is needed. Parallel decoding and lightweight verification address the memory-bound inference phase while preserving edge deployability.

  • The pipeline aims to fine-tune a modest base LLM for competitive complex reasoning while minimizing output tokens under memory and latency constraints.
  • LoRA-based supervised fine-tuning unlocks reasoning while keeping the pretrained backbone frozen and reusable across domains.
  • Reinforcement learning combines budget forcing with answer-based rewards to encourage concise reasoning while preserving correctness.
  • A lightweight Switcher predicts whether reasoning is needed and bypasses adapters for simple queries, reducing latency and KV-cache growth.
  • The complete system includes adapter adaptation, dynamic routing, budget-forced reinforcement learning, parallel test-time scaling, quantization, and on-device execution.

3 LoRA for Modular Reasoning

LoRA adapters enable modular reasoning in Qwen2.5 instruct models, with OT3 providing strong gains and 7B adapters recovering much of dense or distilled performance at lower adaptation cost. The results also expose scale-, task-, and hyperparameter-dependent trade-offs between reasoning specialization, general capability, and efficient deployment.

  • LoRA adaptation: LoRA-based SFT uses teacher-generated reasoning traces to add modular reasoning capabilities while retaining a shared base model.Adapters can be enabled or disabled at runtime for task-specific specialization.
  • Dataset effects: OT3 produces the largest and most consistent reasoning gains across Qwen2.5-3B and Qwen2.5-7B, outperforming the smaller gains from MoT.Evaluations span mathematics, science, and coding benchmarks.
  • 7B performance: 4.24% of parameters are updated by OT3 LoRA rank 128, which reaches performance close to R1-Distill-Qwen-7B on several reasoning benchmarks for Qwen2.5-7B.Increasing rank from 64 to 128 generally narrows the gap to dense OT3 on core benchmarks.
  • Scale dependence: OT3 LoRA rank 128 substantially underperforms dense OT3 for Qwen2.5-3B, showing that adapter capacity or optimization matters more at smaller scale.This motivates further ablations of PEFT settings.
  • Coding trade-offs: SFT improves LiveCodeBench but degrades MBPP and HumanEval, reflecting a trade-off between explicit reasoning and direct-answer code generation.For Qwen2.5-7B, LoRA often retains stronger HumanEval and MBPP performance than dense OT3.
  • Hyperparameters: Rank 128 is a practical edge operating point: it performs strongly for 3B under memory constraints and is best on average for 7B, where rank effects are smaller.For 7B, average score rises from 0.388 at rank 32 to 0.402 at rank 128, while rank 256 reaches 0.397.
  • Hyperparameters: For Qwen2.5-7B, learning rates of 1e−4 and 2e−4 train reliably, whereas 5e−4 can produce unstable or collapsed runs.Batch size has comparatively minor effects, with average accuracy varying only from 0.394 to 0.398.

4 Dynamic LoRA Routing via the Switcher Module

The Switcher dynamically routes queries between a base instruct model and reasoning LoRA adapters, activating reasoning selectively during inference. This routing exposes a controllable accuracy–completion-length trade-off while supporting chunked prefill and adapter-compatible KV-cache reuse.

  • Dynamic routing: The Switcher classifies each prompt and activates reasoning adapters only when reasoning is likely to help.It enables the shared backbone to operate in standard chat or reasoning mode.
  • Switcher design: The Switcher is a lightweight prefilling-stage classification head that averages final-layer hidden states before binary routing.Its compact MLP uses hidden dimension 8, ReLU, and dropout p = 0.2.
  • Switcher design: Chunked prefill is supported by incrementally updating the sequence representation instead of buffering all prompt hidden states.This matches compute-bound prefill on edge devices.
  • KV-cache reuse: The KV-cache reuse strategy avoids re-encoding the prompt when a query is routed from base mode to reasoning mode after prefilling.Standard LoRA training otherwise expects prompt keys and values generated with adapters active.
  • Switcher training: Approximately 2k labeled samples combine easy SQuAD2.0 and MMLU math prompts with harder S1K and StrategyQA prompts for switcher training.Labels 0 and 1 indicate low and high complexity, respectively.
  • Deployment motivation: Routing simple queries to the base model can reduce aggregate token generation, latency, and power consumption while reserving reasoning adapters for harder queries.The design targets day-to-day interactions where most prompts do not require multi-step reasoning.
  • Evaluation: On MATH500, increasing the fraction of queries routed to reasoning raises accuracy smoothly from the base-model baseline toward the reasoning-only upper bound.The right-panel comparison relates average completion length to overall accuracy across switcher thresholds.
  • Evaluation: Higher accuracy requires proportionally greater average completion length, while lower-cost operating points remain available when accuracy demands are modest.The Switcher therefore provides a tunable accuracy–computation trade-off.

5 Budget Forcing and Inference-Time Compute Optimization

The section presents budget forcing as reinforcement-learning control over reasoning length, using a soft barrier and total-generation penalty to reduce verbosity while preserving accuracy. Experiments show substantial compression with minimal performance degradation, while unconstrained models waste tokens on redundant verification.

  • Motivation: CoT prompting improves complex-task performance but often incurs high latency and token costs, while unconstrained models can overthink easy problems.These costs motivate explicit compute and length control for edge inference.
  • Reward design: The method uses discrete budgets of 1000, 3000, 4000, and 6000 tokens with a multiplicative, piecewise-linear soft barrier around target budget B.The reward decays from 1.0 to 0.0 across the exceedance window, with p = 0, Llow = (1 −m)B, and Lhigh = (1 + m)B.
  • Reward design: Budget-forcing RL combines accuracy incentives with a total-response-length penalty to discourage verbose reasoning and prevent reward hacking.Penalizing total length prevents models from moving verbose reasoning after a prematurely closed reasoning block.
  • Results: βKL = 10−3 provides the best reported balance, significantly reducing generation length with negligible performance drops.βKL = 10−4 improves formatting adherence at very short lengths but causes slightly greater regression on larger, unbounded contexts.
  • Results: ∼2.4× average completion-length reduction and up to ∼8× compression are achieved while maintaining performance comparable to the base model.The approach preserves the essential reasoning backbone and human readability while pruning redundant validation loops and verbose parsing.

6 Parallel Test-Time Scaling and Reasoning

Parallel test-time scaling generates multiple reasoning candidates and aggregates them, using a lightweight verifier to select among answers without deploying a separate verifier model. Weighted voting improves accuracy over greedy decoding and standard voting, with benefits increasing as parallelism grows.

  • Parallel scaling: Parallel decoding exploits memory-bound autoregressive inference by generating multiple independent reasoning paths with limited incremental runtime overhead.This approach better uses available compute during decoding, complementing the compute-bound prefill phase.
  • Candidate selection: Majority voting aggregates candidate answers, but verifier scoring is useful when candidates disagree or the answer space is ambiguous.A lightweight verifier can score candidates and break ties that unweighted voting cannot resolve.
  • Verifier design: The verifier reuses the generator’s KV-cache and adds only minimal overhead, avoiding the storage, memory, and latency costs of a separate verifier model.A short verification prompt requires an additional prefill step while preserving reuse from the original generation.
  • Results: 72.7% accuracy with two parallel responses exceeds greedy decoding at 71.0% and standard majority voting at 70.0%.The evaluation uses a 4-bit-weight-quantized Qwen-2.5-7B-Instruct model on MATH500.
  • Results: At eight parallel responses, weighted majority voting improves upon the baselines by 10% and shows slightly lower variance than unweighted voting.Both voting methods gain steadily with increased parallelism, while verifier weighting remains consistently stronger.

7 Quantization

This section describes quantization methods for reducing LLM deployment costs and the challenges of preserving performance under low-bit representations. It presents a W4A16KV8 setup and quantization-aware modular reasoning that substantially recovers reasoning performance.

  • Background: Quantization reduces model footprint, data transfer, compute requirements, latency, and energy consumption by representing weights and activations with low-bit numbers.
  • Background: Post-training quantization is difficult because numerical outliers create a range-precision trade-off that can strongly degrade model performance.
  • Quantization setup: The deployment configuration uses INT4 per-channel quantization for linear-layer weights, INT8 KV-cache and embeddings, and INT16 remaining activations.Weights, KV-cache, and embeddings use symmetric quantization, while activations use asymmetric quantization.
  • Quantization setup: The quantized base-model pipeline trains function-preserving transformations and quantization parameters end-to-end after initialization on DCLM-Edu data.Quantization is simulated using FastForward, following the FPTQuant pipeline.
  • Base-model results: A simple min-max post-training quantization pipeline produces unacceptable accuracy and perplexity degradation, driven mainly by precision loss in 4-bit weights.
  • Base-model results: Function-preserving transformations and improved range initialization recover performance, matching full-precision accuracy on CSR with about 0.4 perplexity loss on WikiText-2.
  • Reasoning results: Quantization-aware modular reasoning recovers substantial MATH500 and GPQA performance, whereas a naïvely quantized base with a full-precision reasoning module is essentially non-functional.
  • Reasoning results: With extended training, the approach reaches performance within roughly 2% of an equivalently trained full-precision reasoning model while remaining more compact and inference-efficient.

8 Discussions and Challenges

The framework combines modular LoRA reasoning, dynamic routing, budget-forced reinforcement learning, quantization-aware adaptation, and parallel test-time scaling to balance accuracy with edge-device efficiency. Results show substantial token savings and accuracy gains, while highlighting open challenges in routing, compression, quantization, and memory reduction.

  • LoRA for modular reasoning: LoRA enabled reasoning in 3B and 7B models, with rank 128 allowing the 7B model to nearly match dense fine-tuning on challenging benchmarks.The 3B model remained more sensitive to adapter-capacity limits, and specialization occasionally degraded zero-shot performance on simpler coding tasks.
  • Dynamic LoRA routing: Dynamic routing preserves base-model speed for standard queries and activates reasoning adapters for complex queries, while masked prefill enables KV-cache reuse without re-encoding prompts.The current switcher uses supervised query-complexity classification; future work proposes reinforcement-learned routing and broader adapter selection.
  • Budget forcing: 2.4× lower average reasoning-token usage, with up to 8× maximum compression, was achieved with minimal task-accuracy degradation through budget forcing.The method uses a multiplicative penalty and a soft barrier to reduce completion length while mitigating reward hacking from formatting exploits.
  • Parallel test-time scaling and reasoning: A lightweight verifier makes parallel test-time scaling practical on resource-constrained devices by combining aggregation with a learned notion of correctness at negligible additional cost.The design could be extended to process-level reward models and interdependent parallel generation.
  • Quantization: Quantization-aware modular reasoning addresses distribution shifts from quantization, while a compact 4-bit quantized 7B model achieved reasoning performance comparable to substantially larger models.The approach combines function-preserving transformations, improved range initialization, and joint tuning of transformation and quantization parameters.
  • Remaining challenges: Reasoning remains memory-bound because token generation dominates, motivating future work below 4-bit weight quantization and semantic-aware token budgets.Current budget forcing assumes every token has equal cost despite differences in semantic utility, while smaller weight footprints could improve efficiency.

9 Conclusion

The framework makes LLM reasoning practical on resource-constrained edge devices through hardware-aware co-design across adaptation, routing, training, decoding, and cache reuse.

  • The end-to-end framework targets practical LLM reasoning under edge-device resource constraints.Its design combines parameter-efficient adaptation, dynamic routing, budget-forced reinforcement learning, parallel test-time scaling, latent verification, KV-cache reuse, and quantization-aware training.

2018. URL https://arxiv.org/abs/1806.03822. (Cited on page 10)

The evaluation uses diverse benchmarks covering mathematics, science, and coding. The suite includes competition mathematics, graduate-level science, coding generation, and Python function-generation tasks.

  • The benchmark suite spans mathematics, science, and coding domains.
  • Mathematics benchmarks: AIME 24/25 contains 30 challenging mathematics competition problems requiring multi-step reasoning with integer answers from 0 to 999.
  • Mathematics benchmarks: MATH500 contains 500 multi-step mathematical questions across algebra, geometry, number theory, precalculus, and probability.
  • Science benchmark: GPQA Diamond contains 198 multiple-choice science questions at the PhD level across physics, chemistry, and biology.
  • Coding benchmarks: LiveCodeBench v2 comprises 511 coding problems, evaluated here only in the Code Generation scenario.
  • Coding benchmarks: HumanEval and HumanEval+ comprise 164 Python function-generation problems, with HumanEval+ increasing verification unit tests 80 times.

B LoRA ablation study

The appendix documents parameter-efficient fine-tuning and budget-forcing evaluation, including ablations, training settings, and qualitative comparisons of reasoning traces.

  • LoRA ablation study: The PEFT ablation varies learning rates, batch sizes, and LoRA ranks for 3B and 7B backbones.Models use a 50,000-entry OpenThoughts3 subset for one epoch and are evaluated on mathematical and scientific benchmarks.
  • LoRA ablation study: Tables 12 and 13 report ablation results for Qwen2.5-3B-Instruct and Qwen2.5-7B-Instruct, marking the best result within each learning-rate subgroup.
  • Budget-forcing training: Table 14 lists the hyperparameters used for budget-forcing reinforcement-learning training.The table states that all experiments share these settings unless otherwise noted.
  • Qualitative comparisons: Figures 9 and 10 qualitatively compare unconstrained baseline reasoning traces with budget-forced traces.
  • Number theory reasoning: The baseline identifies the relevant number-theory property early but performs redundant verification, whereas the forced trace applies the reasoning directly.
  • Pattern recognition: For pattern recognition, the baseline spends nearly 1000 tokens validating alternatives, while the forced trace directly retrieves the formula and computes the requested term.
Loading 2603.16867v2…