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

Binarization, which converts weight parameters to binary values, has emerged as an effective strategy to reduce the size of large language models (LLMs). However, typical binarization techniques significantly diminish linguistic effectiveness of LLMs. To address this issue, we introduce a novel binarization technique called Mixture of Scales (BinaryMoS). Unlike conventional methods, BinaryMoS employs multiple scaling experts for binary weights, dynamically merging these experts for each token to adaptively generate scaling factors. This token-adaptive approach boosts the representational power of binarized LLMs by enabling contextual adjustments to the values of binary weights. Moreover, because this adaptive process only involves the scaling factors rather than the entire weight matrix, BinaryMoS maintains compression efficiency similar to traditional static binarization methods. Our experimental results reveal that BinaryMoS surpasses conventional binarization techniques in various natural language processing tasks and even outperforms 2-bit quantization methods, all while maintaining similar model size to static binarization techniques.

Paper

Similar papers

Peer review

Reviewer iCrR7/10 · confidence 4/52024-07-08

Summary

This paper studies the problem of post-training binarization of large language models. Built upon OneBit, this work propose BinaryMoS, to use a mixture of scaling weights for the linear layer of binarized LLMs. Instead of static scales in OneBit, BinaryMoS employs a set of scaling weight experts and adaptively combines them given the input representation of current tokens. As a result, BinaryMoS offers stronger representation capabilities while incurring minimal computation and memory overhead. Experimental results on various LLM benchmarks show the proposed technique compares favorably to other binarization methods for LLMs.

Strengths

1. BinaryMoS has minimal computational and memory overhead, and keeps the architecture simple without introducing sparse weight matrices. 2. Mixture-of-expert technique could potentially applied to other quantization method. 3. The method works for a wide range of LLMs. 3. The proposed method out-performs other binarization techniques on several LLM benchmarks.

Weaknesses

1. The paper lacks discussion of the cost of post-training adaptations. The training cost seems comparably high as it requires three epochs over the selected dataset. 2. The improvement over OneBit is small, and the improved binarization still perform significantly lower than other less aggressively quantized method.

Questions

1. While post-training binarization performs poorly compared to the original base model, methods like BitNet seem to bridging the gap by training low-bit models from scratch. Can your method extend to the scenario where the model is trained from scratch?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

The authors have adequately addressed the limitations of the work.

Reviewer fpKa3/10 · confidence 4/52024-07-12

Summary

This paper proposes using a mixture of scales for binarizing continuous latent weights. It includes a router implemented by a linear layer, outputting a softmax score which is then used as the combining matrix of the scale basis (called scale experts). The method uses two score basses, one for input $S_{in}$ and one for output $S_{out}$, each includes a number of scale vectors. The author conducted an experimental analysis of the number of experts to be used, and evaluated the proposed method on different language models with results showing performance gains by mixture of scales.

Strengths

The proposed idea is pretty nice, simple, and shows performance gains.

Weaknesses

Weak significance: Although the idea is nice and valid, its scientific and technological significance as well as its originality is low. The paper could be more suitable for a workshop than a major venue like NeurIPS. The experimental analysis of the number of scaling experts shows that the performance is not always improved with the number of experts. This also strengthens the above opinion of low significance. Several minor typos to be corrected, such as to use a unique term for MoS.

Questions

Related to (4) where the utilized scales are obtained by combining the scale basis using $G$, scale experts $S_{in}$ and $S_{out}$ are of primary importance in the method. How are these scale experts are constructed and/or learnt? If the scale experts are learnt, what can be the benefits of learning $S_{in, out}$ then compute $\hat{S}_{in, out}$ instead of learning these scales directly?

Rating

3

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes

Reviewer fpKa2024-08-12

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

Reviewer fpKa2024-08-13

Regarding my rating is lower than that the other reviewers, I read again the paper to make sure that I hadn't missed something important, I didn't find. Instead, I have some following suggestions about math eq (2) -- (5) for the authors to improve the quality of the paper: 1. Notation $\\odot$ seems denoting the elementwise product. First, it should be specified. 2. Suppose that my above interpretation 1. is correct, given that $A \odot B$ requires that $A$ and $B$ must have the same dimensions, eq (2) - (5) need to be corrected mathematically. In particular, eq (2) in its current form cannot read well. 3. The dimension of variables needs to be revised. Take your current notation of variable dimensions: $X: k \times m$, $W_{FP}: n \times m$: - For eq (3), it should be that $W_R: m \times e$ (currently $n \times e$). It result in $G: k \times e$. - Is $S_{in}: m \times e$ and $S_{out}: n \times e$ right? If so, eq (4) has dimension mismatch issue, maybe to be corrected as $hat S_{in} = G S_{in}^T$, $hat S_{out} = G S_{out}^T$? - If correcting as above, then $hat S_{in}: k \times m$, $hat S_{out}: k \times n$. This leads a problem with eq (5): $\odot$ products of the LHS have mismatched dimensions, while its RHS should be corrected as (please double-check my suggestion here) $X \odot hat S_{in} Sign(W^T) \odot hat S_{out}$.

Authorsrebuttal2024-08-14

Thank you for your additional comments. We acknowledge that there are some typos and unclear explanations in the equations presented in our paper. We will address them in this response and revise the manuscript accordingly. **Response to 1 and 2. Notation $\odot$ and Dimension Mismatch in $\odot$ Operation** As you suggested, we will explicitly state in the revised paper that the notation $\odot$ represents element-wise multiplication. Additionally, we want to clarify that we assume element-wise multiplication with broadcasting to a common shape, which functions exactly the same as in current deep learning frameworks. This means that in A $\odot$ B, the dimensions of A and B do not need to match exactly for the operation to proceed. For example, if A has dimensions $n \times m$ and B has dimensions $1 \times m$, B will be broadcasted to a common shape of $n \times m$ and then multiplied by A. Specifically, for the first dimension, the same value of B is multiplied by A. We will also clarify this broadcasting to a common shape in our revised paper. **Response to 2 and 3. Clarification of Eq.(2)-(5)** **Eq.(2):** This equation was adopted from the OneBit operation. To ensure compatibility with the subsequent equations, we will swap the row and column dimensions of $S_{in}$ and $S_{out}$. In the revised version, Eq. (2) will be expressed as $X [S_{in}^T \odot \text{Sign}(W^T_{FP}) \odot S_{out}] = [(X \odot S_{in})\text{Sign}(W^T_{FP})] \odot S_{out}$, where $S_{in}\in \mathbb{R}^{1 \times m}$, $S_{out}\in \mathbb{R}^{1 \times n}$. **Eq.(3):** As you correctly pointed out, $W_R$ should have dimensions $m \times e$, but it was mistakenly written as $n \times e$. We will correct this in the revision. **$S_{in}$ and $S_{out}$ of Eq.(4):** Given that $e$ scaling factors are used, the dimension of $S_{in}$ and $S_{out}$ should be $e \times m$ and $S_{out}$ as $e \times n$, respectively. We will clarify these dimensions in the revision. Meanwhile, as the dimension of $G$ is $k \times e$, the matrix multiplication in Eq. (4) between $G$ and $S_{in/out}$ is correct. **Eq.(5):** As you mentioned, the transpose notation should be removed from $\hat{S_{in}}$ and $\hat{S_{out}}$. We will correct Eq. (5) in the revision to: $[(X \odot \hat{S_{in}})\text{Sign}(W^T_{FP})] \odot \hat{S_{out}}$. We apologize for any confusion these equations may have caused. We will include clearer descriptions and correct typos in the final version of the paper. Lastly, we would like to emphasize the major contribution of this work once again. We propose a new LLM binarization method called BinaryMoS. It introduces the concept of Mixture-of-Scale to adaptively generate scaling factors, thereby enhancing the representational power of binary models. The accuracy evaluation results in our paper and the latency measurements provided in the global rebuttal demonstrate that the proposed method achieves state-of-the-art accuracy while maintaining the latency advantage of binary models. We hope our response has clarified your concerns. Thank you again for your feedback.

Reviewer s3b87/10 · confidence 4/52024-07-12

Summary

This paper proposes a binarization technique for LLMs, inspired by the mixture-of-expert (MOE) model. In the proposed approach, multiple scaling factors for the binarized matrices are available, each treated as an expert just like in MOE. The model infers a weight combination of these scaling factors adaptively at each time step, hence the binarization strategy is different for each input token. The paper shows that this adaptive binarization technique does not increase much memory overhead as the traditional MOE does, but can effectively improve model quality, as compared with existing 1-bit or 2-bit quantization baselines.

Strengths

The proposed approach effectively improved the quality of binarized LLMs, which has always been a challenging task. The introduction of adaptive binarization makes a lot of sense and deserves broader attention and experimentation.

Weaknesses

More detailed analysis in terms of efficiency, scaling behaviors and ablation studies could have been given.

Questions

1. The performance of BinaryMoS depends critically on the choice of static experts S_in, S_out. How robust is the proposed approach is to random initialization of these experts? Does the magnitude of these scaling experts matter a lot? 2. Are there any interesting patterns learned for expert weights? For examples, whenever one expert gets assigned larger weights, do the input sequences contain specific n-grams or syntactic structure? It would be great to find some insights into how and why the model learned to make expert choices, for better interpretability. 3. Related to question 2, if there is no clear interpretable expert assignment patterns that can be found, what if you just determine expert weights randomly without training? This should be a baseline to validate the effectiveness of the learning procedure. 4. How important is the knowledge distillation loss? It seems there is no analysis to the relative importance of distillation. 5. How does the quantization performance scale with model size? If there is a clear scaling law of BinaryMoS which appears to be more efficient than full-precision training, then this approach should become a standard for very large model training and inference.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper addressed limitations in the "Discussion and future work" section.

Reviewer scP17/10 · confidence 3/52024-07-13

Summary

This paper introduces a method to compress large language models (LLMs) by quantizing weight values to 1-bit. The goal is to mitigate performance degradation seen in previous quantization methods applied to LLMs. They utilize ideas from Mixture-of-Experts and introduce token-adaptive scaling factors that help control the effective values of binarized weights. Their approach results in high compression ratios compared to earlier binarization methods, maintaining memory efficiency.

Strengths

- Quantization helps lower the barriers to deploying large language models in compute-constrained environments. While previous approaches have been able to achieve this, it has come at the expense of linguistic utility. BinaryMoS in an attempt to lower these barriers without these costs. - Through experiments on various benchmarks, they demonstrate the effectiveness of BinaryMoS in memory-efficient quantization - The paper offers a clear explanation of existing research work on binarization and this helps with clarity and readability.

Weaknesses

- The motivation for extending the ideas from MoE to One-Bit is not entirely clear. Is the paper focused on applying binarization to MoE-style models? Please clarify this aspect. - The benefits of using 4 experts compared to 2 experts do not seem significant across tasks. Given the additional memory overhead introduced by these experts, wouldn't it be better to advocate for using two experts instead? - The analysis section on token-adaptive scaling factors does not adequately explain why BinaryMoS is preferable to One-Bit. While the variation in gating scores for experts across tokens is shown, there is no context provided about the sentences being analyzed. It seems possible that BinaryMoS is only useful and sensitive to certain domains and tasks and may not apply to every task. - The omission of latency measurements compared to other binarization methods raises questions about the robustness and efficiency of BinaryMoS. Since latency measurements are critical for real-world deployments of large language models, a discussion on this topic would be great

Questions

See weaknesses

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors claim to address the limitations of their work in the "Discussion and Future Work" section, but this discussion is incomplete. They do not mention latency requirements and potential failure cases of BinaryMoS are also omitted and not listed as limitations. Additionally, they should acknowledge that performance loss is still expectedly incurred compared to FP16.

Reviewer scP12024-08-12

Thanks a lot for the clarifications! I will increase my rating to 7!

Authorsrebuttal2024-08-14

Thank you for your feedback! We appreciate the constructive reviews for improving our work.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC