NeuZip: Memory-Efficient Training and Inference with Dynamic Compression of Neural Networks

## Abstract

Abstract The performance of neural networks improves when more parameters are used. However, the model sizes are constrained by the available on-device memory during training and inference. Although applying techniques like quantization can alleviate the constraint, they suffer from performance degradation. In this work, we introduce NeuZip, a new weight compression scheme based on the entropy of floating-point numbers in neural networks. With NeuZip, we are able to achieve memory-efficient training and inference without sacrificing performance. Notably, we significantly reduce the memory footprint of training a Llama-3 8B model from 31GB to less than 16GB, while keeping the training dynamics fully unchanged. In inference, our method can reduce memory usage by more than half while maintaining near-lossless performance. Our code is publicly available. 1 1 1 https://github.com/BorealisAI/neuzip

## 1 Introduction

Deep learning with neural networks has become the backbone of numerous artificial intelligence applications. The search for better-performing networks is a longstanding topic in deep learning. Without modifying the design, scaling up the number of parameters (e.g., number of hidden dimensions or layers) has been demonstrated as an effective practice to boost the performance of neural networks of the same kind (Kaplan et al., ). This idea has been successfully applied to text, image, audio, and multi-modal tasks with a wide range of model architectures (Yu et al., ; Radford et al., ; Brown et al., ). Recently, the number of parameters in the state-of-the-art models has become more than 100 billion or even a trillion parameters. For example, one of the state-of-the-art language models in 2020, GPT-3, has 175B parameters (Brown et al., ), growing by nearly 100 times compared with the largest Transformer architecture in the 2017 paper (Vaswani et al., ).

Despite the growth in the model size, the hardware capacity is not keeping up with the pace: the largest on-device memory of GPUs was 32GB in 2017, and is 80GB to this date in 2024, growing by only 2.5 times. The available hardware supply poses a limitation on the trainable model size, bottlenecking the scaling capacity. Although this problem can be alleviated by using more GPUs and sharding the model in multiple devices (Rajbhandari et al., ), such a practice introduces more communication overheads among GPUs, making large-scale distributed training less efficient. Therefore, saving the total memory usage is critical in scaling up neural networks.

The peak memory usage is dominated by three relatively independent parts: the optimizer, the saved activations for back-propagation, and the model itself. For the optimizer, there are already memory-efficient optimizers achieving a sublinear space complexity (Shazeer and Stern, ; Hao et al., ); for the activations, the memory can be saved by enabling activation checkpointing (Chen et al., ), which saves the storage by recomputing the forward activations during the back-propagation. For the model parameters, there has not been an effective method to save the memory while preserving the ability to train the model. Recently, Dettmers et al. () proposed the quantized low-rank adaptation (QLoRA), which freezes the parameters using a 4-bit data type for the backbone pre-trained model. While significantly saving the memory for the model, it imposes a constraint on the overall change of the model to be low-rank, limiting the model capacity.

In this paper, we propose NeuZip, an algorithm to compress the neural networks while maintaining their full abilities. Specifically, each floating-point number is represented by three parts: the sign bit, the exponent bits, and the mantissa bits. Following the observation that weights are concentrated around zero (Kalamkar et al., ), we demonstrate that this corresponds to the low-entropy nature of the exponent bits. We hence compress the exponent bits using the asymmetric numeral system (ANS Duda ()), a lossless compression algorithm that achieves a high throughput on parallel computing devices like GPUs. Since the compression is lossless, the memory reduction comes without compromising any precision loss and enables full-parameter training.

In addition to lossless compression for training, we also propose a lossy variant of NeuZip for inference that further reduces the memory footprint. Specifically, we control the relative change of each parameter by storing only the top-$k$ significant bits of the mantissa. We empirically show that lossy NeuZip lies at the Pareto frontier of the memory–performance trade-off when compared with several state-of-the-art quantization baselines.

## 2 Our Approach

The Shannon entropy (Shannon, ) is used to measure the “stochasticity” of a random variable with the following definition:

$$ $\displaystyle H(X):=\mathop{\mathbb{E}}_{X\sim p(X)}[-\log_{2}p(X)]$ (1) $$

for a random variable $X$ with probability $p$.
A lower entropy indicates a less stochasticity of a random variable. In fact, the entropy equals the minimum number of bits required, in expectation, to represent a random variable, therefore corresponding to data compressibility. For the non-concentrating random variable with all possible values sharing an equal probability, the entropy of which reaches the maximum value $\log_{2}n$, where $n$ is all possible values $X$ can take. On the other hand, for highly-concentrating (e.g., fully deterministic) random variables, the entropy can be as low as 0.

### 2.1 Low-Entropy Nature of Neural Network Parameters

We argue that the parameters in neural network tend to have low entropy. First, parameters are typically initialized with Gaussian distribution for matrices (Glorot and Bengio, ; He et al., ). This encourages all weights to be centered around zero, effectively reducing the entropy (or randomness). In addition, regularization is also applied for better generalization ability. For example, the weight decay technique reduces the magnitudes of weights at every update iteration. Similarly in Bayesian inference, prior distributions (e.g., Gaussian and Laplace distributions) are often applied, imposing a zero-concentrated preference over the parameters.Even without explicit regularization,
stochastic gradient descent (SGD) or its variants are shown to have the implicit regularization effect on neural networks, meaning the model parameters are implicitly encouraged to have smaller magnitudes during training (Soudry et al., ; Vardi and Shamir, ).
All the above effects and techniques lead to the following observation:

###### Observation 2.1 .

Assuming neural network parameters are i.i.d. random variables, the entropy of the distribution is likely to be low.

Specifically, each parameter is represented is represented by three components: the sign bit, the exponent bits, and the mantissa bits in the IEEE 754 standard (IEEE, ).(^1^11We use BF16 (Kalamkar et al., ) in this paper.) Therefore, we conduct a fine-grained analysis and investigate the distribution of each component of a floating-point number in neural networks.

Figure: Figure 1: The histograms of different components of the parameters of LLama-3 8B model (Dubey et al., ). The $x$-axis is all possible binary values and the $y$-axis represent the frequency of each value.
Refer to caption: x1.png

As shown in Figure 1, the sign bit has a high entropy as it is evenly distributed; hence, it is not compressible. For the exponent bits, there is a clear pattern that they demonstrate a low-entropy nature, carrying only less than 3 bits of information with 8 bits of capacity. For the mantissa bits, they store nearly 7-bit information with 7-bit capacity. In fact, we shown in Appendix A that this is common in deep learning.

This phenomenon suggests that by simply compressing the exponents, we are able to recover the overall optimal compression ratio. In this example, an ideal compression algorithm is able to achieve a ratio as high as 1.501 (the sum of the three entropy values), only marginally below the overall compression ratio 1.505.

###### Observation 2.1 .

### 2.2 Lossless NeuZip: Compressing Exponents for Training

#### Compressed representation.

Based on observation, we see that the number of bits per exponent is largely inflated compared with the information entropy. However, previous research demonstrates that the dynamic range provided by the 8-bit exponents are critical for neural networks (Kalamkar et al., ). We therefore propose to compress the exponent bits in a lossless manner based on the entropy. This practice mainly has three benefits: (1) it increases the throughput of compression as only part of the bits are processed by the compression; (2) it reduces the burden of maintaining the statistics of a large set of symbols (e.g., 256 symbols for 8-bit exponents versus 65,536 symbols for 16-bit representations), enabling a great efficiency of compression algorithms; (3) most importantly, it recovers most of the compressibility as shown in Figure 1.

#### Multi-layer neural networks.

The compression alone does not save any memory for maintaining a single array. This is because, either compression or decompression, requires at least one buffer of the same size as the uncompressed array. In the scope of neural networks, the whole model is prohibitively large and it is infeasible to duplicate the memory. In NeuZip, however, we exploit the multi-layer structure of modern neural networks to avoid creating a large buffer. Without loss of generality, we focus on the linear function as a common building block in neural networks at layer $l$:

$$ $\displaystyle\bm{x}_{l}\leftarrow\bm{W}_{l}\bm{x}_{l-1}+\bm{b}_{l},$ (2) $$

where $\bm{W}_{l}\in\mathbb{R}^{m\times n}$ is the weight matrix, $\bm{b}_{l}\in\mathbb{R}^{m}$ is the bias vector of layer $l$, and $\bm{x}_{l}$ is the input of layer $l$. We propose to modify the compressed forward pass in the following form

$$ $\displaystyle\hat{\bm{W}}\leftarrow$ $\displaystyle\mathrm{decompress}(\bm{c}_{l})$ (3) $\displaystyle\bm{x}_{l}\leftarrow$ $\displaystyle\hat{\bm{W}}\bm{x}_{l-1}+\bm{b}_{l},$ (4) $$

where $\bm{c}_{l}$ is the compressed storage of the matrix $\bm{W}_{l}$. In this way, we only need to store $\bm{c}_{i}$ for each layer, enjoying low-memory usage. During each forward pass, weight matrices stay in the compressed form until the original data is needed, in which case it is decompressed into a temporary space $\hat{\bm{W}}$ for computation. As a result, the entire network is never fully decompressed at any point in time, making the overall forward pass memory efficient. The per-layer procedure is shown in Figure 2.

Note that although we alter the forward pass, the back-propagation for each linear layer is fully unaffected. This is because

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\bm{W}_{l}}=\frac{\partial \mathcal{L}}{\partial\bm{x}_{l}}\frac{\partial\bm{x}_{l}}{\partial\bm{W}_{l}}= \left(\nabla_{\bm{x}_{l}}\mathcal{L}\right)\bm{x}_{l-1}^{\top}.$ (5) $$

Therefore, we are able to obtain the gradient as long as the activations are saved. Similarly, we can also propagate the gradient of inputs with

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\bm{x}_{l-1}}=\frac{\partial \mathcal{L}}{\partial\bm{x}_{l}}\frac{\partial\bm{x}_{l}}{\partial\bm{x}_{l-1} }=\left(\nabla_{\bm{x}_{l}}\mathcal{L}\right)^{\top}\bm{W}_{l},$ (6) $$

where $\bm{W}_{l}$ can be constructed by decompression.
It is worth noting that our NeuZip is compatible with activation checkpointing (Chen et al., ) by recomputing the activations, opening more opportunity for memory saving.

Figure: (a) Vanilla

For weight updates, we decompress the matrix into the original floating-point format and compress the updated matrix again. This procedure is done in a layer-by-layer fashion, similar to LOMO (Lv et al., ). The overall training procedure is described in Appendix B.

#### Compression algorithm.

In our implementation, we choose to use the asymmetric numeral systems (ANS) (Duda, ) as our backbone compression algorithm because it can be easily parallelized and achieves a high throughput with parallel execution, making it an ideal candidate on deep learning accelerators like GPUs. Specifically, ANS encodes a sequence of symbols by treating them as base-$n$ numbers. However, unlike the common numerical system that uses a uniform base for each digit, ANS treats every single digit with a different base $\lceil 1/\hat{p}_{i}\rceil$, where $\hat{\bm{p}}$ is the frequency of symbols. As a result, it achieves a near-optimal compression rate by suing around $1/\hat{p}_{i}$ bits for the $i^{\mathrm{th}}$ symbol.

### 2.3 Lossy NeuZip: Additionally Truncating Mantissa for Inference

In the algorithm above, we show that the training of neural networks can be completely unaffected by lossless compression. On the other hand, inference is known to be less sensitive to precision loss compared with training (Dettmers et al., ; Dettmers and Zettlemoyer, ). This enables further memory reduction of NeuZip by reducing the precision.
In our study, we conduct a pilot experiment that perturbs each weight with a noise proportional to the weight magnitude. We observe that with a small noise ratio there is little or no effect on the overall performance (Appendix C). Motivated by this, we propose a variant of NeuZip that compresses mantissa in a lossy way during inference.

In its core, we simply round and truncate the mantissa to fewer bits. Specifically, we assume the original floating-point number $f$ has an exponent $e$ and mantissa $m$. After rounding, the mantissa is denoted by $\hat{m}$ and the resulting floating-point number is denoted by $\hat{f}$.

The rounding introduces an error expressed as:

$$ $\displaystyle|f-\hat{f}|=\left|2^{e-127}\cdot\frac{m}{2^{7}}-2^{e-127}\cdot \frac{\hat{m}}{2^{7}}\right|=2^{e-134}\cdot|m-\hat{m}|$ (7) $$

where $e-127$ interprets the exponent bits $e$ as an integer, which can be either positive, negative, or 0. In the fraction $m/2^{7}$, $m$ is the significand (an unsigned integer) and $7$ is the precision.
It is straightforward to see that the relative error is given by

$$ $\displaystyle\frac{|f-\hat{f}|}{|f|}=\frac{2^{e-134}\cdot|m-\hat{m}|}{2^{e-134 }\cdot|m|}=\frac{|m-\hat{m}|}{m}.$ (8) $$

Suppose our rounding keeps $k$ most significant bits in $\hat{m}$, the earliest point where $\hat{m}$ could differ from the original number $m$ is at the $(k+1)$th bit. This means that the maximum possible relative change introduced by this rounding is $1/2^{k}$.
Given that the mantissa bits are highly uniform as shown in Figure 1, such a practice resembles the weight perturbation based on relative magnitudes, justifying the rounding trick applied to mantissas.

In our implementation, we store the sign and mantissa bits together as a signed integer to minimize the requests of memory write. Further, given that modern architectures are mostly byte (8-bit) addressable,
we pack multiple such signed integers into a single byte for memory efficiency. To align with an 8-bit byte, we let the precision after rounding to be $\{0,1,3\}$, ensuring that all the bits in a byte are utilized efficiently. We illustrate the process in Figure 8(b).

Lastly, we enable a block-wise normalization technique (Dettmers et al., ), where a block is a chunk of weights that are stored contiguously in memory. Such block-wise normalization makes sure that the weight with the largest magnitude in a block will always be normalized to 1, invariant to mantissa rounding and truncation. The normalization coefficient—which handles mantissa while ignoring the exponent—is stored with 8 bits, and is used for de-normalization during the decompression of the weight. This strategy is based on the observation that larger weights play a more important role in neural networks (Han et al., ).

## 3 Experiments

We empirically verify the effectiveness of NeuZip across different model architectures and datasets. Given the success of large language models, we mainly consider Transformer-based models for our experiments. We choose two designs of Transformer, decoder-only and encoder–decoder models, to show the generality of our method. All experiments are conducted on RTX A6000 GPUs where the uncompressed data type is BFloat16.

### 3.1 Lossless NeuZip for Pre-Training

#### Settings.

We choose decoder-only models to evaluate our method on the pre-training task. We select 3 models with different sizes to study the scaling effect, including GPT-Neo 2.7B (Black et al., ), Llama-3 8B (Dubey et al., ), and LLama-2 13B (Touvron et al., ). For fair comparison, all competing methods are initialized with the same random weights.

For the task, we consider language modeling, which requires the model to predict the next token given the context. We use the Wikitext-2 dataset (Merity et al., ), where each data sample is a fixed-length sequence from an article on Wikipedia. We set the length to 1024 following the common practice (Radford et al., ).

For each experiment, we report the loss (negative log-likelihood) on unseen samples. To study memory saving, we report the peak memory usage for each run during the training process. The numbers are shown in gibibyte (GiB, $1024^{3}$ bytes). We also report the speed by the number of iterations per second to demonstrate the time-efficiency of each method.

We apply the vanilla SGD update to all runs for efficiency. The activation checkpointing technique (Chen et al., ) is enabled by default. It is worth noting that pre-training these large models to the optimal performance is extremely expensive (Rajbhandari et al., ). Given that our NeuZip training method is lossless, we only train the models for 1 epoch to showcase its effectiveness. We use the same hyper-parameters for all runs.

#### Results.

We present the results in Table 1. We first test the vanilla training method, where only the activation checkpointing is applied (shown in Figure 2(b)). As shown, the vanilla training requires the highest amount of memory because it stores the uncompressed weights and gradients for all layers.

We also test the LOMO technique (Lv et al., ), which promptly updates the weights in a layer-by-layer fashion (shown in Figure 2(c)). This allows LOMO to reuse a buffer to store the gradients for each layer. As a result, LOMO approximately reduces the peak memory usage by the size of a model.

Finally, we apply our NeuZip on top of LOMO (shown in Figure 2(d)). For all models, NeuZip additionally reduces more than 20% percentage of memory compared with LOMO, accounting for a total memory reduction of more than 50%. Notably, NeuZip reduces the peak memory of training a Llama-2 13B model to less than 20GB, enabling training a 13B model on consumer-grade GPUs without any precision loss.

**Table 1: Pre-training decoder-only models on the language modeling task. The loss numbers are calculated on the validation set with the cross-entropy loss. Memory is reported in GiB ($1024^{3}$ B). Speed represents the number of iterations per second. The bold numbers represent the top results.**
|  | GPT-Neo-XL 2.7B | Llama-3 8B | LLama-2 13B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Name | Loss | Mem | Speed | Loss | Mem | Speed | Loss | Mem | Speed |
| Vanilla | 8.81 | 11.22 | 0.96 | 8.61 | 30.97 | 0.77 | - | OOM | - |
| LOMO | 8.81 | 6.97 | 0.94 | 8.61 | 19.47 | 0.78 | 9.10 | 26.26 | 0.49 |
| +NeuZip Lossless | 8.81 | 5.54 | 0.70 | 8.61 | 15.25 | 0.45 | 9.10 | 18.58 | 0.28 |

**Table 2: Fine-tuning encoder–decoder models on the SQL generation task. The BLEU scores are calculated with SacreBLEU. Memory is reported in GiB ($1024^{3}$ B). Speed represents the number of iterations per second. The bold numbers represent the top results.**
|  | T5 1B | T5 3B | T5 11B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Name | BLEU | Mem | Speed | BLEU | Mem | Speed | BLEU | Mem | Speed |
| Vanilla | 79.9 | 3.82 | 3.69 | 85.1 | 11.32 | 2.43 | - | OOM | - |
| LOMO | 79.9 | 2.75 | 3.68 | 85.1 | 7.07 | 2.47 | 82.3 | 25.95 | 0.69 |
| + NeuZip Lossless | 79.9 | 2.39 | 2.02 | 85.1 | 5.21 | 1.33 | 82.3 | 20.68 | 0.46 |
| QLoRA INT8 | 70.4 | 5.84 | 1.11 | 72.1 | 11.54 | 1.12 | 63.5 | 33.36 | 0.37 |
| QLoRA FP4 | 70.1 | 3.63 | 1.70 | 72.1 | 7.35 | 1.74 | 63.3 | 22.73 | 0.58 |
| QLoRA FP4^2 | 70.6 | 3.61 | 1.63 | 72.0 | 7.27 | 1.61 | 60.6 | 22.38 | 0.57 |
| QLoRA NF4 | 70.4 | 3.63 | 1.83 | 71.2 | 7.35 | 1.65 | 59.4 | 22.73 | 0.57 |
| QLoRA NF4^2 | 70.5 | 3.61 | 1.64 | 71.2 | 7.07 | 1.57 | 57.9 | 22.38 | 0.57 |

**Table 3: Evaluating lossy NeuZip on different models and tasks. ‘PPL” represents the perplexity values. Memory is reported in GiB. Speed represents the number of iterations per second. The bold numbers represent the top results, whereas the underlined numbers are the second-best ones.**
|  | Llama-3 8B | Llama-2 13B | Yi-1.5 34B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Name | PPL | Mem | Speed | PPL | Mem | Speed | PPL | Mem | Speed |
| Vanilla | 9.89 | 15.08 | 5.07 | 10.87 | 24.36 | 3.59 | - | OOM | - |
| Quant INT8 | 10.07 | 8.63 | 3.54 | 10.97 | 12.74 | 2.27 | 10.87 | 33.41 | 1.13 |
| Quant FP4 | 11.51 | 5.77 | 3.45 | 11.38 | 7.37 | 1.87 | 11.57 | 19.54 | 1.75 |
| Quant NF4 | 10.75 | 5.77 | 3.38 | 11.15 | 7.37 | 1.83 | 11.06 | 19.54 | 1.67 |
| Quant FP4^2 | 11.50 | 5.44 | 3.41 | 11.38 | 6.87 | 1.86 | 11.57 | 18.11 | 1.61 |
| Quant NF4^2 | 10.75 | 5.44 | 3.34 | 11.15 | 6.87 | 1.81 | 11.06 | 18.11 | 1.54 |
| NeuZip 0-bit | 13.64 | 5.24 | 3.44 | 12.46 | 6.30 | 1.87 | 12.06 | 16.20 | 0.94 |
| NeuZip 1-bit | 10.77 | 6.05 | 3.38 | 11.17 | 7.77 | 1.86 | 11.04 | 20.14 | 0.93 |
| NeuZip 3-bit | 9.93 | 7.70 | 3.38 | 10.90 | 10.73 | 1.84 | 10.76 | 27.92 | 0.93 |
| NeuZip 7-bit (lossless) | 9.89 | 10.95 | 3.39 | 10.87 | 16.66 | 1.84 | 10.72 | 43.40 | 0.94 |

### 3.2 Lossless NeuZip for Fine-Tuning

#### Settings.

A benefit of using lossless compression comes from retaining the pre-trained weight without any information loss. We conduct a fine-tuning experiment with encoder–decoder models to test the performance of our NeuZip on broader architectures. In particular, we choose
three T5 models: T5 1B, T5 3B, and T5 11B (Raffel et al., ), where the pre-trained parameters are used for initialization.

The T5 models are pre-trained on the C4 dataset (Lin et al., ), which is filtered to contain natural language only. To avoid data leaks from pre-training, we choose a non-natural language generation dataset for fine-tuning. Specifically, we use a public SQL generation dataset (Zhong et al., ; Yu et al., ) as the test bed. For each sample, the model is required to generate the SQL command from a human question. For example, the question could be “CREATE TABLE head (age INTEGER). How many heads of the departments are older than 56 ?”. The model is expected to generate “SELECT COUNT(*) FROM head WHERE age > 56”. We feed the question and response into the encoder and decoder, respectively. The objective is to minimize the cross-entropy loss on the response.

Similar to the pre-training experiments, we also sweep the learning rate from $10^{-3}$ to $3\times 10^{-1}$ for each run. After fine-tuning, we generate with the model on the validation set with greedy decoding. The generated SQL commands are then compared with the ground truths by SacreBLEU (Post, ), a metric that evaluates the similarity between corpora based on precision scores.

#### Results.

The results are reported in Table 2. All baselines in the pre-training experiment (i.e., the vanilla training, LOMO, and NeuZip) are included in this table. Similar to the results in Section 3.1, they achieve the same BLEU scores for each model. Specifically, our NeuZip is able to train a 11B model within 24GB.

For fine-tuning, it is possible to apply other memory-efficient training techniques. For example, QLoRA (Dettmers et al., ) compresses the pre-trained model by using low-precision data types and train the LoRA modules only (Hu et al., ). In our comparison experiment, we choose the widely used quantization data types for QLoRA, including INT8 (Dettmers et al., ), FP4, and NF4 (Dettmers et al., ). We apply the LoRA modules (Hu et al., ) on all linear layers, where every LoRA rank is set to $8$ to control the memory usage.(^2^22It should be noted that the down-projection matrices in each T5 feed-forward network are not quantized for stability, as otherwise the model performance is seriously jeopardized. See [https://github.com/huggingface/transformers/issues/20287](https://github.com/huggingface/transformers/issues/20287) for more details.) As shown in the second half of Table 2, all quantization methods underperform NeuZip in terms of both generation quality and memory usage. In terms of time efficiency, some quantization methods are slower than others, but in general, they are in the same magnitude as our method. Overall, NeuZip achieves the least memory usage while maintaining the highest performance.
The results strongly suggests the practicality of our NeuZip.

### 3.3 Lossy Compression for Inference

As mentioned in Section 2.3, the inference process is less sensitive in precision loss, which provides an opportunity for compressing mantissa in a lossy fashion during inference. We evaluate the performance of our lossy NeuZip in such scenarios.

#### Settings.

Following the settings in previous sections, we test our approach with both decoder-only and encoder–decoder architectures. For the decoder-only models, we select the LLama-3 8B (Dubey et al., ), LLama-2 13B (Touvron et al., ), and Yi-1.5 34B (Young et al., ). For the encoder–decoder architecture, we use the T5 1B, 3B, and 11B models as in Section 3.2.

Since all decoder-only models are trained for language modeling, we evaluate the performance with language modeling tasks. Specifically, we test all methods on the Wikitext-2 validation set (Merity et al., ) following Section 3.1, where each sequence consists of 1024 tokens. On the other hand, the encoder–decoder models (T5 series) contain multiple tasks in pre-training. Since they excel at zero-shot translation, we evaluate them on the WMT14 En-De translation task (Bojar et al., ), where each source sentence is prepended with “translate from English to German:” based on the pre-training format (Raffel et al., ).

Following the standard evaluation pipeline for lossy compression (Frantar et al., ; Dettmers and Zettlemoyer, ), we evaluate all models with the perplexity metric, which is sensitive to how distorted the compressed model is.

#### Results.

The results for decoder-only and encoder–decoder models are shown in Tables 3(a) and 3(b), respectively.
We see that the vanilla (uncompressed BFloat16) models achieve the best perplexity scores in all experiments at a cost of the excessive memory usage.
For quantization methods, we choose the same INT8 (Dettmers et al., ), FP4, and NF4 (Dettmers et al., ) data types mentioned in Section 3.2. In general, quantization methods suffer from notable perplexity degradation. Although the INT8 variant (Dettmers et al., ) manages to better preserve the perplexity, it uses around 50% more memory compared with other quantization methods.

For our lossy NeuZip, we set three different levels of precision: 0-bit, 1-bit, and 3-bit mantissa preserved. We choose these values because they are aligned in 8-bit byte arrays (discussed in Section 2.3). All these variants use a block size of 512 for normalization. We additionally include the lossless NeuZip (7-bit mantissa) for a full comparison.
As shown in the table, our lossy NeuZip demonstrates a spectrum of memory saving and performance preservation. The 0-bit NeuZip attains the best memory efficiency in all experiments, whereas the lossless 7-bit NeuZip obtains the best perplexity scores. Notably, the 3-bit NeuZip achieves nearly lossless performance in all experiments while using less than 50% memory compared with the uncompressed model. The results confirm the effectiveness of our method.

Figure: Figure 3: The trade-off between memory and performance for different methods.
Refer to caption: x2.png

### 3.4 In-Depth Analyses

#### The memory–performance trade-off.

In Section 3.3, we observe that the performance is generally decreased with less memory usage. We analyze this trade-off of our NeuZip as well as quantization methods in Figure 3. Note that the optimal methods are the ones on the Pareto frontier (Pareto, ), i.e., the more bottom-left, the better. In addition to measuring the perplexity, we also include a preliminary study by evaluating the end-to-end performance on the MMLU dataset (Hendrycks et al., ) in Appendix E.

As shown, three out of four NeuZip variants are on the Pareto frontier, with the remaining one staying fairly close to the frontier. On the other hand, there is only one quantization technique that lies on the Pareto frontier. This result demonstrates that our NeuZip generally achieves a better memory–performance trade-off than quantization.

#### The effect of block size in lossy compression.

As introduced in Section 2, we apply normalization to lossy NeuZip to ensure the weight with the largest absolute value will not be affected by truncation. We show the effect of block size in this experiment with a giant model, Llama-3 70B evaluated on the Wikitext-2 dataset.

**Table 4: The effect of block size.**
|  | Block 32 | Block 64 | Block 128 | Block 256 | Block 512 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Name | PPL | Mem | PPL | Mem | PPL | Mem | PPL | Mem | PPL | Mem |
| NeuZip 0-bit | 6.341 | 35.7 | 6.694 | 34.6 | 6.853 | 34.2 | 7.639 | 33.8 | 7.104 | 33.5 |
| NeuZip 1-bit | - | OOM | 4.611 | 42.7 | 4.662 | 42.2 | 4.640 | 41.8 | 4.649 | 41.4 |

As seen in Table 4, a smaller block size clearly leads to better performance at the cost of compromising memory efficiency due to the overhead of storing normalization coefficients. Therefore, the block-wise normalization provides a more fine-grained trade-off between memory and performance by varying the block size.

#### Throughputs of NeuZip.

Figure: Figure 4: The throughput experiment. (a) Comparison of CPU-offloading, quantization, lossy NeuZip compression, and lossless NeuZip compression. (b) Comparison of GPU-reloading, de-quantization, lossy NeuZip decompression, and lossless NeuZip decompression.
Refer to caption: x3.png

In addition to overall time efficiency presented in Tables 1–3, we analyze the throughput of matrix compression and decompression with our NeuZip, in comparison with the throughput of matrix quantization and de-quantization based on the NF4 data type (Dettmers et al., ) using the popular library bitsandbytes.(^3^33Available at [https://github.com/bitsandbytes-foundation/bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes))
We additionally include the CPU-offloading technique as a baseline, which lowers the GPU memory pressure by transferring data to CPU and reloading them to GPU when needed.
Figure 4 measures the throughput of matrix processing in GiB/s when we vary the matrix size from $10^{5}$ to $10^{8}$ bytes.

We see that CPU-offloading is generally slow across different sizes of matrices. This is due to the bottleneck of CPU–GPU communication through PCIe.
For quantization, the bitsandbytes package has been highly optimized for GPU, and its throughput is one magnitude higher than the CPU-offloading technique when the matrix size is large.
Profoundly, our NeuZip achieves the highest throughput for compression among all methods (Figure 4a), and a high throughput for decompression similar to de-quantization (Figure 4b). The results suggest that our NeuZip, albeit causing overhead compared with uncompressed vanilla models, is still highly efficient in practice.

## 4 Related Work

#### Model compression.

Previous work has explored different techniques to reduce the memory usage of neural networks, including knowledge distillation (Hinton et al., ) and pruning (Kwon et al., ). Most related to our work is the quantization technique, which represents each parameter with fewer bits; common approaches include $k$-means-based quantization (Han et al., ), linear quantization (Han et al., ), and mixed precision quantization (Dettmers et al., , ). When training data are available, one may incorporate the quantization into the training process to improve performance (Xiao et al., ; Frantar et al., ). In this paper, our NeuZip compression is a zero-shot method, and therefore, our experiments consider the widely used zero-shot quantization methods (Dettmers et al., , ) for fair comparison. We leave the utilization of additional data of NeuZip to future work.

#### Memory-efficient optimizers.

The optimizer also occupies a considerable amount of memory during training (Rajbhandari et al., ). To address this, memory-efficient optimizers (Shazeer and Stern, ; Zhao et al., ; Hao et al., ) are developed to reduce the memory footprint of training. Our NeuZip is orthogonal to these optimization techniques, as it can be seamlessly combined with any of these methods for further memory saving. In particular, the lossless NeuZip is expected to have exactly the same results with less memory.

#### Parameter-efficient training.

Another line of research saves memory by training a subset of parameters (Houlsby et al., ; Zaken et al., ) so the optimizer only stores information about a small set of trainable parameters. One notable example is the low-rank adaptation (LoRA (Hu et al., )). However, such a practice restricts the optimization space of parameters, and thus usually leads to significant performance degradation. Moreover, low-rank methods are unsuitable for pre-training.

It is important to mention that memory-efficient optimizers and parameter-efficient training cannot reduce the memory cost during inference. By contrast, our NeuZip is suitable for both training and inference.

## 5 Conclusion

#### Summary.

In this work, we present NeuZip, a novel compression scheme for neural networks that achieves memory-efficient training and inference. By analyzing the floating-point structures, we propose to compress the exponent in a lossless way and to compress the mantissa in a lossy way. The lossless variant of our NeuZip may be applied to both training and inference, while yielding exactly the same result as the uncompressed model. The lossy NeuZip provides additional memory saving for inference, achieving superior memory–performance trade-off.

#### Limitations and future work.

Due to the hardware constraint, the largest model that we consider in this paper is 70B. We would like to verify our NeuZip on even larger models like GPT-3 (Brown et al., ) with more capable hardware. Another limitation of NeuZip is that the throughput is lower than the vanilla model. However, it has a comparable speed with highly optimized quantization methods while achieving significantly better performance. By using NeuZip, we expect to create opportunities for researchers with academic budgets to explore and study large models.

## Acknowledgments

The research is supported in part by the Natural Sciences and Engineering Research Council of Canada (NSERC), a Mitacs Accelerate project, the Amii Fellow Program, the Canada CIFAR AI Chair Program, an Alberta Innovates Program, and the Digital Research Alliance of Canada (alliancecan.ca).

## Appendix A Inspecting the Entropy on More Models

#### Random initialization.

When training from scratch, the parameters are randomly initialized. To verify the compressibility in this case, we check the parameter entropy of a randomly initialized model with the same architecture as Lllama-3. The initialization methods follow the standard procedure provided in the Hugging Face library. The results show a similar pattern to what the released Llama model has, suggesting the compressibility with NeuZip occurs even with random weights.

Figure: Figure 5: The histograms of different floating-point components of the parameters of a randomly initialized Llama-3 8B model.
Refer to caption: x4.png

#### Diffusion.

We also inspect the parameter entropies beyond Transformer models. In Figure 6, we check all four models in a diffusion pipeline. We see that the low-entropy exponents not only occur in Transformer models but other architectures like convolution-based VAE and U-Net models.

Figure: Figure 6: The histograms of the exponent bits in different components of Stable Diffusion 1.5 model. We omit the sign and mantissa bits for simplicity as we do not compress them based on their entropies.
Refer to caption: x5.png

Both experiments show that the occurrence of low-entropy components is a common phenomenon in deep learning.

## Appendix B The Algorithm for Training with Lossless NeuZip

In this section, we describe the forward-backward procedure of NeuZip. First, we compress all the linear layers in the original model and store the compressed information on-device. During the training iterations, we decompress the compressed weights in a layer-by-layer manner for the forward pass. For the backward pass, the input is recalculated again following the forward pass like activation checkpointing (Chen et al., ). A linear operation calculates the gradients for both the weight matrix and input. To do so, we need to decompress the weight again, which is used to calculate the gradient of input. After the gradient is calculated, we directly update the weight without storing it similar to LOMO (Lv et al., ).

Figure: Algorithm 1 Memory-efficient training with NeuZip

## Appendix C The Tolerance of Random Perturbation

In this experiment, we would like to explore the sensitivity of neural network weights to random perturbations. For each parameter, we have two types of magnitudes: absolute and relative magnitudes. The former one represents the actual numerical error, whereas the second one is calculated based on the original value. For example, when the original value is $-1.5$, an absolute magnitude of $0.125$ means the perturbed range is $[-1.5-0.125,-1.5+0.125]$. On the other hand, a relative magnitude of $0.125$ means the perturbed range is $[-1.5*(1+0.125),-1.5*(1-0.125]$.
We conduct such a experiment with the perturbation grid in Figure 7. For each cell, we choose the maximum error between relative error and absolute value for perturbing. A random value is sampled from the perturbed range uniformly as the perturbation. The weight value is then set to the random sample.

Figure: Figure 7: Evaluating the byte-level perplexity with perturbed LLama-3 8B model (Dubey et al., ) on Wikitext-2 (Merity et al., ). Each parameter is perturbed with controlled noises. Both the $x$- and $y$-axes are log-scale with base 2.
Refer to caption: x6.png

As shown in Figure 7, we see a clear pattern that the model tends to tolerate the relative change rather than the absolute change.

## Appendix D The Storage for Lossless and Lossy Compression

In this section, we describe the underlying storage layout for NeuZip in Figure 8.

Figure: (a) Lossless compression scheme.
Refer to caption: x7.png

Essentially, each BFloat16 number is first split into an exponent and a signed mantissa. We group all the exponents in the matrix and perform the lossless compression. The signed mantissa is optionally truncated, depending on the required precision. The signed mantissa is then stored separately in memory.

## Appendix E Evaluating on MMLU

We provide the results on MMLU (Hendrycks et al., ) in Figure 9. Here, the theoretical optimal point should be at the top left corner.

Figure: Figure 9: The memory–performance trade-off on the MMLU dataset.
Refer to caption: x9.png

Similar to the results in Section 3.4, all of our NeuZip variants are on the Pareto frontier, suggesting the optimal trade-off between memory and performance.

## Peer Reviews

confidence: 2
contribution: 1
initial_rating: 3
presentation: 1
questions: Are the experiments of pre-trained being ran for a full epoch? How many A6000s are used and how long it took?
soundness: 2
strengths: The paper made relatively extensive experiments on the modern LLMs, and showed its advantage in reducing memory footprint.
summary: This paper presents a entropy based quantization method to reduce the memory footprint in training. It is claimed that this method could reduce around half of the memory without losing accuracy.
weaknesses: 1. The motivation is unclear. It's not clear that where the entropy for the sign bit (Figure 1) comes from. The author did not mention in details, where the pre-trained model weight is obtained and whether it is an empirical value or theoretical value, making the analysis not persuasive enough.

2. The compression algorithm is not clearly presented.

3. Codes are not released, otherwise we can make clear what the methodology is by checking the implementation.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: - Transformer-type language models such as BERT-type are getting implemented on edge devices such as cell phones. Do you see a potential that this compression helps the edge implementation of LLMs?
- Do you see a potential for Transformer-alternative architectures such as MAMBA, xLSTMs, JAMBA, etc?
soundness: 4
strengths: The paper addresses the important problem of increased size of LLMs with lack of resources for training and inference. The algorithm can be built on top of other compression methods. The authors also show that their lossy compression falls on the Pareto frontier and therefore a good compression candidate overall.
summary: This is an interesting work that uses small entropy in the exponent bit to compress neural networks in two settings: lossless and lossy. In the lossless setting the memory consumption and inference speed are both improved. The authors try their models on LLMs with <70B parameters but the method is general and can be used to compress LLMs of larger size. The paper is properly structured, well-written, and easy to follow.
weaknesses: The paper motivates on linear layer, and their compression requires activations to be saved. This is a restriction specially when small values like GELU type activations are generated in the network. I wonder if the paper suits partial linear activations like ReLU better. Discussion on convutional models, and RNNs will strengthen the paper further.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: - I wonder how the reported speeds for the quantization methods such as Quant INT8 in Table 2 came to be. Especially for large models, inference speed-ups should be expected, see Dettmers et al. 2023 (LLM.int8(), https://arxiv.org/abs/2208.07339), appendix D, p.17ff, where LLM.int8 reports increased inference speed for all models that have more than 13B parameters). 
- Similarly, for QLoRA, which has been shown to make networks more memory efficient in the original work, this paper seems to report contrary results, decreasing or barely reducing the memory efficiency compared to vanilla networks as well as suffering a large performance hit (f.e. T5 1B). A discussion of why this happens could be beneficial to the paper, and if it turns out that QLoRA is not suitable for networks this small, other state-of-the-art methods should rather be used to compare _NeuZip_ with.
soundness: 2
strengths: - The paper is well structured, the experiments are well described and the figures and tables are informative.
- The developed method in the paper is well explained and clear to understand.
- The experiment settings and presentation of results are well chosen to show the effect of proposed method.
summary: The paper presents NeuZip a compression technique for reducing the memory requirements of neural networks for training and inference. NeuZip contains both a lossless and a lossy operating mode, where either only the exponent bits are (losslessly) compressed, or additionally a truncation of the mantissa bits is performed. The performance is tested on various models from the Natural Language Processing community.
weaknesses: - The novelty of the _NeuZip_ method seems poor.  Lossless NeuZip consists of the direct application of a very widely used compression method (ANS) to (parts of) the weight vector. The lossy compression part is a truncation of mantissa bits, which is a well-known method in data compression (f.e. used in the GRIB2 format). Moreover, the paper fails to mention that this truncation itself is equivalent to quantisation to a non-uniform grid (f.e. float32 with 0 mantissa bits corresponds to quantization with the grid $\{\pm 2^{-127}, \pm2^{-126}, \dots \pm2^{128}\}$). Quantization with non-uniform grids is a widely used idea (f.e. in "Data Free Network Quantization", Chikin et al. 2022, or NF4 from QLoRA, Dettmers et al. 2023).
- The methods that the paper compares lossless _NeuZip_ to are not fairly chosen, and this is not immediately made obvious in the experimental descriptions. In Table 2, the quantisation methods that _NeuZip_ compares to contain no entropy coding, while _NeuZip_ contains entropy-coding in the exponent bits (which will be especially significant as the ratio of exponent bits to mantissa bits increases). Including entropy coding into the compared quantisation methods would reduce the speed by a small bit, but increase the compression performance at no perplexity cost (which are the metrics the methods are compared on). Again, as the proposed method is equivalent to non-uniform grid quantisation, I would expect most of the gains that are reported in this table to have come from the included entropy coding.
- A large part of the memory efficiency gains of lossless _NeuZip_ come from the established LOMO method (see Table 1). Adding _NeuZip_ on top often reduces the execution speed by a large margin.

Meta-review: This paper proposes NeuZip, a weight compression technique that takes account of the entropy of floating point numbers in NNs. The authors propose a lossless, and lossy variant, and evaluate these on various language models.

I will not be taking Reviewer nr8v's review into account for this meta review, as it is extremely short and poor quality. I asked the reviewer the day after the review was posted to provide a high quality review ant this was ignored. 

In terms of strengths, Reviewer UCfs was appreciative of the problem tackled, and the experimental results. They also approved of how the algorithm could be combined with existing techniques. Reviewer DiBQ thought the paper was strong on presentation.

In terms of weaknesses, Reviewer UCfs wanted to see applicability to model types beyond transformers (which I note that the authors provide with Mamba experiments). Reviewer DiBQ had major concerns about novelty, and the experimental setup.

After an extensive discussion, Reviewer DiBQ still has concerns that this work does not provide much additional knowledge to the community, and that the experimental setup presented is unfairly chosen to benefit the proposed method. From my perspective, the negatives presented outweigh the positives (it isn't clear what elevates this work from the bulk of existing compression methods), so I am leaning towards rejection. However, I am not fully confident in my decision as I am effectively basing this on 2 reviews, so I will indicate this in my confidence score.

Decision: Reject