M o D ification: Mixture of Depths Made Easy (2025)

Abstract: Long-context efficiency has recently become a trending topic in serving large language models (LLMs). And mixture of depths (MoD) is proposed as a perfect fit to bring down both latency and memory. In this paper, however, we discover that MoD can barely transform existing LLMs without costly training over an extensive number of tokens. To enable the transformations from any LLMs to MoD ones, we showcase top-k operator in MoD should be promoted to threshold-p operator, and refinement to architecture and data should also be crafted along. All these designs form our method termed MoDification. Through a comprehensive set of experiments covering model scales from 3B to 70B, we exhibit MoDification strikes an excellent balance between efficiency and effectiveness. MoDification can achieve up to ~1.2 \times speedup in latency and ~1.8 \times reduction in memory compared to original LLMs especially in long-context applications.

## introduction long-context efficiency is turning to be one of the core concerns in large language model (llm) serving @xcite @xcite . typically, a long context can incur dramatically huge latency and memory overhead either at prefilling or decoding stage @xcite . to address this consumption, strategies like speculative decoding @xcite @xcite and key-value (kv) cache compression @xcite @xcite have already been utilized to enhance the latency and memory, respectively. however, dedicating efforts to separate components is not always ideal. as a consequence, mixture of depths (mod, @xcite has recently been introduced to simultaneously consider both * dawei song is the corresponding author. figure 1 : the sub-optimality of mod due to the use of top-k operator. for efficiency, the efficiency improvement is restricted because top-k operator is not cheap and unimportant computations are preserved by retaining a constant number of tokens at every layer. for effectiveness, the effectiveness is undesired because dominated computations are saved by skipping a constant number of tokens at every layer. the sub-optimality is mainly due to the improper use of top-k operator. the top-k operator is basically employed at each layer to determine which tokens should be skipped in the computation of the layer. unfortunately, for efficiency, the top-k operator is computationally expensive @xcite and forces a constant number of tokens retained, possibly restricting the efficiency improve-ment brought by computation saving at one unimportant layer, as shown in figure 1 (left). for effectiveness, it deviates from the distribution of layer significance @xcite , potentially resulting in undesired computation saving at one dominated layer, as shown in figure 1 (right). these disadvantages largely impede the applications of mod. in this paper, in order to broaden the scope of mod, we target at converting existing llm checkpoints to mod ones with our designed modification. 1 the designs of modification follow two guidelines: 1) alleviating the sub-optimality of mod, and 2) minimizing the training of mod. under the guidelines, we firstly use threshold-p operator as an drop-in replacement of the top-k operator. on one hand, threshold-p operator is not much computationally expensive and allows any number of retained tokens; on the other hand, it can make the layer itself determine the priority of tokens through adaptive gates across layers. we secondly find around 10b tokens with decent diversity are fairly enough in the conversion process. we conduct experiments on model scales ranging from 3b to 70b. the efficiency and effectiveness of modification are verified on such a wide array of scales. with acceptable performance decline, modification can achieve up to ∼1.2× speedup in latency and ∼1.8× reduction in memory compared to original llms especially in long-context applications. on the contrary, mod can unexpectedly lead to slow-down in latency under the same experimental settings. further ablations and discussions validate the trustworthiness of the design choices in modification. while the transformer architecture admits the scaling of pretraining, a critical bottleneck during inference is also revealed by the transformer. in other words, the transformer architecture mainly grants parallelism during training, yet only generates tokens one followed by another and requires kv caches adequately stored during inference. this bottleneck may easily slow down the expected speed and shoot over the provided hardware, especially for long-context circumstances @xcite . to this demand, many approaches have been proposed to either counter the latency or the memory increment @xcite . among these approaches, speculative decoding @xcite is a representative cluster in reducing the latency and kv cache compression @xcite is another symbolic group in mitigating the memory. however, for an integrated inference system, it is might be somehow redundant to conquer the latency and the memory separately. consequently, later solutions to some extent shift the focus to conditional computation @xcite , in which early exiting @xcite and mixture of depths (mod, @xcite are of central roles. ## modification in this work, we put forward such a method, named modification, that is aimed at converting existing llm checkpoints to mod ones with slight training compute. to circumvent the sub-optimality of original mod, two designs are engaged, in which one is associated with architecture and the other is connected with data. the first one is to use thresholdp operator in place of the top-k operator, as shown in figure 2 . threshold-p operator is not much computationally expensive and allows any number of retained tokens. we then instantly discover this threshold-p replacement correlates modification to mixture of experts (moe) @xcite and we can empower modification with cutting-edge techniques that have been widely used. the second one is to enhance the data diversity and constrain the training to the scale of 10b tokens. we luckily unveil this data scale is fairly enough in practice. architecture after applying threshold-p operator to equation 2, we make a further refinement to the formulation. the formula can now be briefly formed as below: in fact, the conditional provided by f i can be interpreted from an moe view as below: the proposed data mixture generally covers domains like webs, wikis, books, codes, maths, etc. in doing so, we expect modification could be trained to approximate original mixtures of any existing llm checkpoints. implementation we train modification and mod with an overall batch size of 1,024 and a sequence length of 4,096. this leads to 4m tokens per optimization step. the learning rate is 3e-5, and weight decay is 1e-1. the learning rate is scheduled with a linear warm-up and a cosine decay, where the warm-up stage takes 1% optimization steps and the decay stage takes all the left optimization steps. the gradients will be properly clipped to keep the norm under 5e-1. the training precision is bfloat16. deepspeed @xcite , flashattention @xcite , and gradient checkpointing @xcite are necessarily enabled to make the training possible with 16 nvidia a100 gpus. the k for the top-k operator is 512 by default, and the p for the threshold-p operator is 0.5 for minima-2-3b as mentioned earlier but 0.55 for llama-2-{7,13,70}b since slightly larger threshold is demanded for larger llms in our trials. the best coefficient α for layer load reducing objective is 0.01 in our developments. results from the main comparison results between mod and modification in one may think that mod is not properly configured in training and thus shows very unstable and almost meaningless results. to demonstrate that the training of mod does not suffer from fluctuations such like loss spikes, we hereby display the training losses as in figure 3 . from the loss curves, we conclude that both mod and modification are rigorously optimized. besides, one may argue that the sparsity of mod influenced by k does not align with the sparsity of modification influenced by p, thus exhibiting unreasonable results. to show that the sparsity of mod is not the key factor, we confirm in to further enhance the plausibility of the comparison, we also compare modification to an alternative efficient method named shortgpt @xcite that is designed to directly drop unimportant layers, therefore decreasing latency and memory, as in figure 4 . the considered shortgpt baseline is configured as preserving 22 out of 32 layers (i.e., number of activated tokens is 90,112). the major conceptual difference between modification and shortgpt lies in that the layer skipping of shortgpt happens at task level. we observe that shortgpt may lead to more apparent efficiency gains, nevertheless, it results in unacceptable performance degradation. this phenomenon implies that modification has considerable advantages over other efficient designs. ablation studies to get an in-depth understanding of the inner working mechanisms of modification, we perform a series of ablation studies on the threshold value p, layer loading reducing coefficient α, etc. threshold value p. we show the impact of the threshold value p in figure 5 , from which we inspect that while the memory consumption is always reduced across all threshold values, the latency and the performance can vary a lot from one threshold to another. disappointingly, though smaller threshold values can cut down memory consumption, they can cause increased latency as a result of limited computation saving. for minima-2-3b in the plot, the most suitable p value is 0.5 regarding the trade-off between efficiency and effectiveness. however, as aforementioned, larger models such as llama-2-70b urge larger p values like 0.55 due to potentially larger structural sparsity. layer load reducing coefficient α. in particularly, the coefficient trades effectiveness for efficiency. for instance, as α increases itself by an order of magnitude, the latency diminishes, so is the performance. in contrast, without the layer load reducing objective (i.e., when the coefficient is 0), modification tends to save very few computations. in this case, the latency is even enlarged. in summary, we assert that α would be better set to 0.01 to reach equilibrium. to confirm the interleaved strategy is the best choice, we compare the interleave startegy to a tactic that applies modification to top half layers and another that applies modification to full layers, as in in figure 6 , we unleash that mod saves computations with the sub-optimality due to the use of topk operator. the forced constraint that a constant number of tokens should be retained may: 1) limit the computation saving (efficiency) when computa-tions of more tokens should be skipped, 2) affect the computation preserving (effectiveness) when computations of fewer tokens should be skipped. conversely, modification permits more flexible computation savings, leading to a better balance between efficiency and effectiveness. ## long-context applications while we have fully justified the usefulness of modification in comparably short-context scenarios, we are obliged to substantiate the potential of modification in long-context applications. we take two long-context applications as testbeds, i.e., longtext generation that requires long decoding and long-context retrieval that requires long prefilling. as a prerequisite, we extend the context capabil-ity of modification-7b from 4k to 32k. we adopt ntk-aware interpolation @xcite and conduct post training on pg19 dataset @xcite . in long-text generation scenario where prompts sampled from proof-pile @xcite , modification-7b can achieve up to ∼1.2× speedup in latency and ∼1.8× reduction in memory in comparison to llama-2-7b. and through longcontext extension, modification is compatible with long-context modeling and consistently behaves in low perplexity without compromise on efficiency. ## related work efficient language models improving efficiency of language models is a long-standing and challenging task @xcite . besides system-level efficiency optimizations, algorithmlevel efficiency optimizations mainly lie in model compression and efficient architecture. model compression approaches @xcite aim to reduce the model size directly. these compression approaches can be grouped into quantization @xcite , pruning @xcite , and distillation @xcite @xcite . efficient architectures @xcite instead aim to reduce the quadratic time complexity of attention to quasi-linear one. mixture of depths @xcite stands at the intersection of model compression and efficient architecture. on one hand, mixture of depths can skip layers, which can be viewed as a type of pruning. on the other hand, mixture of depths can drop tokens, which can be viewed as a sort of complexity reduction in time. conditional computation a third perspective towards improving efficiency of language models would be conditional computation @xcite . it aims to relieve unnecessary computations by conditionals. this dynamic property can to the maximum extend preserve effectiveness while improving efficiency. typical work in this area would be mixture of experts @xcite , early exiting @xcite , and mixture of depths. and mixture of depths is recently viewed as one of the best choices among available conditional computation approaches. ## conclusions it is no doubt that mod is an ideal choice to lift the efficiency of llms. in this paper, we would like to address the concern that mod comes with a must that costly training (from scratch) should be executed. via both architecture and data designs, we make it possible to convert existing llms to mod ones with appropriate data. the designed modification successfully outweighs mod in both efficiency and effectiveness.