Mixture of Scales: Memory-Efficient Token-Adaptive Binarization for Large Language Models

We appreciate the constructive reviews, and here we address your comments in detail:

W1. High training cost

Though training three epochs over the selected dataset may seem costly, the datasets used for fine-tuning LLMs for quantization, particularly C4 and Wikitext2, are generally very small, making the training cost affordable. Our selected dataset consists of a total of 195M tokens, with 192M tokens from the C4 dataset and 3M tokens from the Wikitext2 dataset.

We will highlight the exact size of the selected dataset to clarify the training cost in the revised paper.

W2. Impact of BinaryMoS

Because the accuracy improvement of BinaryMoS over OneBit is not as dramatic as the improvement of OneBit over PB-LLM and BiLLM, there may be concerns that the proposed solution's improvement is marginal. However, we would like to highlight that the improvement of BinaryMoS over OneBit is substantial enough to make binary LLMs more applicable in practice. For example, both PB-LLM and BiLLM fail to generate grammatically correct sentences, while both OneBit and BinaryMoS can generate grammatically correct sentences. However, as presented in Table A4-1, BinaryMoS can generate contextually proper answers, whereas OneBit fails to generate correct answers. This is because BinaryMoS processes each token with token-adaptive scaling factors which contain contextual information.

We will also update the above discussion and Table A4-1 in the Appendix of the revised paper.

Table A4-1. Comparison of generation quality on the LLaMA-1-13B models with BinaryMoS and OneBit.

Q1. Training from scratch.

Thank you for your question and sharing your insight. The main novelty of the proposed BinaryMoS is its ability to increase the representational power of binary models by introducing the concept of Mixture-of-Scale, which can adaptively calculate proper scaling factors based on the context. This BinaryMoS architecture can also be trained from scratch. As you highlighted, post-training binarization inherently performs poorly compared to training-from-scratch approaches like BitNet, as post-training binarization introduces much less training iterations. Hence, if we have enough training facilities, we can further improve accuracy of binary models by adopting training-from-scratch approaches for BinaryMoS.

Thank you again for your valuable feedback. If you have any additional questions, please let us know.

We appreciate the reviews, and hope our rebuttal could convince you and change your stance on our paper.

W1. Weak significance and low originality

In this work, we propose a new binary LLM architecture to increase the representational power of binary models with negligible inference overhead by introducing the concept of Mixture-of-Scale, which can adaptively calculate proper scaling factors based on the context. We would like to highlight that good model architecture design is significant in the field of deep learning, as the model architecture determines the fundamental capability of the models. For this reason, many papers proposing modifications to model architecture are published in major venues. For example, Bi-Real Net, which applied a residual path to every convolution layer to enhance the accuracy of 1-bit convolutional neural networks (CNNs), has been published in ECCV [R3-1].

[R3-1] Liu, Zechun, et al. "Bi-real net: Enhancing the performance of 1-bit cnns with improved representational capability and advanced training algorithm." Proceedings of the European conference on computer vision (ECCV). 2018.

W2. Performance of the number of scale experts

We agree that performance does not always improve as the number of experts increases. Meanwhile, we want to emphasize that the key point of the proposed BinaryMoS is the introduction of multiple scaling experts (e.g., 2, 4, 8) to generate token-adaptive scaling factors, not the specific number of experts. In this way, BinaryMoS can improve the accuracy of binary models compared to OneBit, which uses fixed scaling factors. As shown in Tables 3 and 4 of our paper, BinaryMoS significantly enhances the accuracy of binary models. This highlights the effectiveness of token-adaptive scaling factor generation in BinaryMoS for improving the accuracy of binary models.

Q1. Effectiveness of training scale experts

As you correctly pointed out, the choice of scale experts $S_{in}$ and $S_{out}$ significantly influences the accuracy of BinaryMoS, so it is important to find proper $S_{in}$ and $S_{out}$. If we randomly set $S_{in}$ and $S_{out}$t as you questioned, BinaryMoS fails to achieve sufficient accuracy results. Therefore, during the training process of BinaryMoS, both $S_{in}$ and $S_{out}$ are also trained. Before training, BinaryMoS adopts SVD-based $S_{in}$ and $S_{out}$ initialization methods proposed in OneBit, so that $S_{in}$ and $S_{out}$ are initialized to minimize the binarization error of pre-trained weights, rather than applying random initialization. We will clarify that both $S_{in}$ and $S_{out}$ are also trained and their initialization methods in the revised version of the paper.

If we directly learn $\hat{S_{in}}$ and $\hat{S_{out}}$, it works exactly the same as OneBit. During inference, the scaling factors trained with this approach are fixed regardless of input context. Meanwhile, in the proposed BinaryMoS, on top of training $S_{in}$ and $S_{out}$, the router weight $W_R$ is also trained to generate gating score $G$ based on the input context $X$ so the token-adaptive scaling factors can be generated as described in Eq. (3) and (4) of our paper.

Thank you again for your valuable feedback. If you have any additional questions, please let us know.

Thank the authors for the response that partially answered my concerns, therefore I'm happy to increase my rating by 1.

We appreciate the constructive reviews, and here we address your comments in detail:

Q1. Importance of scale experts

As you correctly pointed out, the choice of static experts $S_{in}$ and $S_{out}$ significantly influences the accuracy of BinaryMoS, so it is important to find proper $S_{in}$ and $S_{out}$. If we randomly set $S_{in}$ and $S_{out}$ as you questioned, BinaryMoS fails to achieve sufficient accuracy results. Therefore, during the training process of BinaryMoS, both $S_{in}$ and $S_{out}$ are also trained.

The term ‘static’ might cause confusion, but $S_{in}$ and $S_{out}$ are static only during the inference stage, not during the training stage. Moreover, before training, BinaryMoS adopts SVD-based $S_{in}$ and $S_{out}$ initialization methods proposed in OneBit, so that $S_{in}$ and $S_{out}$ are initialized to minimize the binarization error of pre-trained weights, rather than applying random initialization.

We will clarify that both $S_{in}$ and $S_{out}$ are also trained and their initialization methods in the revised version of the paper.

Q2. Linguistic pattern of assigning expert

The weights of experts vary across input tokens, and we were not able to identify any interesting patterns learned for expert weights. However, we have included example sentences and the corresponding expert weights assigned to each token in Figure 1 of the attached PDF file of author rebuttal. This is provided in case others might find some interesting patterns. If you have any further insights or opinions on the patterns, please let us know.

Q3. Influence of random expert weights

Thank you for the valuable insight. As you suggested, we measure the accuracy of BinaryMoS by setting expert weights randomly without training as shown in Table A2-1. The evaluation results demonstrate significant accuracy degradation, underscoring the importance of training to fully achieve the advantage of the proposed Mixture-of-Scale scheme. We will update the analysis result in Table A2-1 and the above discussion in the Appendix of the revised paper.

Table A2-1. Accuracy evaluation of BinaryMoS model with random expert weights

Q4. Importance of knowledge distillation

Knowledge distillation (KD) has been widely adopted to improve the accuracy of QAT techniques, and it has been also utilized in previous work such as OneBit. To verify the importance of KD loss, we compare the accuracy of both OneBit and BinaryMoS with and without KD in Table A2-2. The binarized models without KD show worse perplexity and accuracy results compared to those with KD. However, regardless of the application of KD, BinaryMoS consistently outperforms OneBit, demonstrating the effectiveness of the proposed Mixture-of-Scale approach. We will update the analysis result in Table D and the above discussion in the Appendix of the revised paper.

Table A2-2. Perplexity and averaged zero-shot accuracy results of OPT-1.3B model

Q5. Scaling law of BinaryMoS

As you pointed out, if there is a clear scaling law for BinaryMoS, it could become a standard for compressing very large models. However, since BinaryMoS adopts QAT-based strategies, which require training the full weight parameters to adapt the model to quantization, it is very challenging to evaluate the effect of BinaryMoS on very large models, such as LLaMA 30B and 70B, due to high training cost. Hence, there is a difficulty in determining the scaling law of BinaryMoS at this point. Please note that this limitation is not specific to BinaryMoS but is a general limitation of QAT-based strategies, including previous works like OneBit. To scale BinaryMoS to very large models, we need to find a way to integrate the BinaryMoS approach with parameter-efficient fine-tuning techniques (PEFT) such as LoRA [R2-1]. This integration would make the training procedure feasible for very large models. Therefore, we believe that this is an important future direction for BinaryMoS to extend its usability and effectiveness to much larger models.

[R2-1] Yelysei Bondarenko, et al. “Low-Rank Quantization-Aware Training for LLMs”, arxiv:2406.06385.

Thank you again for your valuable feedback. If you have any additional questions, please let us know.

We appreciate the constructive reviews, and here we address your comments in detail:

W1. Motivation of BinaryMoS

The motivation stems from the fact that previous binarization methods, including OneBit, still have low representational power. To push the limits of binarized weights, we propose token-adaptive scaling factors inspired by MoE. Typically, MoE is used not for model compression techniques like quantization but to enhance the original model's capabilities by duplicating the weights of FFN layers according to the number of experts. Simply combining binarization methods with traditional MoE can enhance the model's capabilities, but the memory overhead is large, conflicting with the goal of extreme compression inherent in binarization. Therefore, we propose a novel binarization method that adopts the concept of using multiple experts from MoE and applies it to the scaling factors of binarization. In this way, while OneBit uses fixed scaling factors regardless of context, the proposed BinaryMoS can generate token-adaptive scaling factors which can improve the representation power of the binarized model.

In summary, our BinaryMoS scheme is inspired by MoE-style models, but it does not employ multiple experts with separate weights. Instead, it uses multiple scaling experts to generate token-adaptive scaling factors while fixing binarized weights, thereby retaining the memory efficiency of binarized models.

W2. Benefits of using 4-expert

As the reviewer pointed out, in terms of zero-shot accuracy, the average accuracy of 2-experts and 4-experts is 50.64% and 50.61%, respectively, showing no significant difference. However, in terms of perplexity, the 4-expert configuration achieves remarkable improvements compared to the 2-expert configuration by lowering perplexity from 12.18 to 11.85 (reported in Table 2 of our paper). Hence, we can expect that 4-experts will generally achieve better language modeling capabilities compared to 2-experts. Moreover, due to the low memory cost of adopting scaling experts, as shown in Table A1-1, the model sizes of 2-expert and 4-expert BinaryMoS are similar, while both significantly reduce the model size compared to the original Float16 model. In other words, the cost of adopting the 4-expert configuration is low, while we can expect an improvement in language modeling capabilities. Based on these observations, we chose to employ 4-experts in BinaryMoS.

Table A1-1. Comparison of memory size for 2-expert and 4-expert BinaryMoS.

W3. Relationship between gating score and context

While Figure 3 shows the variation in gating scores for experts across tokens in an exemplary sentence, we observed a similar tendency throughout our experiments for various tasks reported in Section 4.4. Each token is assigned a different scaling factor with our token-adaptive binarization method. Although our experiments are still limited, we believe this trend could extend to other LLM tasks as well, since the concept of applying multiple scaling experts for weight binarization is general.

W4. Latency measurements compared to other binarization methods

Thanks for pointing out the latency measurements. As you highlighted, the latency is a critical component for real-world deployment. However, previous binarization papers, including PB-LLM, BiLLM, and OneBit, have not reported their latency due to the lack of a CUDA kernel for matrix multiplication between FP activation and 1-bit weights. Therefore, we first developed a custom CUDA kernel for 1-bit matrix multiplication by modifying the CUDA kernel for multi-bit matrix multiplication [R1-1, R1-2]. Then, we also developed a custom CUDA kernel for BinaryMoS by fusing the operations of scaling experts and routers on top of the 1-bit matrix multiplication CUDA kernel.

We measured the latency of the linear layers in LLaMA-7B and LLaMA-13B (batch size: 1) and reported the latency results in Table G-1 and G-2 of global response. As illustrated in our paper, PB-LLM and BiLLM require extra matrix multiplications, so they tend to show similar or larger latency results compared to the original FP16 models. OneBit, which employs the simplest binarization scheme, achieves significant improvement over the original FP16 model and shows the minimum latency. Meanwhile, as our BinaryMoS introduces additional operations on processing scaling experts, which require far fewer operations compared to the matrix multiplication, our BinaryMoS also shows similar latency results as OneBit. This demonstrates that the multi-scaling factor module in BinaryMoS improves performance in terms of perplexity and zero-shot accuracy with minimal overhead to latency.

[R1-1] Taesu Kim, et al., “QUICK: Quantization-aware Interleaving and Conflict-free Kernel for efficient LLM inference”, arxiv:2402.10076

[R1-2] https://github.com/SqueezeBits/QUICK

L1. Limitation of BinaryMoS

Thanks for the suggestion.

First, we will include the latency results which we reported in the response to W4 and global response in the revised manuscript. Regarding other limitations, we will include the following parts in the limitation section of the revised manuscript.

1. Our results are still limited to relatively small models up to 13B. We plan to work on larger models such as 30B.

2. The performance loss from the FP16 model is still substantial, so further innovations regarding binarization are needed to reduce the performance gap.

Thank you again for your valuable feedback. If you have any additional questions, please let us know.