FleXOR: Trainable Fractional Quantization

Parameter quantization is a popular model compression technique due to its regular form and high compression ratio . In particular , quantization based on binary codes is gaining attention because each quantized bit can be directly utilized for computations without dequantization using look-up tables . Previous attempts , however , only allow for integer numbers of quantization bits , which ends up restricting the search space for compression ratio and accuracy . Moreover , quantization bits are usually obtained by minimizing quantization loss in a local manner that does not directly correspond to minimizing the loss function . In this paper , we propose an encryption algorithm/architecture to compress quantized weights in order to achieve fractional numbers of bits per weight and new compression configurations further optimize accuracy/compression trade-offs . Decryption is implemented using XOR gates added into the neural network model and described as tanh ( x ) , which enable gradient calculations superior to the straight-through gradient method . We perform experiments using MNIST , CIFAR-10 , and ImageNet to show that inserting XOR gates learns quantization/encrypted bit decisions through training and obtains high accuracy even for fractional sub 1-bit weights . 1 INTRODUCTION . Deep Neural Networks ( DNNs ) demand larger number of parameters and more computations to support various task descriptions all while adhering to ever-increasing model accuracy requirements . Because of abundant redundancy in DNN models ( Han et al. , 2016 ; Frankle & Carbin , 2019 ; Courbariaux et al. , 2015 ) and limited scaling technology of integrated circuits due to physical challenges , numerous model compression techniques are being aggressively studied to expedite inference of DNNs across a range of device scales . As a practical model compression scheme , parameter quantization is a popular choice because compression ratio is high and regular formats after compression enable full memory bandwidth utilization that is crucial in highly parallel computing systems . Quantization schemes based on binary codes are gaining increasing attention since quantized weights follow specific constraints to allow simpler computations during inference . Specifically , using binary codes , a weight vector is represented as ∑q i=1 ( αi ∑v j=1 bi , j ) , where q is the number of quantization bits , v is the vector size , α is a scaling factor ( α ∈ R ) , and bi , j is a binary code ( bi , j ∈ { −1 , +1 } ) . Then , a dot product with activations is conducted as ∑q i=1 ( αi ∑v j=1 ajbi , j ) , where aj is a full-precision activation . Then , the number of multiplications is reduced from v to q ( expensive floating-point multipliers are less required for inference ) . Moreover , even though we do not discuss a new activation quantization method in this paper , if activations are also quantized by using binary codes , then most computations are replaced with bit-wise operations ( using XNOR logic and population counts ) ( Xu et al. , 2018 ; Rastegari et al. , 2016 ) . Consequently , even though representation space is constrained compared with quantization methods based on look-up tables , various inference accelerators can be designed to exploit advantages of binary codes ( Rastegari et al. , 2016 ; Xu et al. , 2018 ) . Since a successful 1-bit weight quantization method has been demonstrated in BinaryConnect ( Courbariaux et al. , 2015 ) , advances in compression-aware training algorithms in the form of binary codes ( e.g. , binary weight networks ( Rastegari et al. , 2016 ) and LQ-Nets ( Zhang et al. , 2018 ) ) produce 1-3 bits for quantization while accuracy drop is modest or negligible . Fundamental investigations on DNN training mechanisms using fewer quantization bits have also been actively reported to improve quantization quality ( Li et al. , 2017 ) . Previously , binary-coding-based quantization has only permitted integer numbers of quantization bits , limiting the compression/accuracy trade-off search space , especially in the range of very low quantization bits . In this paper , we propose a flexible encryption algorithm/architecture ( called “ FleXOR ” ) to enable fractional sub 1-bit numbers to represent each weight while quantizated bits are trained by gradient descent . Even though vector quntization is also a well-known scheme with high compression ratio ( Stock et al. , 2019 ) , we assume the form of binary codes . Note that the number of quantization bits can be different for each layer ( e.g. , Wang et al . ( 2019 ) ) to allow fractional quantization bits on average . FleXOR implies fractional quantization bits for each layer that can be quantized with different bits . The main contributions of this paper are as follows : XOR-based encryption of quantized bits enhances compression ratio : We encrypt quantized weights to further compress quantized DNNs applying basic principles from cryptography and digital communication . XOR-aware training algorithm learns encrypted weights : A binary model of XOR gates is used for forward propagation and model inference . We suggest calculating gradients of XOR gates using tanh functions instead of using a straight-through estimator ( STE ) ( Bengio et al. , 2013 ) for efficient compression-aware training . High model accuracy with sub 1-bit quantization : To the best of our knowledge , our work is the first to explore model accuracy under 1 bit/weight when weights are quantized based on binary codes . Such an exploration plays a significant role for mobile devices where it may be possible to sacrifice modest amounts of accuracy drop in exchange for improved energy efficiency . 2 ENCRYPTING QUANTIZED BITS USING XOR GATES . The main purpose of FleXOR is to compress quantized bits into encrypted bits that can be reconstructed by XOR gates as shown in Figure 1 . Suppose that Nout bits are to be compressed into Nin bits ( Nout > Nin ) . The role of an XOR-gate network is to produce various Nout-bit combinations using Nin bits . In other words , in order to maximize the chance of generating a desirable set of quantized bits , the encryption scheme is designed to seek a particular property where all possible 2Nin outcomes through decryption are evenly distributed in 2Nout space . A linear Boolean function , f ( x ) , maps f : { 0 , 1 } Nin → { 0 , 1 } and has the form of a1x1 ⊕ a2x2 ⊕ · · · ⊕ aNinxNin where aj ∈ { 0 , 1 } ( 1 ≤ j ≤ Nin ) and ⊕ indicates bitwise modulo-2 addition . In Figure 1 , six binary outputs are generated through six Boolean functions using four binary inputs . Let f1 ( x ) and f2 ( x ) be two such linear Boolean functions usingx = ( x1 , x2 , ... , xNin ) ∈ { 0 , 1 } Nin . The Hamming distance between f1 ( x ) and f2 ( x ) is the number of inputs on which f1 ( x ) and f2 ( x ) differ , and defined as dH ( f1 , f2 ) : = wH ( f1 ⊕ f2 ) = # { x ∈ { 0 , 1 } Nin |f1 ( x ) 6= f2 ( x ) } , ( 1 ) where wH ( f ) = # { x ∈ { 0 , 1 } Nin |f ( x ) = 1 } is the Hamming weight of a function and # { } corresponds to the size of a set ( Kahn et al. , 1988 ) . The Hamming distance is a well-known method to express non-linearity between two Boolean functions ( Kahn et al. , 1988 ) and increased Hamming distance between a pair of two Boolean functions results in a variety of outputs produced by XOR gates . Increasing Hamming distance is a required feature for cryptography to derive complicated encryption structure such that inverting encrypted data becomes difficult . For digital communication , the Hamming distance between encoded signals is closely related to the amount of error correction possible . In Figure 1 , y1 is represented as x1 ⊕ x3 ⊕ x4 , or equivalently a vector [ 1 0 1 1 ] denoting which inputs are selected . Concatenating such vectors , a XOR-gate network in Figure 1 can be described as a binary matrixM⊕ ∈ { 0 , 1 } Nout×Nin ( e.g. , the second row ofM⊕ is [ 1 1 0 0 ] and the third row is [ 1 1 1 0 ] ) . Then , decryption through XOR gates is simply represented as y =M⊕x where x and y are the binary inputs and binary outputs of XOR gates , and addition is ‘ XOR ’ and multiplication is ‘ AND ’ ( see Appendix for more details and examples ) . Encrypted weight bits are stored in 1-dimensional vector format and sliced into blocks of Nin-bit size as shown in Figure 2 . Then , decryption of each slice is performed by a XOR-gate network that is shared by all slices ( temporally- or spatially-shared ) . Depending on the quantization scheme and characteristics of layers , quantized bits may need to be scaled by a scaling factor and/or reshaped . 3 FLEXOR TRAINING ALGORITHM FOR QUANTIZATION BITS DECISION . Once the structure of XOR gates has been pre-determined and fixed to increase the Hamming distance of XOR outputs , we find quantized and encrypted bits by adding XOR gates into the model . In other words , we want an optimizer that understands the XOR-gate network structure so as to compute encrypted bits and scaling factors via gradient descent . Note that for inference , we store binary encrypted weights ( converted from real number encrypted weights ) in memory and generate binary quantized weights through Boolean XOR operations . Activation quantization is not discussed in this paper . Similar to the STE method introduced in ( Courbariaux et al. , 2015 ) , Boolean functions need to be described in a differentiable manner to obtain gradients in backward propagation . For two real number inputs x1 and x2 ( x1 , x2 ∈ R to be used as encrypted weights ) , the Boolean version of a XOR gate for forward propagation is described as ( note that 0 is replaced with −1 ) F⊕ ( x1 , x2 ) = ( −1 ) sign ( x1 ) sign ( x2 ) . ( 2 ) For inference , we store sign ( x1 ) and sign ( x2 ) instead of x1 and x2 . On the other hand , a differentiable XOR gate for backward propagation is presented as f⊕ ( x1 , x2 ) = ( −1 ) tanh ( x1 · Stanh ) tanh ( x2 · Stanh ) , ( 3 ) where Stanh is a scaling factor for FleXOR . Note that tanh functions are widely used to approximate Heaviside step functions ( i.e. , y ( x ) = 1 if x > 0 or 0 , otherwise ) in digital signal processing and Stanh can control the steepness ( as explained in Appendix ) . In the case of consecutive XOR operations , the order of inputs to be fed into XOR gates should not affect the computation of partial gradients for XOR inputs . Therefore , as a simple extension of Eq . ( 3 ) , a differentiable XOR gate network with n inputs can be described as f⊕ ( x1 , x2 , . . . , xn ) = ( −1 ) n−1 tanh ( x1 · Stanh ) tanh ( x2 · Stanh ) . . . tanh ( xn · Stanh ) . ( 4 ) Then , a partial derivative of f⊕ with respect to xi ( an encrypted weight ) is given as ∂f⊕ ( x1 , x2 , . . . , xn ) ∂xi = Stanh ( −1 ) n−1 ( 1− tanh2 ( xi · Stanh ) ) ∏n j=1 tanh ( xj · Stanh ) tanh ( xi · Stanh ) ( 5 ) Note that increasing the Hamming distance is associated with more tanh multiplications for each XOR-gate network output . Then , from Eq . ( 5 ) , we may suffer from the vanishing gradient problem since | tanh ( x ) | ≤ 1 . To resolve this problem , we also consider a simplified alternative partial derivative expressed as ∂f⊕ ( x1 , x2 , . . . , xn ) ∂xi ≈ Stanh ( −1 ) n−1 ( 1− tanh2 ( xi · Stanh ) ) ∏ j 6=i sign ( xj ) . ( 6 ) Eq . ( 6 ) shows that when we compute a partial derivative , all XOR inputs other than xi are assumed to be binary , i.e. , the magnitude of a partial derivative is then only determined by xi . We use Eq . ( 6 ) in this paper to calculate custom gradients of encrypted weights due to fast training computations and convergence , and use Eq . ( 2 ) for forward propagation . Algorithm 1 : Pseudo codes of Conv Layer with FleXOR1 we ∈ Rd ( k·k·Cin·Cout ) /Noute·Nin . Encrypted weights M⊕ ∈ { 0 , 1 } Nout×Nin . XOR gates ( shared ) α ∈ RCout . Scaling factors for each output channel Function FleXOR Conv ( input , stride , padding ) : for i← 0 to d ( k · k · Cin · Cout ) /Noute − 1 do for j ← 1 to Nout do wqi·Nout+j ← ( −1 ) · Nin∏ l=1 M⊕j , l=1 Signc ( wei·Nin+l ) · ( −1 ) . Eq . ( 2 ) Wq ← Reshape ( wq , [ k , k , Cin , Cout ] ) return Conv ( input , Wq , α , stride , padding ) . Convolution operation for binary codes Forward Function Signc ( x ) : return sign ( x ) Gradient Function Signc ( x , ∇ ) : return ∇ · ( 1− tanh2 ( x · Stanh ) ) · Stanh . Eq . ( 6 ) By training the whole network including FleXOR components with using custom gradient computation methods described above , encrypted and quantized weights are obtained in a holistic manner . FleXOR operations for convolutional layers are described in Algorithm 1 , where encrypted weights ( inputs of a XOR-gate network ) and quantized weights ( outputs of a XOR-gate network ) are we and Wq . We first verify basic training principles of FleXOR using LeNet-5 on the MNIST dataset . LeNet-5 consists of two convolutional layers and two fully-connected layers ( specifically , 32C5MP2-64C5-MP2-512FC-10SoftMax ) , and each layer is accompanied by an XOR-gate network with Nin binary inputs and Nout binary outputs . The quantization scheme follows 1-bit binary code with full-precision scaling factors that are shared across weights for the same output channel number ( for conv layers ) or output neurons ( for FC layers ) . Encrypted weights are randomly initialized with N ( µ=0 , σ2=0.0012 ) . All scaling factors for quantization are initialized to be 0.1 ( note that if batch normalization layers are immediately followed , then scaling factors for quantization are redundant ) . 1Kernel size is k × k , the number of input channel and output channel are Cin and Cout , respectively Using the Adam optimizer with an initial learning rate of 10-4 and batch size of 50 without dropout , Figure 3 shows training loss and test accuracy when Stanh=100 , elements of M⊕ are randomly filled with 1 or 0 , and for two values of Nout – 10 and 20 . Using the 1-bit internal quantization method and ( Nin , Nout ) encryption scheme , one weight can be represented by ( Nin/Nout ) bits . Hence , Figure 3 represents training results for 0.4 , 0.6 , and 0.8 bits per weight . Note that as for a randomly filled M⊕ , increasing Nout ( and Nin correspondingly for the same compression ratio ) increases the Hamming distance for a pair of any two rows of M⊕ and , hence , offers the chance to produce more diversified outputs . Indeed , as shown in Figure 3 , the results for Nout=20 present improved test accuracy and less variation compared withNout=10 . See Appendix for the distribution of encrypted weights at different training steps .

Summary: The paper proposes a new approach for quantizing neural networks. It looks at binary codes for quantization instead of the usual lookup tables because quantization approaches that rely on lookup tables have the drawback that during inference, the network has to fetch the weights from the lookup table and work with the full precision values, which means that the computation cost is remains the same as the non-quantized network. The paper presents FleXOR gates, a fast and efficient logic gate for mapping bit sequences to binary weights. The benefit of the gate is that the bit sequence can be shorter than the binary weights which means that the code length can be less than 1 bit per weight.