BitNet b1.58 Reloaded: State-of-the-art Performance Also on Smaller Networks

Our BitLinear layer functions as a drop-in replacement for PyTorch’s torch.nn.Linear layer. Figure 1 illustrates BitLinear’s 5-step computation flow:

1. 1.

   The activations are normalized.

2. 2.

   The normalized activations are quantized to k-bit precision.

3. 3.

   The 16-bit shadow weights are quantized to 1.58-bit weights.

4. 4.

   The quantized activations are multiplied with the 1.58-bit weights.

5. 5.

   The result of the multiplication is dequantized by rescaling.

We define the quantized weights $W_{quant}$ (denoted as 1.58-Bit Weights in Figure 1) as:

Having quantized both the activations and the weights, we can apply a kernel with $q_{quant}$ and $w_{quant}$ as inputs:

where $b$ is optional bias. We detach both $x_{quant}$ and $w_{quant}$ from the computation graph to achieve a straight-through estimation of the gradients. The gradients update the “shadow weights”, i.e., the 16-Bit Weights that are quantized by AbsMeasure Quantization.

Finally, we rescale the output $y$ during the Dequantization process:

Comparing to the original BitNet b1.58, there are a number of differences:

• •

  We chose to use a standard layer normalization (LayerNorm) rather than RMS normalization, as the computational overhead is minimal and we observed slightly better performance with the standard layer norm in preliminary experiments.

• •

  We allow the use of both the median and the mean for quantizing weights. Prior works [13, 11] solely employ the mean. We investigate the impact of this choice in Section 3.

• •

  We actually quantize weights and activations to integer values. This means the matrix multiplications are performed between the 1.58-bit weights with integer values $\{-1,0,1\}$ and the 8-bit quantized activations with integer values ${-128,\ldots,127}$. This allows to develop multiplication-free kernels, as multiplication with $-1$ corresponds to the subtraction of an 8-bit integer value, multiplication with $0$ to the disregard of a value, and multiplication with $1$ to the addition of an 8-bit integer value.

  This is in contract to previous work [11], where the quantized weights have floating point values $\{\frac{-1}{w_{scale}},0,\frac{1}{w_{scale}}\}$ while quantized activations have floating point values $\{\frac{-128}{x_{scale}},\ldots,\frac{127}{x_{scale}}\}$ according to the published information about the implementation[1]. Consequently, our BitNet 1.58 Bits Reloaded architecture is more directly amenable to custom software kernels and hardware implementations.

We conduct all experiments with standard networks in small configurations with the torch.nn.Linear layers replaced by our BitLinear layers. The Adam[6] optimizer and a batch-size of 128 are employed. The number of model parameters is slightly higher in the BitLinear setting, as we both have 1.58-bit weights as well as the 16-bit shadow weights. However, this fact does not change the number of trainable/optimized parameters in practice.

For SLMs, we train small Mistral-like models with 4 layers and hidden sizes of $32$, $64$, $128$, and $256$. The number of attention head and key-value cache heads is set to the ceiling of the hidden size divided by 64, i.e., $1$ head for $32$ and $64$ hidden sizes and $2$ and $4$ heads for $128$ and $256$, respectively. The resulting models sizes are 6M, 12M, 24M, and 48M parameters. We use a text corpus of 135M tokens and train from scratch for 10 epochs unless otherwise noted, corresponding to a total of 1.35B tokens for each training. We trained a Byte Pair Encoding tokenizer with a vocabulary size of $8{,}000$. The experiments are conducted with the standard trainer from the Hugging Face transformers library³³3https://github.com/huggingface/transformers.

For vision models, we consider a standard serial implementation of classifier for MNIST and standard CNN-based implementations for CIFAR-10, and CIFAR-100. The model for MNIST is the smallest in this paper with only 100K parameters. The CIFAR-10 and CIFAR-100 models represent the but smallest models with 2.1M and 2.2M, respectively. The difference in model size is explained by CIFAR-10 having 10 classes and CIFAR-100 having 100 classes. The experiments are based on Pytorch Lightning⁴⁴4https://github.com/Lightning-AI/pytorch-lightning and use torchvision’s⁵⁵5https://pytorch.org/vision/stable/index.html versions of the datasets.

The MNIST [2] dataset consists of 60.000 train and 10.000 test samples. The CIFAR10 [7] and CIFAR100 [7] datasets both contains 50.000 train and 10.000 test samples. All models are trained from scratch. We calculate the accuracy as the mean of the percentage of correct batches across the test set.

The first experiment for SLMs is a scaling experiment, where we perform 16-bit and 1.58-bit training on all four model sizes. The second experiment is a hyperparameter tuning for the learning rate and weight decay in a 12M SLM, with a fixed hidden size of 64. We show the results of the first and second experiment in Tables 3.1 and 3.1, respectively.

Both tables show the different configurations and perplexities after 10 epochs. For most configurations, the training has converged or is close to convergence at the end of the experiment. It is important to keep in mind that the reported perplexity is the exponentiation of the entropy, i.e., here the exponentiation of the loss defined via cross-entropy. Thus, minor changes in the loss result in quite discernible changes to the perplexity.

The first two columns give the hidden layer size and the number of parameters. The third column provides the bit-depth and implementation: “16” stands for 16-bit training, “1.58-mean” for our BitLinear implementation with 1.58 bits and AbsMean quantization of weights, and “1.58-median” for our BitLinear implementation with 1.58 bits and AbsMedian quantization of weights.

We show the results of this first experiment in Figure 2(d). Figure 2(a)) show that the 16-bit training scales exactly as expected when the number of hidden layers, and thus the models capacity, increases. We see in Figure 2(b), that 1.58-bit training follows the same trend, albeit with slightly lower performance. In Figure 2(c), we can visually compare the scaling between 16-bit training for models with 32 and 64 hidden sizes and 1.58-bit training for models with 64 and 128 hidden sizes. The observed perplexities suggest that the effective capacities of the models with 1.58-bit weights are around half that of the models with 16-bit weights, i.e., that hidden layers of approximately double size are needed for 1.58-bit models to reach performance comparable with the 16-bit counterparts. Figure 2(d) shows that the median generally converges slower than the mean over the employed weight decays. We discuss this fact in Section 4.

The fourth column shows the learning rate. For 16-bit training, we took a high but stable learning rate of 0.001 (1e-3). For 1.58-bit, we used the same or a larger learning rate of 0.01 (1e-2), as 1.58-bit training has been found to be more robust to higher learning rate in the context of LLMs[11]. The fifth column shows the weight decay. We tried both a small but noticeable decay of 5%, which is pretty prevalent in the pre-training and fine-tuning of LLMs, and no weight decay. Both Tables 3.1 and 3.1 hint that a weight decay of 5% yields the best or similar performance compared to other values of weight decay. This is also visualized in Figures 3(a), 3(b), and 3(c), where trainings with a weight decay of 5% are represented by a red line.

The sixth column provides the perplexity after 10 epochs. For nearly all configurations, after 10 epochs the training had converged. The best perplexities for 16-bit and 1.58-bit training are marked in bold, respectively. The seventh and last column shows the number of epochs, with total training length corresponding to 135M tokens per epoch, i.e., 1.35B tokens per 10-epoch experiment.