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.
Figures & tables
Figure 1: Overview of the proposed efficient reasoning framework for edge devices. (a) The model architecture utilizes parameter-efficient LoRA adapters and a lightweight switcher to dynamically route queries. This design allows the base model and the reasoning-activated mode to seamlessly share a reusable KV cache during prefill. (b) Parallel test-time scaling strategy, generating multiple reasoning streams concurrently to improve accuracy without severe latency penalties. (c) The end-to-end deployment pipeline, illustrating the progression from multi-stage training (SFT and budget-forced RL) to quantization, model export, and final on-device execution.
Figure 2: Architecture of the Hybrid Reasoning Model. The pipeline begins with a compact base LLM, which is specialized for reasoning via LoRA-based supervised fine-tuning (SFT). To enforce concise generation and prevent excessive verbosity, these adapters undergo reinforcement learning (RL) with Budget Forcing. Finally, a lightweight Switcher module is introduced to act as a reasoning-needed classifier, creating a hybrid model that dynamically routes incoming queries to either the fast base model or the specialized reasoning adapters based on task complexity.
FT data
LoRA
AIME24
AIME25
MATH500
GPQA
AMC23
LCB
HE
HE+
MBPP
MBPP+
Qwen-3B
-
-
0.09
0.03
0.66
0.31
0.36
0.23
0.40
0.64
0.72
0.61
MoT
-
0.20
0.22
0.80
0.37
0.61
0.38
0.43
0.38
0.34
0.29
OT3
-
0.51
0.43
0.92
0.40
0.83
0.50
0.48
0.43
0.36
0.30
OT3
rk 128
0.18
0.21
0.79
0.32
0.55
0.23
0.40
0.35
0.39
0.32
Qwen-7B
-
-
0.10
0.17
0.76
0.37
0.60
0.36
0.83
0.75
0.80
0.68
R1∗
-
0.55
0.40
0.92
0.49
0.89
0.59
0.41
0.38
0.42
0.36
Table 1: Results of the main experiments. FT data indicates the finetuning dataset. The LoRA column specifies the adapter rank, or its absence for dense finetuning. LCB refers to LiveCodeBench, and HE denotes HumanEval. R1∗ corresponds to the R1-Distill-Qwen-7B model released by DeepSeek; we evaluated this model following our standardized evaluation protocol.
LR
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
1\mbox{e-}4
0.040
0.025
0.573
0.277
0.281
0.239
2\mbox{e-}4
0.051
0.033
0.574
0.274
0.299
0.246
5\mbox{e-}4
0.051
0.038
0.556
0.280
0.275
0.240
Table 2: Average performance grouped by learning rate for Qwen2.5-3B-Instruct.
Rank
%TP
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
32
1.94%
0.047
0.020
0.568
0.270
0.284
0.238
64
3.88%
0.040
0.023
0.570
0.283
0.284
0.240
128
7.76%
0.054
0.038
0.560
0.272
0.287
0.242
256
15.52%
0.048
0.048
0.573
0.282
0.284
0.247
Table 3: Average performance grouped by LoRA rank for Qwen2.5-3B-Instruct. %TP denotes the percentage of trainable parameters.
Batch size
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
32
0.054
0.030
0.562
0.277
0.288
0.242
64
0.048
0.038
0.566
0.278
0.287
0.243
128
0.040
0.028
0.575
0.277
0.281
0.240
Table 4: Average performance grouped by batch size for Qwen2.5-3B-Instruct.
LR
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
1\mbox{e-}4
0.158
0.132
0.775
0.366
0.533
0.393
2\mbox{e-}4
0.170
0.148
0.779
0.354
0.543
0.399
5\mbox{e-}4
0.148
0.112
0.635
0.350
0.439
0.337
Table 5: Average performance grouped by learning rate for Qwen2.5-7B-Instruct.
Rank
%TP
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
32
1.06%
0.152
0.119
0.776
0.357
0.534
0.388
64
2.12%
0.159
0.150
0.779
0.355
0.535
0.396
128
4.24%
0.178
0.141
0.784
0.367
0.541
0.402
256
8.48%
0.165
0.152
0.771
0.359
0.536
0.397
Table 6: Average performance grouped by LoRA rank for Qwen2.5-7B-Instruct. %TP denotes the percentage of trainable parameters.
Batch size
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
32
0.154
0.148
0.781
0.355
0.530
0.394
64
0.164
0.146
0.774
0.366
0.538
0.398
128
0.172
0.127
0.776
0.358
0.541
0.395
Table 7: Average performance grouped by batch size for Qwen2.5-7B-Instruct.
Figure 3: Impact of the Switcher module on MATH500. Left: Combined model accuracy as the fraction of queries routed to the reasoning adapters. Right: Average completion length versus overall accuracy across different switcher thresholds.
Figure 4: Average Completion Length Distributions. Left: Evaluation with a forced maximum completion length of 4K tokens. Right: Evaluation with a maximum of 6K tokens. Note that distribution tails extending below zero or above the maximum budget are standard artifacts of Kernel Density Estimation (KDE) curve smoothing. The progression from the baseline ( purple ) through the intermediate ( blue ) to the final RL fine-tuned checkpoint ( green ) demonstrates stable, progressive learning of concise generation ( βKL=10−3 ).
Figure 5: Average Completion Length Comparison. Left: C.D.F. of the average completion length for base model, orange curve, and RL fine-tuned one, green curve, with βKL=1.e−3 . We considered a maximum completion length of 6K. Right: Reduction in completion length from the RL fine-tuned model. We use the same models in the left plot. The RL fine tuned achieved average reduction length of 2.38±0.07 .
Model
Accuracy ↑(%)
Budget = 1K
Budget = 2K
Budget = 4K
Budget = 6K
Average (Budget = 32K)
SFT Baseline (r=128)
34
57
73
83
95
BF RL βKL=1.e−3
62
78
85
90
92
BF RL βKL=1.e−4
72
80
84
85
90
Table 8: Performance Results on MATH500. Results for different budget values. For budgets ∈{1K,2K,4K,6K} , we manually force the model to generate the final answer once the limit is hit. In the last column, we do not force the model to generate the answer. We let it generate tokens untill the budget hit.
Figure 6: Qualitative comparison on algebraic simplification. Middle: The Baseline trace correctly identifies the difference of squares strategy immediately but engages in excessive self-verification, re-calculating the result via expansion, direct computation, and alternative factorizations. Bottom: The Budget Forced trace recognizes the nested difference of squares structure and executes the solution linearly without redundant checking.
Figure 7: Qualitative comparison on modular arithmetic. Middle: The Baseline trace correctly computes the sum and remainder immediately but engages in extensive, redundant verification using four different methods (step-by-step addition, digit sum rule, pairing, and re-calculation). Bottom: The Budget Forced trace performs the direct calculation and returns the result without hesitation.
Parallel Responses
1
2
4
6
8
Greedy (baseline)
71.0
-
-
-
-
Majority Vote
69.9 ± 1.3
70.0 ± 1.3
75.1 ± 1.0
76.6 ± 1.0
77.5 ± 0.8
Weighted MV (ours)
69.9 ± 1.3
72.7 ± 1.0
76.1 ± 0.9
77.5 ± 0.8
78.2 ± 0.7
Table 9: Accuracy of our lightweight verifier weighted majority voting compared to majority voting without verifier and greedy decoding on MATH500. The mean and standard deviation are computed from 20 random draws from 16 independent 4bit-weight-quantized Qwen-2.5-7B-Instruct responses.
Figure 8: Function-Preserving Transformations . We use 4 transform types from FPTQuant: scale-and-rotate transform \mathscompmodernTk merged into query and key, a per-channel scaler \mathscompmodernTu merged into up and down projection, and \mathscompmodernTv that consists of invertible matrices per head merged into value and output weights, and a rotation matrix \mathscompmodernTr for rotating residuals (shared across layers). After training of the transforms is complete, the transformation parameters from each ‘merge group’ are merged into the original model weights W .
Method
Bitwidth
Lp
\mathscompmodernT
train
WikiText-2 ( ↓ )
CSR ( ↑ )
MMLU ( ↑ )
Full-precision
BF16
-
-
-
6.85
72.90
74.28
Min-max quantization
W4A16KV8
-
-
-
102.4
51.71
62.35
W4A16KV8
✓
-
-
9.18
65.83
67.59
W4A16KV8
-
✓
-
8.48
67.85
69.06
W4A16KV8
✓
✓
-
7.53
70.68
72.26
FPTQuant∘ (ours)
W4A16KV8
✓
✓
✓
7.26
72.94
72.81
Table 10: Quantized W4A16KV8 base model (Qwen2.5-7B-Instruct) results . We report Wikitext perplexity, average 0-shot CSR, and 5-shot MMLU accuracies. ‘ Lp ’ denotes the use of Lp range initialization, \mathscompmodernT = using the set of mergeable transforms, ‘train’ = end-to-end training of transformation and quantization parameters.
Bitwidth
FPTQuant∘
QAMR
N
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
BF16
-
-
50k
21.8
20.3
82.6
38.6
65.2
45.70
W4A16KV8
-
-
50k
0.0
0.0
0.0
0.0
0.0
0.00
W4A16KV8
-
✓
50k
17.3
6.3
75.6
33.0
64.0
39.25
W4A16KV8
✓
✓
50k
23.3
15.0
79.6
33.7
57.0
41.72
BF16
-
-
1.2M
53.3
33.0
94.0
39.9
82.5
60.54
W4A16KV8
✓
✓
1.2M
46.6
36.6
89.6
37.8
80.0
58.12
Table 11: Quantized W4A16KV8 reasoning model results (Qwen2.5-7B-Instruct base) . We report AIME24/25, MATH500, GPQA, and AMC23 accuracies (higher is better). N = number of OT data examples used for training reasoning module.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
LR
BS
Rank
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
0.0001
32
32
0.05
0.01
0.55
0.25
0.24
0.220
0.0001
32
64
0.03
0.02
0.58
0.29
0.28
0.240
0.0001
32
128
0.08
0.03
0.56
0.28
0.30
0.250
0.0001
32
256
0.03
0.04
0.57
0.28
0.28
0.240
0.0001
64
32
0.05
0.01
0.54
0.27
0.28
0.230
0.0001
64
64
0.03
0.03
0.57
0.29
0.29
0.242
Appendix
Table 12: Ablation results for Qwen2.5-3B-Instruct. LR denotes learning rate and BS stands for batch size. In each learning rate subgroup, the best performance is marked in bold.
LR
BS
Rank
AIME24
AIME25
MATH500
GPQA
AMC23
Avg
0.0001
32
32
0.18
0.12
0.79
0.36
0.50
0.390
0.0001
32
64
0.13
0.16
0.78
0.36
0.51
0.388
0.0001
32
128
0.17
0.11
0.78
0.37
0.55
0.396
0.0001
32
256
0.14
0.14
0.79
0.35
0.54
0.392
0.0001
64
32
0.14
0.12
0.77
0.35
0.54
0.384
0.0001
64
64
0.16
0.17
0.76
0.38
0.56
0.406
Appendix
Table 13: Ablation results for Qwen2.5-7B-Instruct. LR denotes learning rate and BS stands for batch size. In each learning rate subgroup, the best performance is marked in bold.
Hyperparameter
Value
Optimizer
AdamW
Learning Rate
2×10−5
LR Scheduler
Cosine
Warmup Ratio
0.05
Batch Size (Global)
256
Generations per Prompt ( G )
8
Appendix
Table 14: GRPO Training Hyperparameters. All experiments share these settings unless otherwise noted.
Figure 9: Qualitative comparison on number theory reasoning. Text highlighted in red denotes redundant verification and verbal parsing, while bold text identifies essential reasoning steps. We use “ [ … ] ” as a placeholder for brevity. Top: Prompt. Middle: The Baseline trace correctly identifies the property ( p2 ) early on but falls into extensive, redundant verification loops checking composite numbers and re-listing primes (highlighted in red). Bottom: The Budget Forced trace directly applies the prime-square property and computes the result without unnecessary hesitation or syntactic noise.
Figure 10: Qualitative comparison on pattern recognition. Middle: The Baseline trace correctly identifies the formula 2n−1 initially but spends nearly 1000 tokens validating it against alternative arithmetic formulas ( a+(n−1)d ) and hypothetical user errors (confusing term number with value, testing n2 , etc.). Bottom: The Budget Forced trace directly retrieves the formula and computes the specific term requested.
Real-time AI experiences call for on-device large language models (OD-LLMs) optimized for efficient deployment on resource-constrained hardware. The most useful OD-LLMs produce near-real-time responses and exhibit broad hardware compatibility, maximizing user reach. We present a methodology for designing such models using hardware-in-the-loop architecture search under mobile latency constraints. This system is amenable to industry-scale deployment: it generates models deployable without custom kernels and compatible with standard mobile runtimes like Executorch. Our methodology avoids specialized attention mechanisms and instead uses attention skipping for long-context acceleration. Our approach jointly optimizes model architecture (layers, dimensions) and attention pattern. To efficiently evaluate candidates, we treat each as a pruned version of a pretrained backbone with inherited weights, thereby achieving high accuracy with minimal continued pretraining. We leverage the low cost of latency evaluation in a staged process: learning an accurate latency model first, then searching for the Pareto-frontier across latency and quality. This yields MobileLLM-Flash, a family of foundation models (350M, 650M, 1.4B) for efficient on-device use with strong capabilities, supporting up to 8k context length. MobileLLM-Flash delivers up to 1.8x and 1.6x faster prefill and decode on mobile CPUs with comparable or superior quality. Our analysis of Pareto-frontier design choices offers actionable principles for OD-LLM design.
We present IronLLM-0.6B, a 654M-parameter language model designed for efficient on-device inference. IronLLM-0.6B combines a hybrid attention architecture with X-MTP, a lightweight shared-KV multi-token prediction design that eliminates per-depth KV-cache replay and employs a lightweight verification head for rollback-free drafting, achieving a 1.48x decoding speedup. The model is pretrained on approximately 6.2 trillion tokens using a quality-oriented data pipeline and is further post-trained with Multi-Domain On-Policy Distillation to integrate capabilities from domain-specialized teachers. To better meet the low-latency requirements of on-device scenarios, IronLLM-0.6B adopts an Instruct-Only design. Evaluations show that IronLLM-0.6B achieves competitive performance relative to larger models such as Qwen3.5-0.8B and MiniCPM5-1B, while producing more concise responses on many tasks. We further present IronLLM-0.6B-Light, which replaces RMSNorm with Dynamic Tanh and simplifies several computationally expensive components to improve inference and quantization efficiency. Together, the IronLLM models provide an effective performance-efficiency trade-off for resource-constrained deployment.
Large Language Models (LLMs) have demonstrated remarkable capabilities across a range of Natural Language Processing (NLP) tasks, but their high computational and memory demands pose significant challenges for deployment on resource-constrained edge devices. Existing approaches to model compression and optimization often rely on coarse-grained pruning or quantization, which can compromise accuracy or require re-training and fine-tuning. In this work, we introduce SelectInfer, a neuron-level optimization framework that enables efficient LLM inference on edge devices through selective neuron loading and computation. By profiling and identifying both task-specific and general-purpose neurons using an offline LLM profiler, SelectInfer implements two key optimizations: selective loading, which reduces memory footprint by selectively loading a subset of neurons that were identified to be most important during the offline stage, and selective computation, which dynamically computes only the most relevant neurons at runtime. Evaluation across multiple datasets shows that SelectInfer achieves significant reductions in memory footprint and computation while preserving task performance, making it a practical step towards enabling LLM deployment on edge devices
Huzaifa Shaaban Kabakibo, Eric Schniedermeyer, Artem Burchanow +1
Department of Computer Networks, Paderborn University, Paderborn, Germany