Robust Training of Neural Networks at Arbitrary Precision and Sparsity

## Abstract

Abstract The discontinuous operations inherent in quantization and sparsification introduce obstacles to backpropagation. This is particularly challenging when training deep neural networks in ultra-low precision and sparse regimes. We propose a novel, robust, and universal solution: a denoising affine transform that stabilizes training under these challenging conditions. By formulating quantization and sparsification as perturbations during training, we derive a perturbation-resilient approach based on ridge regression. Our solution employs a piecewise constant backbone model to ensure a performance lower bound and features an inherent noise reduction mechanism to mitigate perturbation-induced corruption. This formulation allows existing models to be trained at arbitrarily low precision and sparsity levels with off-the-shelf recipes. Furthermore, our method provides a novel perspective on training temporal binary neural networks, contributing to ongoing efforts to narrow the gap between artificial and biological neural networks.

## 1 Introduction

The recent surge in the size and complexity of generative AI models has elevated computational efficiency to the forefront of AI research [, ]. Among the diverse approaches to achieving efficiency, quantization and sparsification techniques stand out as two classic and widely explored methods. Quantization and sparsity techniques can effectively reduce the computational requirements of large language models (LLMs). Quantization reduces the precision of model weights and activations, thereby decreasing storage requirements. Sparsification, on the other hand, prunes redundant weights, leading to a more compact model. These techniques enable LLMs to be deployed on resource-constrained devices, such as mobile phones and embedded systems, while also improving their speed and memory efficiency. This facilitates the widespread adoption of LLMs by a broader range of individuals and businesses.

Figure: (a)
Refer to caption: x1.png

Despite their promise, quantization and sparsification introduce non-differentiable operations, such as rounding and hard thresholding, which are incompatible with the differentiable design of backpropagation, the cornerstone of neural network training. This incompatibility has plagued training algorithms for decades, hindering progress in the field of efficient neural networks [, , ].

To tackle the discontinuity challenges inherent in training efficient neural networks, algorithms have primarily focused on adapting gradient descent algorithms to work with non-differentiable operations. Empirical techniques such as the straight-through estimator (STE) [] have been employed to define gradients for non-differentiable operations. However, even with the STE, the perturbation introduced by quantization has been observed to disrupt existing training recipes (Fig. 1(a)). Consequently, clipping is commonly applied to limit the signals within a small range to prevent divergence. Despite these techniques, training quantized networks remains restricted to specific precisions [, , , ] or models [, ]. When moving to low precisions, these approaches usually also make changes to model architectures and recipes, such as inserting extra normalizations[, ], changing learning rates[], replacing optimizers [], keeping several layers unquantized [, ] or through fine-tuning []. It is unclear if these techniques are effective and flexible enough in the large generative AI models.

In contrast to the intricate and often empirically tuned techniques of prior methods, we adopt another approach that formulates discontinuous operations as the introduction of perturbations []. We address this challenge directly by suppressing the effects of these perturbations, effectively denoising the signal. Our approach comprises three fundamental steps:

1. Affine Transform for Quantization (Sec. 4.1): An initial affine transform $f$ scales the input signal without introducing additional operations (e.g., clipping).

2. Perturbation Injection (Sec. 4.2): A controlled perturbation $\bm{\delta}$ is injected into the signal, precisely modeling the effect of quantization.

3. A Denoising Affine Transform for Reconstruction (Sec. 4.3): A key innovation of our approach is the introduction of another affine transform $g$ that effectively reconstructs the original signal while mitigating quantization noise.

Our method offers several key advantages: 1. Continuous control over quantization noise: This ensures stable model training and prevents divergence.
2. Graceful degradation: Under extreme noise, our approach seamlessly transitions to a lower-resolution (through averaging) model, guaranteeing a performance lower bound.
3. Compatibility: Our method works with existing architectures and training recipes, eliminating the need for extensive modifications or hyperparameter tuning.

These innovations enable the development of diverse, efficient neural networks for various applications. We demonstrate the effectiveness of our approach by training sparse models and temporal binary neural networks. Our key contributions include:

- •
A simple, robust, and universal solution for training quantized neural networks.
- •
A novel theoretical perspective and practical solution for training networks with discontinuous operations, based on decomposition and denoising.
- •
State-of-the-art results with ultra-low precision models, using standard architectures and training recipes.
- •
Easy adaptation to train sparse neural networks.
- •
Successful training of temporal binary neural networks, demonstrating the robustness of our approach and its potential for bridging the gap between artificial and biological neural networks.

## 2 Motivations

### 2.1 A Comprehensive Picture of Model Efficiency

While extreme quantization and sparsification can significantly reduce storage and computational requirements, they also introduce the risk of quality degradation. Existing studies on quantization and sparsity primarily focus on showcasing the effectiveness of specific implementations through meticulous tuning. While these studies provide valuable insights, they fall short of providing a comprehensive understanding of the performance trade-offs involved in applying these techniques to extreme levels. The absence of a robust, universal algorithm for training quantized and sparse neural networks has hindered an accurate comparison of different approaches and has limited the exploration of the full potential of these techniques. To address these limitations, a more comprehensive and rigorous approach is needed to evaluate the trade-off between efficiency and quality. This necessitates the development of a universal approach that can effectively manage extreme quantization and sparsity levels (Fig. 1(c,d)).

### 2.2 A Biophysical Basis for Sparse Quantization

One of the most significant achievements in 20th-century neural physiology was the development of the Hodgkin–Huxley model for modeling animal neural networks [, ]. This model unveiled the intriguing fact that animal neural networks primarily consist of brief electrical impulses, commonly known as spikes in their activity patterns. Despite this critical insight, the question of how animals efficiently learn and process information through these spike trains continues to pose a profound and unresolved challenge. Since the rise of regular neural networks, researchers have strived to find a more biophysically plausible approach to artificial intelligence through spiking neural networks [, , , ]. However, the lack of differentiability in spiking neural networks prevents the straightforward application of gradient descent algorithms, the foundation of modern neural networks. Consequently, these networks have limited application and popularity. The development of universal training algorithms for quantized and sparse neural networks could open up new avenues in the field of spiking neural networks.

## 3 Related Work

Neural network quantization has become a widely adopted technique for reducing memory footprint and accelerating inference time, enabling efficient deployment on resource-constrained devices []. While full-precision models typically store weights in floating-point format, quantized weights are represented as integers, typically using 8 bits [, , ], 3-4 bits [, , , ], or even 1 bit [, , , , , ]. In addition to quantizing weights, model activations can also be quantized to further enhance computational efficiency [, , ].

Although 8-bit quantization is commonly used as a standard practice in industry [], achieving lower-bit quantization remains challenging and requires specialized techniques to ensure robust training. Several common techniques include: 1. Mixed precision quantization: This approach selectively assigns different bit levels to different weights, aiming to optimize the trade-off between model size and accuracy [, , , ]. 2. Training recipes: These techniques compensate for the discontinuities introduced by quantization by employing strategies such as sharpness-aware minimization [, ], state-aware optimization [], knowledge distillation [], and multi-phase training []. 3. Quantization-friendly network architectures: This approach involves replacing original network layers with alternatives that are more amenable to quantization [].

In contrast to prior work, our method explicitly models quantization discontinuities as perturbations. We decompose the perturbed signal into clean and noisy components, then apply denoising to suppress the noise. This approach leads to a closed-form solution that guarantees training convergence even at extremely low bitwidths. While previous methods have also modeled quantization noise using continuous distributions (e.g., Uniform or Gaussian) for gradient estimation [, ], they do not optimize the reconstruction process itself to enhance training stability.

To further reduce model footprint, researchers have been combining sparsity/pruning and quantization in a unified formulation to further compress neural networks [, ]. In this paper, we extend our noise injection and denoising reconstruction theory to sparsity and argue that instead of pruning small values to zero, moving near-mean values to mean better preserves signal fidelity.

## 4 Methods

We primarily focus on explaining our formulation for quantization, as it is more widely supported by modern hardware and offers a broader range of applications. Subsequently, we demonstrate how our formulation can be extended to sparsification. Our method’s flexibility enables it to address quantization and sparsification individually or together, allowing users to apply it to model weights and activations independently or in combination.

### 4.1 Affine Transform for Quantization

We begin our study by using min-max scaling to move and scale the floating-point vector $\bm{x}$ to the desired range, e.g. $[0,2^{bits}-1]$. We formulate with this range for simplicity. The range can be shifted if signed integers are more compatible with specific hardware. This affine transform is implemented using standard functions in neural network libraries for backpropagation training.

$$ $f(\bm{x})=\frac{\bm{x}-x_{min}}{x_{max}-x_{min}+\epsilon}\cdot(2^{bits}-1)$ (1) $$

### 4.2 Perturbation Injection

Quantization reduces the number of bits used to represent $f(\bm{x})$ by rounding to the nearest integer. This rounding operation introduces discontinuities at half-integer values, rendering it non-differentiable. These discontinuities can lead to training difficulties, as neural networks rely on gradients for learning. Most existing methods for training neural networks with non-smooth operations employ the straight-through estimator (STE) []. The STE provides a technique for defining gradients for non-differentiable operations. However, even with the STE, training algorithms may diverge when moving to lower bits (Fig. 1(a)).

Building upon prior research [, , ], we model the impact of these operations as perturbations within the network, recognizing their role in causing training instability.
Specifically, quantization of $f(\bm{x})$ can be modeled as injecting a bounded perturbation:

$$ $\bm{q}=f(\bm{x})+\bm{\delta},$ (2) $$

where $\bm{\delta}=\mbox{round}(f(\bm{x}))-f(\bm{x}),\bm{\delta}_{i}\in[-0.5,0.5]$. We avoid empirical operations such as clipping, which can introduce larger perturbations and degrade model performance. By directly incorporating the quantization noise into the model, our method maintains signal fidelity and achieves superior results compared to traditional clipping-based approaches. Combining with the first affine transform (Eq. 1), our quantization can be found in Fig. 2.

Figure: Figure 2: JAX reference code for quantization: $\bm{q}=f(\bm{x})+\bm{\delta}$.
Refer to caption: x5.png

### 4.3 Denoising Affine Transform for Reconstruction

A core innovation of our approach is the introduction of an additional affine transformation $g$, designed to be resilient to perturbations. This enhances the robustness of quantized neural networks, and remarkably, we find that standard training methods used for full-precision models can still effectively converge even when perturbations are present.

Quantization algorithms typically introduce a one-dimensional affine or linear transform, often referred to as dequantization, to approximate the original vector $\bm{x}$. To streamline computations, most practical implementations simply invert the scaling involved in Eq. 1 [, , ] or minimize the $L_{2}$ reconstruction error [, ]. However, focusing solely on minimizing this approximation error may overlook the crucial challenge of training neural networks to be robust against perturbations (Fig. 1(a)).

Our approach, counterintuitively, increases the approximation error but naturally resolves this long-standing issue. We formulate the reconstruction as a ridge regression problem, introducing a regularization factor $\lambda$ and solving for two parameters:

$$ $\min_{a,b}\frac{1}{2N}||a\cdot\bm{q}+b-\bm{x}||^{2}+\frac{\lambda}{2}a^{2}$ (3) $$

Here $N$ is the dimension/length of $\bm{x}$. Taking the derivative with respect to $b$, and setting to zero yields the following equation: $\frac{1}{N}\sum_{i}(b+a\bm{q}_{i}-\bm{x}_{i})=0$.

Solving for $b$ yields:

$$ $b=\overline{\bm{x}}-a\overline{\bm{q}}$ (4) $$

And from setting the derivative with respect to $a$ to zero: $\frac{1}{N}\sum_{i}\bm{q}_{i}(a\bm{q}_{i}+b-\bm{x}_{i})+\lambda a=0$. Simplifying the equation, we obtain:

$$ $\overline{\bm{q}^{2}}a+\overline{\bm{q}}b-\overline{\bm{xq}}+\lambda a=0$ (5) $$

Substituting Eq. 4 into Eq. 5: $\overline{\bm{q}^{2}}a+\bar{\bm{q}}\bar{\bm{x}}-\overline{\bm{q}}^{2}a-
\overline{\bm{xq}}+\lambda a=0$.

We arrive at the solution for $a$:

$$ $a=\frac{Cov_{xq}}{Var_{q}+\lambda}$ (6) $$

Substituting $a$ and $b$ back into the ridge regression yields the reconstructed vector, representing the quantized version of $\bm{x}$:

$$ $\bm{r}=g(\bm{q})=a\cdot\bm{q}+b=\frac{Cov_{xq}}{Var_{q}+\lambda}(\bm{q}- \overline{\bm{q}})+\overline{\bm{x}}$ (7) $$

This affine transformation can also be seamlessly implemented using fundamental operations readily available in neural network libraries (Fig. 3). Quantization and reconstruction can be represented as a straightforward addition of $\bm{\delta}$ between the two affine transform operations: $r=g(f(\bm{x})+\bm{\delta})$. During actual computations, $f(\bm{x})+\bm{\delta}$ is cast to the appropriate data type.

#### 4.3.1 A Novel View of the Quantized Signal

The reconstruction from Eq. 7 provides a novel decomposition of the quantized signal $\bm{r}$ into a smooth component, $\overline{\bm{x}}$, and a non-smooth component, $a(\bm{q}-\overline{\bm{q}})$, drawing inspiration from detail enhancement techniques in computer vision [, ]. Importantly, the bounded perturbation from quantization, $\bm{\delta}$, is entirely contained within this non-smooth component. This observation is key, as this non-smoothness directly contributes to training instability. Consequently, we can stabilize the training by suppressing this component using the parameter $\lambda$. In essence, $\lambda$ acts as a control knob, regulating the balance between signal and noise that enters the training process, as visualized in Fig. 1(b).

Figure: Figure 3: The denoising affine transform $\bm{r}=g(\bm{q})$.
Refer to caption: x6.png

Our design exhibits two important properties:

###### Proposition 1 .

By adjusting $\lambda$, the quantized model can be trained to converge if the training algorithm converges on a smaller scale network.

###### Proof.

As $\lambda$ approaches infinity, the perturbation can be completely suppressed, resulting in the mean value $\overline{\bm{x}}$ serving as a fail-safe vector for the reconstruction of $\bm{x}$. This behavior is mathematically expressed as:

$$ $\lim_{\lambda\to\infty}a=0,\lim_{\lambda\to\infty}b=\overline{\bm{x}}.$ (8) $$

∎

The structure of this smaller scale network will be explained in Sec. 4.5. In practice, we find that a small $\lambda$ provides a good trade-off between preserving the original signal and suppressing quantization noise. This leads us to discuss the other extreme case of $\lambda=0$.

###### Proposition 2 .

The regularized model evolves continuously from the original model through the control parameter $\lambda$.

###### Proof.

Setting $\lambda=0$ (no regularization) and $\bm{\delta}=0$ (eliminating perturbations) restores the network to its original form. Intuitively, this behavior arises from the fact that for unperturbed input data $\bm{x}$, $\bm{q}=f(\bm{x})$. The scaling factor $a$, calculated as
$\frac{Cov_{xf(x)}}{Var_{f(x)}}=\frac{x_{max}-x_{min}+\epsilon}{2^{bits}-1}$, effectively inverts the scaling in Eq. 1. Since $a\cdot f(\bm{x})=\bm{x}-x_{min}$, the reconstructed signal $\bm{r}$ is then computed as $\bm{r}=a\cdot(f(\bm{x})-\overline{f(\bm{x})})+\overline{\bm{x}}=\bm{x}-x_{min}
-\overline{\bm{x}}+x_{min}+\overline{\bm{x}}=\bm{x}$. In practice, a very small $\lambda$ is used, the resulting model is a small change from the original network. This property suggests the full precision networks can be easily finetuned to a low precision networks. These properties highlight the robustness and flexibility of our proposed quantization method.
∎

###### Proposition 1 .

###### Proof.

###### Proposition 2 .

###### Proof.

#### 4.3.2 The Sensitivity of the Scaling Factor

By inverting the scaling factor in Equation 1, we can map the maximum perturbation of $0.5$ to a perturbation scaled at a magnitude of $\frac{x_{max}-x_{min}}{2^{bits}}$ in the original training signals. This implies that the perturbation intensity doubles with each reduction in bit precision, leading to a substantially amplified impact on the training process. This heightened perturbation has disrupted traditional training recipes (Fig. 1(a)), necessitating the introduction of a regularization term to stabilize the training process.

The analysis of the solution for $a$ is rooted in well-established principles of perturbation theory, as outlined in Golub’s work [] (Section 2.6).

$$ $(Var_{q}+\lambda)a=Cov_{xq}$ (9) $$

When $\bm{\delta}$ is bounded, the impact of $\bm{\delta}$ on $a$ is proportional to $\kappa=\frac{1}{Var_{q}+\lambda}$  [] (Eq. 2.6.4). $\lambda$ plays a crucial role in establishing an upper bound on $\kappa$. In the absence of $\lambda$, the influence of $\bm{\delta}$ on $a$ may become unbounded. This aligns with our observations that strong perturbations associated with lower precision lead to corrupted training. The introduction of $\lambda$ effectively stabilizes the training process by mitigating the impact of quantization noise and ensuring that the solution to Eq. 9 remains well-conditioned.

### 4.4 Extension to Sparsification

Similar to quantization, sparsification introduces discontinuities through the hard thresholding operation $H$. Consequently, the sparsification process can also be modeled as introducing perturbations, where $\bm{\delta}=H(\bm{x})-\bm{x}$. In our experiment section we show this leads to significant gain in sparsifying the model. Since the perturbation from sparsification rarely leads to divergence, we may set $\lambda=0$ to pass the entire perturbed signal when only sparsification is applied.

When dealing with extreme sparsity or quantization, it is noteworthy that $100\%$ sparsity, under the traditional definition, results in an all-zero vector. To address this limitation, we revisit Equation 7, which defines the reconstructed signal. Equation 7 involves adding back significant deviations to a mean vector of $\overline{\bm{x}}$, ensuring that the mean $\overline{\bm{x}}$ remains the foundation of the signal even at extreme sparsity levels. Therefore, sparsification towards the mean presents itself as the most intuitive approach. In this context, significance is measured by the deviation from the mean $\bm{x}_{i}-\overline{\bm{x}}$, not from zero. Consequently, extreme sparsity results in the mean $\overline{\bm{x}}$ being preserved, rather than zero.

### 4.5 Quantized Matrix Multiplication

In matrix multiplication $Y=XW$, we apply sub-channel quantization by reshaping the last dimension of $X$ and the first dimension of $W$, effectively splitting vectors along the contraction dimension into smaller, manageable chunks. Quantizing each chunk independently leads to a higher overall approximation quality than quantizing the entire vector or matrix at once. This finer-grained approach reduces perturbation, resulting in a more stable model. The resulting “fail-safe” backbone model (obtained as $\lambda\to\infty$) can be visualized as a piecewise constant function (Fig. 1(b)), where each piece is represented by its mean value. This is akin to lowering the granularity of the original model, effectively providing a trainable backbone with reduced complexity.

This sub-channel quantization approach permits low-precision block-wise calculation of the expensive matrix multiplication, followed by summation of partial results, can be implemented through batch matrix multiplication, drastically reducing the overall computational burden. While the provided reference code (Fig. 4) utilizes “fake-quantization” for clarity, the actual implementation involves true quantization and match matrix multiplication to achieve performance gains.

Figure: Figure 4: Reference code for a quantized matrix multiplication.
Refer to caption: x7.png

The memory savings achieved depend on the chosen block size $B$. For instance, when $a$ and $b$ are stored as $16$-bit floats, this results in an additional memory overhead of $32$ bits per block. The effective number of bits per element is reduced to $32/B$. In practice we notice that even $8$-bit floats deliver similar quality results. A typical block size of $128$ is employed to maintain a storage overhead of less than one bit. The block size introduces a trade-off between memory savings and model accuracy. We investigate this trade-off further in our experiments.

## 5 Experiments

To ensure consistency, we maintained a fixed block size of $128$ across all experiments unless otherwise specified. Our experiments revealed that the regularization parameter $\lambda=0.01$ consistently yielded satisfactory results. Quantization was applied to all linear transforms within the model, including convolutional layers. All models were trained from scratch, with architectures and parameters unchanged to enable a fair comparison between full-precision and quantized models. We utilize the notation “AxWy” to represent the quantization configuration of a neural network, where “A” denotes the bitwidth for activations and “W” denotes the bitwidth for weights. For instance, “A4W4” indicates that both activations and weights are quantized to 4 bits.

Perturbations can have a two-sided effect on neural network training. While small perturbations can mimic data augmentation and enhance the training process, they can also introduce unwanted variability into the training signal, especially when dealing with ultra-low bit and extremely sparse representations. This unwanted variability can prevent the model from learning effectively and slow down the model’s convergence. To further assess the performance of the low-precision models, we extended our main experiments by four times the training duration. With extended training our $1$-bit models remain highly competitive without any architectural change or recipe tweaking (Tables 1, 3). This demonstrates the robustness of our method and its ability to maintain model performance over extended training periods.

Our methodology stands out for its simplicity and generality, having undergone extensive testing on a diverse range of models and datasets, from small to large-scale. To showcase its efficacy, we present results on two well-established architectures applied to two widely recognized datasets. First, we demonstrate the effectiveness of our approach on ResNet-50 trained on ImageNet, achieving satisfactory performance without any empirical tuning. Subsequently, we evaluate our method under the same setting on the Transformer model, which is more relevant to generative AI applications. Our
low precision training consistently surpasses the full-precision training results, demonstrating its adaptability to different model architectures and tasks.

### 5.1 Ultra-low Precision Models

#### 5.1.1 ResNet-50 on ImageNet

**Table 1: Top-1 Validation Accuracy of ResNet-50 on the ImageNet Dataset.**
| Precision | 100 Epochs | 400 Epochs |
| --- | --- | --- |
| A32W32 | 76.41 | - |
| A4W4 | 76.45 | - |
| A4W2 | 75.12 | 75.59 |
| A4W1 | 72.04 | 73.97 |

We utilized the Flax framework to train ResNet-50 from scratch on ImageNet, employing stochastic gradient descent with an initial learning rate of 0.1, training the model for 100 epochs with weight decay of 0.0001. As shown in Table 1, the top-1 accuracy from the A4W4 configuration (76.45) surpasses the baseline (76.41) without any hyperparameter tuning. This demonstrates the effectiveness of our method in achieving competitive performance without requiring extensive optimization. We compared our results to previously reported A4W4 quantization results for ResNet-50 trained on ImageNet. Our results compare favorably to existing work, without the need for additional operations such as parameter search, fine-tuning, calibration, clipping, gradient estimation, or reinforcement learning (Table 2).

**Table 2: Comparison of top-1 accuracy for A4W4 ResNet-50. The columns represent: FP32: Full precision baseline. A4W4: Quantizing both activations and weights to 4 bits. GE: Whether gradient estimation is involved. PT: Pretraining/finetuning/calibration required. Clip: Clipping required. LB: The lowest bitwidth reported in the corresponding paper. * estimated from Fig. 1 []**
| Method | FP32 | A4W4 | GE | PT | Clip | LB |
| --- | --- | --- | --- | --- | --- | --- |
| AQT [] | 76.65 | 76.4^∗ | Y | Y | Y | 4 |
| VS-Quant [] | 76.16 | 75.28 | Y | Y | Y | 3 |
| FAQ [] | 76.15 | 76.25 | Y | Y | Y | 4 |
| HAQ [] | 76.15 | 76.14 | N | Y | Y | 4 |
| Ours | 76.41 | 76.45 | N | N | N | 1 |

#### 5.1.2 Transformer on WMT

To evaluate the effectiveness of our method on transformer models, we employed the Flax framework to train the transformer model on two WMT2017 datasets (EN-DE, DE-EN) and subsequently assessed its performance on the corresponding WMT2014 datasets. The training process utilized the AdamW optimizer with weight decay set to 0.1 and a cosine scheduler for 25,000 steps, employing a batch size of 1024. Recognizing the known slow convergence of transformer models, we extended the training duration to 100,000 steps (Table 3). Remarkably, our low-precision results consistently surpass the full-precision baseline.

Given the prevalence of transformers in large language models, extensive research has been dedicated to quantizing transformer models. We compare our findings to other works, and ours stands out as the only method that can surpass the full-precision baseline (Table 4). This achievement highlights the unique strength of our formulation, which not only preserves signal fidelity but also benefits from regularization effects. Several recent works have explored alternative quantization approaches using different datasets, which are not included in this table. One such method is AWQ, a weight-only quantization 4-bit quantization method [], requires retaining $1\%$ of salient weights and all activations unquantized. Their method also involves searching for an optimal scaling factor and a calibration set. Additionally, BitNet [], presents a 1-bit quantization method for transformers. The lowest activation precision achieved in their work is 8 bits, exceeding the highest activation bit in our method. Their method also necessitates clipping, additional normalization, and recipe changes.

**Table 3: BLEU Score of training low precision Transformers on the WMT datasets.**
|  | DE-EN | EN-DE |  |  |
| --- | --- | --- | --- | --- |
| Steps | 25k | 100k | 25k | 100k |
| A32W32 | 33.5 | 33.9 | 29.49 | 29.8 |
| A4W4 | 33.78 | 33.64 | 29.71 | 30.17 |
| A4W2 | 33.45 | 34.04 | 28.58 | 30.03 |
| A4W1 | 32.76 | 33.66 | 27.06 | 28.32 |
| A2W2 | 32.32 | 33.51 | 27.56 | 28.61 |
| A2W1 | 31.39 | 32.51 | 26 | 27.4 |
| A1W1 | 27.4 | 28.27 | 21.42 | 23.64 |

**Table 4: BLEU score comparison of A4W4 transformers. Columns are defined as in Table 2**
| Method | FP32 | A4W4 | GE | PT | Clip | LB |
| --- | --- | --- | --- | --- | --- | --- |
| LSQ+LUQ [] | 27.5 | 27.17 | Y | N | Y | 4 |
| Fixed-Point [] | 28.48 | 26.94 | Y | Y | Y | 4 |
| GradScale [] | 27.5 | 25.9 | Y | N | N | 4 |
| LUQ+SMP [] | 27.5 | 27.25 | N | Y | Y | 4 |
| Ours | 29.49 | 29.71 | N | N | N | 1 |

#### 5.1.3 Binary Transformers

**Table 5: BLEU Score of Transformers with binary activations on the WMT datasets. For the last column we record the drop from full precision models.**
|  | DE-EN | EN-DE |  |  |
| --- | --- | --- | --- | --- |
| Steps | 25k | 100k | 25k | 100k |
| A1W4 | 29.74 | 30.74 | 24.07 | 26.28 (-11.81%) |
| A1W2 | 28.81 | 29.81 | 23.4 | 25 (-16.11%) |
| A1W1 | 27.4 | 28.27 | 21.42 | 23.64 (-20.67%) |
| A1W1 [] | - | - | - | 17.87 (-32.18%) |

In biological neural networks, information is transmitted via electrical impulses called action potentials, or spikes. The complex processes governing spike transmission within the nervous system have been extensively investigated []. While binary signals can effectively represent these spike trains, the learning rule for spikes remains an open question [, ].

Inspired by the temporal nature of the transformer model, we reduced the activation precision to 1-bit for all linear layers, transforming it into a temporal binary network, akin to a quasi-spiking neural network. However, unlike traditional spiking neural networks, our model doesn’t simulate spike generation or consider spiking frequency.

To evaluate our approach, we assigned 1, 2, and 4 bits to the weights (Table 5). While introducing perturbations into spiking neural networks during training using backpropagation is not entirely new, our approach differs from previous attempts in several key aspects. Earlier studies have predominantly focused on introducing perturbations only at the spikes []. In contrast, our formulation introduces perturbations during signal quantization, irrespective of the spikes. This mirrors the inherent noisiness of the learning process and aligns with the biological reality of neural networks, where noise is an intrinsic part of neural signaling. Our results showcase that these converted binary transformers remain highly competitive with full-precision counterparts.

A recent study attempted to binarize transformers []. Their approach included extra normalization layers, clipping, and progressive quantization during training. We compared our method to their A1W1 configuration (Table 1, BMT-6 in their work, trained with 200k steps), achieving significant improvements (Table 5, last column).

### 5.2 Quantization and Sparsification

**Table 6: BLEU score of training Transformers with sparse weights for 25k steps on the WMT datasets. We present the baseline result, our sparsification result, and ours with both sparsification and quantization. Sparsification contributes to regularization, we mark improved results in bold (compared with Table 3).**
|  | Sparsity | Baseline | Ours | + A4W4 | + A4W1 |
| --- | --- | --- | --- | --- | --- |
| DE-EN | 25% | 33.45 | 33.64 | 33.63 | 32.55 |
| 50% | 33.38 | 33.73 | 33.94 | 31.32 |  |
| 75% | 32.08 | 33.4 | 33.37 | 30.28 |  |
| 90% | 29.5 | 31.94 | 32.22 | 29.31 |  |
| EN-DE | 25% | 29.25 | 30.03 | 29.92 | 27.58 |
| 50% | 28.54 | 29.07 | 29.45 | 26.32 |  |
| 75% | 27.25 | 28.98 | 28.73 | 25.05 |  |
| 90% | 20.6 | 27.02 | 26.7 | 23.24 |  |

In this section, we assess the effectiveness of our sparsification through perturbation proposal by comparing it to a baseline approach that employs a multiplicative mask, $\bm{w}_{sparse}=\bm{w}\odot\bm{1}_{|\bm{w}|>threshold}$, obtaining significant improvements. Additionally, we evaluate the performance of combining the proposed sparsity technique with quantization. Our findings indicate that moderate levels of sparsification perturbations introduce beneficial regularization effects during training, leading to improved BLEU scores. These findings are summarized in Table 6. Our supplementary section 8 further explores the integration of structured sparsity, demonstrating how its combination with binarization techniques yields sub-1 bit models that achieve competitive performance.

Quantization and sparsification are both valuable approaches for compressing neural networks, but both can also introduce performance degradation. Therefore, carefully balancing these techniques is crucial to achieve the desired trade-off between model size and accuracy. Sparsification is performed within each block (of size $128$) before quantization. This choice of order is made because quantized values have limited sparsity levels, making the reverse process less well-defined and potentially less effective. Striking a balance between compression and accuracy is paramount when applying quantization and sparsification techniques. Excessive application can lead to substantial performance degradation. Our experiments demonstrate that low-level sparsification perturbations are beneficial. However, as quantization levels increase, the tolerable level of sparsity decreases (Table 6).

#### 5.2.1 Block Size and λ 𝜆 \lambda italic_λ

In addition to precision and sparsity, block size also presents a trade-off between accuracy and efficiency. We observe that its influence is more pronounced at extremely low precision levels, such as 1-bit. While using a smaller block size can improve performance, the effective bits per element can easily exceed the original design. We provide some comparisons here. Considering this trade-off, lower precision models may not always be more efficient (Table 7). Firstly, the accuracy achieved with smaller blocks in an A1W1 models may not surpass the accuracy achieved using A2W2. The perturbations introduced during lower precision training often remain high, hindering the achievement of high quality. Optimal selection needs to be made based on the underlying hardware support and the problem of interest. Our work facilitates a comprehensive study of this trade-off.

**Table 7: BLEU Score comparison of adjusting the block size when training the A1W1 Transformers for 100k steps on the WMT DE-EN Dataset.**
| Block | 32 | 128 | 512 |
| --- | --- | --- | --- |
| A1W1 | 29.71 | 28.27 | 27.14 |

The parameter $\lambda$ suppresses the impact of perturbations and prevents training explosion, especially during the transition to 1-bit quantization. Standard quantization approaches often omit the use of $\lambda$ in an attempt to minimize quantization error. However, this can lead to training divergence in the early stages, as demonstrated in Fig. 1(a). We compared with the classic inverse scaling-based quantization implementation in the AQT library [] to observe this effect. In our experiments, we observe that a wide range of $\lambda$ values yield satisfactory results (Table 8). However, our preference set on the safer side, and use $\lambda=0.01$ for all settings. For higher precision ($\geq 4$-bits), smaller $\lambda$ values, such as 0.0001, can be safely used to allow more signals to pass through without causing training instability. Extremely small $\lambda$ can cause numerical instability, particularly exacerbated by the inherent numerical noise introduced by neural network operations.

**Table 8: BLEU Score comparison of adjusting the $\lambda$ when training the A1W1 Transformers for 25k steps on the WMT EN-DE Dataset.**
| $\lambda$ | 1.0 | 0.01 | 0.0001 | 0 |
| --- | --- | --- | --- | --- |
| A1W1 | 5.83 | 21.42 | 20.08 | NaN |

#### 5.2.2 Limitations

As demonstrated in Section 5.2.1(Table 7), quantization algorithms generally achieve the highest accuracy with small sub-channel quantization blocks. However, the data reshaping required for these blocks can create a performance bottleneck, especially in low-precision settings where integer matrix multiplications are relatively fast. Furthermore, the ridge regression computations add additional overhead. To achieve an efficient implementation in practice, careful low-level optimization and consideration of the underlying hardware capabilities will be crucial.

## 6 Summary

Discontinuous operations such as quantization and sparsification pose a significant challenge to backpropagation training, as they introduce non-differentiability into the learning process. This non-differentiability has long been considered the Achilles heel of backpropagation, hindering the development of efficient and accurate neural network training algorithms. We address this challenge by re-framing these discontinuities as perturbations. This allows us to introduce a novel, continuous control mechanism that seamlessly handles non-differentiability during training. Our approach is expected to significantly facilitate the deployment of large-scale neural network models on resource-constrained devices such as mobile phones, enabling the widespread adoption of deep learning for mobile applications. Additionally, our technique holds promise for the development of biophysically plausible neural networks, which have the potential to revolutionize artificial intelligence and machine learning.

## 7 Acknowledgements

We extend our sincere gratitude to Abhijit Ogale, Jeff Dean, Jian Li, Kyuyeuan Kim, Mark Sandler, Rasmus Larsen, Sameer Agarwal, Sanjiv Kumar, Tammo Splank for their insights and supports.

## 8 Ternary Weights via Structured Sparsity

Our sparsification method flexibly handles M:N structured sparsity, enabling a class of ultra-low precision models (even below 1-bit) by combining quantization and sparsity.

First, we introduce perturbations to enforce an M:N sparsity constraint (for simplicity, we set $N=4$ and $M\in{1,2,3}$).
Non-zero values are then quantized to ${-1,1}$ by taking their sign. This structured sparsity effectively introduces an extra bin of 0, resulting in ternary weight encoding.

Since model weights typically have a close to 0 0 mean, during weights reconstruction, we reformulating the ridge regression without the bias term $b$:

$$ $\min_{a}\frac{1}{2N}||a\cdot\bm{q}-\bm{x}||^{2}+\frac{\lambda}{2}a^{2}$ (10) $$

And for each quantization block of $128$, we solving for the scaling factor only:

$$ $a=\frac{\overline{\bm{qx}}}{\overline{\bm{q^{2}}}+\lambda}$ (11) $$

We applied this method to a transformer model trained on the WMT DE-EN dataset with 4-bit activations and ternary weights. Our results show that this 2:4 ternary implementation achieves comparable performance (Table 9) with dense binary weights.

Regarding storage efficiency, encoding 1:4 structured sparsity necessitates only 2 bits for every group of 4 elements (due to one non-zero position within each group). 2:4 sparsity demands 4 bits per 4-element block, resulting in storage requirements on par with dense binary weights. Importantly, the 1:4 structured sparsity design yields a remarkably efficient 0.5-bit per element model that maintains competitive performance (as shown in Table 9). Moreover, the additional zero-bin inherent to structured sparsity contributes to both diminished storage needs and enhanced outcomes for this class of models (compared with Table 6).

**Table 9: BLEU Score comparison of Transformer on WMT DE-EN using A4W1 with/without structured sparsity. Ignoring the storage overhead of the reconstruction coefficients, we report the average number of bits required to store each weight parameter in the final column.**
| Sparsity | Precision | BLEU | Bits |
| --- | --- | --- | --- |
| Dense | Binary | 32.76 | 1 |
| 25% | Binary | 32.55 | 1 |
| 1:4 | Ternary | 32.1 | 0.5 |
| 2:4 | Ternary | 32.66 | 1 |
| 3:4 | Ternary | 33.1 | 1.5 |

## 9 Lower Precision Floats

In practical implementations, our scaling factor and bias demonstrate resilience to lower precision representations. Empirical evidence confirms that utilizing float8 (E5M2) precision does not adversely affect the accuracy of our results. This robustness can be interpreted as introducing minor perturbations to the scaling and bias values.

Leveraging this property, we can further reduce the subchannel block size to 32 elements while maintaining a storage overhead of less than 1 bit per element. This approach strikes a compelling balance between memory efficiency and model performance.

In addition to quantizing to integer vectors, our proposed method also supports quantization to low-precision floats. In this scenario, the quantization vector $\bm{q}$ can take on values in FP4 or FP8 formats.

## 10 Connection with the Straight-Through Estimator

Unlike traditional straight-through estimators (STEs) that rely on defining backward gradients [, ], our novel approach directly incorporates discontinuous operations into the forward pass. This is effectively a forward implementation of the STE. Our denoising reconstruction process explicitly mitigates the disruptive effects of these discontinuities.

We experimented with replacing the additive noise term (Eq. 2) with a controlled multiplicative activation function:
$\delta(f(\bm{x}))=\bm{s}\cdot f(\bm{x})$ where $\bm{s}=stop_gradient(\frac{\bm{q}}{f(\bm{x})})$ represents element-wise precomputed scaling values. This approach aimed to eliminate explicit gradient scaling during backpropagation while incorporating quantization directly into the activation function. However, our experiments did not reveal any significant performance gains resulting from this modification.

## 11 Implementation through Integer Matrix Multiplication

Consider quantizing the matrix multiplication: $Y_{n\times o}=X_{n\times i}\cdot W_{i\times o}$. We use $\cdot$ for the regular dot product, and $\odot$ for the Hadamard (element-wise) product. Quantized matrix multiplication computes:

$$ $\displaystyle Y^{Q}_{n\times o}$ $\displaystyle=[(a^{X}_{n\times 1}\cdot 1^{X}_{1\times i})\odot Q^{X}_{n\times i }+b^{X}_{n\times 1}\cdot 1^{X}_{1\times i}]$ (12) $\displaystyle\cdot[Q^{W}_{i\times o}\odot(1^{W}_{i\times 1}\cdot a^{W}_{1 \times o})+1^{W}_{i\times 1}\cdot b^{W}_{1\times o}]$ $\displaystyle=(a^{X}_{n\times 1}\cdot a^{W}_{1\times o})\odot(Q^{X}_{n\times i }\cdot Q^{W}_{i\times o})$ $\displaystyle+b^{X}_{n\times 1}\cdot[(1^{X}_{1\times i}\cdot Q^{W}_{i\times o} )\odot a^{W}_{1\times o}]$ $\displaystyle+[a^{X}_{n\times 1}\odot(Q^{X}_{n\times i}\cdot 1^{W}_{i\times 1} )]\cdot b^{W}_{1\times o}$ $\displaystyle+b^{X}_{n\times 1}\cdot(1^{X}_{1\times i}\cdot 1^{W}_{i\times 1}) \cdot b^{W}_{1\times o}$ $$

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: See weaknesses above.
soundness: 2
strengths: - A unified framework that applies to both quantization and sparsification.
- Denoising affine transform for training stability is novel. By framing quantization as a perturbation and introducing a controlled denoising mechanism, the authors provide a theoretically grounded way to mitigate the instability associated with ultra-low precision training.
- A solid mathematical foundation, especially in using ridge regression to stabilize training; analysis of the parameter sensitivity; shortcut for quantized matrix multiplication and the sub-channel quantization.
- Compatibility with standard architectures and training methods, which makes the approach practical.
- Competitive performance on ResNet-50 and Transformer, better results than the baselines (including the full precision training).
summary: This paper proposes a universal framework for training neural networks at arbitrary precision and sparsity by treating quantization and sparsification as perturbations and applying a denoising affine transform to stabilize training. The method is compatible with standard architectures and recipes, achieving competitive results on ResNet-50 and Transformer models at low precision. It provides theoretical grounding through ridge regression and introduces efficient matrix multiplication for quantized models.
weaknesses: - The proposed denoising technique is applied to both quantization and sparsification, but sparsity deserves special attention. Sparsity has unique challenges, such as gradient instability and poor convergence behavior in sparse regimes. By framing sparsification as a perturbation, the method simplifies the impact of sparsity on training. Sparse models typically exhibit unique instability issues, such as dead weights that don’t recover during training. The paper does not address how the proposed denoising affine transform could prevent or manage this issue. More dedicated experiments and ablations would be helpful to understand the unique contribution of the proposed solution to sparse training. The BLEU scores provided in Tab 6 are insufficient to validate the claim that the proposed approach robustly supports sparse training across different tasks and architectures. The method’s utility in extreme sparsity settings, such as >90%, is not tested.
- Certain layers, like the first and last layers or specific convolutional layers, may require lower sparsity levels / higher precision than intermediate layers due to their critical role in feature extraction and final prediction. The paper does not discuss how the proposed approach could be adapted to handle these layer-specific sparsity requirements, which are essential for maintaining performance.
- How does this approach work in the fine-tuning settings, where small shifts in weights might accumulate errors?
- The authors demonstrate results on ResNet-50 and Transformer models, without testing across a broader range of architectures, such as MobileNet and EfficientNet, which are widely used in low-power / mobile applications. Also consider providing performance results on architectures that have very different structural and computational properties, eg RNNs.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. What is the 'Flax framework'?
soundness: 2
strengths: The paper proposes a novel affine transform to perform dequantization step of quantizer by solving a regularized least squares problem with respect to the scaling factor.
summary: The paper introduce an affine transform scheme for the quantization and sparsification of neural networks.
weaknesses: 1. The authors appear to overstate their contributions. The primary challenge in training quantized neural networks is addressing the zero-gradient issue associated with discrete loss functions. However, if I understand correctly, the main contribution here is a de-quantization scheme used to compute the scaling factor for the quantizer, which does not directly address the zero gradient issue as what straight-through estimator does.

2. In the Abstract, it is stated that "Our solution employs a piecewise constant backbone model to ensure a performance lower bound and features an inherent noise reduction mechanism to mitigate perturbation-induced corruption." Could you clarify what the "piecewise constant backbone model" and the "performance lower bound" refer to specifically? I was unable to find details on these terms in the main text.

3. It is unclear whether the authors used quantization-aware training or post-training quantization in the experiments. Could you please clarify?

4. All experiments are conducted for group-wise quantization. To better demonstrate the advantages of the proposed method, I recommend including results for channel-wise quantization as well as testing on a broader range of neural architectures, such as vision transformers and MobileNet.

5. It would be beneficial to provide more technical details on how this method is extend to sparsification.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. In Table 8, why there is a NaN when $\lambda=0$?
2. It is surprising that your A4W4 results are better than A32W32 results in tables 1/2/3. (Low-precision training is better than the conventional float32-precision training). Can you please make an explanation? Usually, low-precision computation reduces the accuracy during the inference.
3. How is your method's actual running memory consumption and run-time compared to other methods?
soundness: 3
strengths: 1. The approach achieves state-of-the-art results with models trained at very low precisions, demonstrating robustness and effectiveness.
2. Demonstrated effectiveness on both ResNet and Transformer architectures across multiple datasets, suggesting that the method is adaptable to a variety of deep learning tasks.
summary: This paper proposes a denoising affine transformation that stabilizes training under these conditions. It introduces perturbation modeling and noise reduction strategies using a piecewise constant backbone model, which ensures performance stability across varying precisions and sparsities.
weaknesses: 1. In Eq (1), $\epsilon$ is not defined. How do you choose it?
2. In line 162, why do you use $\delta_i$ instead of $\delta$? Is it a typo?
3. Run-time analysis is missing. This is important for low-precision networks.
4. Actual running memory consumption is missing. This is important for low-precision networks.
5. The added denoising transformation and affine operations might introduce computational overhead, especially in cases where ridge regression is used for reconstruction, possibly slowing down training on certain hardware.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: **Ambiguous and Inaccurate Mathematical Expressions**  
The readability is low, and I do not believe it is at a level suitable for acceptance.

Line 161:  
Although I understand the actual algorithm, I believe the mathematical formulation of the perturbation is inappropriate. Specifically, the author defines $\delta$ by $\delta = \text{round}(f(x)) - f(x)$, and substituting this into eq.(2) makes $q = \text{round}(f(x))$. So, as a mathematical operation, adding $\delta$ is meaningless for $q$. In other words, rather than adding a perturbation, it simply appears to be stating $q = \text{round}(f(x))$.

eqs.(3-8):  
Clarify the dimensionality of each variable, such as $x \in \mathcal{R}^N$. It is strange that $x$ is in bold font while $b$ is in regular font, implying it is a scalar, yet the authors use $\mathbf{x} - b$, subtracting a scalar from a vector. Additionally, there is no clear definition of $\text{Cov}_{xq}$. Is this a scalar?


**Memory efficiencty**  

Compared to existing methods, is the proposed method more memory efficient? Since it introduces continuous variables $a$ and $b$, it appears that training cannot be performed on low-precision hardware alone.
soundness: 2
strengths: - The proposed method is quite simple and implementation-friendly. It can potentially be combined with various other architectures and  algorithms.

- It empirically works even for high quantization cases such as A1W1.
summary: This work proposes a denoising affine transform method for quantization and sparsification of deep learning. This is based on ridge regression between the original vector and its quantization, which reconstructs the original signal. The proposed method enables us to train deep nets with ultra-low precision and high sparsity levels. In particular, it provides training on temporal binary networks, which has been considered challenging in previous works.
weaknesses: **Lack of Comparative Experimental Validation**  

The paper demonstrates the behavior of the proposed method through various experiments, but it barely shows how much better it is compared to existing methods. Starting from Table 2, the experiments only present results in terms of accuracy, with very little insight into the stability of the learning process, leaving only the final accuracy for comparison. In particular:  
- Tables 2 and 4 seem to compare results, but only for the A4W4 case, where the precision is not extremely low. Moreover, the difference in accuracy with existing methods seems negligible although there is no standard deviation provided and its significance is unclear.  As long as  the metric of performance is limited to the test accuracy, I must say that the contribution is minimal.  
- The comparison for ultra-low precision (A1W1) is limited to a single cell in Table 5.

Thus, while the method's functionality is evident, it is unclear in what aspects it excels.

As I am not an expert in quantization, I cannot make a definitive judgment, but it seems that there is little prior research on precision lower than A4W4. If this is true, then the main claim of the paper should not be the proposal of a method, but rather the experimental validation at precision levels lower than A4W4, which has rarely been tested before. The paper claims state-of-the-art performance in line 94, but if there are no previous studies, it would be more appropriate to describe the setup as novel.  

Furthermore, without a comprehensive evaluation of factors beyond accuracy, such as the stability against hyperparameters (e.g., learning rate), changes in training time, and memory efficiency, it seems difficult to demonstrate the advantages of the proposed method.

Meta-review: This paper introduces a novel perspective by treating quantization and sparsification as noise and proposing affine transformations to "denoise" and stabilize training. The reviewers acknowledged the originality of this idea and its subsequent implementation. However, they also raised substantial concerns regarding ambiguous expressions, unclear statements, and unconvincing experimental validation.

Given these significant issues, the overall ratings are consistently negative, and I recommend rejection for ICLR2025. I hope the feedback provided will assist the authors in refining their paper and further developing this promising method to ensure its potential can be fully realized.

Decision: Reject