Stepped MoE: Segment-Level Routing with Configurable Inference Complexity

Stepped MoE. Segment-Level Routing with Configurable Inference Complexity. By Arnav Kundu and five others, from Apple. Published on the fifth of October. No code repository is linked to it.

Training a large language model costs a great deal, and once it is trained it still has to be adapted for deployments that differ in how much computation they can afford. Two families of approach already address parts of that. Elastic architectures make deployment flexible, letting one model be served at different sizes. Sparsely activated models make computation input-adaptive, spending effort according to what the input needs. The trouble is that existing approaches treat these two dimensions independently, as if a model had to choose between adapting to its hardware and adapting to its input. On top of that, models meant for on-device edge inference must fit the memory and compute limits of the device that serves them. For those devices, fitting is not optional.

The proposal is a unified framework that combines elastic structures with sparsely gated architectures, so that a single model adapts at once to deployment constraints and to task requirements. The backbone is conditioned on two things together: the context it is reading, and a target efficiency specification. That second input is what turns the accuracy and efficiency trade-off into something chosen at inference time, with fine-grained control, rather than something fixed when the model was built.

Inside the model, the sub-networks are elastically nested, one capacity sitting within the next. The model learns to activate the parameters relevant to the task within those nested sub-networks. That is how one set of weights can span several capacity points while still routing according to the input. Choosing a smaller capacity does not switch off the input-adaptive behaviour. It narrows the space in which that behaviour operates.

The experiments produce one model that can be used at four capacity points, starting from one billion parameters and rising in steps of one billion. At those sizes it is more accurate than the corresponding dense models, by two to five percent on knowledge-intensive benchmarks. Against their static versions it performs at par. And it delivers latency metrics similar to the dense models, so the added flexibility is not paid for in speed.

The practical case is made in three parts. Sharing parameters across capacity points saves disk space on the device, because there is one model rather than four. Serving can flex according to the DRAM and compute available at the moment. And the results are more accurate than the dense alternatives. Two kinds of adaptivity that used to live in separate models now live in one, and the device picks how much of it to use.