[论文解读] The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits

Shuming Ma, Hongyu Wang+8 more|arXiv (Cornell University)|Feb 27, 2024

Natural Language Processing Techniques被引用 37

一句话总结

本论文介绍 BitNet b1.58，一种 1.58 位 LLM，其每个参数均为三值 {-1, 0, 1}。在相近的尺寸下，它的性能与 FP16/LLaMA 相当，同时在速度、内存和能耗方面实现显著提升，推动新的放缩规律以及对 1 位 LLM 的潜在硬件优化。

ABSTRACT

Recent research, such as BitNet, is paving the way for a new era of 1-bit Large Language Models (LLMs). In this work, we introduce a 1-bit LLM variant, namely BitNet b1.58, in which every single parameter (or weight) of the LLM is ternary {-1, 0, 1}. It matches the full-precision (i.e., FP16 or BF16) Transformer LLM with the same model size and training tokens in terms of both perplexity and end-task performance, while being significantly more cost-effective in terms of latency, memory, throughput, and energy consumption. More profoundly, the 1.58-bit LLM defines a new scaling law and recipe for training new generations of LLMs that are both high-performance and cost-effective. Furthermore, it enables a new computation paradigm and opens the door for designing specific hardware optimized for 1-bit LLMs.

研究动机与目标

• 在不牺牲性能的前提下，推动降低 LLM 的推理成本和能耗。
• 提出一种包含零态的 1.58 位权重量化方案，以提升建模能力。
• 证明 BitNet b1.58 从 3B 尺度开始在困惑度和端任务指标上可达到全精度基线的水平。
• 展示硬件与能耗影响，并讨论对更大模型及边缘设备的可扩展性。

提出的方法

• 引入 BitNet b1.58，一种 Transformer，其中线性层被 BitLinear 替换，权重量化为 {-1, 0, 1}。
• 采用 absmean 量化来约束权重，gamma 定义为权重绝对值的平均值。
• 采用类似 LLaMA 的架构组件（RMSNorm、SwiGLU、旋转嵌入），不使用偏置。
• 在 RedPajama 数据上，使用 1.58 位权重和 8 位激活从头开始训练。
• 在模型规模 (700M–3.9B) 范围内，对比 FP16 LLaMA 基线，评估困惑度和零-shot 端任务性能。
• 使用 FasterTransformer 及 2 位/ Ladder 内核测量内存、延迟、吞吐量和估计能耗；分析可扩展性与能耗节省。

实验结果

研究问题

• RQ1 在相同模型尺寸与训练设置下，1.58 位 LLM 是否能在困惑度和端任务上达到全精度 LLM 的性能？
• RQ2 随着模型规模增加，1.58 位 LLM 在内存、延迟和能耗方面的影响是什么？
• RQ3 在不同尺寸上，1.58 位 LLM 是否在成本和精度上相对于 FP16 基线呈现帕累托改进？
• RQ4 用 2T tokens 训练能否使 BitNet b1.58 实现或改进在标准基准上的泛化能力？

主要发现

模型	大小	内存 (GB)	延迟 (ms)	PPL	模型（端任务）均值	ARC-E	ARC-C	HS	BQ	OQ	PQ	WGe	均值
LLaMA LLM	700M	2.08	1.18	12.33	45.5	54.7	23.0	37.0	60.0	20.2	68.9	54.8	45.5
BitNet b1.58	700M	0.80	0.96	12.87	44.3	51.8	21.4	35.1	58.2	20.0	68.1	55.2	44.3
LLaMA LLM	1.3B	3.34	1.62	11.25	46.2	56.9	23.5	38.5	59.1	21.6	70.0	53.9	46.2
BitNet b1.58	1.3B	1.14	0.97	11.29	45.4	54.9	24.2	37.7	56.7	19.6	68.8	55.8	45.4
LLaMA LLM	3B	7.89	5.07	10.04	49.7	62.1	25.6	43.3	61.8	24.6	72.1	58.2	49.7
BitNet b1.58	3B	2.22	1.87	9.91	50.2	61.4	28.3	42.9	61.5	26.6	71.5	59.3	50.2
BitNet b1.58	3.9B	2.38	2.11	9.62	51.2	64.2	28.7	44.2	63.5	24.2	73.2	60.5	51.2

• BitNet b1.58 在相同配置下，从 3B 尺度开始在困惑度上达到 FP16 LLaMA 基线。
• BitNet b1.58 在 3B 时相比 FP16 LLaMA 提升 2.71x 速度，GPU 内存减少 3.55x。
• 在 3.9B 时，BitNet b1.58 提升 2.4x 速度、内存效率高 3.32x，同时在端任务上优于 LLaMA 3B。
• 在不同尺寸上，BitNet b1.58 缩小了与 FP16 基线的性能差距，从 3B 起达到可比的端任务结果。
• BitNet b1.58 70B 相对于 LLaMA 70B 最多实现 4.1x 的加速，在两块 A100 GPU 上吞吐量最高提升 11x，且内存降低。
• 能耗分析表明，与 FP16 基线相比，在 7nm 芯片上矩阵乘法的算术能量节省约 71.4 倍，且对更大模型优势日益明显。