SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our results for DeiT-Tiny on CIFAR10 and ImageNet can be seen in Table 2. The PAM based matrix multiplication training matches the baseline accuracy for both datasets. We compare to Adder Attention [30] which also replaces multiplications by cheaper addition based operations. They do not replace the batched matrix multiplication and use standard matrix multiplication for the first and last transformer layer on ImageNet and also use multiplications to smooth the gradients during the backwards pass. Adder Attention transformers can roughly match their reported baseline accuracy for CIFAR10, but do suffer from a sizable degradation on ImageNet.
For IWSLT14 we obtain a baseline test BLEU score of 34.37. Replacing all matrix multiplications on the forward and backwards passes with PAM based matrix multiplications gives a test BLEU score of 34.22. Both scores are averaged over three runs. The observed difference of 0.15 is slightly larger than the run-to-run variance we observed, but fairly small overall. We also note that we use standard hyperparameters originally obtained from FairSeq and tuning the hyperparameters may improve performance.
In Appendix C we show similar experiments for several non-transformer models including convolutional neural and mixer architectures on CIFAR-10. We find little to no degradation across all networks tested.
3.3 Replacing Other Operations
In Table 3 we measure the impact of using piecewise affine versions of different network components. We isolate the effect of each component and experiment with the two different derivative types described in Section 2. MATMUL refers to the linear layers and the batched matrix multiplications that we replaced in Section 3.2. Using exact derivatives for the matrix multiplications performs notably worse than the approximate ones. The same holds for piecewise linear layer normalization, with approximate derivatives we observe little to no degradation but using the exact ones cause a notable one. ATTENTION SOFTMAX refers to the softmax operation performed in the attention computation and excludes the softmax for the loss. Here we find that the using the exact derivatives cause a training instability when using the default hyperparameters whereas the with the approximate backwards pass the performance impact is negligible. The Loss function, softmax cross-entropy with label smoothing, is the only operation where we find the approximate BWD pass to perform
Table 3: BLEU Test Scores on IWSLT14 DE-EN translation for modified Transformer-Small training where certain operations are replaced with piecewise affine approximations using either the exact or approximate backward pass. Cumulative experiments replace all previously listed operations using the better performing derivative version for each one. The brackets denote the change from the baseline BLEU score of  $34.4 \pm 0.1$ . Each reported score is the average±std over three runs with different seeds.
PA OPERATION(S)	EXACT BWD	MIMIC BWD	CUMULATIVE
MATMUL	33.4[-1.0] ± 0.1	34.2[-0.2] ± 0.2	34.2[-0.2] ± 0.2
ATTENTION SOFTMAX	3.0[-31.4] ± 0.9	34.3[-0.1] ± 0.7	34.2[-0.2] ± 0.2
LAYER NORM	33.7[-0.7] ± 0.2	34.4[-0.0] ± 0.1	34.4[-0.0] ± 0.1
LOSS	33.7[-0.7] ± 0.2	33.0[-1.4] ± 0.9	33.7[-0.7] ± 0.2
OPTIMIZER	34.2[-0.2] ± 0.1	33.5[-0.9] ± 0.3
worse than the exact one. It is also the only network component where both types of derivatives have a notable performance loss.
Finally we experiment with using piecewise affine operations for the optimizer update. The AdamW optimizer uses multiplication, division and square root. Replacing all of them only causes a minor degradation of 0.2 BLEU compared to the standard version.
3.4 Fully Multiplication-Free Training
Here we combine all the piecewise affine replacements from the previous section one by one. The resulting BLEU scores are displayed in the last column of Table 3. We chose the derivative type that performed better for each operation which is the approximate one for everything except the loss. Each transformer block contains a gain that we replace with PAM when we replace the attention softmax.
Overall there is little to no performance impact from replacing the matrix multiplications, softmax and layer normalization. Replacing all of them together does not seem to cause any unforeseen issues. Consistent with our previous observations, replacing the loss with a piecewise affine function causes a small BLUE score degradation.
The last entry in Table 3 covers every multiplication and non-linearity in the training process. This includes the forward pass, the backwards pass and the parameter updates by the optimizer. The resulting performance impact is only 0.9 BLEU points and hyperparameter tuning could potentially reduce it even further.
As mentioned in the introduction, most related work falls into the two classes of approaches, depending if the original neural network architecture is preserved or not.
For the first type, fine-grained changes of operations are made to individual multiplications or groups thereof, aiming to closely approximate the original multiplications. Examples for this are mixed precision training [23], and more generally block-floating point schemes, where several weights are grouped together to share the same floating point exponent, enabling inference and training [18, 8]. Another class of related approaches is quantization, which is widely used for inference but introduces additional challenges for training [20]. All mentioned approaches have in common that some multiplications are still needed, thus might not be viable on multiplication-free hardware.
Approaches of the second type allow more drastic changes to the architecture, thus offering more promise to new types of multiplication-free operations. A prominent example are AdderNets [4, 30] or the related EF-operator [38] which rely on absolute differences instead of inner products, and show these can serve as replacement layers for convolutions or also parts of self-attention layers [30], though training can not be made fully multiplication-free while maintaining performance. Another type of fundamentally different architectures can arise by replacing classical arithmetic with tropical algebra, where all multiplications are replaced by additions and additions are replaced by maximum. While numerically very efficient, performance was not shown to be competitive with standard neural networks [22, 9].
To the best of our knowledge, the approach of Mogami [25]—which serves as the basis of our work here—is currently the only work which allows to preserve the neural network architecture while at the same time paving the way for training on fully multiplication-free hardware, if applied to all elements of the training process.
The main goal of this work is to demonstrate the applicability of piecewise affine function approximations to end-to-end transformer training. In our experiments we showed that all operations during training can be replaced by such functions. We believe this is the first fully multiplication-free training of a transformer model and potentially any neural network. While PAM approximates real multiplication to a certain extent, it does not aim to do so as closely as possible, unlike floating point multiplication in finite precision. Instead PAM uses line segments with slopes that are powers of 2. By doing so it becomes significantly simpler, providing the same hardware benefit as replacing all multiplications with e.g. addition, like some other approaches to multiplication-free networks aim to
do. At the same time PAM does roughly approximate standard multiplications overall, allowing it to function as a plug-in replacement across various usecases. This includes the optimizer update and the weighted average in attention, which would be hard to replace with something like addition. PAM therefore seems to strike a good balance between cost and utility.
Our main inspiration is the work of Mogami [25] which applied piecewise affine approximations to convolutional networks, including convolutional layers, linear layers, batch normalization and exponents. Depending on the implementation, this may have been sufficient to make the forward and backwards pass fully multiplication-free but it is unfortunately unclear from the paper and no code is publicly available. We focus on transformers making both the forward and backwards passes multiplication-free and fully piecewise affine. We also extend this to the optimizer eliminating all multiplications in the training process, allowing it to run on entirely on hardware without multiplication support.
Although piecewise affine operations are logically simpler than their standard equivalents, special hardware is required to fully reap the theoretical advantages. Currently available hardware accelerators, such as GPUs and TPUs, have arithmetic units that are specialized for floating point multiplication but lack the comparatively simpler logic circuits that would enable fast piecewise affine multiplication. In practice this means that our present GPU implementation has to perform multiple instructions for every standard multiply-and-accumulate operation causing it to run significantly slower, even with custom kernels. Field-Programmable Gate Arrays (FPGAs) are one type of hardware that could immediately benefit from PAM. Since their logic circuits are fully configurable, they don't have the "sunken cost" of standard multiplication arithmetic and can instead spend it on PAM circuits.
As described in Section 2.4, replacing all multiplications and non-linearities in a neural network with piecewise affine operations results in a network that is fully jointly piecewise affine in both inputs and weights. This is an interesting theoretical property since it differs significantly from standard transformers.
Surprisingly, replacing all operations in the network itself doesn't seem to have a large impact on the performance in our setup. This differs from the observation by Mogami 2020 which found that linear layers did not work well with approximate derivatives. We observe a similar phenomenon for the loss function which could indicate that training is sensitive to slight errors early in backpropagation.
This project focused on standard neural networks without any modifications. However it would be interesting to explore modifying the network architecture to be even better suited to piecewise affine approximations. One potential way of doing this is to use  $\log_2$  and  $\exp_2$  throughout, for instance in the softmax and loss computation. These operations are much easier to approximate than their base- $e$  equivalents. Exploring the effects of error both from quantization as the mantissa is decreased and from magnitude bias is another future research direction.
In this work we explored the use of hardware-efficient piecewise affine operations for transformer training. We experiment on CIFAR10 [19] and ImageNet-1k [6] with DeiTtiny [32] as well as on IWSLT14 German to English translation [3] using a small transformer network. We show that floating point multiplications within matrix multiplications can be replaced by cost-effective piecewise affine multiplications with minor or no performance degradation for all three datasets, and without additional hyperparameter tuning.
We further show that we can extend this to replace all multiplications, division and non-linear functions in the small transformer, the loss function and the optimizer with piecewise affine functions. The resulting performance impact is minor and could potentially be reduced further by hyperparameter tuning. This eliminates all multiplications in the entire training process, for fully multiplication-free training which we believe has not been demonstrated before. This also makes the network fully and jointly piecewise affine in both inputs and weights, a property that could be of theoretical interest.
We publicly release our code, including the kernels that made our experiments computationally viable. This could help further research in this area and into other multiplication-free methods.
Figure 3: Piecewise affine multiplication, division, square and square root along with their standard equivalents. Each row corresponds to one function. Left: A function (solid line) with its piecewise affine alternative (dashed line). Left middle: The relative error  $\frac{\hat{f}(x) - f(x)}{|f(x)|}$  when the piecewise affine function  $\hat{f}$  is viewed as an approximation of the standard function  $f$ . Right middle: The derivatives of  $f$  (solid black), the exact derivative for  $\hat{f}$  (dotted blue) and the approximate derivative of  $\hat{f}$  (dashed red). We assume  $\delta_Y = 1.25$  in the equations in Section 2.5 which emphasizes the PAM operations used, compared to  $\delta_Y = 1$  where PAM matches standard multiplication. Right: The relative error  $\frac{\hat{f}'(x) - f'(x)}{|f'(x)|}$  for the exact derivative  $\hat{f}'$  (dotted blue) and the approximate derivative  $\hat{f}'$  (dashed red) when viewed as an approximation of  $f'(x)$ , the derivative of  $f(x)$ .
Figure 4: The equivalent of Figure 3 for exponentials and logarithms, showing piecewise affine functions with their standard counterparts. Piecewise affine exponentials and logarithms along with their standard equivalents. Each row corresponds to one function. Left: A function (solid line) with its piecewise affine alternative (dashed line). Left middle: The relative error  $\frac{\hat{f}(x) - f(x)}{|f(x)|}$  when the piecewise affine function  $\hat{f}$  is viewed as an approximation of the standard function  $f$ . Right middle: The derivatives of  $f$  (solid black), the exact derivative for  $\hat{f}$  (dotted blue) and the approximate derivative of  $\hat{f}$  (dashed red). We assume  $\delta_Y = 1.25$  in the equations in Section 2.5 which emphasizes the PAM operations used, compared to  $\delta_Y = 1$  where PAM matches standard multiplication. Right: The relative error  $\frac{\hat{f}'(x) - f'(x)}{|f'(x)|}$  for the exact derivative  $\hat{f}'$  (dotted blue) and the approximate derivative  $\hat{f}'$  (dashed red) when viewed as an approximation of  $f'(x)$ , the derivative of  $f(x)$ .
Table 4: Hardware costs of arithmetic operations in different precision from [12, 10].
FORMAT	ADDITION	MULTIPLICATION
ENERGY [pJ]	AREA [μm2]	ENERGY [pJ]	AREA [μm2]
INT32	0.1	137	3.1	3495
INT16	0.05	67
INT8	0.03	36	0.2	282
FLOAT32	0.9	4184	3.7	7700
FLOAT16	0.4	1360	1.1	1640
B Hardware Cost of Piecewise Affine Multiplication
It is important to note that our objective in this work is to investigate the machine learning characteristics of training with piecewise affine operations, rather than designing a hardware implementation. We demonstrate that such training is viable, with final model performance and convergence speed roughly identical to standard training. Without such algorithmic evidence to motivate their development, the development of efficient hardware implementations is unlikely. Although the details of the integrated circuit design should not affect e.g. convergence, they can drastically affect the cost of the operation in terms of energy or area.
Such hardware considerations include which aspects of standard floating point numbers are supported, e.g. denormal numbers, different types of NaNs and so on, as well as whether to use a custom encoding of the exponent for instance. An actual hardware implementation would also have to consider factors such as latency, as well as memory characteristics like bandwidth and read costs. The design could either focus on individual operands or vectors like tensor cores and systolic arrays. Operating directly on vectors or matrices can allow significant amortization of some of the costs involved.
Ignoring most of these considerations we can arrive at a very rough estimate of the costs. We assume that the cost of multiply-and-accumulate operations is the sum of the costs provided in Table 4 for multiplication and addition. A PAM operation can be performed with one full int32 addition and one int8 addition for the exponent which also needs to check for under / overflow. When implemented as a logic circuit, we estimate the cost of this could be comparable to two int32 additions. Using these numbers float32 PAM uses roughly  $\frac{2\cdot0.1}{3.7} = 5.4\%$  of the energy and  $\frac{2\cdot137}{7700} = 3.6\%$  of the area compared to float32 multiplication and  $\frac{2\cdot0.1}{1.1} = 18\%$  of the energy and  $\frac{2\cdot137}{1640} = 17\%$  of the area compared to float16 multiplication. PAM is likely to work with smaller numbers such as float16 which would reduce the cost further by a factor of around  $2\times$ . In Appendix D we find that mantissas as narrow as 4 bits may work well, significantly less than the 10 and 7 bits used in float16 and bfloat16.
PAM reduces the cost of multiplications by roughly an order of magnitude but does not affect accumulation. The average savings for multiply-and-accumulate (MAC) operations will therefore be lower. For float32-multiply float32-accumulate PAM could have a cost of roughly  $\frac{2\cdot 0.1 + 0.9}{3.7 + 0.9} = 24\%$  of the energy and  $\frac{2\cdot 137 + 4184}{7700 + 4184} = 38\%$  of the area. For standard mixed precision float16-multiply float32-accumulate the costs are  $\frac{2\cdot 0.1 + 0.9}{1.1 + 0.9} = 55\%$  of the energy and  $\frac{2\cdot 137 + 4184}{1640 + 4184} = 77\%$  of the area. As we have eliminated the majority of the cost associated with the multiplication, the accumulation accounts for the majority of the remaining arithmetic costs of dot products. These costs can likely be reduced with methods such as chunked accumulation in lower precision [29, 35], combining multiple operations as done in tensor cores, or the Kulisch accumulation method used in [13]. We have not explored the use of these methods here but conjecture they are likely complementary to our approach, just as they have been for standard multiplication.
C Additional Architectures
The main focus of this work is on transformer networks. However we have found that piecewise affine matrix multiplications work well across a range of convolutional architectures. Table 5 shows the results of training five different non-transformer networks on CIFAR-10 using piecewise affine matrix multiplications with the approximate bwd pass, replacing all linear and convolutional layers.
Table 5: Top-1 Test Accuracy on CIFAR-10 for different networks when trained with standard matrix multiplications vs piecewise affine matrix multiplications (using approximate bwd). Each value is the average±std for three runs with different seeds. Brackets show the difference from the baseline.
NETWORK	BASELINE	PA-MATMUL
VGG-13	92.9 ± 0.3%	92.9[+0.0] ± 0.2%
RESNet-20	92.1 ± 0.3%	92.0[-0.1] ± 0.2%
RESNet-110	94.2 ± 0.1%	94.1[-0.1] ± 0.2%
RESNext-20 4x16	93.7 ± 0.2%	93.8[+0.1] ± 0.2%
CONVMixer-256/8	94.8 ± 0.2%	94.8[-0.0] ± 0.1%
In each case we roughly match the baseline (within the run-to-run variance), without any adjustments to the hyperparameters.
We selected the five models to demonstrate variety. Here are some of the ways in which they differ:
- VGG-13 [31]: Plain CNN without skip connection or normalization layers. Contains several fully connected layers at the end.
- ResNet-20 [11]: Model with skip connections and normalization.
- ResNet-110 [11]: Model of significant depth.
- ResNeXt-20 4x16 [37]: Uses bottleneck blocks and grouped convolutions.
- ConvMixer-256/8 [33]: Has a mixer style architecture and uses GeLU activations. The training setup also uses heavier data augmentation than for the other networks.
The exact details of the training setup used in each case can be found in our code.
D PAM with Narrower Mantissas
In this section we explore whether PAM can work with narrower mantissas which would decrease memory read costs and improve effective memory bandwidth. We simulate training with piecewise affine matrix multiplications (approximate bwd) using numerical formats with fewer mantissa bits than the 23 used in float32. We implement this by rounding the inputs and masking the extra mantissa bits, but do not change any other aspects of the training setup (i.e. no tuning or special low precision tricks). The internal accumulation is left unchanged similar to standard fp16-fp32 mixed precision. We focus on the training for two tasks, the IWSLT transformer from Section 3.2 and the VGG-13 CIFAR-10 network from Appendix C. The results can be seen in Table 6 below.
For both networks, we observe minimal to no discrepancy when training with 7 bits (equivalent to bfloat16). Mantissas as narrow as 4 bits work well, providing a comfortable margin for bfloat16 and offering promise for extensions into very narrow formats like 8-bit formats (although these might necessitate additional techniques to accommodate the narrow exponent). However, a 3-bit mantissa noticeably impairs the transformer's performance and may marginally impact the VGG training.
Table 6: The impact of the mantissa size of the PAM inputs on VGG-13 CIFAR-10 Test Accuracy and IWSLT14 BLEU Score for a small transformer. Each value is the average±std for three runs with different seeds.
MATMUL TYPE	CIFAR-10 TEST ACCURACY	IWSLT14 BLEU SCORE
FLOAT32	92.9 ± 0.3%	34.4 ± 0.1
PAM FLOAT32	92.9 ± 0.2%	34.2 ± 0.2
PAM BFLOAT	92.9 ± 0.2%	34.4 ± 0.2
PAM 4 BIT MANTISSA	92.9 ± 0.5%	34.2 ± 0.1
PAM 3 BIT MANTISSA	92.8 ± 0.2%	29.4 ± 0.5
E Runtime
Since PAM is not natively supported by the GPUs we use for our experiments, we require multiple instructions for each operation and can not make use of special hardware features such as the tensor
cores. We can also not use the highly optimized cuDNN and cuBLAS libraries at all, having to write the kernels required ourselves. Overall this can easily slow training down significantly compared to well optimized baselines that fully benefit from the GPU hardware including tensor cores. Here we want to emphasize that this is not a limitation of the method or our implementation but rather the lack of hardware support for the operation. We would incur a similar slowdown if we wanted to simulate training in a new custom floating point format, like some variant of float16, using only int and float32 operations. A significant effort went into enabling our experiments to run at the scale we report.
IWSLT14 DE-EN: The baseline runtime is around 1 hour on a V100 GPU. With PAM approximate matrix multiplications this number is around 4.5 hours. Replacing all operations results in a runtime of around 5.5 hours.
DeiT i1k and CIFAR-10: The CIFAR-10 baseline runs in about 4 hours on 4xV100 GPUs. The PAM approximate matrix multiplication version takes around 15 hours. The ImageNet version took 34 hours for the baseline vs 131 hours for PAM matrix multiplications on 6xA100 GPUs.
Convolutional Networks on CIFAR-10: The runtimes (baseline vs PAM matmuls) are roughly:
VGG-13: 20 min vs 4.5 hours
- ResNet-20: 20 min vs 7 hours
- ResNet-110: 45 min vs 10 hours
- ResNeXt-20 4x16: 25 min vs 10 hours
- ConvMixer-256/8: 20 min vs 5 hours
We note that we did not focus on speeding up convolutions, electing to simply perform them as matrix multiplications using relatively inefficient folding operations.
The scale of typical deep learning problem makes it hard to investigate new low-level hardware algorithms in this field. The existing hardware accelerators that enable common large workloads are also highly specialized and do not necessarily perform well when performing slightly different operations. In this study the scale of the problems we can study is somewhat limited due to the high cost of training. We have tried to address this by writing custom kernels that allow us to run larger problems at a reduced cost.
Ideally, we could show direct hardware benefit from the piecewise-affine operations. However, we do not have the expertise to design an efficient hardware implementation that fully reaps the hardware advantages. Without significant investments, such an implementation would also be unlikely to allow training at the scale we demonstrate here. Instead, we have tried to estimate the potential hardware benefits and view this work as algorithmic evidence to inspire further work in this area.