Rethinking Output Alignment For 1-bit Post-Training Quantization of Large Language Models

TL;DR

This work tackles the challenge of 1-bit post-training quantization for large language models, where naive output alignment often underperforms due to inter-layer error accumulation and disruption of attention patterns. It introduces a data-aware, selective layer-wise output matching framework and Attention Matrix Preservation (AMP) to explicitly align the true target output $WX$ with the quantized output while preserving token interactions. The method derives closed-form updates for quantization parameters and employs a block-wise strategy that applies output matching only to the last layer of each block, while using weight alignment for other layers. Empirical results on OPT and LLaMA demonstrate consistent gains over state-of-the-art 1-bit PTQ baselines with modest overhead, enabling more reliable deployment of 1-bit quantized LLMs.

Abstract

Large Language Models (LLMs) deliver strong performance across a wide range of NLP tasks, but their massive sizes hinder deployment on resource-constrained devices. To reduce their computational and memory burden, various compression techniques have been proposed, including quantization, pruning, and knowledge distillation. Among these, post-training quantization (PTQ) is widely adopted for its efficiency, as it requires no retraining and only a small dataset for calibration, enabling low-cost deployment. Recent advances for post-training quantization have demonstrated that even sub-4-bit methods can maintain most of the original model performance. However, 1-bit quantization that converts floating-point weights to \(\pm\)1, remains particularly challenging, as existing 1-bit PTQ methods often suffer from significant performance degradation compared to the full-precision models. Specifically, most of existing 1-bit PTQ approaches focus on weight alignment, aligning the full-precision model weights with those of the quantized models, rather than directly aligning their outputs. Although the output-matching approach objective is more intuitive and aligns with the quantization goal, naively applying it in 1-bit LLMs often leads to notable performance degradation. In this paper, we investigate why and under what conditions output-matching fails, in the context of 1-bit LLM quantization. Based on our findings, we propose a novel data-aware PTQ approach for 1-bit LLMs that explicitly accounts for activation error accumulation while keeping optimization efficient. Empirical experiments demonstrate that our solution consistently outperforms existing 1-bit PTQ methods with minimal overhead.

Figures (3)

• Figure 1: Comparison of block-level loss under ARB (weight alignment) versus ARB-X (layer-wise output alignment) for each layer of LLaMA-2-7B. Bar color denotes which method incurs lower block loss: blue for ARB, red for ARB-X. Several layers demonstrate that minimizing layer-level loss via ARB-X does not necessarily achieve lower block-level loss than ARB, indicating the limitations of naive layer-wise output alignment.
• Figure 2: Accumulated quantization error in LLaMA-2-7B under ARB-X. The top plot reports cosine similarity and MSE between block outputs of the quantized and full-precision models. While ARB-X maintains low cosine similarity loss between pre- and post-quantization outputs, the MSE relative to the full-precision baseline grows with depth, indicating error accumulation. The bottom plot reports the MSE loss for token similarity matrices, computed as pairwise cosine similarities between token representations within each sample. These matrices also drift from the full-precision baseline after quantization, suggesting that ARB-X's output alignment may degrade the learned attention mask in subsequent layers.
• Figure 3: Block-wise MSE reconstruction error between quantized and full-precision attention score in LLaMA-2-7B. Two curves represent the errors of quantized models with and without the AMP mask, computed over C4 calibration data.