D-LLM: A Token Adaptive Computing Resource Allocation Strategy for Large Language Models

We truly appreciate the reviewer for the valuable feedback. We have carefully considered the comments and suggestions and would like to address each of the concerns raised.

Q1: Token Granularity

We follow the tokenization released by corresponding LLMs. Today's LLMs utilize their own tokenization respectively. For example, Llama 2[1] employs a BytePair Encoding (BPE) algorithm with a vocabulary of 32k tokens, while Llama 3[2] uses the tiktoken tokenizer, expanding the vocabulary size up to 128k tokens. At the same time, LLMs are trained based on their specific tokenizations. The binding relationship between tokenization and LLM means that if we simply change the tokenization, the LLM will lose its ability. If we want to compare different levels of token granularity, training extra LLMs from scratch is necessary but unaffordable. Our method has been verified on llama2 and llama3, which demonstrates that D-LLM is applicable and effective across various tokenizers.

[1] Llama 2: Open Foundation and Fine-Tuned Chat Models

[2] The Llama 3 Herd of Models

Q2: Decision Module Criteria

• The decision module outputs a probability distribution with dimensionality 2, denoted as $b=[b_1,b_2]$ where $b_1+b_2=1$. The transformer layer is decided to be executed or skipped according to the max arguments in $b$. Specifically, if $b_1<b_2$, the layer should be executed; if $b_1\geq b_2$, the layer should be skipped.
• Decision-making depends on the max argument in vector $b$, therefore we don't have to set or optimize any thresholds.

Q3: Hyper-Parameters

Hyper-parameters in our method include target acceleration rate $\Omega$, reserved initial tokens $m$, and the weight factor $\alpha$ of $\mathcal{L}\_{rate}$. We provide additional descriptions and details for these hyper-parameters as follows.

• Target acceleration rate $\Omega$ is user-defined for expected computation overheads. We elaborate the performance under different $\Omega$ in Figure 2 (line 198-199). More details are available in Table 4 in A.2 (line 532-540).

• The importance of initial tokens is pointed out in recent works for LLMs' generalization to long contexts[1,2]. Initial tokens encode strong absolute position information and collect significant attention scores, which makes it necessary to be kept when pruning context. Our experiments in Table 3 (line 270-271) demonstrate that initial tokens reserved in the proposed eviction policy lead to better performance.

• We conduct additional analysis on hyper-parameters $\alpha$ in Eq.14. (line 213). Details are available in global rebuttal #2.

Q4: Compatibility with Different LLM Architectures

• Our proposed D-LLM is designed for decoder-only transformers, which is the mainstream architecture used by current LLMs. With respect to other architectures like encoder-decoder models, our method needs further adaptation, especially the strategy of KV-cache eviction.

• We are currently focused on the decoder-only architecture but are open to adapting our method to other LLM architectures in the future.

Q5: Impact on Latency

• We make a detailed analysis on latency and other overheads taken by decision modules in global rebuttal #1. In real-time applications, D-LLM with actual acceleration rate 40% raises the generating speed from 32 token/s to about 41 token/s. The latency is negligible compared to the gain of generating speed.
• We also have considered several optimization directions to minimize the latency, for example, reducing the complexity of decision module; utilizing global decision module to predict the routes for all transformer layers; reducing the frequency of utilizing decision module when generating; and so on.

Q6: Comprehensive Evaluation

• In terms of benchmark selection, we mainly follow the popular evaluation datasets used by the mainstream LLMs[1]. Currently, D-LLM is designed and verified on specific domain tasks, like commonsense reasoning and math solving. We used a considerable amount of benchmarks compared with other research on LLMs[2,3]. Certainly, we plan to expand our scale of training of D-LLM in the future, including large-scale common datasets for LLMs' training and the size of training model's parameters. We also plan to assess D-LLM's generalizability and robustness across a broader range of tasks, for example, coding, reading comprehension, and so on.

[1] Llama 2: Open Foundation and Fine-Tuned Chat Models

[2] DoRA: Weight-Decomposed Low-Rank Adaptation

[3] Not all Layers of LLMs are Necessary during Inference

Q7: Implementation Details

• We have released our implementation code and the web address is provided in the checklist (line 690). The framework of our proposed D-LLM can be viewed in the directory ./llama. We also provide codes for training and inference in the repository.

• There are no particular challenges or restrictions to implement D-LLM on different hardware or software environments. We provide the necessary library list to train and infer our D-LLM in the file ./requirements.txt from the repository.

We truly appreciate the reviewer for the valuable feedback. We have carefully considered the comments and suggestions and would like to address each of the concerns raised.

W1: While the experimental results are impressive, the paper lacks a deep theoretical analysis of the dynamic decision module's behavior and its impact on model performance.

We understand your expectations of theoretical analysis and would like to make following clarifications:

• Theoretically, dynamic inference is a conditional execution paradigm, similar to mixture of experts and sparsity algorithms. The key idea is to reduce network redundancy and customize a network topology for each sample.

• We formulate the inference of dynamic inference as a sequential decision problem. Each decision module decides whether a transformer block should be executed or not.

• We provide detailed analysis on dynamic decision module's behavior related to the property of grammatical terms (line 279-293), the relationship between FLOPs in A.1 (line 522-531), and instance complexity in A.5 (line 563-573).

In future work, we will explore more theoretical tools and methods to provide a more in-depth explanation of the model's behavior.

W2: The computational overhead introduced by the dynamic decision module itself is not sufficiently analyzed.

We provide computational overhead analysis taken by the decision module of D-LLM in global rebuttal #1.

W3: The manuscript requires a deeper examination of the impact of various components of the loss function. Specifically, a thorough analysis is needed to investigate how varying the value of 𝛼 affects the performance of model.

We perform a comprehensive analysis on hyper-parameter $\alpha$ and details are available in global rebuttal #2.

W4: Although the paper demonstrates the effectiveness of the overall approach, it would benefit from more detailed ablation studies to isolate the impact of different components of the dynamic decision module and eviction policy.

We would like to kindly remind you that we have performed ablation studies to demonstrate the impact of the dynamic decision module and eviction policy in our paper (line 254-270). We select two baselines for comparison: 'LoRA finetune' only utilizes LoRA modules to finetune the LLM, and 'D-LLM w/o Evl. Str.' utilizes dynamic decision module without eviction policy. The results show that decision module can effectively achieve acceleration by layer skipping while maintaining comparable performance on benchmarks. With the help of eviction policy, D-LLM can further reduce the computation overhead due to fewer KV calculations and lower computations in attention operation.

Additionally, we consider that you might be interested in components of dynamic decision module for finer granularity. We utilize one layer network 'Linear & Softmax' as dynamic module for ablation study. The results in following table demonstrate that the module designed in our D-LLM has non-linear fitting capability to help the model control the actual acceleration rate more precisely and achieve better performance.

Values in table represent (FLOPs$\downarrow$, Accuracy$\uparrow$).

W5: The manuscript necessitates a more elaborate description of the evaluation metrics. For instance, a detailed explanation of the computation methodology for FLOPs is required to provide readers with a comprehensive understanding of the evaluation process.

We give a detailed description of all metrics mentioned in our paper as follows:

• PPL measures the divergence between the model's predicted probability distribution and ground truth. Given a sentence $X=(x_0,x_1,\cdots,x_n)$ and predicted probability distribution $p_\theta$ from model, PPL can be computed as $\text{PPL}(X)=\exp\{-\frac{1}{t}\sum_{i}^t\log p_{\theta}(x_i|x_{<i})\}$.

• Accuracy: For math-solving tasks, the answers are marked as correct only if the numerical difference between results in prediction and ground truth is less than $10^{-5}$. For commonsense reasoning tasks, answers are marked as correct only if the output matches the correct options in lowercase.

• FLOPs: FLOPs, short for floating point operations, measures the complexity of algorithms/models[2]. We calculated the flops of fully connected layers and attention operations in transformer. For fully connected layers without bias, we compute FLOPs as $\text{FLOPs}=(2I-1)O$, where $I$ is the input dimensionality and $O$ is the output dimensionality. FLOPs of attention operations are proportional to the cache size, so we calculated the average FLOPs when generating the 1st token to the 1K-th token. The FLOPs of baselines are normalized to 1 for better comparisons.

We will add the description in the revised version.

[1] Perplexity—a measure of the difficulty of speech recognition tasks, 1977

[2] Pruning Convolutional Neural Networks for Resource Efficient Inference, 2017

We truly appreciate the reviewer for the valuable feedback. We would like to address each of concerns raised as follows.

W1 & Q4: Additional overheads from the decision module training

We provide an analysis of overheads taken by decision modules in global rebuttal #1.

W2: Baselines do not include “dynamic inference mechanisms”

We would like to address the choice of baselines as follows:

• D-LLM w/o eviction strategy shown in Table 2 (line 260) can be regarded as a baseline for dynamic inference mechanisms on layer skipping. We propose KV-Cache eviction strategy on layer skipping for better adaption to LLM and prove its effectiveness in our ablation study.

• Dynamic inference mechanisms share similar methods with LLMs acceleration. AdaInfer, one of our baselines, introduces an early-exit mechanism to stop inference at intermediate layers, which is an important approach of dynamic inference mechanisms.

W3: Add references or clear evidence for preserving initial tokens

The importance of initial tokens is pointed out in recent works for LLMs' generalization to long contexts[1,2]. Initial tokens encode strong absolute position information and collect significant attention scores, which makes it necessary to be kept when pruning context. We will include the references in revision.

[1] LM-Infinite: Simple on-the-fly length generalization for large language models.

[2] Efficient streaming language models with attention sinks.

W4: Writing needs revision. Duplicated citations.

We will correct the duplicate citations and carefully check over other mistakes in revision.

Q1: Difference between 'dynamic pruning' (line 99) and 'dynamic inference mechanism' (line 105)

We discuss related works in other fields like computer vision in the context of dynamic inference mechanisms. 'Dynamic pruning methods' like Adainfer and MoD are designed specifically for LLMs. Although some dynamic inference mechanisms are shared in computer vision and LLMs, we want to discuss the LLMs acceleration methods separately to emphasize the domain of our work. We will make it clearer in revision.

Q2-3: Specific use of LoRA

LoRA primarily helps LLM adapt to the layer-skipping computation mode and specific domain tasks. We perform ablation study on D-LLM with and w/o LoRA under $\Omega=0.5$. The following table shows that the accuracy significantly decreases when LoRA modules are not trained together. The model also exhibits less control over acceleration rate.

Values in table represent (FLOPs$\downarrow$, Accuracy$\uparrow$).

Q5: D-LLM on a completely new dataset

Currently, we conduct experiments on specific tasks with few-shot training, To infer on a completely new dataset, further training of D-LLM is required to adapt model to domain knowledge and task templates. The table below shows that D-LLM maintains stable acceleration rates but has decreasing performance on new datasets, which also exists in PEFT methods (LoRA). We will research on generalization of D-LLM with more common datasets in the future.

Values in table represent (FLOPs$\downarrow$, Accuracy$\uparrow$).

Q6: Difference between target and actual acceleration rate

We give comparison between actual acceleration rate and target in the following table. Results show that larger target rate brings the larger difference. The difference is also influenced by the difficulty of tasks and hyper-parameter $\alpha$ in Eq.14 (line 213). For example, actual rate of OBQA is closer to the target because making choices is easier than summarization (SAMSum). Analysis on $\alpha$ is available in global rebuttal #2.

Q7: Why is the average acceleration rate calculated with b^tilde rather than using b? (Eq.11)

We use $\tilde{b}\_l$ instead of $b_l$ for gradients backward propagation. Utilizing $b_l$ from $\arg\max$ operation, which is not differentiable, makes it fail to optimize $\mathcal{L}\_{rate}$ in Eq.12.

In practical implementation, $\tilde{b}\_l$ can be approximated to one-hot vectors in 2 ways:

• Descend parameter $\tau$ to 0, then the softmax becomes 'sharp' and $\tilde{b}\_{l,1}$ numerically close to $b_l$.

• Use straight-through estimator as $\tilde{b}\_{l,1}\leftarrow(\tilde{b}\_{l,1}-b_l)\text{.detach()}+b_l$. The $\text{.detach()}$ operation means 0 gradients for backward propagation.

Both ways are viable in our framework. We will clarify details in revision.

Q8: Can FLOPS be interpreted for various purposes of evaluation?

In addition to computation, FLOPs can essentially measure the rate of memory overhead in KV-cache, which is exactly the same as executed-layers rate since only executed layers store the KV-cache in D-LLM. We provide the exact values on SAMSum in the following table to show the slight difference between FLOPs and other metrics, which is caused by extra few computation of decision modules.

We truly appreciate the reviewer for the valuable feedback. We have carefully considered the comments and suggestions and would like to address each of the concerns raised.

W1: Highlighting how D-LLMs differ significantly from or improve upon existing methods like AdaInfer and SkipNet would strengthen the novelty claim. Providing a more detailed comparison with these methods in terms of theoretical underpinnings or empirical performance could clarify the unique contributions of D-LLMs.

We make a detailed description of existing comparison methods in our paper as follows.

• ShortenLlaMA prunes unimportant layers based on designed metrics like Talylor and PPL and then apply LoRA to finetune pruned LLM on specific tasks. The pruning is static, which is fundamentally different from our dynamic inference method.

• AdaInfer is an early-exit method to stop the inference at intermediate layers. It underperforms on complex tasks such as Q&A.

• Mixture-of-Depth selects top-k tokens for calculation only at specific layers. Our method produces a more comprehensive dynamic inference mechanism and performs better at various benchmarks.

• Dynamic inference mechanisms in computer vision like SkipNet utilize layer skipping. Building on a similar mechanism, our method introduces the KV-cache eviction strategy for further overhead reduction in memory and computation for LLM inference.

We will incorporate these details in the revision to provide a clearer clarification of our contributions.

W2: The ablation on hyper-parameter \alpha in Eq. 14 is not provided. Please discuss its influence on task performance and computation reduction.

We analyze hyper-parameter $\alpha$ and discuss its influence on task performance and computation reduction. Details are available in global rebuttal #2.

W3: The font in Fig. 2-4 is too small for visualization.

We have increased the font size in Fig. 2-4 of our paper to enhance readability. The revised figures can be seen in Figure 2-4 of the PDF in the global rebuttal. We will update these figures in the revision.