Stepped moe adapts large language models to device limits
Stepped MoE: Segment-Level Routing with Configurable Inference Complexity
Machine LearningArtificial IntelligenceComputation and Language
Summary
Training very large language models takes a lot of computer power, and running them on different devices with limited memory or speed is hard. The authors created a new kind of model that can change its size and which parts it uses based on what the device can handle and the task needs. This model can act like smaller or bigger versions depending on the situation without losing accuracy. Their tests show it can be more accurate than fixed-size models while using less memory on devices like phones or computers.
What this means in practice
- •For mobile app developers: Deploy language models that dynamically adjust size and memory use to fit phones and tablets with different hardware limits.
- •For cloud service engineers: Provide scalable language model APIs that adapt computation and capacity based on client device and request complexity.
Authors
Arnav Kundu, Zhaoyang Xu, Bairu Hou, Chang Gao, Reed Li, Tao Lei
Abstract
Training large language models (LLMs) is resource-intensive, and adapting them for diverse deployment scenarios with varying computational constraints remains challenging. While elastic architectures enable flexible model deployment and sparsely activated models allow input-adaptive computation, existing approaches treat these dimensions independently. Moreover, models catered towards on-device edge inference need to conform to the memory and compute limitations of the serving devices. In this paper, we introduce a unified framework that combines elastic structures with sparsely gated architectures to create models that adapt simultaneously to both deployment constraints and task requirements. Our approach employs a model backbone that conditions on both the context and target efficiency specifications, enabling fine-grained control over the accuracy-efficiency trade-off at inference time. The model learns to activate task-relevant parameters within elastically-nested sub-networks, allowing a single model to span multiple capacity points while maintaining input-adaptive routing. Through experiments we demonstrate that we can create a model that allows the flexibility to use 1,2,3,4 billion parameters while being more accurate than their dense counter-parts (2-5\% on knowledge-intensive benchmarks) and at par with their static versions while delivering similar latency metrics as dense models. Overall, we save on device disk space by sharing the model parameters, allow flexibility of serving based on DRAM and compute available while delivering more accurate results.