Bridging On-Device and Cloud LLMs for Collaborative Reasoning: A Unified Methodology for Local Routing and Post-Training
Abstract
Device-cloud collaboration holds promise for deploying large language models (LLMs), leveraging lightweight on-device models for efficiency while relying on powerful cloud models for superior reasoning. A central challenge in this setting is determining, for each incoming query, whether it should be processed locally or offloaded to the cloud. Existing approaches typically rely on external routers, which often struggle to determine difficulty from the prompt itself, especially for tasks involving complex reasoning. Motivated by this limitation, we propose enabling on-device LLMs to decide internally whether to invoke cloud assistance at inference time, with this capability instilled through reinforcement learning based post-training. Casting on-device LLM post-training as a reward maximization problem, we design hierarchical rewards to encourage local problem solving and judicious cloud offloading. To solve the resulting problem, we develop an algorithm featuring a group-level policy gradient that stabilizes optimization, together with adaptive prompt filtering that provides complementary learning signals to mitigate policy collapse (i.e., exclusive local execution or exclusive cloud offloading). Extensive experiments on on-device-scale LLaMA and Qwen models across multiple reasoning benchmarks show that our method consistently outperforms baselines and significantly narrows the gap to full cloud LLMs.
Lay Summary
Today's best AI assistants often depend on cloud servers to answer difficult questions. Smaller AI models that run directly on phones or laptops are faster and cheaper, but they are not as good at complex reasoning. This creates a practical question: when should a device answer on its own, and when should it ask the cloud for help? Our work helps the smaller on-device model learn to make that choice by itself. Instead of relying on a separate decision system, we train the model to both improve its own problem-solving skills and recognize when a harder question needs help from a stronger cloud model. It learns to solve easier questions locally and use the cloud only when necessary. This leads to a better balance between speed, cost, and quality. In our experiments, the method worked better than existing approaches and brought small on-device models closer to the performance of much larger cloud systems.